codex-app-server/cloud/sub2api.py

"""sub2api 网关接入：系统用管理 API 为每位用户建号发 key。

设计（2026-10-27 独立部署 sub2api 后）：
- 系统侧只用管理 API 管用户：管理员登录 -> 建 sub2api 账号 ->
  以该账号登录 -> 发 key -> key 存我方库，任务注入该用户的 key。
- key 创建是用户侧端点 POST /api/v1/keys（不是 /api/v1/api-keys，
  后者只有管理员改组）；管理员不能代用户建/删 key，所以 refresh_token
  也要存下来：以后轮换/吊销 key 时靠它换用户 access token。
- key 明文存 PG：与 github_app_config 的 App 私钥同一信任边界
  （库只在宿主 loopback 可达，服务端 env 同机）；at-rest 加密是
  后续项，不在本批假装已做。
- 零第三方依赖，只用标准库；HTTP 可注入 request 可调用对象，
  便于单测（见 cloud/tests/test_sub2api.py）。
"""

import collections
import json
import os
import random
import secrets
import threading
import time
import urllib.error
import urllib.request

GATEWAY_TABLE_SQL = (
    "CREATE TABLE IF NOT EXISTS user_gateway_keys("
    "user_id TEXT PRIMARY KEY, "
    "sub2api_user_id BIGINT NOT NULL, "
    "api_key TEXT NOT NULL, "
    "refresh_token TEXT, "
    "created_at DOUBLE PRECISION, "
    "updated_at DOUBLE PRECISION)")


class Sub2apiError(RuntimeError):
    def __init__(self, message, status=None):
        super().__init__(message)
        self.status = status


def _default_request(method, url, body=None, token=None):
    """urllib 实现的 request：返回 (status, parsed_json_dict)。"""
    headers = {"Content-Type": "application/json"}
    if token:
        headers["Authorization"] = "Bearer " + token
    req = urllib.request.Request(
        url, method=method,
        data=json.dumps(body).encode() if body is not None else None,
        headers=headers)
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            raw = resp.read() or b"{}"
            return resp.status, json.loads(raw)
    except urllib.error.HTTPError as e:
        try:
            data = json.loads(e.read() or b"{}")
        except ValueError:
            data = {}
        return e.code, data


def gateway_origin():
    """sub2api origin（不带 /v1）：未配则返回 None。"""
    raw = os.environ.get("SUB2API_BASE_URL", "").strip()
    return raw.rstrip("/") or None


def gateway_name():
    raw = os.environ.get("CLOUD_GATEWAY", "xcmapi").strip().lower()
    return raw if raw in ("xcmapi", "sub2api") else "xcmapi"


def gateway_group_name():
    return os.environ.get("SUB2API_GROUP_NAME", "codex-system").strip() \
        or "codex-system"


def opening_balance():
    """开通时授的初始余额（美元）：默认 10，显式值优先。

    新用户余额为 0 时网关鉴权层直接 403（余额门在任何上游路由
    之前），所以开通必须同时授余额，否则发出的 key 立刻不可用。
    显式配 0 表示不授（调用方自理）；非法/负值回落默认 10。
    """
    raw = os.environ.get(
        "CLOUD_SUB2API_OPENING_BALANCE_USD", "").strip()
    if not raw:
        return 10.0
    try:
        value = float(raw)
    except ValueError:
        return 10.0
    return value if value >= 0 else 10.0


def ensure_gateway_keys_table(conn):
    conn.execute(GATEWAY_TABLE_SQL)


def get_user_gateway_key(conn, user_id):
    """只读：返回该用户的网关行 dict 或 None，绝不触发开通。"""
    ensure_gateway_keys_table(conn)
    cur = conn.execute(
        "SELECT user_id, sub2api_user_id, api_key, refresh_token "
        "FROM user_gateway_keys WHERE user_id = ?",
        (user_id,))
    row = cur.fetchone()
    if row is None:
        return None
    try:
        keys = ("user_id", "sub2api_user_id", "api_key", "refresh_token")
        return dict(zip(keys, row))
    except (TypeError, ValueError):
        return None


