proxy.py

#!/usr/bin/env python3
"""OpenAI-compatible proxy for AtomCode CodingPlan.

Features:
- OpenAI-compatible /v1/models and /v1/chat/completions
- model alias mapping for DeepSeek, Qwen3-VL and GLM-5.2
- token loading from env, ~/.atomcode/auth.toml or ~/.codingplan/token.json
- JSON refresh-token flow
- /health, /status, /usage and /claim helpers
- upstream 429 retry with Retry-After support
"""

import argparse
import http.server
import json
import os
import ssl
import sys
import time
import urllib.error
import urllib.request


UPSTREAM = "https://api.gitcode.com/api/v5"
OAUTH_REFRESH = "https://acs.atomgit.com/oauth/refresh"
UA = "atomcode/4.25.7"
TOKEN_JSON = os.path.expanduser("~/.codingplan/token.json")
TOKEN_TXT = os.path.expanduser("~/.codingplan/token.txt")
ATOMCODE_AUTH = os.path.expanduser("~/.atomcode/auth.toml")

MODEL_MAP = {
    "deepseek-v4-flash": "deepseek-ai/DeepSeek-V4-Flash",
    "deepseek-v4": "deepseek-ai/DeepSeek-V4-Flash",
    "deepseek-r1": "deepseek-ai/DeepSeek-R1",
    "qwen-vl": "Qwen/Qwen3-VL-8B-Instruct",
    "qwen3-vl": "Qwen/Qwen3-VL-8B-Instruct",
    "qwen3-vl-8b": "Qwen/Qwen3-VL-8B-Instruct",
    "Qwen/Qwen3-VL-8B-Instruct": "Qwen/Qwen3-VL-8B-Instruct",
    "glm-5": "GLM-5.2",
    "GLM-5.2": "GLM-5.2",
}

TOKEN_CACHE = {"token": "", "refresh": "", "expires_at": 0, "ts": 0}


def read_atomcode_auth():
    data = {}
    try:
        with open(ATOMCODE_AUTH, encoding="utf-8") as f:
            for line in f:
                if "=" not in line:
                    continue
                key, value = line.split("=", 1)
                key = key.strip()
                value = value.strip().strip('"')
                if key in ("access_token", "refresh_token", "expires_in", "created_at"):
                    data[key] = value
    except OSError:
        pass
    return data


def read_token_json():
    try:
        with open(TOKEN_JSON, encoding="utf-8") as f:
            data = json.load(f)
            return {
                "token": data.get("access_token") or data.get("token") or "",
                "refresh": data.get("refresh_token") or "",
                "expires_at": int(data.get("expires_at") or 0),
            }
    except Exception:
        return {"token": "", "refresh": "", "expires_at": 0}


def write_token_json(token, refresh="", expires_at=0):
    os.makedirs(os.path.dirname(TOKEN_JSON), exist_ok=True)
    data = {"access_token": token, "refresh_token": refresh, "expires_at": expires_at}
    with open(TOKEN_JSON, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)


def load_tokens(force=False):
    now = time.time()
    if not force and TOKEN_CACHE["token"] and now - TOKEN_CACHE["ts"] < 60:
        return TOKEN_CACHE.copy()

    token = os.environ.get("CODINGPLAN_TOKEN", "")
    refresh = os.environ.get("CODINGPLAN_REFRESH", "")
    expires_at = int(os.environ.get("CODINGPLAN_EXPIRES_AT", "0") or 0)

    if not token:
        saved = read_token_json()
        token = saved["token"]
        refresh = refresh or saved["refresh"]
        expires_at = expires_at or saved["expires_at"]

    if not token:
        atom = read_atomcode_auth()
        token = atom.get("access_token", "")
        refresh = refresh or atom.get("refresh_token", "")
        try:
            created = int(atom.get("created_at", "0") or 0)
            expires = int(atom.get("expires_in", "0") or 0)
            expires_at = expires_at or created + expires
        except ValueError:
            pass

    if not token:
        try:
            with open(TOKEN_TXT, encoding="utf-8") as f:
                token = f.read().strip()
        except OSError:
            pass

    TOKEN_CACHE.update({"token": token, "refresh": refresh, "expires_at": expires_at, "ts": now})
    return TOKEN_CACHE.copy()


