From 4a37d7ef4e4332aa887166768c5760f8bbe22af3 Mon Sep 17 00:00:00 2001 From: Codex Proxy Lab Date: Sun, 19 Jul 2026 22:17:56 +0800 Subject: [PATCH 1/6] fix: bind response continuity to physical websocket --- src/auth/session-affinity.ts | 122 ++++++++++++- src/proxy/codex-types.ts | 18 +- src/proxy/ws-pool.ts | 121 ++++++++++++- src/proxy/ws-transport.ts | 89 +++++++--- src/routes/shared/non-streaming-handler.ts | 5 +- src/routes/shared/non-streaming-helpers.ts | 5 +- src/routes/shared/proxy-handler.ts | 67 +++++++- .../shared/proxy-implicit-resume-lifecycle.ts | 5 + .../shared/proxy-implicit-resume-request.ts | 4 - .../shared/proxy-request-preparation.ts | 7 +- src/routes/shared/proxy-retry-classifier.ts | 9 +- src/routes/shared/proxy-retry-recovery.ts | 32 ++++ src/routes/shared/proxy-session-context.ts | 22 ++- src/routes/shared/proxy-ws-context.ts | 19 +- src/routes/shared/streaming-handler.ts | 5 +- .../proxy-handler-recovery.test.ts | 98 ++++++++--- tests/integration/proxy-handler.test.ts | 10 +- .../ws-physical-continuity.test.ts | 162 ++++++++++++++++++ tests/unit/auth/session-affinity.test.ts | 67 ++++++-- tests/unit/proxy/ws-pool.test.ts | 141 +++++++++++++++ .../proxy/ws-transport-early-error.test.ts | 18 ++ tests/unit/proxy/ws-transport.test.ts | 13 +- .../shared/non-streaming-affinity.test.ts | 2 +- .../proxy-implicit-resume-lifecycle.test.ts | 12 +- .../proxy-implicit-resume-request.test.ts | 19 +- .../shared/proxy-request-preparation.test.ts | 8 +- .../shared/proxy-retry-classifier.test.ts | 10 ++ .../shared/proxy-session-context.test.ts | 2 - .../shared/proxy-ws-context-boundary.test.ts | 6 +- .../routes/shared/proxy-ws-context.test.ts | 25 ++- .../routes/shared/streaming-handler.test.ts | 2 +- 31 files changed, 961 insertions(+), 164 deletions(-) create mode 100644 tests/integration/ws-physical-continuity.test.ts diff --git a/src/auth/session-affinity.ts b/src/auth/session-affinity.ts index 9b04f139..5bae16ac 100644 --- a/src/auth/session-affinity.ts +++ b/src/auth/session-affinity.ts @@ -9,10 +9,24 @@ import { createHash } from "crypto"; +export interface ChainAdvanceTicket { + conversationId: string; + variantHash?: string; + generation: number; + expectedParentResponseId: string | null; +} + +interface ChainHead { + conversationId: string; + variantHash?: string; + responseId: string | null; + generation: number; + updatedAt: number; +} + interface AffinityEntry { entryId: string; conversationId: string; - turnState?: string; /** SHA-256 hex of the instructions string. Stored as hash to bound memory usage. */ instructionsHash?: string; inputTokens?: number; @@ -31,6 +45,7 @@ const CLEANUP_INTERVAL_MS = 10 * 60 * 1000; // 10 minutes export class SessionAffinityMap { private map = new Map(); + private chainHeads = new Map(); private ttlMs: number; private cleanupTimer: ReturnType | null = null; @@ -39,7 +54,25 @@ export class SessionAffinityMap { this.cleanupTimer = setInterval(() => this.cleanup(), CLEANUP_INTERVAL_MS); } - /** Record that a response was created by a specific account in a conversation. */ + /** Capture the current chain generation before dispatching an upstream request. */ + captureChainAdvance( + conversationId: string, + variantHash?: string, + expectedParentResponseId?: string | null, + ): ChainAdvanceTicket { + const head = this.getChainHead(conversationId, variantHash); + return { + conversationId, + variantHash, + generation: head?.generation ?? 0, + expectedParentResponseId: + expectedParentResponseId === undefined + ? head?.responseId ?? null + : expectedParentResponseId, + }; + } + + /** Record response metadata and conditionally advance the implicit chain head. */ record( responseId: string, entryId: string, @@ -49,19 +82,45 @@ export class SessionAffinityMap { inputTokens?: number, functionCallIds?: string[], variantHash?: string, - ): void { + chainTicket?: ChainAdvanceTicket, + ): boolean { + void turnState; + const now = Date.now(); this.map.set(responseId, { entryId, conversationId, - turnState, instructionsHash: instructions !== undefined ? createHash("sha256").update(instructions ?? "").digest("hex") : undefined, inputTokens, functionCallIds: functionCallIds ? [...functionCallIds] : undefined, variantHash, - createdAt: Date.now(), + createdAt: now, + }); + + const key = this.chainKey(conversationId, variantHash); + const current = this.getChainHead(conversationId, variantHash); + if (current?.responseId === responseId) return true; + + if (chainTicket) { + if ( + chainTicket.conversationId !== conversationId || + chainTicket.variantHash !== variantHash || + (current?.generation ?? 0) !== chainTicket.generation || + (current?.responseId ?? null) !== chainTicket.expectedParentResponseId + ) { + return false; + } + } + + this.chainHeads.set(key, { + conversationId, + variantHash, + responseId, + generation: (current?.generation ?? 0) + 1, + updatedAt: now, }); + return true; } /** Look up which account created a given response. */ @@ -88,6 +147,20 @@ export class SessionAffinityMap { variantHash?: string, ): string | null { const now = Date.now(); + if (variantHash !== undefined) { + const head = this.getChainHead(conversationId, variantHash); + if (head) { + if (!head.responseId) return null; + const entry = this.getEntry(head.responseId); + if (!entry) { + this.invalidateHead(head); + return null; + } + if (maxAgeMs !== undefined && now - entry.createdAt > maxAgeMs) return null; + return head.responseId; + } + } + let latestResponseId: string | null = null; let latestCreatedAt = -1; for (const [responseId, entry] of this.map) { @@ -104,10 +177,12 @@ export class SessionAffinityMap { return latestResponseId; } - /** Look up the upstream turn-state token for a given response. */ - lookupTurnState(responseId: string): string | null { - const entry = this.getEntry(responseId); - return entry?.turnState ?? null; + /** + * turnState is scoped to a single Codex turn and must never be restored from + * cross-turn affinity. Kept for compatibility with older callers/tests. + */ + lookupTurnState(_responseId: string): null { + return null; } lookupInstructionsHash(responseId: string): string | null { @@ -134,6 +209,9 @@ export class SessionAffinityMap { /** Drop a response ID — called after upstream rejects it as not-found. */ forget(responseId: string): void { this.map.delete(responseId); + for (const head of this.chainHeads.values()) { + if (head.responseId === responseId) this.invalidateHead(head); + } } /** Drop every response recorded for a conversation, optionally scoped to a @@ -150,9 +228,29 @@ export class SessionAffinityMap { this.map.delete(responseId); dropped++; } + for (const head of this.chainHeads.values()) { + if (head.conversationId !== conversationId) continue; + if (variantHash !== undefined && head.variantHash !== variantHash) continue; + this.invalidateHead(head); + } return dropped; } + + private chainKey(conversationId: string, variantHash?: string): string { + return JSON.stringify([conversationId, variantHash ?? null]); + } + + private getChainHead(conversationId: string, variantHash?: string): ChainHead | null { + return this.chainHeads.get(this.chainKey(conversationId, variantHash)) ?? null; + } + + private invalidateHead(head: ChainHead): void { + head.responseId = null; + head.generation += 1; + head.updatedAt = Date.now(); + } + private getEntry(responseId: string): AffinityEntry | null { const entry = this.map.get(responseId); if (!entry) return null; @@ -171,6 +269,11 @@ export class SessionAffinityMap { this.map.delete(key); } } + for (const [key, head] of this.chainHeads) { + if (now - head.updatedAt > this.ttlMs) { + this.chainHeads.delete(key); + } + } } get size(): number { @@ -183,6 +286,7 @@ export class SessionAffinityMap { this.cleanupTimer = null; } this.map.clear(); + this.chainHeads.clear(); } } diff --git a/src/proxy/codex-types.ts b/src/proxy/codex-types.ts index fcd2fac7..344b81a5 100644 --- a/src/proxy/codex-types.ts +++ b/src/proxy/codex-types.ts @@ -216,9 +216,23 @@ export class CodexApiError extends Error { } } -/** previous_response_id 只能通过 WebSocket 安全续链，失败后不能降级为 HTTP delta-only。 */ +export type PreviousResponseContinuityReason = + | "busy" + | "dead" + | "expired" + | "missing_owner" + | "account_mismatch" + | "disabled" + | "no_key" + | "no_context" + | "transport"; + +/** previous_response_id can only continue on its owning physical WebSocket. */ export class PreviousResponseWebSocketError extends CodexApiError { - constructor(public readonly causeMessage: string) { + constructor( + public readonly causeMessage: string, + public readonly continuityReason?: PreviousResponseContinuityReason, + ) { super( 0, JSON.stringify({ diff --git a/src/proxy/ws-pool.ts b/src/proxy/ws-pool.ts index bf43e0ee..2c4d39a5 100644 --- a/src/proxy/ws-pool.ts +++ b/src/proxy/ws-pool.ts @@ -146,15 +146,30 @@ function isTerminalWsEvent(type: string): boolean { } function isEarlyMetadataWsEvent(type: string): boolean { - return type === "response.created" || type === "response.in_progress"; + return type === "response.created" || + type === "response.in_progress" || + type === "response.metadata" || + type === "codex.response.metadata"; +} + +function completedResponseId(msg: Record, type: string): string | null { + if (type !== "response.completed") return null; + const response = typeof msg.response === "object" && msg.response !== null + ? msg.response as Record + : null; + const id = response?.id ?? msg.response_id; + return typeof id === "string" && id.length > 0 ? id : null; } // ── PersistentWs ─────────────────────────────────────────────────── export interface PersistentWsHooks { /** Called when this WS becomes unusable (close, error, eviction). - * The pool uses this to remove the entry from its map. */ + * The pool uses this to remove the entry and all response owners. */ onDead(): void; + /** Called only after a response.completed frame establishes the newest + * connection-local previous-response anchor. */ + onResponseCompleted?(responseId: string): void; } /** Default keepalive cadence. 25s sits comfortably under the typical 30-60s @@ -456,6 +471,8 @@ export class PersistentWs { if (isTerminalWsEvent(type)) { sess.sawTerminalEvent = true; + const responseId = completedResponseId(msg, type); + if (responseId) this.hooks.onResponseCompleted?.(responseId); queueMicrotask(() => this.releaseAfterTerminalFrame()); } } else { @@ -566,12 +583,26 @@ export interface AcquireResult { reused: boolean; } -export type AcquireBypassReason = "busy" | "cap" | "dead" | "disabled" | "no_key"; +export type AcquireBypassReason = + | "busy" + | "cap" + | "dead" + | "expired" + | "disabled" + | "no_key" + | "missing_owner" + | "account_mismatch"; export interface AcquireBypass { bypass: AcquireBypassReason; } +export type ResponseOwnerBypassReason = Exclude; + +export interface ResponseOwnerBypass { + bypass: ResponseOwnerBypassReason; +} + export interface PersistentWsFactory { /** Called when the pool needs a new WS. The factory must construct a * PersistentWs whose `hooks.onDead` callback maps back to the pool. */ @@ -581,6 +612,10 @@ export interface PersistentWsFactory { export class WsConnectionPool { private readonly map = new Map(); private readonly byEntry = new Map>(); + /** Response IDs are valid only on the physical WS that completed them. */ + private readonly ownerByResponse = new Map(); + /** The upstream keeps only the most recent response per physical WS. */ + private readonly responseByPoolKey = new Map(); private readonly config: WsPoolConfig; private gcInterval: NodeJS.Timeout | undefined; private shuttingDown = false; @@ -616,7 +651,8 @@ export class WsConnectionPool { } let existing = this.map.get(poolKey); - if (existing && !existing.isAlive()) { + if (existing && (!existing.isAlive() || existing.isExpired(this.config.maxAgeMs))) { + existing.closeGracefully(); this.removeEntry(existing); existing = undefined; } @@ -633,16 +669,26 @@ export class WsConnectionPool { return { bypass: "cap" }; } + let freshRef: PersistentWs | undefined; const fresh = await factory({ entryId, poolKey, hooks: { onDead: () => { - // Pool-side cleanup. PersistentWs already marked itself dead. - this.removeEntryByKey(poolKey); + // A same-key connection may have won the factory race. Never let a + // discarded fresh connection remove that winner from the pool. + if (freshRef && this.map.get(poolKey) === freshRef) { + this.removeEntryByKey(poolKey); + } + }, + onResponseCompleted: (responseId) => { + if (freshRef && this.map.get(poolKey) === freshRef) { + this.registerResponseOwner(poolKey, responseId); + } }, }, }); + freshRef = fresh; // Race: another acquire for the same key may have completed during // factory() await. If so, prefer the one already in the map. @@ -674,6 +720,50 @@ export class WsConnectionPool { return { ws: fresh, reused: false }; } + /** Atomically acquire the physical WS that owns `previousResponseId`. + * This method never creates a connection: a response ID must not cross a + * physical WebSocket boundary when store=false. */ + acquireForResponse(entryId: string, previousResponseId: string): AcquireResult | ResponseOwnerBypass { + if (!this.config.enabled || this.shuttingDown) return { bypass: "disabled" }; + if (!entryId || !previousResponseId) return { bypass: "no_key" }; + + const poolKey = this.ownerByResponse.get(previousResponseId); + if (!poolKey) return { bypass: "missing_owner" }; + const owner = this.map.get(poolKey); + if (!owner) { + this.forgetResponseOwner(previousResponseId); + return { bypass: "missing_owner" }; + } + if (owner.entryId !== entryId) return { bypass: "account_mismatch" }; + if (!owner.isAlive()) { + this.removeEntry(owner); + return { bypass: "dead" }; + } + if (owner.isExpired(this.config.maxAgeMs)) { + owner.closeGracefully(); + this.removeEntry(owner); + return { bypass: "expired" }; + } + if (!owner.tryAcquire()) return { bypass: "busy" }; + return { ws: owner, reused: true }; + } + + /** Test/diagnostic helper: return the owning physical WS id. */ + ownerWsId(previousResponseId: string): string | null { + const poolKey = this.ownerByResponse.get(previousResponseId); + return poolKey ? this.map.get(poolKey)?.id ?? null : null; + } + + /** Remove a stale response owner without evicting an otherwise healthy WS. */ + forgetResponseOwner(previousResponseId: string): void { + const poolKey = this.ownerByResponse.get(previousResponseId); + if (!poolKey) return; + this.ownerByResponse.delete(previousResponseId); + if (this.responseByPoolKey.get(poolKey) === previousResponseId) { + this.responseByPoolKey.delete(poolKey); + } + } + /** Evict every WS for the given entryId. Used when the account is * rate-limited / banned / disabled / refreshed (token rotated). */ evictByEntryId(entryId: string): void { @@ -715,6 +805,8 @@ export class WsConnectionPool { // acquires would fail the disabled check anyway. this.map.clear(); this.byEntry.clear(); + this.ownerByResponse.clear(); + this.responseByPoolKey.clear(); } /** Periodic sweep: drop dead/expired idle entries. Skips busy ones. */ @@ -727,6 +819,18 @@ export class WsConnectionPool { } } + private registerResponseOwner(poolKey: string, responseId: string): void { + if (!this.map.has(poolKey)) return; + const previous = this.responseByPoolKey.get(poolKey); + if (previous && previous !== responseId) this.ownerByResponse.delete(previous); + const previousPoolKey = this.ownerByResponse.get(responseId); + if (previousPoolKey && previousPoolKey !== poolKey) { + this.responseByPoolKey.delete(previousPoolKey); + } + this.responseByPoolKey.set(poolKey, responseId); + this.ownerByResponse.set(responseId, poolKey); + } + private removeEntry(ws: PersistentWs): void { this.removeEntryByKey(ws.poolKey); } @@ -735,6 +839,11 @@ export class WsConnectionPool { const ws = this.map.get(poolKey); if (!ws) return; this.map.delete(poolKey); + const ownedResponse = this.responseByPoolKey.get(poolKey); + if (ownedResponse) { + this.responseByPoolKey.delete(poolKey); + this.ownerByResponse.delete(ownedResponse); + } const entryKeys = this.byEntry.get(ws.entryId); if (entryKeys) { entryKeys.delete(poolKey); diff --git a/src/proxy/ws-transport.ts b/src/proxy/ws-transport.ts index 29675971..3878430a 100644 --- a/src/proxy/ws-transport.ts +++ b/src/proxy/ws-transport.ts @@ -20,8 +20,9 @@ import type { CodexInputItem } from "./codex-api.js"; import type { ParsedRateLimit } from "./rate-limit-headers.js"; import { parseRateLimitsEvent } from "./rate-limit-headers.js"; -import { CodexApiError } from "./codex-types.js"; +import { CodexApiError, PreviousResponseWebSocketError } from "./codex-types.js"; import { getProxyUrl } from "../tls/proxy.js"; +import { isPreviousResponseNotFoundError } from "./error-classification.js"; import { PersistentWs, WsReusedConnectionError, @@ -85,7 +86,10 @@ function isTerminalWsEvent(type: string): boolean { } function isEarlyMetadataWsEvent(type: string): boolean { - return type === "response.created" || type === "response.in_progress"; + return type === "response.created" || + type === "response.in_progress" || + type === "response.metadata" || + type === "codex.response.metadata"; } const WS_CONNECTING = 0; @@ -249,9 +253,41 @@ export async function createWebSocketResponse( onRateLimits?: (rl: ParsedRateLimit) => void, poolCtx?: WsPoolContext, ): Promise { + const previousResponseId = request.previous_response_id; + + if (previousResponseId) { + if (!poolCtx) { + throw new PreviousResponseWebSocketError( + "No pooled WebSocket context is available for previous_response_id", + "no_context", + ); + } + const acquired = poolCtx.pool.acquireForResponse(poolCtx.entryId, previousResponseId); + if (!("ws" in acquired)) { + poolCtx.onDecision?.({ kind: "bypass", reason: acquired.bypass }); + throw new PreviousResponseWebSocketError( + `Owning WebSocket is unavailable (${acquired.bypass})`, + acquired.bypass, + ); + } + poolCtx.onDecision?.({ kind: "reuse", wsId: acquired.ws.id }); + try { + return await acquired.ws.send({ request, signal, onRateLimits, reused: true }); + } catch (err) { + if (isPreviousResponseNotFoundError(err)) { + poolCtx.pool.forgetResponseOwner(previousResponseId); + } + if (err instanceof WsReusedConnectionError) { + throw new PreviousResponseWebSocketError(err.message, "transport"); + } + throw err; + } + } + if (poolCtx) { + let acquired; try { - const acquired = await poolCtx.pool.acquire( + acquired = await poolCtx.pool.acquire( poolCtx.entryId, poolCtx.poolKey, (deps) => @@ -264,33 +300,36 @@ export async function createWebSocketResponse( hooks: deps.hooks, }), ); - if ("ws" in acquired) { - poolCtx.onDecision?.({ - kind: acquired.reused ? "reuse" : "new", - wsId: acquired.ws.id, - }); - try { - return await acquired.ws.send({ request, signal, onRateLimits, reused: acquired.reused }); - } catch (err) { - if (err instanceof WsReusedConnectionError) { - // Stale-reuse: open a fresh one-shot WS for this single request. - // The pool's onDead hook has already evicted the dead entry. - poolCtx.onDecision?.({ kind: "retry-after-stale-reuse", wsId: acquired.ws.id }); - return openOneShotWs(wsUrl, headers, request, signal, proxyUrl, onRateLimits); - } - throw err; - } - } - // Bypass (busy / cap / dead / no_key / disabled) → fall through to one-shot. - poolCtx.onDecision?.({ kind: "bypass", reason: acquired.bypass }); } catch (err) { - // Pool itself failed (e.g. factory could not connect). Don't punish the - // caller — fall back to the legacy one-shot path. The error is still - // visible in the one-shot path if the underlying issue persists. + // Only connection construction/acquisition errors reach this fallback. const msg = err instanceof Error ? err.message : String(err); console.warn(`[ws-pool] acquire failed, using one-shot fallback: ${msg}`); poolCtx.onDecision?.({ kind: "bypass", reason: "factory_error" }); + return openOneShotWs(wsUrl, headers, request, signal, proxyUrl, onRateLimits); } + + if ("ws" in acquired) { + poolCtx.onDecision?.({ + kind: acquired.reused ? "reuse" : "new", + wsId: acquired.ws.id, + }); + try { + return await acquired.ws.send({ request, signal, onRateLimits, reused: acquired.reused }); + } catch (err) { + // With full input and no previous_response_id, a pre-response failure + // on a reused WS is safe to replay once on a fresh one-shot. + if (err instanceof WsReusedConnectionError) { + poolCtx.onDecision?.({ kind: "retry-after-stale-reuse", wsId: acquired.ws.id }); + return openOneShotWs(wsUrl, headers, request, signal, proxyUrl, onRateLimits); + } + // Real upstream errors must propagate. They are not pool acquisition + // failures and must never trigger a cross-WS replay. + throw err; + } + } + + // No previous_response_id: a full-input one-shot is safe on pool bypass. + poolCtx.onDecision?.({ kind: "bypass", reason: acquired.bypass }); } return openOneShotWs(wsUrl, headers, request, signal, proxyUrl, onRateLimits); diff --git a/src/routes/shared/non-streaming-handler.ts b/src/routes/shared/non-streaming-handler.ts index fc811c63..bad7b0eb 100644 --- a/src/routes/shared/non-streaming-handler.ts +++ b/src/routes/shared/non-streaming-handler.ts @@ -6,7 +6,7 @@ import type { AccountPool } from "../../auth/account-pool.js"; import type { CookieJar } from "../../proxy/cookie-jar.js"; import type { ProxyPool } from "../../proxy/proxy-pool.js"; import { EmptyResponseError, UpstreamPrematureCloseError } from "../../translation/codex-event-extractor.js"; -import type { SessionAffinityMap } from "../../auth/session-affinity.js"; +import type { ChainAdvanceTicket, SessionAffinityMap } from "../../auth/session-affinity.js"; import type { FormatAdapter, ProxyRequest, UsageHint } from "./proxy-handler-types.js"; import { retryNonStreamingEmptyResponse, @@ -48,6 +48,7 @@ export interface HandleNonStreamingOptions { buildPoolCtx?: (forEntryId: string) => WsPoolContext | undefined; setActiveAccount?: (entryId: string, api: CodexApi) => void; variantHash?: string; + chainAdvanceTicket?: ChainAdvanceTicket; } export async function handleNonStreaming(options: HandleNonStreamingOptions): Promise { @@ -72,6 +73,7 @@ export async function handleNonStreaming(options: HandleNonStreamingOptions): Pr buildPoolCtx, setActiveAccount, variantHash, + chainAdvanceTicket, } = options; let currentEntryId = initialEntryId; let currentApi = initialApi; @@ -108,6 +110,7 @@ export async function handleNonStreaming(options: HandleNonStreamingOptions): Pr inputTokens: result.usage.input_tokens, responseFunctionCallIds, variantHash, + chainAdvanceTicket, }); if (result.responseId && conversationId && variantHash && reasoningReplayItems.length > 0) { getReasoningReplayCache().record({ diff --git a/src/routes/shared/non-streaming-helpers.ts b/src/routes/shared/non-streaming-helpers.ts index 6f541f60..479f9143 100644 --- a/src/routes/shared/non-streaming-helpers.ts +++ b/src/routes/shared/non-streaming-helpers.ts @@ -1,6 +1,6 @@ import type { Context } from "hono"; import type { StatusCode } from "hono/utils/http-status"; -import type { SessionAffinityMap } from "../../auth/session-affinity.js"; +import type { ChainAdvanceTicket, SessionAffinityMap } from "../../auth/session-affinity.js"; import type { AccountPool } from "../../auth/account-pool.js"; import { clearCfChallengeCooldown } from "../../auth/cf-challenge-cooldown.js"; import type { CodexApi, WsPoolContext } from "../../proxy/codex-api.js"; @@ -36,6 +36,7 @@ export interface RecordNonStreamingSuccessAffinityOptions { inputTokens: number; responseFunctionCallIds: Iterable; variantHash?: string; + chainAdvanceTicket?: ChainAdvanceTicket; } export function recordNonStreamingSuccessAffinity( @@ -51,6 +52,7 @@ export function recordNonStreamingSuccessAffinity( inputTokens, responseFunctionCallIds, variantHash, + chainAdvanceTicket, } = options; if (!responseId || !affinityMap || !conversationId) return false; @@ -64,6 +66,7 @@ export function recordNonStreamingSuccessAffinity( inputTokens, Array.from(new Set(responseFunctionCallIds)), variantHash, + chainAdvanceTicket, ); return true; } diff --git a/src/routes/shared/proxy-handler.ts b/src/routes/shared/proxy-handler.ts index 3419e572..bf96d855 100644 --- a/src/routes/shared/proxy-handler.ts +++ b/src/routes/shared/proxy-handler.ts @@ -22,7 +22,7 @@ * - non-streaming-handler.ts — collect / retry response lifecycle */ -import { CodexApi, CodexApiError } from "../../proxy/codex-api.js"; +import { CodexApi, CodexApiError, PreviousResponseWebSocketError } from "../../proxy/codex-api.js"; import { toQuota } from "../../auth/quota-utils.js"; import { acquireAccount, releaseAccount } from "./account-acquisition.js"; import { handleCodexApiError } from "./proxy-error-handler.js"; @@ -52,12 +52,13 @@ import { applyProxyRetryRecoveryDecision, applyCascadingBanDefense, buildProxyRetryRecoveryDecision, + invalidateRejectedPreviousResponse, } from "./proxy-retry-recovery.js"; import { classifyRetryAction } from "./proxy-retry-classifier.js"; import { buildProxySessionContext } from "./proxy-session-context.js"; import { staggerIfNeeded } from "./proxy-stagger.js"; import { sendProxyUpstreamAttempt } from "./proxy-upstream-attempt.js"; -import { buildWsPoolContext } from "./proxy-ws-context.js"; +import { buildWsPoolContext, forgetWsResponseOwner } from "./proxy-ws-context.js"; import { containsInvalidEncryptedContentSignal, getReasoningReplayCache, @@ -72,12 +73,15 @@ export async function handleProxyRequest(options: HandleProxyRequestOptions): Pr ensureProxyRequestInputArray(req); const originalRequestState = captureImplicitResumeRequestState(req); const sessionContext = buildProxySessionContext({ request: req, affinityMap }); + let chainAdvanceTicket = sessionContext.chainAdvanceTicket; + let recoveryWsKeySuffix: string | undefined; + let continuityRecoveryCount = 0; - // Turn state: sticky routing token from upstream, echoed back on subsequent requests + // turnState is scoped to one Codex turn. Preserve only a value supplied by + // the current client request; never restore it from cross-turn affinity. applyProxyRequestForwardingDefaults({ request: req, promptCacheKey: sessionContext.promptCacheKey, - explicitTurnState: sessionContext.explicitTurnState, }); const released = new Set(); @@ -246,6 +250,7 @@ export async function handleProxyRequest(options: HandleProxyRequestOptions): Pr variantHash: sessionContext.variantHash, requestId, tag: fmt.tag, + poolKeySuffix: recoveryWsKeySuffix, }); for (;;) { @@ -282,6 +287,7 @@ export async function handleProxyRequest(options: HandleProxyRequestOptions): Pr turnState: upstreamTurnState, usageHint: implicitResume.getUsageHint(), variantHash: sessionContext.variantHash, + chainAdvanceTicket, implicitResumeActive: implicitResume.isActive(), }); } @@ -312,8 +318,15 @@ export async function handleProxyRequest(options: HandleProxyRequestOptions): Pr if (!triedEntryIds.includes(nextEntryId)) triedEntryIds.push(nextEntryId); }, variantHash: sessionContext.variantHash, + chainAdvanceTicket, }); } catch (err) { + invalidateRejectedPreviousResponse({ + err, + previousResponseId: req.codexRequest.previous_response_id, + affinityMap, + forgetResponseOwner: forgetWsResponseOwner, + }); if (containsInvalidEncryptedContentSignal(err)) { reasoningReplayCache.evictByIdentity({ entryId, @@ -323,7 +336,13 @@ export async function handleProxyRequest(options: HandleProxyRequestOptions): Pr } const retryAction = classifyRetryAction( err, - { stripAndRetryDone, modelRetried, implicitResumeActive: implicitResume.isActive(), previousResponseId: req.codexRequest.previous_response_id }, + { + stripAndRetryDone, + modelRetried, + implicitResumeActive: implicitResume.isActive(), + previousResponseId: req.codexRequest.previous_response_id, + explicitPreviousResponseId: Boolean(sessionContext.explicitPrevRespId), + }, (e) => implicitResume.canReplayAfterError(e), ); @@ -332,9 +351,33 @@ export async function handleProxyRequest(options: HandleProxyRequestOptions): Pr releaseAccount(accountPool, entryId, annotateImageGenOutcome(undefined, req.expectsImageGen), released); throw err; - case "implicit_resume_replay": - implicitResume.replayFullInputAfterError(err); + case "implicit_resume_replay": { + if (!implicitResume.replayFullInputAfterError(err)) throw err; + stripAndRetryDone = true; + const staleId = sessionContext.implicitPrevRespId; + const continuityReason = err instanceof PreviousResponseWebSocketError + ? err.continuityReason + : undefined; + + // A busy owner is a live sibling branch. Keep the parent head and + // original ticket so only the first sibling completion advances it. + // Missing/dead owners are stale: invalidate and rebuild from a root. + if (staleId && continuityReason !== "busy") { + affinityMap.forget(staleId); + forgetWsResponseOwner(staleId); + chainAdvanceTicket = affinityMap.captureChainAdvance( + sessionContext.chainConversationId, + sessionContext.variantHash, + null, + ); + } + + // A unique pooled key avoids both the busy canonical connection and + // one-shot fallback while establishing a new response owner. + recoveryWsKeySuffix = + `recovery-${requestId.slice(0, 8)}-${++continuityRecoveryCount}`; continue; + } case "strip_and_retry": { stripAndRetryDone = true; @@ -348,6 +391,16 @@ export async function handleProxyRequest(options: HandleProxyRequestOptions): Pr affinityMap, restoreImplicitResumeRequest: implicitResume.restore, }); + if (decision.action === "retry" && decision.staleId) { + forgetWsResponseOwner(decision.staleId); + chainAdvanceTicket = affinityMap.captureChainAdvance( + sessionContext.chainConversationId, + sessionContext.variantHash, + null, + ); + recoveryWsKeySuffix = + `recovery-${requestId.slice(0, 8)}-${++continuityRecoveryCount}`; + } continue; } diff --git a/src/routes/shared/proxy-implicit-resume-lifecycle.ts b/src/routes/shared/proxy-implicit-resume-lifecycle.ts index c8784dcf..907d1027 100644 --- a/src/routes/shared/proxy-implicit-resume-lifecycle.ts +++ b/src/routes/shared/proxy-implicit-resume-lifecycle.ts @@ -112,6 +112,11 @@ export function createImplicitResumeLifecycle( if (!shouldReplayFullInputAfterImplicitResumeError(err, active)) return false; warn(`[${tag}] 隐式续链 WebSocket 失败，回退为完整历史重放：${err.causeMessage}`); restore(); + // Rebuild a response-owner chain on a pooled WS. If WS connection setup + // itself fails, CodexApi may still fall back to HTTP with full input. + request.codexRequest.useWebSocket = true; + request.codexRequest.previous_response_id = undefined; + request.codexRequest.turnState = snapshot.turnState; return true; }, restore, diff --git a/src/routes/shared/proxy-implicit-resume-request.ts b/src/routes/shared/proxy-implicit-resume-request.ts index 4dec0103..f57567e1 100644 --- a/src/routes/shared/proxy-implicit-resume-request.ts +++ b/src/routes/shared/proxy-implicit-resume-request.ts @@ -1,7 +1,6 @@ import type { ProxyRequest, UsageHint } from "./proxy-handler-types.js"; export interface ImplicitResumeAffinityLookup { - lookupTurnState(responseId: string): string | null; lookupInputTokens(responseId: string): number | null; } @@ -55,9 +54,6 @@ export function applyImplicitResumeRequest( ...reasoningReplayItems, ...request.codexRequest.input.slice(continuationInputStart), ]; - const implicitTurnState = affinityMap.lookupTurnState(implicitPrevRespId); - if (implicitTurnState) request.codexRequest.turnState = implicitTurnState; - return { reusedInputTokensUpperBound: affinityMap.lookupInputTokens(implicitPrevRespId) ?? undefined, }; diff --git a/src/routes/shared/proxy-request-preparation.ts b/src/routes/shared/proxy-request-preparation.ts index fffc63df..ca7ece12 100644 --- a/src/routes/shared/proxy-request-preparation.ts +++ b/src/routes/shared/proxy-request-preparation.ts @@ -3,7 +3,6 @@ import type { ProxyRequest } from "./proxy-handler-types.js"; export interface ApplyProxyRequestForwardingDefaultsOptions { request: ProxyRequest; promptCacheKey: string; - explicitTurnState: string | null; } export function ensureProxyRequestInputArray(request: ProxyRequest): void { @@ -15,14 +14,10 @@ export function ensureProxyRequestInputArray(request: ProxyRequest): void { export function applyProxyRequestForwardingDefaults( options: ApplyProxyRequestForwardingDefaultsOptions, ): void { - const { request, promptCacheKey, explicitTurnState } = options; + const { request, promptCacheKey } = options; request.codexRequest.prompt_cache_key = promptCacheKey; - if (explicitTurnState) { - request.codexRequest.turnState = explicitTurnState; - } - if (request.codexRequest.reasoning && !request.codexRequest.include?.length) { request.codexRequest.include = ["reasoning.encrypted_content"]; } diff --git a/src/routes/shared/proxy-retry-classifier.ts b/src/routes/shared/proxy-retry-classifier.ts index 60bd0a9e..4615356d 100644 --- a/src/routes/shared/proxy-retry-classifier.ts +++ b/src/routes/shared/proxy-retry-classifier.ts @@ -19,6 +19,9 @@ export interface RetryState { modelRetried: boolean; implicitResumeActive: boolean; previousResponseId: string | undefined; + /** True when the downstream client supplied previous_response_id explicitly. + * The proxy has no guaranteed full transcript for this request. */ + explicitPreviousResponseId?: boolean; } export type RetryAction = @@ -53,8 +56,10 @@ export function classifyRetryAction( return { type: "implicit_resume_replay" }; } - // Priority 2: strip stale previous_response_id (only once) - if (!state.stripAndRetryDone) { + // Priority 2: strip stale implicit previous_response_id (only once). + // Explicit Responses continuations may contain delta-only input, so stripping + // their ID would silently lose history; fail closed through the error handler. + if (!state.stripAndRetryDone && !state.explicitPreviousResponseId) { if (isPreviousResponseNotFoundError(err)) { return { type: "strip_and_retry", kind: "previous_response_not_found" }; } diff --git a/src/routes/shared/proxy-retry-recovery.ts b/src/routes/shared/proxy-retry-recovery.ts index 64b07d54..03690428 100644 --- a/src/routes/shared/proxy-retry-recovery.ts +++ b/src/routes/shared/proxy-retry-recovery.ts @@ -19,6 +19,24 @@ export type ProxyRetryRecoveryDecision = } | { action: "none" }; +export interface InvalidateRejectedPreviousResponseOptions { + err: unknown; + previousResponseId: string | undefined; + affinityMap: Pick; + forgetResponseOwner: (responseId: string) => void; +} + +/** Clear both logical and physical ownership after upstream explicitly rejects an ID. */ +export function invalidateRejectedPreviousResponse( + options: InvalidateRejectedPreviousResponseOptions, +): boolean { + const { err, previousResponseId, affinityMap, forgetResponseOwner } = options; + if (!previousResponseId || !isPreviousResponseNotFoundError(err)) return false; + forgetResponseOwner(previousResponseId); + affinityMap.forget(previousResponseId); + return true; +} + export interface BuildProxyRetryRecoveryDecisionOptions { err: unknown; tag: string; @@ -94,6 +112,9 @@ export function applyProxyRetryRecoveryDecision( restoreImplicitResumeRequest(); request.codexRequest.previous_response_id = undefined; request.codexRequest.turnState = undefined; + // This path now applies only to an implicit chain, for which the proxy owns + // a full request snapshot. Rebuild an owner on WebSocket when possible. + request.codexRequest.useWebSocket = true; return true; } @@ -135,6 +156,17 @@ export function applyCascadingBanDefense({ return false; } + // An explicit Responses continuation may contain only the current delta. + // Keep the ID intact: the physical-owner guard will fail closed before any + // cross-account WebSocket send instead of silently dropping history. + if (explicitPrevRespId) { + console.warn( + `[${tag}] Account switched from explicit response owner ${preferredEntryId} to ${acquiredEntryId}; ` + + `preserving previous_response_id for fail-closed continuity enforcement`, + ); + return false; + } + if (!preferredStatus || !BAN_RISK_STATUSES.has(preferredStatus)) { return false; } diff --git a/src/routes/shared/proxy-session-context.ts b/src/routes/shared/proxy-session-context.ts index 2a34ebc7..16404174 100644 --- a/src/routes/shared/proxy-session-context.ts +++ b/src/routes/shared/proxy-session-context.ts @@ -1,4 +1,4 @@ -import type { SessionAffinityMap } from "../../auth/session-affinity.js"; +import type { ChainAdvanceTicket, SessionAffinityMap } from "../../auth/session-affinity.js"; import type { ProxyRequest } from "./proxy-handler-types.js"; import { computeVariantHash } from "./variant-hash.js"; import { @@ -35,7 +35,7 @@ export interface ProxySessionContext { requiredFunctionCallOutputIds: string[]; implicitStoredFunctionCallIds: string[]; preferredEntryId: string | null; - explicitTurnState: string | null; + chainAdvanceTicket: ChainAdvanceTicket; resumeEvaluationInput: ProxyResumeEvaluationInput; } @@ -87,7 +87,21 @@ export function buildProxySessionContext( : implicitPrevRespId && hashInstructions(currentInstructions) === implicitStoredInstructionsHash ? affinityMap.lookup(implicitPrevRespId) : null; - const explicitTurnState = explicitPrevRespId ? affinityMap.lookupTurnState(explicitPrevRespId) : null; + const currentVariantHead = affinityMap.lookupLatestResponseIdByConversationId( + chainConversationId, + undefined, + variantHash, + ); + const expectedChainParent = explicitPrevRespId + ? currentVariantHead === null + ? null + : explicitPrevRespId + : implicitPrevRespId ?? undefined; + const chainAdvanceTicket = affinityMap.captureChainAdvance( + chainConversationId, + variantHash, + expectedChainParent, + ); return { currentInstructions, @@ -105,7 +119,7 @@ export function buildProxySessionContext( requiredFunctionCallOutputIds, implicitStoredFunctionCallIds, preferredEntryId, - explicitTurnState, + chainAdvanceTicket, resumeEvaluationInput: { implicitPrevRespId, continuationInputStart, diff --git a/src/routes/shared/proxy-ws-context.ts b/src/routes/shared/proxy-ws-context.ts index cfd5a71b..1c36ca72 100644 --- a/src/routes/shared/proxy-ws-context.ts +++ b/src/routes/shared/proxy-ws-context.ts @@ -9,6 +9,8 @@ export interface BuildWsPoolContextOptions { variantHash: string; requestId: string; tag: string; + /** Distinguishes a continuity-recovery WS from a busy canonical chain. */ + poolKeySuffix?: string; } export interface BuildWsPoolContextDeps { @@ -21,6 +23,11 @@ const defaultDeps: BuildWsPoolContextDeps = { log: (line) => console.log(line), }; +/** Remove a response-to-physical-WS owner through the pool boundary. */ +export function forgetWsResponseOwner(responseId: string): void { + getWsPool().forgetResponseOwner(responseId); +} + /** Build a per-request WS pool context only when the WS path has a stable chain id. */ export function buildWsPoolContext( options: BuildWsPoolContextOptions, @@ -33,10 +40,14 @@ export function buildWsPoolContext( const entryId = options.entryId; return { pool: (deps.getWsPool ?? defaultDeps.getWsPool)(), - // Physical WS affinity must follow the upstream response chain. Explicit - // previous_response_id continuations can legitimately change instructions - // (and therefore variantHash), but upstream still requires the same WS. - poolKey: `${entryId}:${options.conversationId}`, + // Full-input chains are isolated by variant. Explicit continuations do + // not use this key; response-owner lookup selects their physical WS. + poolKey: [ + entryId, + options.conversationId, + options.variantHash, + options.poolKeySuffix, + ].filter((part): part is string => Boolean(part)).join(":"), entryId, onDecision: (decision) => { const ridShort = options.requestId.slice(0, 8); diff --git a/src/routes/shared/streaming-handler.ts b/src/routes/shared/streaming-handler.ts index 553b44de..7f0b2936 100644 --- a/src/routes/shared/streaming-handler.ts +++ b/src/routes/shared/streaming-handler.ts @@ -2,7 +2,7 @@ import type { Context } from "hono"; import { stream } from "hono/streaming"; import type { AccountPool } from "../../auth/account-pool.js"; import { clearCfChallengeCooldown } from "../../auth/cf-challenge-cooldown.js"; -import type { SessionAffinityMap } from "../../auth/session-affinity.js"; +import type { ChainAdvanceTicket, SessionAffinityMap } from "../../auth/session-affinity.js"; import type { CodexApi } from "../../proxy/codex-api.js"; import { recordStreamCloseEvent } from "../../logs/stream-close-event.js"; import type { UsageInfo } from "../../translation/codex-event-extractor.js"; @@ -31,6 +31,7 @@ export interface HandleStreamingOptions { turnState?: string; usageHint?: UsageHint; variantHash: string; + chainAdvanceTicket: ChainAdvanceTicket; /** Whether this attempt was sent with an implicit-resume * `previous_response_id`. Needed to break the dead-chain retry loop: * if the upstream stream ends without response.completed while resume was @@ -58,6 +59,7 @@ export function handleStreaming(options: HandleStreamingOptions): Response { turnState, usageHint, variantHash, + chainAdvanceTicket, implicitResumeActive = false, } = options; @@ -109,6 +111,7 @@ export function handleStreaming(options: HandleStreamingOptions): Response { usageInfo?.input_tokens, Array.from(metadataCollector.responseFunctionCallIds), variantHash, + chainAdvanceTicket, ); if (!metadataCollector.invalidReasoningReplay && metadataCollector.reasoningReplayItems.length > 0) { reasoningReplayCache.record({ diff --git a/tests/integration/proxy-handler-recovery.test.ts b/tests/integration/proxy-handler-recovery.test.ts index 23bdb8f9..6224eee0 100644 --- a/tests/integration/proxy-handler-recovery.test.ts +++ b/tests/integration/proxy-handler-recovery.test.ts @@ -114,7 +114,7 @@ describe("proxy-handler recovery & defense", () => { vi.clearAllMocks(); }); - it("recovers from previous_response_not_found by stripping ID and retrying", async () => { + it("fails closed for explicit previous_response_not_found without retrying delta-only input", async () => { const notFoundBody = JSON.stringify({ error: { type: "invalid_request_error", @@ -150,10 +150,11 @@ describe("proxy-handler recovery & defense", () => { const { app } = buildTestApp({ accountPool, fmt, req }); const res = await app.request("/test", { method: "POST" }); - expect(res.status).toBe(200); - expect(createCount).toBe(2); - expect(seenPrevIds[0]).toBe("resp_0e2e6e7917486cfd0069eec8532d988194a3da6379c70abe68"); - expect(seenPrevIds[1]).toBeUndefined(); + expect(res.status).toBe(400); + expect(createCount).toBe(1); + expect(seenPrevIds).toEqual([ + "resp_0e2e6e7917486cfd0069eec8532d988194a3da6379c70abe68", + ]); }); it("replays full original input after implicit previous-response WebSocket failure", async () => { @@ -221,7 +222,7 @@ describe("proxy-handler recovery & defense", () => { { input: [{ role: "user", content: "continue" }], previousResponseId: "resp_implicit_ws", - turnState: "turn-implicit", + turnState: "turn-original", useWebSocket: true, }, { @@ -232,7 +233,7 @@ describe("proxy-handler recovery & defense", () => { ], previousResponseId: undefined, turnState: "turn-original", - useWebSocket: false, + useWebSocket: true, }, ]); expect(accountPool.acquire).toHaveBeenCalledTimes(1); @@ -242,7 +243,64 @@ describe("proxy-handler recovery & defense", () => { }); }); - it("recovers when collectTranslator raises previous_response_not_found", async () => { + it("retries implicit continuity recovery at most once", async () => { + const req: ProxyRequest = { + ...createDefaultRequest(), + codexRequest: { + ...createDefaultRequest().codexRequest, + prompt_cache_key: "thread-retry-once", + input: [ + { role: "user", content: "first" }, + { role: "assistant", content: "ok" }, + { role: "user", content: "continue" }, + ], + }, + }; + const identity = resolvePromptCacheIdentity(req.codexRequest, req.clientConversationId); + const variantHash = computeVariantHash( + req.codexRequest.instructions, + req.codexRequest.tools, + buildVariantIdentity(req.codexRequest, identity), + ); + getSessionAffinityMap().record( + "resp_retry_once", + "e1", + "thread-retry-once", + undefined, + "You are helpful", + undefined, + undefined, + variantHash, + ); + + let createCount = 0; + const seen: Array<{ previousResponseId?: string; inputLength: number; useWebSocket?: boolean }> = []; + mockCreateResponse = (request) => { + createCount++; + seen.push({ + previousResponseId: request.previous_response_id, + inputLength: request.input.length, + useWebSocket: request.useWebSocket, + }); + return Promise.reject(new PreviousResponseWebSocketError("still unavailable", "transport")); + }; + + const { app } = buildTestApp({ + accountPool: createMockAccountPool(), + fmt: createMockFormatAdapter(), + req, + }); + const res = await app.request("/test", { method: "POST" }); + + expect(res.status).toBe(502); + expect(createCount).toBe(2); + expect(seen).toEqual([ + { previousResponseId: "resp_retry_once", inputLength: 1, useWebSocket: true }, + { previousResponseId: undefined, inputLength: 3, useWebSocket: true }, + ]); + }); + + it("fails closed when explicit collect raises previous_response_not_found", async () => { const notFoundBody = JSON.stringify({ error: { type: "invalid_request_error", @@ -283,17 +341,12 @@ describe("proxy-handler recovery & defense", () => { const { app } = buildTestApp({ accountPool, fmt, req }); const res = await app.request("/test", { method: "POST" }); - expect(res.status).toBe(200); + expect(res.status).toBe(400); - expect(createCount).toBe(2); - expect(collectCount).toBe(2); - expect(seenPrevIds[0]).toBe("resp_collect_stale"); - expect(seenPrevIds[1]).toBeUndefined(); + expect(createCount).toBe(1); + expect(collectCount).toBe(1); + expect(seenPrevIds).toEqual(["resp_collect_stale"]); expect(accountPool.release).toHaveBeenCalledTimes(1); - expect(accountPool.release).toHaveBeenCalledWith("e1", { - input_tokens: 7, - output_tokens: 3, - }); }); it("does not loop forever when previous_response_not_found persists after strip", async () => { @@ -320,10 +373,10 @@ describe("proxy-handler recovery & defense", () => { const res = await app.request("/test", { method: "POST" }); expect(res.status).toBe(400); - expect(createCount).toBe(2); + expect(createCount).toBe(1); }); - it("recovers from unanswered function_call by stripping ID and retrying", async () => { + it("fails closed for explicit unanswered function_call without delta-only retry", async () => { const unansweredBody = JSON.stringify({ error: { type: "invalid_request_error", @@ -352,10 +405,9 @@ describe("proxy-handler recovery & defense", () => { const { app } = buildTestApp({ accountPool, fmt, req }); const res = await app.request("/test", { method: "POST" }); - expect(res.status).toBe(200); - expect(createCount).toBe(2); - expect(seenPrevIds[0]).toBe("resp_unanswered_chain"); - expect(seenPrevIds[1]).toBeUndefined(); + expect(res.status).toBe(400); + expect(createCount).toBe(1); + expect(seenPrevIds).toEqual(["resp_unanswered_chain"]); }); it("returns descriptive error when banned and remaining accounts disabled/expired", async () => { diff --git a/tests/integration/proxy-handler.test.ts b/tests/integration/proxy-handler.test.ts index 10625fa4..ab89ff9f 100644 --- a/tests/integration/proxy-handler.test.ts +++ b/tests/integration/proxy-handler.test.ts @@ -318,7 +318,7 @@ describe("proxy-handler integration", () => { const affinityMap = getSessionAffinityMap(); expect(affinityMap.lookup("resp_meta")).toBe("e1"); expect(affinityMap.lookupConversationId("resp_meta")).toBe("thread-collect"); - expect(affinityMap.lookupTurnState("resp_meta")).toBe("turn-success"); + expect(affinityMap.lookupTurnState("resp_meta")).toBeNull(); expect(affinityMap.lookupInstructionsHash("resp_meta")).toBe("58d0189aa8572b25a2e4ba09928df2c3d924d07f53de9aeb94ffe7f6f2a1de2b"); expect(affinityMap.lookupInputTokens("resp_meta")).toBe(33); expect(affinityMap.lookupFunctionCallIds("resp_meta")).toEqual(["call_a", "call_b"]); @@ -1084,8 +1084,8 @@ describe("proxy-handler integration", () => { expect(accountPool.acquire).toHaveBeenCalledTimes(1); }); - // 19. Cascading Ban Defense — strips only when preferred is banned - it("strips previous_response_id and turnState when preferred account is banned (cascading ban defense)", async () => { + // 19. Explicit continuation account fallback — preserve state and fail closed in transport + it("preserves explicit continuation state on banned-owner fallback for fail-closed continuity", async () => { const affinityMap = getSessionAffinityMap(); affinityMap.record( "resp_preferred", @@ -1125,9 +1125,9 @@ describe("proxy-handler integration", () => { expect(res.status).toBe(200); expect(capturedRequest).toBeDefined(); - expect(capturedRequest?.previous_response_id).toBeUndefined(); + expect(capturedRequest?.previous_response_id).toBe("resp_preferred"); expect(capturedRequest?.turnState).toBeUndefined(); - expect(affinityMap.lookup("resp_preferred")).toBeNull(); + expect(affinityMap.lookup("resp_preferred")).toBe("e_preferred"); }); // 19b. Cascading Ban Defense — does NOT strip for quota exhaustion diff --git a/tests/integration/ws-physical-continuity.test.ts b/tests/integration/ws-physical-continuity.test.ts new file mode 100644 index 00000000..90b3658d --- /dev/null +++ b/tests/integration/ws-physical-continuity.test.ts @@ -0,0 +1,162 @@ +import { WebSocketServer, type WebSocket } from "ws"; +import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { createWebSocketResponse, type WsCreateRequest } from "@src/proxy/ws-transport.js"; +import { PreviousResponseWebSocketError } from "@src/proxy/codex-types.js"; +import { WsConnectionPool } from "@src/proxy/ws-pool.js"; + +interface SeenRequest { + connection: number; + payload: WsCreateRequest; + socket: WebSocket; +} + +function request(marker: string, previousResponseId?: string): WsCreateRequest { + return { + type: "response.create", + model: "gpt-test", + instructions: "test", + input: [{ role: "user", content: marker }], + ...(previousResponseId ? { previous_response_id: previousResponseId } : {}), + }; +} + +async function drain(response: Response): Promise { + return response.text(); +} + +async function waitFor(predicate: () => boolean): Promise { + for (let i = 0; i < 100; i++) { + if (predicate()) return; + await new Promise((resolve) => setTimeout(resolve, 5)); + } + throw new Error("condition timeout"); +} + +describe("physical WebSocket response continuity", () => { + let server: WebSocketServer; + let url: string; + let pool: WsConnectionPool; + let connectionCount: number; + let seen: SeenRequest[]; + let onRequest: (seenRequest: SeenRequest) => void; + + beforeEach(async () => { + connectionCount = 0; + seen = []; + onRequest = () => undefined; + server = new WebSocketServer({ host: "127.0.0.1", port: 0 }); + server.on("connection", (socket) => { + const connection = ++connectionCount; + socket.on("message", (raw) => { + const seenRequest = { + connection, + payload: JSON.parse(raw.toString()) as WsCreateRequest, + socket, + }; + seen.push(seenRequest); + onRequest(seenRequest); + }); + }); + await new Promise((resolve) => server.once("listening", resolve)); + const address = server.address(); + if (!address || typeof address === "string") throw new Error("invalid test server address"); + url = `ws://127.0.0.1:${address.port}`; + pool = new WsConnectionPool({ enabled: true, maxAgeMs: 60_000, maxPerAccount: 8 }, { startGc: false }); + }); + + afterEach(async () => { + await pool.shutdown(); + await new Promise((resolve) => server.close(() => resolve())); + }); + + function context(poolKey = "entry:conv:variant") { + return { pool, entryId: "entry", poolKey }; + } + + function complete(item: SeenRequest, responseId: string): void { + item.socket.send(JSON.stringify({ type: "response.created", response: { id: responseId } })); + item.socket.send(JSON.stringify({ type: "response.completed", response: { id: responseId } })); + } + + it("continues the newest response only on the owning physical connection", async () => { + let sequence = 0; + onRequest = (item) => complete(item, `resp_${++sequence}`); + + await drain(await createWebSocketResponse(url, {}, request("first"), undefined, null, undefined, context())); + expect(pool.ownerWsId("resp_1")).not.toBeNull(); + + await drain(await createWebSocketResponse(url, {}, request("second", "resp_1"), undefined, null, undefined, context())); + expect(connectionCount).toBe(1); + expect(seen.map((item) => item.connection)).toEqual([1, 1]); + expect(seen[1].payload.previous_response_id).toBe("resp_1"); + expect(pool.ownerWsId("resp_1")).toBeNull(); + expect(pool.ownerWsId("resp_2")).not.toBeNull(); + + await expect( + createWebSocketResponse(url, {}, request("stale", "resp_1"), undefined, null, undefined, context()), + ).rejects.toBeInstanceOf(PreviousResponseWebSocketError); + expect(connectionCount).toBe(1); + expect(seen).toHaveLength(2); + }); + + it("fails closed when the owning connection is busy without opening a one-shot", async () => { + const held = new AbortController(); + onRequest = (item) => { + const marker = (item.payload.input[0] as { content?: string }).content; + if (marker === "first") complete(item, "resp_1"); + }; + + await drain(await createWebSocketResponse(url, {}, request("first"), undefined, null, undefined, context())); + const pending = createWebSocketResponse( + url, {}, request("held", "resp_1"), held.signal, null, undefined, context(), + ); + pending.catch(() => undefined); + await waitFor(() => seen.length === 2); + + await expect( + createWebSocketResponse(url, {}, request("concurrent", "resp_1"), undefined, null, undefined, context()), + ).rejects.toMatchObject({ name: "PreviousResponseWebSocketError" }); + expect(connectionCount).toBe(1); + expect(seen).toHaveLength(2); + held.abort(); + }); + + it("fails closed after the owner dies without opening a replacement carrying the old ID", async () => { + onRequest = (item) => complete(item, "resp_1"); + await drain(await createWebSocketResponse(url, {}, request("first"), undefined, null, undefined, context())); + server.clients.forEach((socket) => socket.terminate()); + await waitFor(() => pool.ownerWsId("resp_1") === null); + + await expect( + createWebSocketResponse(url, {}, request("after-death", "resp_1"), undefined, null, undefined, context()), + ).rejects.toBeInstanceOf(PreviousResponseWebSocketError); + expect(connectionCount).toBe(1); + expect(seen).toHaveLength(1); + }); + + it("keeps metadata behind the barrier and rejects a following 400 before returning Response", async () => { + onRequest = (item) => { + const marker = (item.payload.input[0] as { content?: string }).content; + if (marker === "first") { + complete(item, "resp_1"); + return; + } + item.socket.send(JSON.stringify({ type: "response.created", response: { id: "resp_failed" } })); + item.socket.send(JSON.stringify({ type: "response.in_progress", response: { id: "resp_failed" } })); + item.socket.send(JSON.stringify({ type: "codex.response.metadata", headers: { "x-test": "1" } })); + item.socket.send(JSON.stringify({ + type: "error", + status: 400, + error: { code: "previous_response_not_found", message: "not found" }, + })); + }; + + await drain(await createWebSocketResponse(url, {}, request("first"), undefined, null, undefined, context())); + await expect( + createWebSocketResponse(url, {}, request("second", "resp_1"), undefined, null, undefined, context()), + ).rejects.toMatchObject({ status: 400 }); + expect(pool.ownerWsId("resp_1")).toBeNull(); + expect(connectionCount).toBe(1); + expect(seen).toHaveLength(2); + }); +}); diff --git a/tests/unit/auth/session-affinity.test.ts b/tests/unit/auth/session-affinity.test.ts index ec245f8e..f817639d 100644 --- a/tests/unit/auth/session-affinity.test.ts +++ b/tests/unit/auth/session-affinity.test.ts @@ -233,37 +233,68 @@ describe("SessionAffinityMap", () => { }); }); - // turnState tracking - describe("turnState tracking", () => { - it("lookupTurnState returns recorded turnState", () => { + describe("concurrent chain advancement", () => { + it("keeps the first completed sibling as head when an older branch finishes late", () => { map = new SessionAffinityMap(); - map.record("resp_1", "entry_1", "conv_1", "ts_abc"); - expect(map.lookupTurnState("resp_1")).toBe("ts_abc"); + map.record("resp_parent", "entry_1", "conv_1", undefined, "system", 10, [], "vh"); + + const ticketA = map.captureChainAdvance("conv_1", "vh", "resp_parent"); + const ticketB = map.captureChainAdvance("conv_1", "vh", "resp_parent"); + + expect(map.record( + "resp_B", "entry_1", "conv_1", undefined, "system", 10, [], "vh", ticketB, + )).toBe(true); + expect(map.record( + "resp_A", "entry_1", "conv_1", undefined, "system", 10, [], "vh", ticketA, + )).toBe(false); + + expect(map.lookup("resp_A")).toBe("entry_1"); + expect(map.lookup("resp_B")).toBe("entry_1"); + expect(map.lookupLatestResponseIdByConversationId("conv_1", undefined, "vh")).toBe("resp_B"); }); - it("lookupTurnState returns null when no turnState was recorded", () => { + it("advances different variants independently", () => { map = new SessionAffinityMap(); - map.record("resp_1", "entry_1", "conv_1"); - expect(map.lookupTurnState("resp_1")).toBeNull(); + const mainTicket = map.captureChainAdvance("conv_1", "main"); + const subTicket = map.captureChainAdvance("conv_1", "sub"); + + expect(map.record( + "resp_main", "entry_1", "conv_1", undefined, "system", 10, [], "main", mainTicket, + )).toBe(true); + expect(map.record( + "resp_sub", "entry_1", "conv_1", undefined, "system", 10, [], "sub", subTicket, + )).toBe(true); + + expect(map.lookupLatestResponseIdByConversationId("conv_1", undefined, "main")).toBe("resp_main"); + expect(map.lookupLatestResponseIdByConversationId("conv_1", undefined, "sub")).toBe("resp_sub"); }); - it("turnState expires along with entry", () => { - map = new SessionAffinityMap(50); - map.record("resp_1", "entry_1", "conv_1", "ts_abc"); - expect(map.lookupTurnState("resp_1")).toBe("ts_abc"); + it("invalidates an outstanding ticket when the parent is forgotten", () => { + map = new SessionAffinityMap(); + map.record("resp_parent", "entry_1", "conv_1", undefined, "system", 10, [], "vh"); + const ticket = map.captureChainAdvance("conv_1", "vh", "resp_parent"); + + map.forget("resp_parent"); + expect(map.record( + "resp_late", "entry_1", "conv_1", undefined, "system", 10, [], "vh", ticket, + )).toBe(false); + expect(map.lookup("resp_late")).toBe("entry_1"); + expect(map.lookupLatestResponseIdByConversationId("conv_1", undefined, "vh")).toBeNull(); + }); + }); - const start = Date.now(); - while (Date.now() - start < 60) { - // busy wait - } + describe("turnState isolation", () => { + it("never returns a turnState recorded by a previous user turn", () => { + map = new SessionAffinityMap(); + map.record("resp_1", "entry_1", "conv_1", "ts_abc"); expect(map.lookupTurnState("resp_1")).toBeNull(); }); - it("turnState is updated on re-record", () => { + it("remains null after re-recording a response", () => { map = new SessionAffinityMap(); map.record("resp_1", "entry_1", "conv_1", "ts_old"); map.record("resp_1", "entry_1", "conv_1", "ts_new"); - expect(map.lookupTurnState("resp_1")).toBe("ts_new"); + expect(map.lookupTurnState("resp_1")).toBeNull(); }); }); }); diff --git a/tests/unit/proxy/ws-pool.test.ts b/tests/unit/proxy/ws-pool.test.ts index 09c28367..c5f8eab2 100644 --- a/tests/unit/proxy/ws-pool.test.ts +++ b/tests/unit/proxy/ws-pool.test.ts @@ -588,6 +588,147 @@ describe("WsConnectionPool", () => { expect(factory).toHaveBeenCalledTimes(2); }); + it("registers a response owner only after response.completed", async () => { + const { factory, created } = makeFactory(); + const acquired = await pool.acquire("entry-A", "entry-A:conv-1:variant-A", factory); + if (!("ws" in acquired)) throw new Error("expected acquire success"); + const send = acquired.ws.send({ + request: { type: "response.create", model: "m", instructions: "", input: [] }, + signal: undefined, + onRateLimits: undefined, + reused: false, + }); + const mock = created[0]["ws"] as unknown as MockWs; + mock.pushMessage({ type: "response.created", response: { id: "resp_A" } }); + expect(pool.ownerWsId("resp_A")).toBeNull(); + mock.pushMessage({ type: "response.completed", response: { id: "resp_A" } }); + await send; + expect(pool.ownerWsId("resp_A")).toBe(acquired.ws.id); + }); + + it("keeps only the most recent response owner per physical WS", async () => { + const { factory, created } = makeFactory(); + const first = await pool.acquire("entry-A", "entry-A:conv-1:variant-A", factory); + if (!("ws" in first)) throw new Error("expected acquire success"); + const mock = created[0]["ws"] as unknown as MockWs; + const firstSend = first.ws.send({ + request: { type: "response.create", model: "m", instructions: "", input: [] }, + signal: undefined, + onRateLimits: undefined, + reused: false, + }); + mock.pushMessage({ type: "response.completed", response: { id: "resp_A" } }); + await firstSend; + await nextTick(); + + const second = await pool.acquire("entry-A", "entry-A:conv-1:variant-A", factory); + if (!("ws" in second)) throw new Error("expected acquire success"); + const secondSend = second.ws.send({ + request: { type: "response.create", model: "m", instructions: "", input: [] }, + signal: undefined, + onRateLimits: undefined, + reused: true, + }); + mock.pushMessage({ type: "response.completed", response: { id: "resp_B" } }); + await secondSend; + + expect(pool.ownerWsId("resp_A")).toBeNull(); + expect(pool.ownerWsId("resp_B")).toBe(second.ws.id); + }); + + it("acquireForResponse fails closed on missing, busy, dead, and account-mismatched owners", async () => { + const { factory, created } = makeFactory(); + expect(pool.acquireForResponse("entry-A", "resp_missing")).toEqual({ bypass: "missing_owner" }); + + const acquired = await pool.acquire("entry-A", "entry-A:conv-1:variant-A", factory); + if (!("ws" in acquired)) throw new Error("expected acquire success"); + const mock = created[0]["ws"] as unknown as MockWs; + const send = acquired.ws.send({ + request: { type: "response.create", model: "m", instructions: "", input: [] }, + signal: undefined, + onRateLimits: undefined, + reused: false, + }); + mock.pushMessage({ type: "response.completed", response: { id: "resp_A" } }); + await send; + await nextTick(); + + expect(pool.acquireForResponse("entry-B", "resp_A")).toEqual({ bypass: "account_mismatch" }); + const owner = pool.acquireForResponse("entry-A", "resp_A"); + expect("ws" in owner).toBe(true); + expect(pool.acquireForResponse("entry-A", "resp_A")).toEqual({ bypass: "busy" }); + mock.pushClose(1006, "gone"); + expect(pool.acquireForResponse("entry-A", "resp_A")).toEqual({ bypass: "missing_owner" }); + }); + + it("does not let a discarded same-key factory loser remove the winning connection", async () => { + const pending: Array<{ + deps: { entryId: string; poolKey: string; hooks: PersistentWsHooks }; + resolve: (ws: PersistentWs) => void; + }> = []; + const factory = vi.fn((deps: { entryId: string; poolKey: string; hooks: PersistentWsHooks }) => + new Promise((resolve) => pending.push({ deps, resolve })), + ); + + const firstPromise = pool.acquire("entry-A", "entry-A:conv-race", factory); + const secondPromise = pool.acquire("entry-A", "entry-A:conv-race", factory); + await nextTick(); + expect(pending).toHaveLength(2); + + const winner = new PersistentWs({ + ws: new MockWs(), + entryId: pending[1].deps.entryId, + poolKey: pending[1].deps.poolKey, + hooks: pending[1].deps.hooks, + }); + pending[1].resolve(winner); + const second = await secondPromise; + expect("ws" in second && second.ws).toBe(winner); + + const loser = new PersistentWs({ + ws: new MockWs(), + entryId: pending[0].deps.entryId, + poolKey: pending[0].deps.poolKey, + hooks: pending[0].deps.hooks, + }); + pending[0].resolve(loser); + expect(await firstPromise).toEqual({ bypass: "busy" }); + await nextTick(); + + expect(pool.size()).toBe(1); + expect(pool.countByEntryId("entry-A")).toBe(1); + }); + + it("expires a response owner before allowing continuation", async () => { + let now = 0; + const expiringPool = new WsConnectionPool({ maxAgeMs: 100 }, { startGc: false }); + const mock = new MockWs(); + const factory = vi.fn(async (deps: { entryId: string; poolKey: string; hooks: PersistentWsHooks }) => + new PersistentWs({ ...deps, ws: mock, now: () => now }), + ); + try { + const acquired = await expiringPool.acquire("entry-A", "entry-A:conv-expire", factory); + if (!("ws" in acquired)) throw new Error("expected acquire success"); + const sent = acquired.ws.send({ + request: { type: "response.create", model: "m", instructions: "", input: [] }, + signal: undefined, + onRateLimits: undefined, + reused: false, + }); + mock.pushMessage({ type: "response.completed", response: { id: "resp_expired" } }); + await sent; + await nextTick(); + expect(expiringPool.ownerWsId("resp_expired")).toBe(acquired.ws.id); + + now = 101; + expect(expiringPool.acquireForResponse("entry-A", "resp_expired")).toEqual({ bypass: "expired" }); + expect(expiringPool.ownerWsId("resp_expired")).toBeNull(); + expect(expiringPool.size()).toBe(0); + } finally { + await expiringPool.shutdown(); + } + }); + it("evictByEntryId closes all connections for that entry and frees the cap", async () => { const capped = new WsConnectionPool({ maxPerAccount: 2 }, { startGc: false }); const { factory } = makeFactory(); diff --git a/tests/unit/proxy/ws-transport-early-error.test.ts b/tests/unit/proxy/ws-transport-early-error.test.ts index 703bb228..100289fd 100644 --- a/tests/unit/proxy/ws-transport-early-error.test.ts +++ b/tests/unit/proxy/ws-transport-early-error.test.ts @@ -149,6 +149,24 @@ describe("createWebSocketResponse — early-stream error rejection", () => { } }); + it.each(["codex.response.metadata", "response.metadata"])( + "keeps %s behind the early barrier so a following previous_response_not_found rejects", + async (metadataType) => { + const promise = createWebSocketResponse("wss://test/ws", {}, BASE_REQUEST); + promise.catch(() => undefined); + const ws = await waitForOpen(); + ws.emit("message", JSON.stringify({ type: "response.created", response: { id: "resp_new" } })); + ws.emit("message", JSON.stringify({ type: "response.in_progress", response: { id: "resp_new" } })); + ws.emit("message", JSON.stringify({ type: metadataType, headers: { "x-test": "1" } })); + ws.emit("message", JSON.stringify({ + type: "error", + status: 400, + error: { code: "previous_response_not_found", message: "not found" }, + })); + await expect(promise).rejects.toMatchObject({ status: 400 }); + }, + ); + it("rejects with CodexApiError(402) when first frame is response.failed quota_exhausted", async () => { const promise = createWebSocketResponse("wss://test/ws", {}, BASE_REQUEST); promise.catch(() => { /* asserted below */ }); diff --git a/tests/unit/proxy/ws-transport.test.ts b/tests/unit/proxy/ws-transport.test.ts index 37e00226..5861cfe8 100644 --- a/tests/unit/proxy/ws-transport.test.ts +++ b/tests/unit/proxy/ws-transport.test.ts @@ -290,19 +290,16 @@ describe("createWebSocketResponse", () => { ws.close(); }); - it("passes previous_response_id without store/stream fields", async () => { + it("fails closed before opening a one-shot when previous_response_id has no pool owner", async () => { const req: WsCreateRequest = { ...BASE_REQUEST, previous_response_id: "resp_prev_123", }; - const { ws } = await startConnect(req); - const sent = JSON.parse(ws.sentMessages[0]); - expect(sent.previous_response_id).toBe("resp_prev_123"); - expect(sent.store).toBeUndefined(); - expect(sent.stream).toBeUndefined(); - - ws.close(); + await expect(createWebSocketResponse("wss://test/ws", {}, req)).rejects.toMatchObject({ + name: "PreviousResponseWebSocketError", + }); + expect(wsInstances).toHaveLength(0); }); it("preserves message ordering", async () => { diff --git a/tests/unit/routes/shared/non-streaming-affinity.test.ts b/tests/unit/routes/shared/non-streaming-affinity.test.ts index e1bd34cb..d23a1192 100644 --- a/tests/unit/routes/shared/non-streaming-affinity.test.ts +++ b/tests/unit/routes/shared/non-streaming-affinity.test.ts @@ -29,7 +29,7 @@ describe("recordNonStreamingSuccessAffinity", () => { expect(recorded).toBe(true); expect(affinityMap.lookup("resp-ns")).toBe("entry-1"); expect(affinityMap.lookupConversationId("resp-ns")).toBe("conversation-1"); - expect(affinityMap.lookupTurnState("resp-ns")).toBe("turn-1"); + expect(affinityMap.lookupTurnState("resp-ns")).toBeNull(); // null instructions → hash of empty string (sha256("")) expect(affinityMap.lookupInstructionsHash("resp-ns")).toBe("e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"); expect(affinityMap.lookupInputTokens("resp-ns")).toBe(0); diff --git a/tests/unit/routes/shared/proxy-implicit-resume-lifecycle.test.ts b/tests/unit/routes/shared/proxy-implicit-resume-lifecycle.test.ts index d008f7ea..2c405cf5 100644 --- a/tests/unit/routes/shared/proxy-implicit-resume-lifecycle.test.ts +++ b/tests/unit/routes/shared/proxy-implicit-resume-lifecycle.test.ts @@ -29,11 +29,9 @@ function makeProxyRequest(): ProxyRequest { } function makeAffinityLookup(options: { - turnState?: string | null; inputTokens?: number | null; } = {}): ImplicitResumeAffinityLookup { return { - lookupTurnState: vi.fn(() => options.turnState ?? null), lookupInputTokens: vi.fn(() => options.inputTokens ?? null), }; } @@ -122,7 +120,7 @@ describe("implicit resume lifecycle", () => { const lifecycle = createImplicitResumeLifecycle({ request, snapshot, - affinityMap: makeAffinityLookup({ turnState: "turn-implicit", inputTokens: 123 }), + affinityMap: makeAffinityLookup({ inputTokens: 123 }), tag: "Test", implicitPrevRespId: "resp_implicit", continuationInputStart: 2, @@ -146,7 +144,7 @@ describe("implicit resume lifecycle", () => { expect(lifecycle.resumeReasonForAttempt()).toBeNull(); expect(lifecycle.getUsageHint()).toEqual({ reusedInputTokensUpperBound: 123 }); expect(request.codexRequest.previous_response_id).toBe("resp_implicit"); - expect(request.codexRequest.turnState).toBe("turn-implicit"); + expect(request.codexRequest.turnState).toBe("turn-original"); expect(request.codexRequest.useWebSocket).toBe(true); expect(request.codexRequest.input).toEqual([{ role: "user", content: "continue" }]); @@ -169,7 +167,7 @@ describe("implicit resume lifecycle", () => { const lifecycle = createImplicitResumeLifecycle({ request, snapshot, - affinityMap: makeAffinityLookup({ turnState: "turn-implicit", inputTokens: 123 }), + affinityMap: makeAffinityLookup({ inputTokens: 123 }), tag: "Test", implicitPrevRespId: "resp_implicit", continuationInputStart: 2, @@ -192,9 +190,9 @@ describe("implicit resume lifecycle", () => { expect(warn).toHaveBeenCalledWith("[Test] 隐式续链 WebSocket 失败，回退为完整历史重放：ws down"); expect(lifecycle.isActive()).toBe(false); expect(lifecycle.getUsageHint()).toBeUndefined(); - expect(request.codexRequest.previous_response_id).toBe("explicit-prev"); + expect(request.codexRequest.previous_response_id).toBeUndefined(); expect(request.codexRequest.turnState).toBe("turn-original"); - expect(request.codexRequest.useWebSocket).toBe(false); + expect(request.codexRequest.useWebSocket).toBe(true); expect(request.codexRequest.input).toBe(snapshot.input); }); diff --git a/tests/unit/routes/shared/proxy-implicit-resume-request.test.ts b/tests/unit/routes/shared/proxy-implicit-resume-request.test.ts index 35ad6717..8448e441 100644 --- a/tests/unit/routes/shared/proxy-implicit-resume-request.test.ts +++ b/tests/unit/routes/shared/proxy-implicit-resume-request.test.ts @@ -29,11 +29,9 @@ function makeProxyRequest(): ProxyRequest { } function makeAffinityLookup(options: { - turnState?: string | null; inputTokens?: number | null; } = {}): ImplicitResumeAffinityLookup { return { - lookupTurnState: vi.fn(() => options.turnState ?? null), lookupInputTokens: vi.fn(() => options.inputTokens ?? null), }; } @@ -54,10 +52,7 @@ describe("implicit resume request state helpers", () => { it("applies implicit resume by using the previous response id, WebSocket, sliced input, and usage hint", () => { const request = makeProxyRequest(); - const affinityMap = makeAffinityLookup({ - turnState: "turn-implicit", - inputTokens: 123, - }); + const affinityMap = makeAffinityLookup({ inputTokens: 123 }); const usageHint = applyImplicitResumeRequest({ request, @@ -68,12 +63,11 @@ describe("implicit resume request state helpers", () => { expect(request.codexRequest.previous_response_id).toBe("resp_implicit"); expect(request.codexRequest.useWebSocket).toBe(true); - expect(request.codexRequest.turnState).toBe("turn-implicit"); + expect(request.codexRequest.turnState).toBe("turn-original"); expect(request.codexRequest.input).toEqual([ { role: "user", content: "continue" }, ]); expect(usageHint).toEqual({ reusedInputTokensUpperBound: 123 }); - expect(affinityMap.lookupTurnState).toHaveBeenCalledWith("resp_implicit"); expect(affinityMap.lookupInputTokens).toHaveBeenCalledWith("resp_implicit"); }); @@ -84,10 +78,7 @@ describe("implicit resume request state helpers", () => { request, implicitPrevRespId: "resp_implicit", continuationInputStart: 2, - affinityMap: makeAffinityLookup({ - turnState: "turn-implicit", - inputTokens: 123, - }), + affinityMap: makeAffinityLookup({ inputTokens: 123 }), reasoningReplayItems: [ { type: "reasoning", id: "rs_replay", summary: [], encrypted_content: "encrypted" }, { type: "function_call", call_id: "call_1", name: "read_file", arguments: "{}" }, @@ -96,7 +87,7 @@ describe("implicit resume request state helpers", () => { expect(request.codexRequest.previous_response_id).toBe("resp_implicit"); expect(request.codexRequest.useWebSocket).toBe(true); - expect(request.codexRequest.turnState).toBe("turn-implicit"); + expect(request.codexRequest.turnState).toBe("turn-original"); expect(request.codexRequest.input).toEqual([ { type: "reasoning", id: "rs_replay", summary: [], encrypted_content: "encrypted" }, { type: "function_call", call_id: "call_1", name: "read_file", arguments: "{}" }, @@ -126,7 +117,7 @@ describe("implicit resume request state helpers", () => { request, implicitPrevRespId: "resp_implicit", continuationInputStart: 2, - affinityMap: makeAffinityLookup({ turnState: "turn-implicit", inputTokens: 123 }), + affinityMap: makeAffinityLookup({ inputTokens: 123 }), reasoningReplayItems: [{ type: "reasoning", id: "rs_replay", summary: [], encrypted_content: "encrypted" }], }); request.codexRequest.instructions = "mutated-system"; diff --git a/tests/unit/routes/shared/proxy-request-preparation.test.ts b/tests/unit/routes/shared/proxy-request-preparation.test.ts index 67231477..2d02a9c7 100644 --- a/tests/unit/routes/shared/proxy-request-preparation.test.ts +++ b/tests/unit/routes/shared/proxy-request-preparation.test.ts @@ -38,17 +38,16 @@ describe("proxy request preparation", () => { expect(request.codexRequest.input).toEqual([]); }); - it("applies prompt cache key and explicit turn state for upstream forwarding", () => { + it("applies the prompt cache key without synthesizing cross-turn state", () => { const request = makeProxyRequest(); applyProxyRequestForwardingDefaults({ request, promptCacheKey: "cache-key-1", - explicitTurnState: "turn-explicit", }); expect(request.codexRequest.prompt_cache_key).toBe("cache-key-1"); - expect(request.codexRequest.turnState).toBe("turn-explicit"); + expect(request.codexRequest.turnState).toBeUndefined(); }); it("does not clear an existing turn state when no explicit turn state is available", () => { @@ -58,7 +57,6 @@ describe("proxy request preparation", () => { applyProxyRequestForwardingDefaults({ request, promptCacheKey: "cache-key-1", - explicitTurnState: null, }); expect(request.codexRequest.turnState).toBe("turn-existing"); @@ -71,7 +69,6 @@ describe("proxy request preparation", () => { applyProxyRequestForwardingDefaults({ request, promptCacheKey: "cache-key-1", - explicitTurnState: null, }); expect(request.codexRequest.include).toEqual(["reasoning.encrypted_content"]); @@ -85,7 +82,6 @@ describe("proxy request preparation", () => { applyProxyRequestForwardingDefaults({ request, promptCacheKey: "cache-key-1", - explicitTurnState: null, }); expect(request.codexRequest.include).toEqual(["custom.include"]); diff --git a/tests/unit/routes/shared/proxy-retry-classifier.test.ts b/tests/unit/routes/shared/proxy-retry-classifier.test.ts index f6622fa2..d6dde849 100644 --- a/tests/unit/routes/shared/proxy-retry-classifier.test.ts +++ b/tests/unit/routes/shared/proxy-retry-classifier.test.ts @@ -74,6 +74,16 @@ describe("classifyRetryAction", () => { }); }); + it("fails closed instead of stripping an explicit previous_response_id", () => { + const result = classifyRetryAction(prevRespNotFound, { + ...defaultState, + previousResponseId: "resp_explicit", + explicitPreviousResponseId: true, + }, neverReplayable); + + expect(result).toEqual({ type: "error_handler_decides" }); + }); + describe("priority 3: error handler", () => { it("delegates 429 to error handler", () => { const result = classifyRetryAction(rateLimitErr, defaultState, neverReplayable); diff --git a/tests/unit/routes/shared/proxy-session-context.test.ts b/tests/unit/routes/shared/proxy-session-context.test.ts index 53803f7d..c434bb17 100644 --- a/tests/unit/routes/shared/proxy-session-context.test.ts +++ b/tests/unit/routes/shared/proxy-session-context.test.ts @@ -84,7 +84,6 @@ describe("buildProxySessionContext", () => { expect(context.implicitPrevRespId).toBeNull(); expect(context.prevRespId).toBe("resp_prev"); expect(context.preferredEntryId).toBe("entry-prev"); - expect(context.explicitTurnState).toBe("turn-prev"); expect(context.continuationInputStart).toBe(0); expect(context.resumeEvaluationInput.implicitPrevRespId).toBeNull(); }); @@ -150,7 +149,6 @@ describe("buildProxySessionContext", () => { expect(context.implicitPrevRespId).toBe("resp_implicit"); expect(context.prevRespId).toBe("resp_implicit"); expect(context.preferredEntryId).toBe("entry-implicit"); - expect(context.explicitTurnState).toBeNull(); expect(context.implicitStoredInstructionsHash).toBe(sha256("system")); expect(context.implicitStoredFunctionCallIds).toEqual(["call_a"]); expect(context.requiredFunctionCallOutputIds).toEqual(["call_a"]); diff --git a/tests/unit/routes/shared/proxy-ws-context-boundary.test.ts b/tests/unit/routes/shared/proxy-ws-context-boundary.test.ts index bb1303b3..a2d7da8c 100644 --- a/tests/unit/routes/shared/proxy-ws-context-boundary.test.ts +++ b/tests/unit/routes/shared/proxy-ws-context-boundary.test.ts @@ -2,7 +2,10 @@ import { readFileSync } from "node:fs"; import { resolve } from "node:path"; import * as ts from "typescript"; import { describe, expect, it } from "vitest"; -import { buildWsPoolContext } from "@src/routes/shared/proxy-ws-context.js"; +import { + buildWsPoolContext, + forgetWsResponseOwner, +} from "@src/routes/shared/proxy-ws-context.js"; const ROOT = process.cwd(); const WS_CONTEXT_MODULE = "src/routes/shared/proxy-ws-context.ts"; @@ -48,6 +51,7 @@ function importsNamedBinding(content: string, moduleSuffix: string, bindingName: describe("proxy websocket context boundary", () => { it("exports websocket pool context construction from its own module", () => { expect(buildWsPoolContext).toBeTypeOf("function"); + expect(forgetWsResponseOwner).toBeTypeOf("function"); const wsContext = source(WS_CONTEXT_MODULE); expect(importsNamedBinding(wsContext, "ws-pool.js", "getWsPool", WS_CONTEXT_MODULE)).toBe(true); diff --git a/tests/unit/routes/shared/proxy-ws-context.test.ts b/tests/unit/routes/shared/proxy-ws-context.test.ts index 9f678929..7f6e8882 100644 --- a/tests/unit/routes/shared/proxy-ws-context.test.ts +++ b/tests/unit/routes/shared/proxy-ws-context.test.ts @@ -71,24 +71,37 @@ describe("buildWsPoolContext", () => { expect(context).toBeDefined(); expect(context?.pool).toBe(pool); expect(context?.entryId).toBe("entry-A"); - expect(context?.poolKey).toBe("entry-A:conv-1"); + expect(context?.poolKey).toBe("entry-A:conv-1:vh-123"); expect(deps.getPoolCalls).toBe(1); }); - it("keeps the same physical WebSocket across explicit previous-response variant changes", () => { + it("isolates full-input variants on different physical WebSocket keys", () => { const deps = createDeps(); - const first = buildWsPoolContext( + const main = buildWsPoolContext( { ...baseOptions(), variantHash: "main-turn" }, deps.deps, ); - const continuation = buildWsPoolContext( + const subagent = buildWsPoolContext( { ...baseOptions(), variantHash: "changed-instructions" }, deps.deps, ); - expect(first?.poolKey).toBe("entry-A:conv-1"); - expect(continuation?.poolKey).toBe(first?.poolKey); + expect(main?.poolKey).toBe("entry-A:conv-1:main-turn"); + expect(subagent?.poolKey).toBe("entry-A:conv-1:changed-instructions"); + expect(subagent?.poolKey).not.toBe(main?.poolKey); + }); + + it("uses a distinct pooled key for a full-replay continuity recovery", () => { + const deps = createDeps(); + const canonical = buildWsPoolContext(baseOptions(), deps.deps); + const recovery = buildWsPoolContext( + { ...baseOptions(), poolKeySuffix: "recovery-rid-1" }, + deps.deps, + ); + + expect(canonical?.poolKey).toBe("entry-A:conv-1:vh-123"); + expect(recovery?.poolKey).toBe("entry-A:conv-1:vh-123:recovery-rid-1"); }); it("logs pool decisions with the route tag and shortened request id", () => { diff --git a/tests/unit/routes/shared/streaming-handler.test.ts b/tests/unit/routes/shared/streaming-handler.test.ts index ff780494..d269a619 100644 --- a/tests/unit/routes/shared/streaming-handler.test.ts +++ b/tests/unit/routes/shared/streaming-handler.test.ts @@ -116,7 +116,7 @@ describe("handleStreaming", () => { }); expect(affinityMap.lookup("resp_stream")).toBe("entry-stream"); expect(affinityMap.lookupConversationId("resp_stream")).toBe("conversation-stream"); - expect(affinityMap.lookupTurnState("resp_stream")).toBe("turn-stream"); + expect(affinityMap.lookupTurnState("resp_stream")).toBeNull(); expect(affinityMap.lookupInstructionsHash("resp_stream")).toBe("58d0189aa8572b25a2e4ba09928df2c3d924d07f53de9aeb94ffe7f6f2a1de2b"); expect(affinityMap.lookupInputTokens("resp_stream")).toBe(10_001); expect(affinityMap.lookupFunctionCallIds("resp_stream")).toEqual(["call_stream"]); From d725c49bd3a50a4dcc0711ef69debdd418a0a864 Mon Sep 17 00:00:00 2001 From: shiitin Date: Mon, 20 Jul 2026 00:41:10 +0800 Subject: [PATCH 2/6] fix: bound websocket response start wait --- src/proxy/ws-pool.ts | 43 +++++++++++++++++++-- src/proxy/ws-transport.ts | 20 ++++++++++ tests/unit/proxy/ws-pool.test.ts | 50 +++++++++++++++++++++++++ tests/unit/proxy/ws-transport.test.ts | 54 +++++++++++++++++++++++++++ 4 files changed, 164 insertions(+), 3 deletions(-) diff --git a/src/proxy/ws-pool.ts b/src/proxy/ws-pool.ts index 2c4d39a5..a619ebbd 100644 --- a/src/proxy/ws-pool.ts +++ b/src/proxy/ws-pool.ts @@ -82,6 +82,7 @@ interface InFlightSession { abortListener: (() => void) | null; signal: AbortSignal | undefined; streamClosed: boolean; + responseStartTimer: ReturnType | undefined; } /** Subset of the `ws` module's WebSocket interface that PersistentWs needs. @@ -183,6 +184,11 @@ export const DEFAULT_PING_INTERVAL_MS = 25_000; * and re-using it would cost a real-request cache miss. */ export const DEFAULT_LIVENESS_TIMEOUT_MULTIPLIER = 2.5; +/** Maximum time to wait for the first non-metadata response event. This keeps + * an upstream that sends only provisional metadata from occupying a pooled + * connection indefinitely. */ +export const DEFAULT_WS_RESPONSE_START_TIMEOUT_MS = 180_000; + export class PersistentWs { readonly id: string; readonly entryId: string; @@ -301,6 +307,7 @@ export class PersistentWs { signal: AbortSignal | undefined; onRateLimits: ((rl: ParsedRateLimit) => void) | undefined; reused: boolean; + responseStartTimeoutMs?: number; }): Promise { if (!this.busy) { throw new Error("PersistentWs.send called without prior tryAcquire"); @@ -335,8 +342,19 @@ export class PersistentWs { abortListener: null, signal: opts.signal, streamClosed: false, + responseStartTimer: undefined, }; + const responseStartTimeoutMs = + opts.responseStartTimeoutMs ?? DEFAULT_WS_RESPONSE_START_TIMEOUT_MS; + if (responseStartTimeoutMs > 0) { + this.currentSession.responseStartTimer = setTimeout( + () => this.handleResponseStartTimeout(responseStartTimeoutMs), + responseStartTimeoutMs, + ); + this.currentSession.responseStartTimer.unref?.(); + } + if (opts.signal) { const listener = () => this.handleAbort(); opts.signal.addEventListener("abort", listener, { once: true }); @@ -379,9 +397,12 @@ export class PersistentWs { this.pingTimer = undefined; } try { this.ws.close(1000, reason.slice(0, 120)); } catch { /* already closing */ } - if (this.currentSession && !this.currentSession.streamClosed) { - try { this.currentSession.controller.close(); } catch { /* already closed */ } - this.currentSession.streamClosed = true; + if (this.currentSession) { + this.clearResponseStartTimer(this.currentSession); + if (!this.currentSession.streamClosed) { + try { this.currentSession.controller.close(); } catch { /* already closed */ } + this.currentSession.streamClosed = true; + } } this.detachAbortListener(); this.busy = false; @@ -389,6 +410,20 @@ export class PersistentWs { try { this.hooks.onDead(); } catch { /* hook errors must not propagate */ } } + private handleResponseStartTimeout(timeoutMs: number): void { + const sess = this.currentSession; + if (!sess || sess.earlyDecisionMade) return; + sess.earlyDecisionMade = true; + sess.reject(new Error(`WebSocket response start timeout after ${timeoutMs}ms`)); + this.markDead("response start timeout"); + } + + private clearResponseStartTimer(sess: InFlightSession): void { + if (!sess.responseStartTimer) return; + clearTimeout(sess.responseStartTimer); + sess.responseStartTimer = undefined; + } + private detachAbortListener(): void { const sess = this.currentSession; if (sess?.signal && sess.abortListener) { @@ -414,6 +449,7 @@ export class PersistentWs { private resolveSessionResponse(sess: InFlightSession): void { if (sess.earlyDecisionMade) return; sess.earlyDecisionMade = true; + this.clearResponseStartTimer(sess); sess.resolveResponse(); for (const chunk of sess.earlyMetadataChunks.splice(0)) { this.enqueueSessionChunk(sess, chunk); @@ -557,6 +593,7 @@ export class PersistentWs { * treat early errors as account-level and keep the WS open only if the * error wasn't connection-fatal. */ private releaseAfterEarlyError(): void { + if (this.currentSession) this.clearResponseStartTimer(this.currentSession); this.detachAbortListener(); this.currentSession = null; this.busy = false; diff --git a/src/proxy/ws-transport.ts b/src/proxy/ws-transport.ts index 3878430a..2bfeeb8b 100644 --- a/src/proxy/ws-transport.ts +++ b/src/proxy/ws-transport.ts @@ -24,6 +24,7 @@ import { CodexApiError, PreviousResponseWebSocketError } from "./codex-types.js" import { getProxyUrl } from "../tls/proxy.js"; import { isPreviousResponseNotFoundError } from "./error-classification.js"; import { + DEFAULT_WS_RESPONSE_START_TIMEOUT_MS, PersistentWs, WsReusedConnectionError, type PersistentWsHooks, @@ -362,6 +363,7 @@ async function openOneShotWs( let expectedCloseBeforeOpen = false; const earlyMetadataChunks: Uint8Array[] = []; let pingTimer: ReturnType | undefined; + let responseStartTimer: ReturnType | undefined; // Open timeout: if the WS handshake never completes, reject after 20s. const openTimer = setTimeout(() => { @@ -373,8 +375,15 @@ async function openOneShotWs( } }, 20_000); + function clearResponseStartTimer() { + if (!responseStartTimer) return; + clearTimeout(responseStartTimer); + responseStartTimer = undefined; + } + function cleanupTimers() { clearTimeout(openTimer); + clearResponseStartTimer(); if (pingTimer) { clearInterval(pingTimer); pingTimer = undefined; @@ -420,6 +429,7 @@ async function openOneShotWs( function resolveResponse() { if (earlyDecisionMade) return; earlyDecisionMade = true; + clearResponseStartTimer(); resolve(buildResponse()); for (const chunk of earlyMetadataChunks.splice(0)) { enqueueChunk(chunk); @@ -472,6 +482,16 @@ async function openOneShotWs( console.log(`[WS-Open] 🟢 WebSocket successfully opened for request. wsUrl: ${wsUrl}`); clearTimeout(openTimer); ws.send(JSON.stringify(request)); + responseStartTimer = setTimeout(() => { + if (earlyDecisionMade) return; + earlyDecisionMade = true; + cleanupTimers(); + closeWs(1000, "response start timeout"); + reject(new Error( + `WebSocket response start timeout after ${DEFAULT_WS_RESPONSE_START_TIMEOUT_MS}ms`, + )); + }, DEFAULT_WS_RESPONSE_START_TIMEOUT_MS); + responseStartTimer.unref?.(); pingTimer = setInterval(() => { try { ws.ping(); } catch { /* ws already closed */ } }, 25_000); diff --git a/tests/unit/proxy/ws-pool.test.ts b/tests/unit/proxy/ws-pool.test.ts index c5f8eab2..eef02f6e 100644 --- a/tests/unit/proxy/ws-pool.test.ts +++ b/tests/unit/proxy/ws-pool.test.ts @@ -157,6 +157,56 @@ describe("PersistentWs", () => { expect(onDead).toHaveBeenCalled(); }); + describe("response start timeout", () => { + beforeEach(() => { vi.useFakeTimers(); }); + afterEach(() => { vi.useRealTimers(); }); + + it("rejects and evicts when upstream sends only provisional metadata", async () => { + const { ws, persistent, onDead } = newPersistentWs({ pingIntervalMs: 0 }); + persistent.tryAcquire(); + const promise = persistent.send({ + request: { type: "response.create", model: "m", instructions: "", input: [] }, + signal: undefined, + onRateLimits: undefined, + reused: false, + responseStartTimeoutMs: 1_000, + }); + promise.catch(() => undefined); + await vi.advanceTimersByTimeAsync(0); + ws.pushMessage({ type: "response.created", response: { id: "resp_waiting" } }); + ws.pushMessage({ type: "codex.response.metadata", headers: { "x-test": "1" } }); + + await vi.advanceTimersByTimeAsync(1_000); + + await expect(promise).rejects.toThrow("WebSocket response start timeout after 1000ms"); + expect(persistent.isAlive()).toBe(false); + expect(ws.closeReason).toBe("response start timeout"); + expect(onDead).toHaveBeenCalledTimes(1); + }); + + it("clears the deadline after the first client-visible event", async () => { + const { ws, persistent } = newPersistentWs({ pingIntervalMs: 0 }); + persistent.tryAcquire(); + const promise = persistent.send({ + request: { type: "response.create", model: "m", instructions: "", input: [] }, + signal: undefined, + onRateLimits: undefined, + reused: false, + responseStartTimeoutMs: 1_000, + }); + await vi.advanceTimersByTimeAsync(0); + ws.pushMessage({ type: "response.created", response: { id: "resp_started" } }); + ws.pushMessage({ type: "response.output_text.delta", delta: "started" }); + const response = await promise; + + await vi.advanceTimersByTimeAsync(5_000); + + expect(persistent.isAlive()).toBe(true); + ws.pushMessage({ type: "response.completed", response: { id: "resp_started" } }); + await expect(response.text()).resolves.toContain("response.completed"); + }); + }); + it("errors the response stream when the WS closes after a visible frame without terminal event", async () => { const { ws, persistent, onDead } = newPersistentWs(); persistent.tryAcquire(); diff --git a/tests/unit/proxy/ws-transport.test.ts b/tests/unit/proxy/ws-transport.test.ts index 5861cfe8..31c6109c 100644 --- a/tests/unit/proxy/ws-transport.test.ts +++ b/tests/unit/proxy/ws-transport.test.ts @@ -266,6 +266,60 @@ describe("createWebSocketResponse", () => { await expect(promise).rejects.toThrow("WebSocket closed before terminal event"); }); + it("rejects and closes a one-shot WS that sends only provisional metadata", async () => { + vi.useFakeTimers(); + try { + const promise = createWebSocketResponse("wss://test/ws", {}, BASE_REQUEST); + promise.catch(() => undefined); + await vi.advanceTimersByTimeAsync(0); + const ws = lastWs(); + ws.emit("message", JSON.stringify({ + type: "response.created", + response: { id: "resp_waiting" }, + })); + ws.emit("message", JSON.stringify({ + type: "codex.response.metadata", + headers: { "x-test": "1" }, + })); + + await vi.advanceTimersByTimeAsync(180_000); + + await expect(promise).rejects.toThrow("WebSocket response start timeout after 180000ms"); + expect(ws.readyState).toBe(3); + } finally { + vi.useRealTimers(); + } + }); + + it("clears the one-shot deadline after the first client-visible event", async () => { + vi.useFakeTimers(); + try { + const promise = createWebSocketResponse("wss://test/ws", {}, BASE_REQUEST); + await vi.advanceTimersByTimeAsync(0); + const ws = lastWs(); + ws.emit("message", JSON.stringify({ + type: "response.created", + response: { id: "resp_started" }, + })); + ws.emit("message", JSON.stringify({ + type: "response.output_text.delta", + delta: "started", + })); + const response = await promise; + + await vi.advanceTimersByTimeAsync(360_000); + + expect(ws.readyState).toBe(1); + ws.emit("message", JSON.stringify({ + type: "response.completed", + response: { id: "resp_started" }, + })); + await expect(readStream(response)).resolves.toContain("response.completed"); + } finally { + vi.useRealTimers(); + } + }); + it("ignores codex.rate_limits for early response resolution", async () => { let rateLimitCalled = false; const onRateLimits = () => { rateLimitCalled = true; }; From 9f98ff4cbb8c9b81e4c2bfaa3c1260c00f32643f Mon Sep 17 00:00:00 2001 From: shiitin Date: Mon, 20 Jul 2026 00:57:24 +0800 Subject: [PATCH 3/6] fix: bypass proxy for container healthcheck --- docker-healthcheck.sh | 4 +- tests/unit/ci/docker-healthcheck.test.ts | 58 ++++++++++++++++++++++++ 2 files changed, 61 insertions(+), 1 deletion(-) create mode 100644 tests/unit/ci/docker-healthcheck.test.ts diff --git a/docker-healthcheck.sh b/docker-healthcheck.sh index cbab987f..16b006d7 100644 --- a/docker-healthcheck.sh +++ b/docker-healthcheck.sh @@ -1,4 +1,6 @@ #!/bin/sh # Read server port from config, fallback to 8080 PORT=$(grep -A5 '^server:' /app/config/default.yaml 2>/dev/null | grep 'port:' | head -1 | awk '{print $2}') -curl -fs "http://localhost:${PORT:-8080}/health" || exit 1 +# Local health checks must never inherit an upstream HTTP proxy. +curl --noproxy '*' --fail --silent --show-error --max-time 3 \ + "http://127.0.0.1:${PORT:-8080}/health" >/dev/null || exit 1 diff --git a/tests/unit/ci/docker-healthcheck.test.ts b/tests/unit/ci/docker-healthcheck.test.ts new file mode 100644 index 00000000..58049762 --- /dev/null +++ b/tests/unit/ci/docker-healthcheck.test.ts @@ -0,0 +1,58 @@ +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import { spawn } from "child_process"; +import { createServer, type Server } from "http"; + +let server: Server; +let responseStatus = 200; + +beforeAll(async () => { + server = createServer((_req, res) => { + res.statusCode = responseStatus; + res.setHeader("content-type", "application/json"); + res.end('{"status":"ok"}'); + }); + await new Promise((resolve, reject) => { + server.once("error", reject); + server.listen(8080, "127.0.0.1", () => resolve()); + }); +}); + +afterAll(async () => { + await new Promise((resolve, reject) => { + server.close((err) => err ? reject(err) : resolve()); + }); +}); + +function runHealthcheck(): Promise { + return new Promise((resolve, reject) => { + const child = spawn("/bin/sh", ["docker-healthcheck.sh"], { + cwd: process.cwd(), + env: { + ...process.env, + HTTP_PROXY: "http://127.0.0.1:9", + HTTPS_PROXY: "http://127.0.0.1:9", + ALL_PROXY: "http://127.0.0.1:9", + http_proxy: "http://127.0.0.1:9", + https_proxy: "http://127.0.0.1:9", + all_proxy: "http://127.0.0.1:9", + NO_PROXY: "", + no_proxy: "", + }, + stdio: "ignore", + }); + child.once("error", reject); + child.once("close", (code) => resolve(code)); + }); +} + +describe("docker-healthcheck.sh", () => { + it("bypasses proxy variables for the local health endpoint", async () => { + responseStatus = 200; + await expect(runHealthcheck()).resolves.toBe(0); + }); + + it("returns a failure when the local health endpoint is unhealthy", async () => { + responseStatus = 503; + await expect(runHealthcheck()).resolves.not.toBe(0); + }); +}); From 20716aad5ebd4cab989a4fc4bf04cc468f4bdbab Mon Sep 17 00:00:00 2001 From: shiitin Date: Mon, 20 Jul 2026 01:00:44 +0800 Subject: [PATCH 4/6] test: cover continuation response-start timeout --- tests/unit/proxy/ws-pool.test.ts | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/tests/unit/proxy/ws-pool.test.ts b/tests/unit/proxy/ws-pool.test.ts index eef02f6e..b1c94d2c 100644 --- a/tests/unit/proxy/ws-pool.test.ts +++ b/tests/unit/proxy/ws-pool.test.ts @@ -184,6 +184,29 @@ describe("PersistentWs", () => { expect(onDead).toHaveBeenCalledTimes(1); }); + it("classifies a continuation timeout as a reused-connection failure", async () => { + const { ws, persistent } = newPersistentWs({ pingIntervalMs: 0 }); + persistent.tryAcquire(); + const promise = persistent.send({ + request: { type: "response.create", model: "m", instructions: "", input: [] }, + signal: undefined, + onRateLimits: undefined, + reused: true, + responseStartTimeoutMs: 1_000, + }); + promise.catch(() => undefined); + await vi.advanceTimersByTimeAsync(0); + ws.pushMessage({ type: "response.created", response: { id: "resp_waiting" } }); + + await vi.advanceTimersByTimeAsync(1_000); + + await expect(promise).rejects.toMatchObject({ + name: "WsReusedConnectionError", + message: "WebSocket response start timeout after 1000ms", + }); + expect(persistent.isAlive()).toBe(false); + }); + it("clears the deadline after the first client-visible event", async () => { const { ws, persistent } = newPersistentWs({ pingIntervalMs: 0 }); persistent.tryAcquire(); From dead7dc8f18418ce67e548d0c7ba2aa367e02b75 Mon Sep 17 00:00:00 2001 From: shiitin Date: Mon, 20 Jul 2026 07:27:44 +0800 Subject: [PATCH 5/6] fix: release one-shot timeout resources --- src/proxy/ws-transport.ts | 9 +++++---- tests/unit/proxy/ws-transport.test.ts | 14 +++++++++++--- 2 files changed, 16 insertions(+), 7 deletions(-) diff --git a/src/proxy/ws-transport.ts b/src/proxy/ws-transport.ts index 2bfeeb8b..697b0db8 100644 --- a/src/proxy/ws-transport.ts +++ b/src/proxy/ws-transport.ts @@ -484,12 +484,13 @@ async function openOneShotWs( ws.send(JSON.stringify(request)); responseStartTimer = setTimeout(() => { if (earlyDecisionMade) return; + const timeoutError = new Error( + `WebSocket response start timeout after ${DEFAULT_WS_RESPONSE_START_TIMEOUT_MS}ms`, + ); earlyDecisionMade = true; - cleanupTimers(); + errorStream(timeoutError); closeWs(1000, "response start timeout"); - reject(new Error( - `WebSocket response start timeout after ${DEFAULT_WS_RESPONSE_START_TIMEOUT_MS}ms`, - )); + reject(timeoutError); }, DEFAULT_WS_RESPONSE_START_TIMEOUT_MS); responseStartTimer.unref?.(); pingTimer = setInterval(() => { diff --git a/tests/unit/proxy/ws-transport.test.ts b/tests/unit/proxy/ws-transport.test.ts index 31c6109c..871b8d07 100644 --- a/tests/unit/proxy/ws-transport.test.ts +++ b/tests/unit/proxy/ws-transport.test.ts @@ -266,10 +266,17 @@ describe("createWebSocketResponse", () => { await expect(promise).rejects.toThrow("WebSocket closed before terminal event"); }); - it("rejects and closes a one-shot WS that sends only provisional metadata", async () => { + it("rejects, detaches abort handling, and closes a metadata-only one-shot WS", async () => { vi.useFakeTimers(); try { - const promise = createWebSocketResponse("wss://test/ws", {}, BASE_REQUEST); + const abortController = new AbortController(); + const removeAbortListener = vi.spyOn(abortController.signal, "removeEventListener"); + const promise = createWebSocketResponse( + "wss://test/ws", + {}, + BASE_REQUEST, + abortController.signal, + ); promise.catch(() => undefined); await vi.advanceTimersByTimeAsync(0); const ws = lastWs(); @@ -282,8 +289,9 @@ describe("createWebSocketResponse", () => { headers: { "x-test": "1" }, })); - await vi.advanceTimersByTimeAsync(180_000); + vi.advanceTimersByTime(180_000); + expect(removeAbortListener).toHaveBeenCalledWith("abort", expect.any(Function)); await expect(promise).rejects.toThrow("WebSocket response start timeout after 180000ms"); expect(ws.readyState).toBe(3); } finally { From 930dbd45d55f6f4f275d06a1caea338e81ebff12 Mon Sep 17 00:00:00 2001 From: shiitin Date: Mon, 20 Jul 2026 20:07:52 +0800 Subject: [PATCH 6/6] fix: harden healthcheck and websocket pool limits --- docker-healthcheck.sh | 7 +-- src/proxy/ws-pool.ts | 60 ++++++++++++++++-------- tests/unit/ci/docker-healthcheck.test.ts | 8 +++- tests/unit/proxy/ws-pool.test.ts | 45 ++++++++++++++++++ 4 files changed, 96 insertions(+), 24 deletions(-) diff --git a/docker-healthcheck.sh b/docker-healthcheck.sh index 16b006d7..e5c2ac07 100644 --- a/docker-healthcheck.sh +++ b/docker-healthcheck.sh @@ -1,6 +1,7 @@ #!/bin/sh -# Read server port from config, fallback to 8080 -PORT=$(grep -A5 '^server:' /app/config/default.yaml 2>/dev/null | grep 'port:' | head -1 | awk '{print $2}') +# Prefer the runtime port override, then config, then the default. +CONFIG_PORT=$(grep -A5 '^server:' /app/config/default.yaml 2>/dev/null | grep 'port:' | head -1 | awk '{print $2}') +HEALTHCHECK_PORT=${PORT:-${CONFIG_PORT:-8080}} # Local health checks must never inherit an upstream HTTP proxy. curl --noproxy '*' --fail --silent --show-error --max-time 3 \ - "http://127.0.0.1:${PORT:-8080}/health" >/dev/null || exit 1 + "http://127.0.0.1:${HEALTHCHECK_PORT}/health" >/dev/null || exit 1 diff --git a/src/proxy/ws-pool.ts b/src/proxy/ws-pool.ts index a619ebbd..6a748d13 100644 --- a/src/proxy/ws-pool.ts +++ b/src/proxy/ws-pool.ts @@ -649,6 +649,8 @@ export interface PersistentWsFactory { export class WsConnectionPool { private readonly map = new Map(); private readonly byEntry = new Map>(); + /** In-progress factories count against the per-account cap. */ + private readonly pendingCreatesByEntry = new Map(); /** Response IDs are valid only on the physical WS that completed them. */ private readonly ownerByResponse = new Map(); /** The upstream keeps only the most recent response per physical WS. */ @@ -700,32 +702,40 @@ export class WsConnectionPool { return { bypass: "busy" }; } - // Miss: enforce per-account cap before creating. + // Miss: reserve capacity before awaiting the factory so concurrent + // different-key acquires cannot all pass the same per-account cap check. const keys = this.byEntry.get(entryId); - if (keys && keys.size >= this.config.maxPerAccount) { + const pendingCreates = this.pendingCreatesByEntry.get(entryId) ?? 0; + if ((keys?.size ?? 0) + pendingCreates >= this.config.maxPerAccount) { return { bypass: "cap" }; } + this.pendingCreatesByEntry.set(entryId, pendingCreates + 1); let freshRef: PersistentWs | undefined; - const fresh = await factory({ - entryId, - poolKey, - hooks: { - onDead: () => { - // A same-key connection may have won the factory race. Never let a - // discarded fresh connection remove that winner from the pool. - if (freshRef && this.map.get(poolKey) === freshRef) { - this.removeEntryByKey(poolKey); - } - }, - onResponseCompleted: (responseId) => { - if (freshRef && this.map.get(poolKey) === freshRef) { - this.registerResponseOwner(poolKey, responseId); - } + let fresh: PersistentWs; + try { + fresh = await factory({ + entryId, + poolKey, + hooks: { + onDead: () => { + // A same-key connection may have won the factory race. Never let a + // discarded fresh connection remove that winner from the pool. + if (freshRef && this.map.get(poolKey) === freshRef) { + this.removeEntryByKey(poolKey); + } + }, + onResponseCompleted: (responseId) => { + if (freshRef && this.map.get(poolKey) === freshRef) { + this.registerResponseOwner(poolKey, responseId); + } + }, }, - }, - }); - freshRef = fresh; + }); + freshRef = fresh; + } finally { + this.releasePendingCreate(entryId); + } // Race: another acquire for the same key may have completed during // factory() await. If so, prefer the one already in the map. @@ -842,6 +852,7 @@ export class WsConnectionPool { // acquires would fail the disabled check anyway. this.map.clear(); this.byEntry.clear(); + this.pendingCreatesByEntry.clear(); this.ownerByResponse.clear(); this.responseByPoolKey.clear(); } @@ -856,6 +867,15 @@ export class WsConnectionPool { } } + private releasePendingCreate(entryId: string): void { + const pendingCreates = this.pendingCreatesByEntry.get(entryId); + if (!pendingCreates || pendingCreates <= 1) { + this.pendingCreatesByEntry.delete(entryId); + return; + } + this.pendingCreatesByEntry.set(entryId, pendingCreates - 1); + } + private registerResponseOwner(poolKey: string, responseId: string): void { if (!this.map.has(poolKey)) return; const previous = this.responseByPoolKey.get(poolKey); diff --git a/tests/unit/ci/docker-healthcheck.test.ts b/tests/unit/ci/docker-healthcheck.test.ts index 58049762..224a7c4c 100644 --- a/tests/unit/ci/docker-healthcheck.test.ts +++ b/tests/unit/ci/docker-healthcheck.test.ts @@ -1,9 +1,11 @@ import { afterAll, beforeAll, describe, expect, it } from "vitest"; import { spawn } from "child_process"; import { createServer, type Server } from "http"; +import type { AddressInfo } from "net"; let server: Server; let responseStatus = 200; +let serverPort: number; beforeAll(async () => { server = createServer((_req, res) => { @@ -13,7 +15,10 @@ beforeAll(async () => { }); await new Promise((resolve, reject) => { server.once("error", reject); - server.listen(8080, "127.0.0.1", () => resolve()); + server.listen(0, "127.0.0.1", () => { + serverPort = (server.address() as AddressInfo).port; + resolve(); + }); }); }); @@ -37,6 +42,7 @@ function runHealthcheck(): Promise { all_proxy: "http://127.0.0.1:9", NO_PROXY: "", no_proxy: "", + PORT: String(serverPort), }, stdio: "ignore", }); diff --git a/tests/unit/proxy/ws-pool.test.ts b/tests/unit/proxy/ws-pool.test.ts index b1c94d2c..beafdb44 100644 --- a/tests/unit/proxy/ws-pool.test.ts +++ b/tests/unit/proxy/ws-pool.test.ts @@ -646,6 +646,51 @@ describe("WsConnectionPool", () => { await capped.shutdown(); }); + it("counts pending factories against the per-account cap", async () => { + const capped = new WsConnectionPool({ maxPerAccount: 1 }, { startGc: false }); + let resolveFactory: ((ws: PersistentWs) => void) | undefined; + const factory = vi.fn((deps: { entryId: string; poolKey: string; hooks: PersistentWsHooks }) => + new Promise((resolve) => { + resolveFactory = resolve; + }), + ); + + const firstPromise = capped.acquire("entry-A", "entry-A:conv-1", factory); + await nextTick(); + const second = await capped.acquire("entry-A", "entry-A:conv-2", factory); + + expect(second).toEqual({ bypass: "cap" }); + expect(factory).toHaveBeenCalledTimes(1); + + const firstDeps = factory.mock.calls[0][0]; + resolveFactory?.(new PersistentWs({ + ws: new MockWs(), + entryId: firstDeps.entryId, + poolKey: firstDeps.poolKey, + hooks: firstDeps.hooks, + })); + await expect(firstPromise).resolves.toMatchObject({ reused: false }); + expect(capped.countByEntryId("entry-A")).toBe(1); + await capped.shutdown(); + }); + + it("releases pending capacity when a factory fails", async () => { + const capped = new WsConnectionPool({ maxPerAccount: 1 }, { startGc: false }); + const failedFactory = vi.fn(async () => { + throw new Error("connect failed"); + }); + await expect( + capped.acquire("entry-A", "entry-A:conv-failed", failedFactory), + ).rejects.toThrow("connect failed"); + + const { factory } = makeFactory(); + await expect( + capped.acquire("entry-A", "entry-A:conv-retry", factory), + ).resolves.toMatchObject({ reused: false }); + expect(factory).toHaveBeenCalledTimes(1); + await capped.shutdown(); + }); + it("dead connection is treated as a miss on next acquire", async () => { const factories: MockWs[] = []; const factory = vi.fn(async (deps: { entryId: string; poolKey: string; hooks: PersistentWsHooks }) => {