# 进程内共享管理员会话：(origin, email) -> (token, refresh, expires_at)。
# 只复用 token 元组，不共用实例（实例不跨线程）；注册与首任务开通
# 每次新建实例都要登录一次（实测约 111ms，占开通链路约 1/3），
# 共享后同一进程内第二次起免登录。401/刷新失败会回落到重登录，
# 密码轮换后下一次登录即用新密码刷新缓存，不会卡死。
_SHARED_ADMIN = {}
_SHARED_ADMIN_LOCK = threading.RLock()


def _store_shared_session(admin):
    key = getattr(admin, "_shared_key", None)
    if key is None or not admin._token:
        return
    with _SHARED_ADMIN_LOCK:
        _SHARED_ADMIN[key] = (admin._token, admin._refresh,
                              admin._expires_at)


def _evict_shared_session(admin):
    key = getattr(admin, "_shared_key", None)
    if key is None:
        return
    with _SHARED_ADMIN_LOCK:
        _SHARED_ADMIN.pop(key, None)


def shared_admin_session(request=None):
    """带进程内 token 复用的管理员会话（生产开通路径默认用它）。

    缓存命中且未过期时直接播种 token，免一次管理员登录；
    未命中/已过期时在锁内登录一次再返回：突发冷缓存下
    N 个并发开通也只登录一次（其余线程等锁后直接播种），
    不会各登各的把网关刷出 N 次 bcrypt。调用方显式传
    admin 的单测不受影响（不走这条工厂）。
    """
    admin = Sub2apiAdmin.from_env(request=request)
    key = (admin.origin, admin.email)
    admin._shared_key = key
    with _SHARED_ADMIN_LOCK:
        cached = _SHARED_ADMIN.get(key)
        if cached and cached[2] > time.time():
            admin._token, admin._refresh, admin._expires_at = cached
            return admin
        # 缓存未命中：在锁内登录（_login 成功后经
        # _store_shared_session 写回同一缓存，RLock 可重入）。
        admin._login()
        return admin


def _is_rate_limited_error(exc):
    if getattr(exc, "status", None) == 429:
        return True
    return "too many requests" in str(exc).lower()


# 登录重试：网关登录端点 20 次/分钟/IP（固定窗口），突发注册
# 必然撞限；退避 2/4/8/16/32 秒（+0-2 秒抖动）可跨过 60 秒窗口。
_LOGIN_RETRY_ATTEMPTS = 6
_LOGIN_RETRY_BASE_SECONDS = 2.0


def _call_with_login_retry(fn):
    """登录类调用：先取进程内节流名额，再遇 429 退避重试。

    节流名额保证单进程内登录不超过预算（网关 20 次/分钟/IP，
    留 5 次余量给其它进程与人工登录）；重试兜底跨进程竞争。
    其它错误与耗尽次数后原错抛出。
    """
    _acquire_login_slot()
    delay = _LOGIN_RETRY_BASE_SECONDS
    for attempt in range(_LOGIN_RETRY_ATTEMPTS):
        try:
            return fn()
        except Sub2apiError as exc:
            if not _is_rate_limited_error(exc) \
                    or attempt == _LOGIN_RETRY_ATTEMPTS - 1:
                raise
            time.sleep(delay + random.uniform(0, 2))
            delay *= 2


# 进程内登录节流：60 秒窗口内最多 _LOGIN_BUDGET_PER_MINUTE 次。
# 只协调本进程；多进程突发仍靠 429 重试兜底（硬保证需共享
# Redis 节流，控制面目前没有 Redis 客户端，未做）。
_LOGIN_BUDGET_PER_MINUTE = 15
_LOGIN_WINDOW_SECONDS = 60.0
_login_slots = collections.deque()
_login_slots_lock = threading.Lock()