def refresh_tokens_if_needed(force=False):
    tokens = load_tokens(force=True)
    token = tokens["token"]
    refresh = tokens["refresh"]
    expires_at = int(tokens["expires_at"] or 0)
    if not token or not refresh:
        return tokens
    if not force and expires_at and time.time() < expires_at - 24 * 3600:
        return tokens

    payload = json.dumps({"refresh_token": refresh}).encode()
    req = urllib.request.Request(OAUTH_REFRESH, data=payload, method="POST")
    req.add_header("User-Agent", UA)
    req.add_header("Content-Type", "application/json")
    ctx = ssl.create_default_context()
    try:
        with urllib.request.urlopen(req, timeout=20, context=ctx) as resp:
            data = json.loads(resp.read())
        new_token = data.get("access_token") or token
        new_refresh = data.get("refresh_token") or refresh
        new_expires_at = int(time.time()) + int(data.get("expires_in") or 604800)
        write_token_json(new_token, new_refresh, new_expires_at)
        TOKEN_CACHE.update({"token": new_token, "refresh": new_refresh, "expires_at": new_expires_at, "ts": time.time()})
        print("[token] refreshed", file=sys.stderr, flush=True)
    except Exception as exc:
        print(f"[token] refresh failed: {exc}", file=sys.stderr, flush=True)
    return TOKEN_CACHE.copy()


def remap_model(name):
    mapped = MODEL_MAP.get(name, name)
    if mapped != name:
        print(f"[model] {name} -> {mapped}", file=sys.stderr, flush=True)
    return mapped


def upstream_request(path, method="GET", body=None, token=None, accept="application/json", timeout=120):
    if token is None:
        token = refresh_tokens_if_needed()["token"]
    data = json.dumps(body).encode() if body is not None else None
    req = urllib.request.Request(f"{UPSTREAM}{path}", data=data, method=method)
    req.add_header("Authorization", f"Bearer {token}")
    req.add_header("User-Agent", UA)
    req.add_header("Accept", accept)
    if body is not None:
        req.add_header("Content-Type", "application/json")
    return urllib.request.urlopen(req, timeout=timeout, context=ssl.create_default_context())


def summarize_status(status):
    usage = status.get("current_usage") or {}
    windows = status.get("rate_limit_windows") or []
    active_window = windows[0] if windows else {}
    return {
        "plan_name": (status.get("codingplan_free") or {}).get("plan_name") or status.get("plan_type"),
        "plan_type": status.get("plan_type"),
        "expires_at": status.get("expires_at") or (status.get("codingplan_free") or {}).get("expires_at"),
        "window_hours": usage.get("window_hours") or active_window.get("window_hours"),
        "window_limit": usage.get("window_token_limit") or active_window.get("call_limit"),
        "used": usage.get("window_tokens_used") or active_window.get("calls_used"),
        "usage_percent": usage.get("usage_percent") or active_window.get("usage_percent"),
        "reset_at": usage.get("reset_at_display") or active_window.get("reset_at_display"),
        "quota_exhausted": status.get("window_quota_exhausted") or active_window.get("quota_exhausted"),
    }


