#!/usr/bin/env python3 """ ocg_router.py — OpenCode Go 多 key 路由代理 =========================================== OpenAI 兼容代理,背后池化 6 个 OpenCode Go API key。 基于 usage_stats.json 的剩余 quota 智能路由,402/429/5xx 自动故障切换。 端点: GET /v1/models — 合并所有 key 的模型列表 POST /v1/chat/completions — 转发(支持 stream) GET /health — 健康检查 GET /api/status — 代理状态(监控用) 架构: ┌─ OMO (Windows) ──────► ocg_router (:19878) ─► key5 (最空闲) ─► opencode.ai │ │ key6 (备用) │ │ key4 (备用) │ ▼ ... │ usage_stats.json └─ dashboard (:5803) ◄── /api/status """ import os, sys, json, time, logging, threading from http.server import ThreadingHTTPServer, BaseHTTPRequestHandler from urllib.request import Request, urlopen, HTTPError from urllib.error import URLError from datetime import datetime, timezone from pathlib import Path # ── 路径(平台自适应)────────────────────────────────────────── _SCRIPT_DIR = Path(__file__).resolve().parent # gateway/scripts/ _GATEWAY_DIR = _SCRIPT_DIR.parent # gateway/ _PROJECT_DIR = _GATEWAY_DIR.parent # AgentsMeeting/ sys.path.insert(0, str(_SCRIPT_DIR)) from proc_guard import guard # ── 配置 ──────────────────────────────────────────────────── LISTEN_HOST = "0.0.0.0" LISTEN_PORT = 19878 OPENCODE_BASE = "https://opencode.ai/zen/go/v1" USAGE_STATS_FILE = _GATEWAY_DIR / "temp" / "usage_stats.json" ACCOUNTS_FILE = _SCRIPT_DIR / "usage_monitor" / "accounts.json" # 故障切换:同一个 key 失败后冷却时间(秒) FAIL_COOLDOWN_SEC = 300 # 5 分钟(key 自身故障) UPSTREAM_COOLDOWN_SEC = 30 # 上游模型级故障(Router.Unavailable)短冷却:换 key 无意义 # 用量数据刷新间隔 USAGE_REFRESH_INTERVAL = 120 # 2 分钟 # 最大重试次数(所有 key 耗尽) MAX_RETRIES = 6 # ── 日志 ──────────────────────────────────────────────────── LOG_DIR = _GATEWAY_DIR / "logs" LOG_DIR.mkdir(parents=True, exist_ok=True) logging.basicConfig( level=logging.INFO, format="%(asctime)s [ocg_router] %(levelname)s: %(message)s", handlers=[ logging.FileHandler(str(LOG_DIR / "ocg_router.log"), encoding="utf-8"), logging.StreamHandler(), ], ) log = logging.getLogger("ocg_router") # ── 全局状态 ───────────────────────────────────────────────── _keys: list[dict] = [] # [{key_id, label, api_key, workspace_id}] _key_usage: dict[str, dict] = {} # key_id → {rolling_pct, weekly_pct, monthly_pct} _key_failures: dict[str, float] = {} # key_id → fail_until_timestamp _key_failure_reason: dict[str, str] = {} # key_id → 最近失败原因(冷却跳过时用于报错) _route_stats: dict[str, int] = {} # key_id → hit_count _total_hits = 0 _current_key_id = "" _state_lock = threading.Lock() # ═══════════════════════════════════════════════════════════════ # 用量数据加载 # ═══════════════════════════════════════════════════════════════ def load_accounts(): """从 accounts.json 读取所有 OCG key。enabled=false 的 key 直接跳过(不在路由池中)。""" keys = [] if not ACCOUNTS_FILE.exists(): log.error("accounts.json not found: %s", ACCOUNTS_FILE) return keys try: with open(str(ACCOUNTS_FILE), "r", encoding="utf-8") as f: data = json.load(f) for a in data.get("accounts", []): if a.get("provider", "").startswith("kimi"): continue # 跳过 Kimi kid = a.get("key_id", "") if not kid: continue # 显式停用的 key 不加入路由池 if a.get("enabled", True) is False: log.info("key %s disabled in accounts.json, skipped", kid) continue keys.append({ "key_id": kid, "label": a.get("label", kid), "api_key": a.get("key", ""), "workspace_id": a.get("workspace_id", ""), # 续订取消/显式标记 → 视为不可用(采集失败时也能正确跳过) "subscribed": a.get("subscribed", a.get("renewal") != "cancelled"), }) log.info("loaded %d OCG keys from accounts.json", len(keys)) except Exception as e: log.error("failed to load accounts.json: %s", e) return keys def load_usage(): """从 usage_stats.json 读取各 key 用量。""" usage = {} if not USAGE_STATS_FILE.exists(): log.warning("usage_stats.json not found: %s", USAGE_STATS_FILE) return usage try: with open(str(USAGE_STATS_FILE), "r", encoding="utf-8") as f: data = json.load(f) for a in data.get("accounts", []): kid = a.get("key_id", "") rolling = a.get("rolling") or {} weekly = a.get("weekly") or {} monthly = a.get("monthly") or {} has_error = bool(a.get("error")) sub_active = a.get("subscribed", True) usage[kid] = { "rolling_pct": rolling.get("usage_percent", 0) if rolling else 0, "weekly_pct": weekly.get("usage_percent", 0) if weekly else 0, "monthly_pct": monthly.get("usage_percent", 0) if monthly else 0, "error": has_error, "subscribed": sub_active, "last_update": a.get("last_update_iso", ""), } log.info("loaded usage for %d keys", len(usage)) except Exception as e: log.error("failed to load usage_stats.json: %s", e) return usage def refresh_state(): """刷新 key 池和用量数据。调用此函数前必须先获取 _state_lock。""" global _keys, _key_usage _keys = load_accounts() _key_usage = load_usage() # 合并 accounts.json 的订阅标记:采集失败时 usage 缺失该信息, # 以 accounts.json 的 renewal/subscribed 为准 for k in _keys: kid = k["key_id"] if not k.get("subscribed", True): _key_usage.setdefault(kid, {})["subscribed"] = False log.info("state refreshed: %d keys, %d with usage data", len(_keys), len(_key_usage)) # ═══════════════════════════════════════════════════════════════ # Key 选择(核心路由逻辑) # ═══════════════════════════════════════════════════════════════ def _key_health_score(kid): """ 计算 key 健康度(越低越好): - 已取消订阅的 key 排最末(score = 999) - 在冷却期内的 key 排倒数第二(score = 888) - 用量采集失败(error)仅降权(score = 500),key 仍可用 - 否则按 rolling% → weekly% → monthly% 加权 """ usage = _key_usage.get(kid, {}) # 订阅取消 → 不可用 if not usage.get("subscribed", True): return 999 # 用量采集失败 → 降权但可用(不因采集问题误杀 key) if usage.get("error"): return 500 # 故障冷却期内 → 低优先级 now = time.time() fail_until = _key_failures.get(kid, 0) if now < fail_until: return 888 # 正常 key:按用量打分(用量越低分越低=越优先) r = usage.get("rolling_pct", 0) or 0 w = usage.get("weekly_pct", 0) or 0 m = usage.get("monthly_pct", 0) or 0 return r * 0.6 + w * 0.3 + m * 0.1 # 滚动窗口占比最大 def pick_key(): """选择当前最空闲的 key。返回 key dict 或 None。""" if not _keys: return None # 按健康度排序 ranked = sorted(_keys, key=lambda k: _key_health_score(k["key_id"])) best = ranked[0] score = _key_health_score(best["key_id"]) if score >= 888: # 所有 key 都不可用 log.warning("ALL keys unavailable — best score=%.1f for %s", score, best["key_id"]) return best def mark_key_failed(kid, reason="", cooldown=FAIL_COOLDOWN_SEC): """标记 key 故障,进入冷却期。记录失败原因供冷却期报错。""" until = time.time() + cooldown with _state_lock: _key_failures[kid] = until if reason: _key_failure_reason[kid] = reason log.warning("key %s marked failed until %s (%ds): %s", kid, datetime.fromtimestamp(until).strftime("%H:%M:%S"), cooldown, reason) def record_route(kid): """记录一次成功路由。""" global _total_hits, _current_key_id with _state_lock: _route_stats[kid] = _route_stats.get(kid, 0) + 1 _total_hits += 1 _current_key_id = kid # ═══════════════════════════════════════════════════════════════ # HTTP 代理核心 # ═══════════════════════════════════════════════════════════════ class RouterHandler(BaseHTTPRequestHandler): """HTTP 请求路由处理器。""" def log_message(self, format, *args): pass # 禁用标准库日志 def do_GET(self): self._handle() def do_POST(self): self._handle() def do_OPTIONS(self): self.send_response(200) self.send_header("Access-Control-Allow-Origin", "*") self.send_header("Access-Control-Allow-Methods", "GET,POST,OPTIONS") self.send_header("Access-Control-Allow-Headers", "*") self.end_headers() def _handle(self): path = self.path.rstrip("/") or "/" # ── 健康检查 ── if path == "/health": self._json_response({"ok": True, "time": datetime.now(timezone.utc).isoformat()}) return # ── 代理状态 ── if path == "/api/status": self._serve_status() return # ── Key 启用/停用 ── if path == "/api/keys/toggle" and self.command == "POST": self._toggle_key() return # ── 立即刷新状态(accounts.json 被外部修改后调用)── if path == "/api/reload" and self.command == "POST": with _state_lock: refresh_state() self._json_response({"ok": True, "keys": len(_keys)}) return # ── 模型列表 ── if path == "/v1/models": self._serve_models() return # ── Chat Completions ── if path.endswith("/v1/chat/completions") or path == "/v1/chat/completions": self._proxy_chat() return # ── 其他请求尝试透传 ── self._proxy_any() # ── Key 启用/停用 ── def _toggle_key(self): """POST /api/keys/toggle {key_id, enabled} → 写 accounts.json 并立即 reload。""" try: content_length = int(self.headers.get("Content-Length", 0)) body = self.rfile.read(content_length) if content_length > 0 else b"" req_data = json.loads(body) if body else {} kid = req_data.get("key_id", "") enabled = bool(req_data.get("enabled", True)) if not kid: self._error_response(400, "key_id required") return if not ACCOUNTS_FILE.exists(): self._error_response(500, "accounts.json not found") return with open(str(ACCOUNTS_FILE), "r", encoding="utf-8") as f: data = json.load(f) found = False for a in data.get("accounts", []): if a.get("key_id") == kid: a["enabled"] = enabled found = True break if not found: self._error_response(404, f"key {kid} not found") return with open(str(ACCOUNTS_FILE), "w", encoding="utf-8") as f: json.dump(data, f, ensure_ascii=False, indent=2) # 立即 reload,使配置生效 with _state_lock: refresh_state() log.info("key %s enabled=%s (via API), pool now %d keys", kid, enabled, len(_keys)) self._json_response({"ok": True, "key_id": kid, "enabled": enabled, "pool_size": len(_keys)}) except Exception as e: log.error("toggle key failed: %s", e) self._error_response(500, str(e)) # ── Status API ── def _serve_status(self): with _state_lock: pool = [] for k in _keys: kid = k["key_id"] usage = _key_usage.get(kid, {}) fail_until = _key_failures.get(kid, 0) pool.append({ "key_id": kid, "label": k.get("label", kid), "rolling_pct": usage.get("rolling_pct", None), "weekly_pct": usage.get("weekly_pct", None), "monthly_pct": usage.get("monthly_pct", None), "healthy": usage.get("subscribed", True) and (time.time() >= fail_until), "in_cooldown": time.time() < fail_until, "fail_reason": _key_failure_reason.get(kid, ""), "error": usage.get("error", False), "hits": _route_stats.get(kid, 0), "last_update": usage.get("last_update", ""), }) status = { "ok": True, "uptime_seconds": time.time() - _start_time, "current_key": _current_key_id, "total_requests": _total_hits, "pool_size": len(_keys), "pool_healthy": sum(1 for p in pool if p["healthy"]), "pool": pool, "last_refresh": _last_refresh.isoformat() if _last_refresh else None, } self._json_response(status) # ── Models ── def _serve_models(self): """合并所有 key 的模型列表。""" seen = set() models = [] with _state_lock: keys_snapshot = list(_keys) for k in keys_snapshot: kid = k["key_id"] api_key = k.get("api_key", "") if not api_key: continue # 跳过不可用 key(采集 error 不跳过,仅降权) usage = _key_usage.get(kid, {}) if not usage.get("subscribed", True): continue if time.time() < _key_failures.get(kid, 0): continue try: req = Request( f"{OPENCODE_BASE}/models", headers={ "Authorization": f"Bearer {api_key}", "User-Agent": "Mozilla/5.0", }, method="GET", ) with urlopen(req, timeout=10) as resp: data = json.loads(resp.read().decode("utf-8")) for m in data.get("data", []): mid = m.get("id", "") if mid not in seen: seen.add(mid) models.append(m) break # 只问第一个健康的 key 就够了(模型列表相同) except Exception: continue if not models: # fallback:硬编码已知模型列表 models = [ {"id": "deepseek-v4-pro", "object": "model"}, {"id": "deepseek-v4-flash", "object": "model"}, {"id": "kimi-k2.7-code", "object": "model"}, {"id": "glm-5.2", "object": "model"}, {"id": "qwen3.7-plus", "object": "model"}, {"id": "qwen3.7-max", "object": "model"}, ] self._json_response({"object": "list", "data": models}) # ── Chat Completions Proxy ── def _proxy_chat(self): start = time.time() content_length = int(self.headers.get("Content-Length", 0)) body = self.rfile.read(content_length) if content_length > 0 else b"" is_stream = False model_name = "unknown" if body: try: req_data = json.loads(body) is_stream = req_data.get("stream", False) model_name = req_data.get("model", "unknown") log.info("chat request: model=%s stream=%s body_len=%d", model_name, is_stream, len(body)) except (json.JSONDecodeError, UnicodeDecodeError) as e: log.warning("failed to parse request body (len=%d): %s", len(body), e) # 尝试所有 key(按用量排序),直到成功 with _state_lock: usage_snapshot = dict(_key_usage) failures_snapshot = dict(_key_failures) ranked_keys = sorted( list(_keys), key=lambda k: _key_health_score(k["key_id"]) ) last_error = "" success = False for attempt, k in enumerate(ranked_keys): kid = k["key_id"] api_key = k.get("api_key", "") if not api_key: continue # 检查 key 是否可用(采集 error 不跳过,仅降权) usage = usage_snapshot.get(kid, {}) if not usage.get("subscribed", True): continue if time.time() < failures_snapshot.get(kid, 0): continue log.info("→ %s via %s (attempt %d/%d)", model_name, kid, attempt + 1, len(ranked_keys)) try: result = self._forward_request(api_key, body, is_stream) if result["ok"]: record_route(kid) elapsed = time.time() - start log.info("✓ %s via %s OK (%.2fs)", model_name, kid, elapsed) success = True break else: last_error = result.get("error", "unknown") log.warning("⚠ %s via %s FAILED: %s", model_name, kid, last_error) # 错误分类:只有真正的 key 级错误才冷却 key # - 400 客户端错误:不重试、不冷却(换 key 无意义),直接返回给调用方 # - 401/403 认证失败:key 级故障 → 冷却 300s # - 402/429 配额/限流:key 级 → 冷却 60s # - 5xx/Router.Unavailable:上游服务问题 → 不冷却 key(误杀元凶),立即失败 status = result.get("status", 0) if status == -1: # 客户端已断开:直接终止,不重试不冷却 log.info("client disconnected during %s via %s, aborting", model_name, kid) success = False break if status == 400: log.warning("client error 400, returning to caller without retry") success = False break if status in (401, 403): mark_key_failed(kid, last_error, cooldown=FAIL_COOLDOWN_SEC) elif status in (402, 429): mark_key_failed(kid, last_error, cooldown=60) elif status >= 500: # 上游服务端错误(Internal server error / Router.Unavailable): # 不杀 key。立即失败返回,避免"全部 key 被误杀进冷却"的雪崩。 log.warning("upstream %d (service-level) via %s — NOT cooling key, failing request", status, kid) success = False break except Exception as e: last_error = str(e) log.warning("⚠ %s via %s exception: %s", model_name, kid, e) mark_key_failed(kid, last_error) if not success: elapsed = time.time() - start if not last_error: # 所有 key 在冷却中被跳过 → 组装原因+最早到期时间 now = time.time() cooling = [(kid, until) for kid, until in failures_snapshot.items() if now < until] if cooling: earliest = min(until for _, until in cooling) reasons = {_key_failure_reason.get(kid, "unknown") for kid, _ in cooling} reason_str = "; ".join(sorted(r for r in reasons if r and r != "unknown")) last_error = (f"all keys in cooldown until " f"{datetime.fromtimestamp(earliest).strftime('%H:%M:%S')}" f"{' (' + reason_str + ')' if reason_str else ''}") log.error("✗ %s ALL KEYS EXHAUSTED (%.2fs): %s", model_name, elapsed, last_error) self._error_response(502, f"All OCG keys exhausted: {last_error}") def _forward_request(self, api_key, body, is_stream): """转发单次请求到指定 key。返回 {ok, status, error}。""" try: req = Request( f"{OPENCODE_BASE}/chat/completions", data=body, headers={ "Authorization": f"Bearer {api_key}", "Content-Type": "application/json", "User-Agent": "Mozilla/5.0", }, method="POST", ) with urlopen(req, timeout=180) as resp: status = resp.status if status >= 400: err_body = resp.read().decode("utf-8", errors="replace")[:500] return {"ok": False, "status": status, "error": f"HTTP {status}: {err_body}"} # 透传响应 resp_body = resp.read() self._send_response(status, dict(resp.headers), resp_body) return {"ok": True} except HTTPError as e: err_body = e.read().decode("utf-8", errors="replace")[:500] return {"ok": False, "status": e.code, "error": f"HTTP {e.code}: {err_body}"} except (BrokenPipeError, ConnectionResetError, ConnectionAbortedError) as e: # 客户端断开:不是 key 问题,标记特殊错误码让调用方不重试不冷却 return {"ok": False, "status": -1, "error": f"client_disconnected: {e}"} except URLError as e: return {"ok": False, "status": 0, "error": f"URLError: {e.reason}"} except Exception as e: return {"ok": False, "status": 0, "error": str(e)} # ── 通用代理(其他路径) ── def _proxy_any(self): with _state_lock: key = pick_key() if not key: self._error_response(503, "No available OCG keys") return api_key = key["api_key"] self._forward_to_upstream(api_key) def _forward_to_upstream(self, api_key): content_length = int(self.headers.get("Content-Length", 0)) body = self.rfile.read(content_length) if content_length > 0 else b"" method = self.command url = f"{OPENCODE_BASE}{self.path}" excluded = {"host", "connection", "keep-alive", "transfer-encoding", "content-length"} headers = { k: v for k, v in self.headers.items() if k.lower() not in excluded } headers["Authorization"] = f"Bearer {api_key}" headers["User-Agent"] = "Mozilla/5.0" try: req = Request(url, data=body or None, headers=headers, method=method) with urlopen(req, timeout=120) as resp: resp_body = resp.read() self._send_response(resp.status, dict(resp.headers), resp_body) except HTTPError as e: err_body = e.read().decode("utf-8", errors="replace")[:500] self._error_response(e.code, err_body) except Exception as e: self._error_response(502, str(e)) # ── 工具方法 ── def _json_response(self, data): body = json.dumps(data, ensure_ascii=False).encode("utf-8") self.send_response(200) self.send_header("Content-Type", "application/json") self.send_header("Access-Control-Allow-Origin", "*") self.send_header("Content-Length", str(len(body))) self.end_headers() self.wfile.write(body) def _error_response(self, code, message): data = json.dumps({"error": str(message)}, ensure_ascii=False).encode("utf-8") self.send_response(code) self.send_header("Content-Type", "application/json") self.send_header("Access-Control-Allow-Origin", "*") self.send_header("Content-Length", str(len(data))) self.end_headers() self.wfile.write(data) def _send_response(self, status, headers, body): self.send_response(status) allowed = { "content-type", "content-encoding", "cache-control", "x-request-id", "x-ratelimit-remaining", "x-ratelimit-reset", "access-control-allow-origin", } # 注意:不转发 transfer-encoding: chunked — 代理已缓存整个响应,不是 chunked for k, v in headers.items(): if k.lower() in allowed: self.send_header(k, v) self.send_header("Content-Length", str(len(body))) self.send_header("Access-Control-Allow-Origin", "*") self.end_headers() self.wfile.write(body) self.wfile.flush() # ═══════════════════════════════════════════════════════════════ # 后台线程:周期性刷新状态 # ═══════════════════════════════════════════════════════════════ _start_time = time.time() _last_refresh: datetime | None = None def _bg_refresh(): """后台线程:每 USAGE_REFRESH_INTERVAL 秒刷新用量数据。""" global _last_refresh while True: time.sleep(USAGE_REFRESH_INTERVAL) try: with _state_lock: refresh_state() _last_refresh = datetime.now(timezone.utc) # 清空已过冷却期的 key now = time.time() with _state_lock: expired = [k for k, v in _key_failures.items() if now >= v] for k in expired: del _key_failures[k] _key_failure_reason.pop(k, None) log.info("key %s cooldown expired, restored", k) log.info("bg refresh: %d keys, healthy pool", len(_keys)) except Exception as e: log.error("bg refresh failed: %s", e) # ═══════════════════════════════════════════════════════════════ # 启动 # ═══════════════════════════════════════════════════════════════ def main(): # PID 锁 g = guard("ocg_router") if not g.ok: log.error("ocg_router 已有实例在运行 (PID %s),退出", g.pid) sys.exit(1) # 初始加载 with _state_lock: refresh_state() global _last_refresh _last_refresh = datetime.now(timezone.utc) # 启动后台刷新线程 refresh_thread = threading.Thread(target=_bg_refresh, name="ocg_router_refresh", daemon=True) refresh_thread.start() # 启动 HTTP 服务(ThreadingHTTPServer — 多线程,避免大请求阻塞 /api/status 等监控端点) # daemon_threads=True:客户端断开/线程卡住不会阻塞服务退出 server = ThreadingHTTPServer((LISTEN_HOST, LISTEN_PORT), RouterHandler) server.daemon_threads = True # 限制最大并发连接数(每个连接一个线程;信号量在 handler 层控制) MAX_CONCURRENT = 16 log.info("ocg_router 启动 → http://%s:%d (max_parallel=%d)", LISTEN_HOST, LISTEN_PORT, MAX_CONCURRENT) log.info("keys loaded: %d (%d healthy)", len(_keys), sum(1 for k in _keys if _key_health_score(k["key_id"]) < 888)) log.info("fail cooldown: %ds, usage refresh: %ds", FAIL_COOLDOWN_SEC, USAGE_REFRESH_INTERVAL) try: server.serve_forever() except KeyboardInterrupt: log.info("收到中断信号,关闭...") server.shutdown() if __name__ == "__main__": main()