def _acquire_login_slot():
    """等一个登录名额：窗口满了就等最早一次滑出窗口再进。

    等待时长由排队位置决定（突发 100 路约 7 分钟）；注册是
    best-effort 语义，失败仍有 worker 首任务懒开通兜底，故不设
    提前放弃——放弃会把有序排队打回同步重试波。
    """
    while True:
        with _login_slots_lock:
            now = time.time()
            while _login_slots \
                    and _login_slots[0] <= now - _LOGIN_WINDOW_SECONDS:
                _login_slots.popleft()
            if len(_login_slots) < _LOGIN_BUDGET_PER_MINUTE:
                _login_slots.append(now)
                return
            wait = _login_slots[0] + _LOGIN_WINDOW_SECONDS - now
        time.sleep(max(wait, 0))


class Sub2apiAdmin:
    """管理员会话：登录、建/删用户、查组。token 自动刷新一次。

    每次实例只在进程内复用 token 元组（见 shared_admin_session）：
    实例本身不跨线程共用，只复用 (token, refresh, expires_at)，
    401 时走已有刷新->重登录链路自愈，密码轮换后下一次登录
    即用新密码并刷新缓存，不会卡在旧 token 上。
    """

    def __init__(self, origin, email, password, request=None):
        if not origin:
            raise Sub2apiError("SUB2API_BASE_URL 未配置")
        if not email or not password:
            raise Sub2apiError("SUB2API 管理员邮箱/密码未配置")
        self.origin = origin
        self.email = email
        self.password = password
        self.request = request or _default_request
        self._token = None
        self._refresh = None
        self._expires_at = 0
        self._shared_key = None

    @classmethod
    def from_env(cls, request=None):
        return cls(gateway_origin(),
                   os.environ.get("SUB2API_ADMIN_EMAIL", ""),
                   os.environ.get("SUB2API_ADMIN_PASSWORD", ""),
                   request=request)

    def _login(self):
        def _once():
            st, data = self.request(
                "POST", self.origin + "/api/v1/auth/login",
                {"email": self.email, "password": self.password})
            payload = (data or {}).get("data") or {}
            if st != 200 or not payload.get("access_token"):
                raise Sub2apiError(
                    "sub2api 管理员登录失败: %s"
                    % ((data or {}).get("message") or st), status=st)
            return payload

        try:
            payload = _call_with_login_retry(_once)
        except Sub2apiError:
            _evict_shared_session(self)
            raise
        self._token = payload["access_token"]
        self._refresh = payload.get("refresh_token")
        try:
            ttl = int(payload.get("expires_in") or 3600)
        except (TypeError, ValueError):
            ttl = 3600
        self._expires_at = time.time() + max(60, ttl - 60)
        _store_shared_session(self)

    def _refresh_token(self):
        if not self._refresh:
            self._login()
            return
        st, data = self.request(
            "POST", self.origin + "/api/v1/auth/refresh",
            {"refresh_token": self._refresh})
        payload = (data or {}).get("data") or {}
        if st == 200 and payload.get("access_token"):
            self._token = payload["access_token"]
            if payload.get("refresh_token"):
                self._refresh = payload["refresh_token"]
            try:
                ttl = int(payload.get("expires_in") or 3600)
            except (TypeError, ValueError):
                ttl = 3600
            self._expires_at = time.time() + max(60, ttl - 60)
            _store_shared_session(self)
            return
        self._login()

    def _call(self, method, path, body=None):
        if not self._token or time.time() >= self._expires_at:
            self._login()
        st, data = self.request(method, self.origin + path, body,
                                token=self._token)
        if st == 401:
            self._refresh_token()
            st, data = self.request(method, self.origin + path, body,
                                    token=self._token)
        return st, data

    def create_user(self, email, password, username):
        st, data = self._call("POST", "/api/v1/admin/users", {
            "email": email, "password": password,
            "username": username, "role": "user"})
        payload = (data or {}).get("data") or {}
        if st not in (200, 201) or payload.get("id") is None:
            raise Sub2apiError(
                "sub2api 建用户失败: %s"
                % ((data or {}).get("message") or st), status=st)
        return payload["id"]

    def delete_user(self, sub_user_id):
        st, _data = self._call(
            "DELETE", "/api/v1/admin/users/%s" % sub_user_id)
        return st in (200, 204)

    def set_balance(self, sub_user_id, amount):
        st, data = self._call(
            "POST", "/api/v1/admin/users/%s/balance" % sub_user_id,
            {"balance": amount, "operation": "set",
             "notes": "codex provisioning opening balance"})
        if st != 200:
            raise Sub2apiError(
                "sub2api 授余额失败: %s"
                % ((data or {}).get("message") or st), status=st)

    def list_groups(self):
        st, data = self._call("GET", "/api/v1/admin/groups")
        payload = (data or {}).get("data")
        if st != 200:
            raise Sub2apiError("sub2api 查组失败: %s" % st, status=st)
        if isinstance(payload, list):
            return payload
        if isinstance(payload, dict):
            for key in ("items", "list", "groups"):
                if isinstance(payload.get(key), list):
                    return payload[key]
        return []

    def resolve_group_id(self):
        raw = os.environ.get("SUB2API_GROUP_ID", "").strip()
        if raw:
            try:
                return int(raw)
            except ValueError:
                raise Sub2apiError("SUB2API_GROUP_ID 非法: %r" % raw)
        want = gateway_group_name()
        for grp in self.list_groups():
            if isinstance(grp, dict) and grp.get("name") == want:
                if grp.get("id") is not None:
                    return grp["id"]
        raise Sub2apiError("sub2api 组不存在: %s" % want)


