feat(proxy): WebSocket connection pool fixes prompt-cache hit-rate jitter (#440)

* feat(proxy): add WS connection pool module (no wiring yet)

新增 PersistentWs + WsConnectionPool + WsReusedConnectionError，准备解决
上游 WS gateway 按连接 ID 哈希路由导致的 prompt cache 命中率抖动
（5%~99% bimodal）。本 commit 纯模块新增 + 26 个单测，未接线，零回归风险。

设计要点：
- Pool key = `${entryId}:${conversationId}`，覆盖显式 + 隐式续链
- 单 WS strict 串行（codex 协议要求），busy 时 acquire 返回 bypass 让 caller
  走旧路径（不排队，避免死锁）
- 无 idle TTL；max_age=55min 留 5min 缓冲（server 60min 硬限）
- 死连/abort/账号状态变化级联清理
- 复用失败（pre-response）抛 WsReusedConnectionError，caller 可单次 retry；
  流中段失败走 controller.error 不 retry（client 已收到部分数据）

config schema 新增 ws_pool: { enabled: true, max_age_ms, max_per_account }

* feat(proxy): wire WS connection pool through ws-transport + proxy-handler

接线 ws-pool 到主请求链路。

ws-transport.ts:
- 抽出 openOneShotWs 保持原有 one-shot 语义（向后兼容，旧调用方零变更）
- createWebSocketResponse 增加 poolCtx? 参数；带 ctx 时先尝试 pool.acquire
  → 命中复用 PersistentWs.send；遇 WsReusedConnectionError 单次回退 one-shot
- pool 自身故障（factory 抛错）也会 fallback 到 one-shot，不污染调用方

codex-api.ts:
- createResponse / createResponseViaWebSocket 透传 poolCtx 到 createWebSocketResponse
- HTTP 路径完全不受影响

proxy-handler.ts:
- buildPoolCtx() 根据 useWebSocket + chainConversationId 生成 poolKey
  = `${entryId}:${chainConversationId}`，仅 WS 路径生效
- 主流程 + handleNonStreaming 的 empty-response retry 都用同一个 builder

account-pool.ts:
- markStatus(non-active) / markRateLimited / removeAccount / updateToken（refresh
  完成）级联调 evictByEntryId 关闭该账号所有池中 WS。理由：refresh 后老 WS
  携带的 access_token 已失效；其他状态变化下账号本身不可用，留池只浪费
- 用 dynamic import 隔离 ws-pool，避免 account-pool 单测必须拉 proxy 层

src/index.ts: SIGINT/SIGTERM shutdown 钩子追加 wsPool.shutdown() 优雅关闭

集成测：tests/integration/ws-pool-reuse.test.ts 起本地 ws.Server 验证
- 5 turn 同 conv → server 仅 1 次 connection
- 不同 conv → 各自 1 个连接
- server 主动 close → 池立即驱逐，下次 acquire 新建
- enabled:false → 退化为 one-shot（每 turn 新连接）
- evictByEntryId → 池清空 + 后续重建
- 不传 poolCtx → 行为完全等价于今天

测试：1702 passed (+1 skipped)，零回归

* feat(proxy): observability for WS pool decisions + CHANGELOG

ws-transport: 新增 WsDispatchDecision 类型 + WsPoolContext.onDecision 回调，
四种决定 (reuse / new / bypass:<reason> / retry-after-stale-reuse) 在 dispatch
时刻一次性 emit 给 caller。

proxy-handler: buildPoolCtx 装上 onDecision listener，对每个 WS 请求 emit
一行 `[fmt] Account E | rid=R | ws=reuse:abc` 之类的日志，便于直接观察
池命中率 + 抖动归因。

CHANGELOG: Unreleased → Fixed 加完整条目，写明问题、原因、配置、回滚。

* fix(proxy): wire ws_pool config into singleton (self-review of #440)

PR #440 加了 config schema 字段 ws_pool 但忘了真正读取它 —— `getWsPool()`
永远用 DEFAULT_WS_POOL_CONFIG，用户改 `ws_pool.enabled: false` 完全无效，
回滚策略破坏。

修复：startServer() 在加载 cfg 后立即调 setWsPoolConfig() 用用户配置替换
单例。同时把 shutdown 钩子的 dynamic import 一起改成 static（同模块复用）。

新增 3 个测试覆盖单例 wiring 链路：
- enabled:false → acquire 永远 bypass(disabled)，factory 不调
- 默认 getWsPool() 工作正常
- setWsPoolConfig 后调覆盖前调

测试：1705 passed (+1 skipped)

---------

Co-authored-by: icebear0828 <icebear0828@users.noreply.github.com>

Author: icebear0828

CHANGELOG.md

@@ -15,6 +15,18 @@
 - Usage history `five_min` granularity（5 分钟桶）+ Dashboard 新增「5 min」粒度选项与「Last 1h / 6h」时间窗：snapshot 默认 5 分钟一记，新粒度等同于一桶一快照，方便排查刚发生的请求；旧的 hourly/daily 不变，按 granularity 自动收敛兼容窗口（`src/auth/usage-stats.ts`、`src/routes/admin/usage-stats.ts`、`shared/hooks/use-usage-stats.ts`、`web/src/pages/UsageStats.tsx`）
 - 共享纯函数 `formatHitRate` / `sumWindow` / `formatUsageNumber` 抽到 `shared/utils/usage-stats.ts`，配套 vitest 单测覆盖边界（input=0 → "—"、<0.01% 截断、windowed 求和等），UsageChart 与 UsageStats 复用同一份格式化逻辑（`shared/utils/usage-stats.ts`、`shared/utils/__tests__/usage-stats.test.ts`）
 
+### Fixed
+
+- **WebSocket 连接池**（`src/proxy/ws-pool.ts` + `src/proxy/ws-transport.ts` + `src/routes/shared/proxy-handler.ts`）：上游 chatgpt.com 的 WS gateway 按"连接 ID"做负载均衡 hash，过去 codex-proxy 对每个 WS 请求都 `new WebSocket(url)`，导致同一会话同一账号的 prompt cache 命中率在 5%~99% 之间剧烈抖动（同一逻辑会话被路由到不同 backend，每个 backend 各自缓存了不同长度的前缀；实测 cached_tokens 反复出现 1920/2432/24448/40320/47488 等离散"checkpoint"）。引入 per-`(entryId, conversationId)` 的持久 WS 连接池：
+  - 单 WS 上 strict request/response 串行（codex 协议要求），busy 时旁路开新一次性 WS 而非排队（避免死锁）
+  - 无 idle TTL，连接保持开放直到自然死亡 / `max_age_ms`（默认 55 min，留 5 min 缓冲，比 server 60 min 硬限制提前关）/ 账号状态变化（`evictByEntryId`）级联清理
+  - 复用失败（pre-response close）抛 `WsReusedConnectionError`，自动单次 fallback 到一次性新连接；流中段失败保持原语义抛给客户端（不重试，client 已收到部分数据）
+  - account-pool 在 `markRateLimited` / `markStatus(non-active)` / `removeAccount` / `updateToken`（refresh 完成）时级联 `evictByEntryId`，避免老 WS 携带的 access_token 被复用
+  - 新增配置 `ws_pool: { enabled: true, max_age_ms: 3300000, max_per_account: 8 }`；可 `enabled: false` + 重启回滚到旧行为
+  - SIGTERM/SIGINT 进程退出钩子追加 `wsPool.shutdown()` 优雅关闭所有池中连接
+  - 入口日志加 `ws=reuse:<id>` / `ws=new:<id>` / `ws=bypass(<reason>)` / `ws=retry-after-stale-reuse:<id>` 字段，配合 `rid` 可对照 cache 命中率
+  - 集成测：`tests/integration/ws-pool-reuse.test.ts` 起本地 `ws.Server` 验证 5 turn 同会话只触发 1 次 `connection`
+
 ### Changed
 
 - `src/routes/shared/proxy-handler.ts` 入口与 Usage 日志补充诊断字段：入口行新增 `rid` / `conv` / `key` / `prev=<src>:<tail8>` / `tools=N` / `resume=on|off:<reason>`（reason 含 `no_pref_entry`/`acct_mismatch`/`instr_diff`/`missing_tool_calls`/`cont_start_eq_len`），Usage 行带 `rid` 与 `hit=X.X%`，便于对照 prompt-cache 命中率为何偏低、或同一会话请求是否落到同一 cache key

src/auth/account-pool.ts

@@ -114,11 +114,30 @@ export class AccountPool {
 
   removeAccount(id: string): boolean {
     this.lifecycle.clearLock(id);
+    this.evictWsPool(id);
     return this.registry.removeAccount(id);
   }
 
   updateToken(entryId: string, newToken: string, refreshToken?: string): void {
     this.registry.updateToken(entryId, newToken, refreshToken);
+    // The new access_token doesn't take effect on already-open WebSocket
+    // sessions (the upstream auth header is captured at handshake), so any
+    // pooled WS for this entry is now using a stale credential. Evict so the
+    // next request opens a fresh WS with the refreshed token.
+    this.evictWsPool(entryId);
+  }
+
+  /** Drop any pooled WebSocket connections for `entryId`. Used by status
+   *  mutations and token refresh to prevent in-flight reuse from carrying
+   *  stale auth or routing into a backend the account is no longer welcome
+   *  on. Lazy-imports ws-pool so this module doesn't pull the proxy layer
+   *  into bootstrap when the pool isn't otherwise reachable. */
+  private evictWsPool(entryId: string): void {
+    // Avoid hard import: account-pool is also exercised in unit tests that
+    // never touch the WS layer, and dynamic resolution keeps that contract.
+    void import("../proxy/ws-pool.js")
+      .then((mod) => mod.getWsPool().evictByEntryId(entryId))
+      .catch(() => { /* pool unavailable in this build/test context — ignore */ });
   }
 
   setLabel(entryId: string, label: string | null): boolean {
@@ -135,6 +154,10 @@ export class AccountPool {
   markStatus(entryId: string, status: AccountEntry["status"]): void {
     if (this.registry.markStatus(entryId, status)) {
       this.lifecycle.clearLock(entryId);
+      // Status transitions to expired/banned/disabled make the account
+      // unusable; reusing a pooled WS would just hit the same wall on the
+      // upstream side. Evict so the pool doesn't hold a doomed connection.
+      if (status !== "active") this.evictWsPool(entryId);
     }
     if (status === "expired" && this._onExpired) {
       this._onExpired(entryId);
@@ -147,6 +170,7 @@ export class AccountPool {
   ): void {
     if (this.registry.markRateLimited(entryId, this.rateLimitBackoffSeconds, options)) {
       this.lifecycle.clearLock(entryId);
+      this.evictWsPool(entryId);
     }
   }
 

src/config-schema.ts

@@ -72,6 +72,18 @@ export const ConfigSchema = z.object({
     auto_download: z.boolean().default(false),
     allow_prerelease: z.boolean().default(false),
   }).default({}),
+  /** WebSocket connection pool — pins same (entryId, conversationId) to the
+   *  same physical WS so the upstream LB keeps prompt cache warm across
+   *  turns. See `src/proxy/ws-pool.ts` for the rationale. */
+  ws_pool: z.object({
+    enabled: z.boolean().default(true),
+    /** Hard upper bound per connection. Server enforces a 60-min cap; we
+     *  close 5 min early to avoid disrupting in-flight requests. */
+    max_age_ms: z.number().int().positive().default(3_300_000),
+    /** Cap on concurrent pooled connections per account, to bound memory
+     *  when a user opens many parallel conversations. */
+    max_per_account: z.number().int().positive().default(8),
+  }).default({}),
   ollama: z.object({
     enabled: z.boolean().default(false),
     host: z.string().default("127.0.0.1"),

src/index.ts

@@ -21,6 +21,7 @@ import { createModelRoutes } from "./routes/models.js";
 import { createWebRoutes } from "./routes/web.js";
 import { CookieJar } from "./proxy/cookie-jar.js";
 import { ProxyPool } from "./proxy/proxy-pool.js";
+import { setWsPoolConfig, getWsPool } from "./proxy/ws-pool.js";
 import { createProxyRoutes } from "./routes/proxies.js";
 import { createResponsesRoutes } from "./routes/responses.js";
 import { startUpdateChecker, stopUpdateChecker } from "./update-checker.js";
@@ -110,6 +111,15 @@ export async function startServer(options?: StartOptions): Promise<ServerHandle>
 
   // Build upstream router from config
   const cfg = getConfig();
+
+  // Wire WS connection pool to user config (defaults to enabled). Without
+  // this call `getWsPool()` would always use DEFAULT_WS_POOL_CONFIG and
+  // ignore `ws_pool.enabled: false` overrides — breaking the rollback path.
+  setWsPoolConfig({
+    enabled: cfg.ws_pool.enabled,
+    maxAgeMs: cfg.ws_pool.max_age_ms,
+    maxPerAccount: cfg.ws_pool.max_per_account,
+  });
   const adapters = new Map<string, UpstreamAdapter>();
   if (cfg.providers.openai) {
     adapters.set(
@@ -299,8 +309,11 @@ async function main() {
     }, 10_000);
     if (forceExit.unref) forceExit.unref();
 
-    handle.close().then(() => {
+    handle.close().then(async () => {
       getTransport().destroy?.();
+      try {
+        await getWsPool().shutdown();
+      } catch { /* never throws today, but defend against future regressions */ }
       console.log("[Shutdown] Server closed, cleanup complete.");
       clearTimeout(forceExit);
       process.exit(0);

src/proxy/codex-api.ts

@@ -15,9 +15,11 @@ import {
   buildHeaders,
   buildHeadersWithContentType,
 } from "../fingerprint/manager.js";
-import { createWebSocketResponse, type WsCreateRequest } from "./ws-transport.js";
+import { createWebSocketResponse, type WsCreateRequest, type WsPoolContext } from "./ws-transport.js";
 import type { ParsedRateLimit } from "./rate-limit-headers.js";
 import { getInstallationId } from "./installation-id.js";
+
+export type { WsPoolContext };
 import { parseSSEBlock, parseSSEStream } from "./codex-sse.js";
 import { fetchUsage } from "./codex-usage.js";
 import { fetchModels, probeEndpoint as probeEndpointFn } from "./codex-models.js";
@@ -175,10 +177,11 @@ export class CodexApi {
     request: CodexResponsesRequest,
     signal?: AbortSignal,
     onRateLimits?: (rl: ParsedRateLimit) => void,
+    poolCtx?: WsPoolContext,
   ): Promise<Response> {
     if (request.useWebSocket) {
       try {
-        return await this.createResponseViaWebSocket(request, signal, onRateLimits);
+        return await this.createResponseViaWebSocket(request, signal, onRateLimits, poolCtx);
       } catch (err) {
         // Real upstream API errors classified by ws-transport (e.g.
         // usage_limit_reached → CodexApiError(429)) must reach the
@@ -211,6 +214,7 @@ export class CodexApi {
     request: CodexResponsesRequest,
     signal?: AbortSignal,
     onRateLimits?: (rl: ParsedRateLimit) => void,
+    poolCtx?: WsPoolContext,
   ): Promise<Response> {
     const baseUrl = this.resolveBaseUrl();
     const wsUrl = baseUrl.replace(/^https?:/, "wss:") + "/codex/responses";
@@ -246,7 +250,7 @@ export class CodexApi {
       "x-codex-installation-id": installationId,
     };
 
-    return createWebSocketResponse(wsUrl, headers, wsRequest, signal, this.proxyUrl, onRateLimits);
+    return createWebSocketResponse(wsUrl, headers, wsRequest, signal, this.proxyUrl, onRateLimits, poolCtx);
   }
 
   /**