class ProxyHandler(http.server.BaseHTTPRequestHandler):
    def log_message(self, fmt, *args):
        print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] {fmt % args}", file=sys.stderr)

    def send_raw(self, code, data=b"", headers=None):
        self.send_response(code)
        self.send_header("Access-Control-Allow-Origin", "*")
        if headers:
            for key, value in headers.items():
                self.send_header(key, value)
        self.end_headers()
        if data:
            self.wfile.write(data)

    def send_json(self, code, obj):
        self.send_raw(code, json.dumps(obj, ensure_ascii=False).encode(), {"Content-Type": "application/json"})

    def do_OPTIONS(self):
        self.send_raw(200, headers={"Access-Control-Allow-Methods": "GET, POST, OPTIONS", "Access-Control-Allow-Headers": "*"})

    def do_GET(self):
        if self.path == "/v1/models":
            models = [{"id": key, "object": "model", "created": 0, "owned_by": "codingplan"} for key in MODEL_MAP]
            self.send_json(200, {"object": "list", "data": models})
            return
        if self.path in ("/", "/health"):
            tokens = load_tokens()
            self.send_json(200, {"status": "ok", "upstream": UPSTREAM, "authenticated": bool(tokens["token"]), "token_expires_at": tokens["expires_at"]})
            return
        if self.path in ("/status", "/usage"):
            self.handle_status()
            return
        self.send_json(404, {"error": "not found"})

    def do_POST(self):
        length = int(self.headers.get("Content-Length", 0))
        raw = self.rfile.read(length) if length else b"{}"
        try:
            body = json.loads(raw)
        except json.JSONDecodeError:
            self.send_json(400, {"error": "invalid JSON"})
            return
        if self.path == "/v1/chat/completions":
            self.handle_chat(body)
            return
        if self.path == "/claim":
            self.handle_claim(body)
            return
        if self.path == "/refresh":
            tokens = refresh_tokens_if_needed(force=True)
            self.send_json(200, {"refreshed": bool(tokens["token"]), "token_expires_at": tokens["expires_at"]})
            return
        self.send_json(404, {"error": f"unknown: {self.path}"})

    def handle_status(self):
        token = refresh_tokens_if_needed()["token"]
        if not token:
            self.send_json(401, {"error": "no auth token"})
            return
        try:
            with upstream_request("/coding-plan/status", token=token, timeout=15) as resp:
                status = json.loads(resp.read())
            self.send_json(200, {"summary": summarize_status(status), "raw": status})
        except urllib.error.HTTPError as exc:
            self.send_json(exc.code, parse_http_error(exc))
        except Exception as exc:
            self.send_json(502, {"error": {"message": str(exc), "type": "connection_error"}})

    def handle_claim(self, body):
        token = refresh_tokens_if_needed()["token"]
        if not token:
            self.send_json(401, {"error": "no auth token"})
            return
        plan = body.get("plan_type") or body.get("plan") or "Pro"
        try:
            with upstream_request("/coding-plan/claim-v2", method="POST", body={"plan_type": plan}, token=token, timeout=20) as resp:
                self.send_raw(200, resp.read(), {"Content-Type": "application/json"})
        except urllib.error.HTTPError as exc:
            self.send_json(exc.code, parse_http_error(exc))
        except Exception as exc:
            self.send_json(502, {"error": {"message": str(exc), "type": "connection_error"}})

    def handle_chat(self, body):
        token = refresh_tokens_if_needed()["token"]
        if not token:
            self.send_json(401, {"error": "no auth token"})
            return
        body = dict(body)
        body["model"] = remap_model(body.get("model", ""))
        stream = bool(body.get("stream"))
        retries = int(os.environ.get("CODINGPLAN_RETRIES", "3"))
        last_error = None
        for attempt in range(retries + 1):
            try:
                with upstream_request("/chat/completions", method="POST", body=body, token=token, accept="text/event-stream" if stream else "application/json") as resp:
                    if stream:
                        self.send_response(200)
                        self.send_header("Content-Type", "text/event-stream")
                        self.send_header("Access-Control-Allow-Origin", "*")
                        self.send_header("Cache-Control", "no-cache")
                        self.end_headers()
                        for line in resp:
                            self.wfile.write(line)
                            self.wfile.flush()
                    else:
                        self.send_raw(200, resp.read(), {"Content-Type": "application/json"})
                    return
            except urllib.error.HTTPError as exc:
                if exc.code == 429 and attempt < retries:
                    wait = retry_wait_seconds(exc, attempt)
                    print(f"[429] retry in {wait:.1f}s", file=sys.stderr, flush=True)
                    time.sleep(wait)
                    continue
                self.send_json(exc.code if exc.code else 500, parse_http_error(exc))
                return
            except Exception as exc:
                last_error = exc
                if attempt < retries:
                    time.sleep(min(2 ** attempt, 8))
                    continue
        self.send_json(502, {"error": {"message": str(last_error), "type": "connection_error"}})


def retry_wait_seconds(exc, attempt):
    retry_after = exc.headers.get("Retry-After") if exc.headers else None
    try:
        return max(0.2, float(retry_after))
    except (TypeError, ValueError):
        return min(2 ** attempt, 8)


def parse_http_error(exc):
    raw = exc.read().decode(errors="replace")
    try:
        data = json.loads(raw)
    except Exception:
        data = {"message": raw[:500]}
    return {"error": {"message": extract_message(data), "type": "upstream_error", "code": exc.code, "raw": data}}


def extract_message(data):
    try:
        return data.get("choices", [{}])[0].get("message", {}).get("content") or data.get("error_message") or data.get("message") or str(data)[:500]
    except Exception:
        return str(data)[:500]


def main():
    parser = argparse.ArgumentParser(description="CodingPlan OpenAI-compatible proxy")
    parser.add_argument("--port", type=int, default=int(os.environ.get("PROXY_PORT", "18999")))
    parser.add_argument("--host", default=os.environ.get("PROXY_HOST", "127.0.0.1"))
    args = parser.parse_args()
    tokens = load_tokens(force=True)
    print(f"proxy: http://{args.host}:{args.port}/v1", file=sys.stderr)
    print(f"auth: {'OK' if tokens['token'] else 'MISSING'}", file=sys.stderr)
    print(f"models: {', '.join(sorted(MODEL_MAP))}", file=sys.stderr)
    http.server.ThreadingHTTPServer((args.host, args.port), ProxyHandler).serve_forever()


if __name__ == "__main__":
    main()