def user_login(origin, email, password, request=None):
    """以 sub2api 用户身份登录：返回 (access_token, refresh_token)。

    遇网关登录限流（429）时退避重试，见 _call_with_login_retry。
    """
    req = request or _default_request

    def _once():
        st, data = req("POST", origin + "/api/v1/auth/login",
                       {"email": email, "password": password})
        payload = (data or {}).get("data") or {}
        if st != 200 or not payload.get("access_token"):
            raise Sub2apiError(
                "sub2api 用户登录失败: %s"
                % ((data or {}).get("message") or st), status=st)
        return payload["access_token"], payload.get("refresh_token")

    return _call_with_login_retry(_once)


def user_create_key(origin, user_token, name, group_id, request=None):
    """以用户身份发 key：返回 (key_id, key_value)。"""
    req = request or _default_request
    body = {"name": name}
    if group_id is not None:
        body["group_id"] = group_id
    st, data = req("POST", origin + "/api/v1/keys", body, token=user_token)
    payload = (data or {}).get("data") or {}
    value = payload.get("key") or payload.get("api_key") \
        or payload.get("value")
    if st not in (200, 201) or not value:
        raise Sub2apiError(
            "sub2api 发 key 失败: %s"
            % ((data or {}).get("message") or st), status=st)
    return payload.get("id"), value


def user_delete_key(origin, user_token, key_id, request=None):
    req = request or _default_request
    st, _data = req("DELETE", origin + "/api/v1/keys/%s" % key_id,
                    token=user_token)
    return st in (200, 204)


def user_refresh_access(origin, refresh_token, request=None):
    """refresh_token 换 access_token：返回 (新 access, 新 refresh|None)。"""
    req = request or _default_request
    st, data = req("POST", origin + "/api/v1/auth/refresh",
                   {"refresh_token": refresh_token})
    payload = (data or {}).get("data") or {}
    if st != 200 or not payload.get("access_token"):
        raise Sub2apiError("sub2api refresh 失败: %s" % st, status=st)
    return payload["access_token"], payload.get("refresh_token")


def list_gateway_models(origin, api_key, request=None):
    """网关 /v1/models：返回模型 id 列表（key 无模型时为空列表）。"""
    req = request or _default_request
    st, data = req("GET", origin + "/v1/models", token=api_key)
    if st != 200:
        raise Sub2apiError("sub2api 查模型失败: %s" % st, status=st)
    payload = (data or {}).get("data", data)
    if isinstance(payload, dict):
        payload = payload.get("data") or []
    out = []
    if isinstance(payload, list):
        for item in payload:
            if isinstance(item, dict) and item.get("id"):
                out.append(item["id"])
            elif isinstance(item, str):
                out.append(item)
    return out


def gateway_key_health(origin, api_key, request=None):
    """真健康检查：用已有 key 调一次 /v1/models。

    返回 True（200，可用）/ False（401/403：key 已失效，或余额
    为 0——网关在余额为 0 时鉴权层直接 403，两者都算死 key）/
    "unknown"（其它状态或网络异常：网关本身不可用时不据此轮换，
    避免把一次抖动当成 key 死亡、反复删号重建）。
    """
    req = request or _default_request
    try:
        st, _data = req("GET", origin + "/v1/models", token=api_key)
    except Exception:
        return "unknown"
    if st == 200:
        return True
    if st in (401, 403):
        return False
    return "unknown"


def _create_fresh_gateway_key(conn, user_id, admin=None, request=None):
    """建号->登录->发 key->授余额->入库，返回 (key, sub_id, key_id)。

    入库失败时尽力清理刚建的 key 与账号（避免孤儿 key），再抛原错；
    授余额失败同样走这条清理（余额为 0 的 key 在网关侧直接 403，
    留着没有意义）。同一用户并发开通时建号撞 email 已存在，
    则轮询复用胜出者入库的行。管理员工单默认走进程内共享会话
    （同一进程内第二次起免管理员登录，见 shared_admin_session）。
    调用方负责 conn 的提交语义：本函数在入库后
    commit（开通是独立事务，不跟调用方的业务事务绑在一起）。
    调用方须保证调用前本用户无有效行（ensure/rotate 先处理
    旧行），否则会主键冲突（并发撞号除外，走轮询复用）。
    """
    origin = gateway_origin()
    if not origin:
        raise Sub2apiError("SUB2API_BASE_URL 未配置，无法开通网关账号")
    admin = admin or shared_admin_session(request=request)
    req = request or _default_request
    email = "codex-%s@codex.local" % user_id
    username = "codex-" + str(user_id)[:12]
    password = secrets.token_hex(16)
    try:
        sub_id = admin.create_user(email, password, username)
    except Sub2apiError as exc:
        if "already exists" not in str(exc).lower():
            raise
        # 并发开通竞态：同一用户的另一路已建号（email 唯一），
        # 等它入库后复用那一行；等不到则原错抛出（任务失败后
        # 可重跑，孤儿号场景下重跑仍撞号需管理员介入清理）。
        for _ in range(5):
            time.sleep(2)
            row = get_user_gateway_key(conn, user_id)
            if row is not None:
                # 合并修复（P1 rebase 到 d254240）：main 的 #92 竞态
                # 分支原在单体 ensure 里只回一个 key 字符串；P1
                # 抽出 _create_fresh 后契约是 3 元组，4 个调用点
                # 都按 3 元组解包，裸字符串会被当成字符序列解包
                # 而 ValueError。这里按契约回 3 元组，key_id 未知
                # 故为 None（与 validate_or_rotate 的 ok 路一致）。
                return (row["api_key"], row["sub2api_user_id"], None)
        raise
    key_id = None
    user_token = None
    try:
        user_token, refresh = user_login(origin, email, password,
                                         request=req)
        group_id = admin.resolve_group_id()
        key_id, key = user_create_key(origin, user_token,
                                      "codex-task-key", group_id,
                                      request=req)
        amount = opening_balance()
        if amount > 0:
            admin.set_balance(sub_id, amount)
        now = time.time()
        conn.execute(
            "INSERT INTO user_gateway_keys(user_id, sub2api_user_id, "
            "api_key, refresh_token, created_at, updated_at) "
            "VALUES(?, ?, ?, ?, ?, ?)",
            (user_id, sub_id, key, refresh, now, now))
        conn.commit()
    except Exception:
        if key_id is not None and user_token is not None:
            try:
                user_delete_key(origin, user_token, key_id, request=req)
            except Exception:
                pass
        try:
            admin.delete_user(sub_id)
        except Exception:
            pass
        raise
    return key, sub_id, key_id


def ensure_user_gateway_key(conn, user_id, admin=None, request=None):
    """幂等开通：已有行直接返回 key；否则走 _create_fresh 开通。

    只保证“行存在”，不保证 key 仍有效：要校验有效性用
    validate_or_rotate_gateway_key（死 key 会触发轮换重开）。
    """
    ensure_gateway_keys_table(conn)
    row = get_user_gateway_key(conn, user_id)
    if row is not None:
        return row["api_key"]
    key, _sub_id, _key_id = _create_fresh_gateway_key(
        conn, user_id, admin=admin, request=request)
    return key


def validate_or_rotate_gateway_key(conn, user_id, admin=None, request=None):
    """懒开通/健康检查合一：返回 (status, key, sub_id, key_id)。

    status：ok（旧 key 经 /v1/models 实测可用）/ provisioned
    （无行，新开通）/ rotated（旧 key 经实测已死：401/403，
    含余额为 0 的 403；删旧行后重开，死 key 不再被当成已开通）。
    网关本身不可用（健康检查 unknown）时抛 Sub2apiError，调用方
    应回 502、保留旧行不轮换；绝不在此时把“待开通”当成可用。
    """
    ensure_gateway_keys_table(conn)
    row = get_user_gateway_key(conn, user_id)
    origin = gateway_origin()
    if not origin:
        raise Sub2apiError("SUB2API_BASE_URL 未配置，无法开通网关账号")
    if row is None or not row.get("api_key"):
        key, sub_id, key_id = _create_fresh_gateway_key(
            conn, user_id, admin=admin, request=request)
        return "provisioned", key, sub_id, key_id
    health = gateway_key_health(origin, row["api_key"], request=request)
    if health is True:
        return "ok", row["api_key"], row.get("sub2api_user_id"), None
    if health is False:
        conn.execute(
            "DELETE FROM user_gateway_keys WHERE user_id = ?",
            (user_id,))
        conn.commit()
        key, sub_id, key_id = _create_fresh_gateway_key(
            conn, user_id, admin=admin, request=request)
        return "rotated", key, sub_id, key_id
    raise Sub2apiError("sub2api 网关健康检查暂不可用，保留现有 key")


def rotate_user_gateway_key(conn, user_id, admin=None, request=None):
    """仅限本人调用的重新开通/轮换：删旧行后重开。

    无旧行时直接走开通（与 validate 的 provisioned 等价）；返回
    (status, key, sub_id, key_id)，调用方记审计时只可用 sub_id /
    key_id，绝不记 key 明文。旧 key 即使只是“疑似失效”也会被
    换掉：轮换是用户显式动作，不以健康检查为前提。
    """
    ensure_gateway_keys_table(conn)
    row = get_user_gateway_key(conn, user_id)
    if row is not None:
        conn.execute(
            "DELETE FROM user_gateway_keys WHERE user_id = ?",
            (user_id,))
        conn.commit()
    key, sub_id, key_id = _create_fresh_gateway_key(
        conn, user_id, admin=admin, request=request)
    return ("rotated" if row is not None else "provisioned",
            key, sub_id, key_id)