diff --git a/docs-site/src/content/docs/guides/codex-integration.md b/docs-site/src/content/docs/guides/codex-integration.md index d4f611149c4..82cf7aaf446 100644 --- a/docs-site/src/content/docs/guides/codex-integration.md +++ b/docs-site/src/content/docs/guides/codex-integration.md @@ -427,6 +427,14 @@ provider advertises `supports_websockets = true` only when `"websockets": true`; built-in provider may try WebSocket first, and a disabled proxy returns `426` so Codex falls back to HTTP/SSE. +The proxy admits up to 128 client-facing Responses WebSocket connections. When a new connection +finds all 128 slots occupied, it closes the least recently used idle connection to make room. Connections +with an active response, pending client tool or approval results, a native steering/injection +session, or buffered output are protected. If no idle connection can be released, the new +connection receives the existing `503` capacity response. A reclaimed client can reconnect for +later requests using the existing continuation handling. This happens only under connection +pressure, needs no setting, and does not change voice connections or upstream connection reuse. + If a canonical ChatGPT forward continuation references expired or missing local replay state, opencodex returns `previous_response_not_found` before sending anything upstream. Codex's WebSocket client recognizes this error and can reconnect with its full retained context, diff --git a/scripts/test-layout/layout.json b/scripts/test-layout/layout.json index ef6d6d5c638..01958ae32f9 100644 --- a/scripts/test-layout/layout.json +++ b/scripts/test-layout/layout.json @@ -722,6 +722,7 @@ "codex-warmup.test.ts": "codex-integration", "codex-web-search-switch.test.ts": "codex-integration", "codex-websocket-registry.test.ts": "codex-integration", + "websocket-capacity-activity.test.ts": "codex-integration", "codex-write-lock.test.ts": "codex-integration", "coding-agent-tool-result-images.test.ts": "adapters", "cold-spawn-warmup.test.ts": "ci-workflows", @@ -1919,6 +1920,8 @@ "winsw.test.ts": "service", "workflow-budget.test.ts": "lib", "ws-ambiguous-resend.test.ts": "responses", "ws-pre-output-refusal.test.ts": "responses", "responses-pool-overload-failover.test.ts": "responses", "compact-transient-alternate.test.ts": "responses", + "ws-capacity-reclamation.test.ts": "responses", + "ws-capacity-tool-ownership.test.ts": "responses", "ws-endpoint.test.ts": "responses", "ws-failure-stage.test.ts": "responses", "ws-native-injection.test.ts": "responses", diff --git a/src/codex/websocket-activity.ts b/src/codex/websocket-activity.ts new file mode 100644 index 00000000000..5173cec7067 --- /dev/null +++ b/src/codex/websocket-activity.ts @@ -0,0 +1,161 @@ +import { createHash } from "node:crypto"; +import { isAuthorizedProviderExecutedCall, type ProviderExecutedCallType } from "../server/responses-undeclared-tool-guard"; + +// Logical recency avoids clock adjustments and ties between messages in one millisecond. +let activityOrder = 0; +const MAX_PENDING_RESULTS = 1024; +const MAX_RESULT_ID_LENGTH = 1024; +const resultTypes: Readonly> = { + function_call: "function_call_output", + custom_tool_call: "custom_tool_call_output", + local_shell_call: "function_call_output", + shell_call: "shell_call_output", + computer_call: "computer_call_output", + apply_patch_call: "apply_patch_call_output", + tool_search_call: "tool_search_output", + mcp_approval_request: "mcp_approval_response", +}; + +function record(value: unknown): value is Record { + return value !== null && typeof value === "object" && !Array.isArray(value); +} + +function resultKey(type: string, id: string): string { + return createHash("sha256").update(JSON.stringify([type, id])).digest("hex"); +} + +/** Content-free activity owned by one admitted incoming Responses socket. */ +export class CodexWebSocketActivity { + lastUsed = ++activityOrder; + closing = false; + private turns = 0; + private pendingResults = new Set(); + private returnedResults = new Set(); + // Hashed identities only; null marks a route-authorized provider-owned call. + private argumentCalls = new Map(); + // Unknown/oversized identities cannot safely prove that client work finished. + private untrackedResults = false; + + get idle(): boolean { + return !this.closing && this.turns === 0 && this.pendingResults.size === 0 && !this.untrackedResults; + } + + touch(): void { + this.lastUsed = ++activityOrder; + } + + /** Each turn retains its own busy reference, including superseded async cleanup. */ + beginTurn(input: unknown): () => void { + this.turns += 1; + this.touch(); + this.acceptResults(input); + // Ordinary creates start a new response even when its upstream omits response.created. + // Prior results settle old waits, but cannot suppress new calls reusing their IDs. + this.returnedResults.clear(); + this.argumentCalls.clear(); + let finished = false; + return () => { + if (finished) return; + finished = true; + this.turns -= 1; + if (this.turns === 0) this.touch(); + }; + } + + /** Called only after admission or native control accepts the client's results. */ + acceptResults(input: unknown): void { + if (this.untrackedResults) return; + if (Array.isArray(input)) { + for (const item of input) { + if (!record(item) || typeof item.type !== "string") continue; + const id = item.type === "mcp_approval_response" ? item.approval_request_id : item.call_id; + if (typeof id === "string" && id.length <= MAX_RESULT_ID_LENGTH) { + // Ordinary Codex shell results use function_call_output; native steering + // also accepts its explicit local_shell_call_output continuation stub. + const type = item.type === "local_shell_call_output" ? "function_call_output" : item.type; + const key = resultKey(type, id); + if (this.pendingResults.delete(key)) { + if (this.returnedResults.size >= MAX_PENDING_RESULTS) { this.protectUntracked(); return; } + else this.returnedResults.add(key); + } + } + } + } + } + + /** Preserve tool waits even when the response terminal omits its output array. */ + observePayload(payload: string, providerExecuted?: ReadonlySet): void { + if (this.untrackedResults) return; + let frame: unknown; + try { frame = JSON.parse(payload); } catch { return; } + if (!record(frame)) return; + // Native injection can return a result before the same terminal repeats its call. + if (frame.type === "response.created") { + this.returnedResults.clear(); + this.argumentCalls.clear(); + } + if (frame.type === "response.output_item.added") this.rememberArgumentCall(frame.item, providerExecuted); + if (frame.type === "response.function_call_arguments.done" || frame.type === "response.custom_tool_call_input.done") { + const type = frame.type === "response.function_call_arguments.done" ? "function_call" : "custom_tool_call"; + const call = typeof frame.item_id === "string" && frame.item_id.length <= MAX_RESULT_ID_LENGTH + ? this.argumentCalls.get(resultKey("item", frame.item_id)) : undefined; + if (!call || call.type !== type) this.protectUntracked(); + else if (call.result !== null) this.observeResultKey(call.result); + } + if (frame.type === "response.output_item.done") this.observeOutput(frame.item, providerExecuted); + if (["response.completed", "response.incomplete", "response.failed"].includes(String(frame.type)) + && record(frame.response) && Array.isArray(frame.response.output)) { + for (const item of frame.response.output) this.observeOutput(item, providerExecuted); + } + } + + private rememberArgumentCall(item: unknown, providerExecuted?: ReadonlySet): void { + if (!record(item) || (item.type !== "function_call" && item.type !== "custom_tool_call")) return; + if (typeof item.id !== "string" || !item.id.length || item.id.length > MAX_RESULT_ID_LENGTH) return; + const key = resultKey("item", item.id); + const hosted = providerExecuted && isAuthorizedProviderExecutedCall(item, providerExecuted); + if (!hosted && (typeof item.call_id !== "string" || !item.call_id.length || item.call_id.length > MAX_RESULT_ID_LENGTH)) { + if (this.argumentCalls.has(key)) this.protectUntracked(); + return; + } + const result = hosted ? null : resultKey(resultTypes[item.type]!, item.call_id as string); + const previous = this.argumentCalls.get(key); + if (previous) { + if (previous.type !== item.type || previous.result !== result) this.protectUntracked(); + return; + } + if (this.argumentCalls.size >= MAX_PENDING_RESULTS) { this.protectUntracked(); return; } + // An unfinished argument stream alone does not prove the client started work. + this.argumentCalls.set(key, { type: item.type, result }); + } + + private observeOutput(item: unknown, providerExecuted?: ReadonlySet): void { + if (this.untrackedResults) return; + if (!record(item) || typeof item.type !== "string") return; + // Hosted execution may publish both the call and its result in the output stream. + this.acceptResults([item]); + if (!Object.hasOwn(resultTypes, item.type)) return; + if (providerExecuted && isAuthorizedProviderExecutedCall(item, providerExecuted)) return; + if (item.type === "tool_search_call" && item.execution !== "client") return; + const id = item.type === "mcp_approval_request" ? item.id + : item.type === "local_shell_call" || item.type === "tool_search_call" ? item.call_id ?? item.id : item.call_id; + if (typeof id !== "string" || !id.length || id.length > MAX_RESULT_ID_LENGTH) { + this.protectUntracked(); + return; + } + this.observeResultKey(resultKey(resultTypes[item.type]!, id)); + } + + private observeResultKey(key: string): void { + if (this.pendingResults.has(key) || this.returnedResults.has(key)) return; + if (this.pendingResults.size >= MAX_PENDING_RESULTS) this.protectUntracked(); + else this.pendingResults.add(key); + } + + private protectUntracked(): void { + this.untrackedResults = true; + this.pendingResults.clear(); + this.returnedResults.clear(); + this.argumentCalls.clear(); + } +} diff --git a/src/codex/websocket-registry.ts b/src/codex/websocket-registry.ts index 91af86aecb4..405179972b0 100644 --- a/src/codex/websocket-registry.ts +++ b/src/codex/websocket-registry.ts @@ -1,23 +1,47 @@ import type { ServerWebSocket } from "bun"; import type { WsData } from "../server/ws-bridge"; import { createAdmissionGate, type AdmissionMetrics, type AdmissionReservation } from "../lib/admission"; +import { CodexWebSocketActivity } from "./websocket-activity"; export const MAX_TRACKED_CODEX_WEBSOCKETS = 128; const websocketGate = createAdmissionGate("codex_websockets", MAX_TRACKED_CODEX_WEBSOCKETS); +const admittedSockets = new Set>(); + +function reclaimOldestIdleSocket(): void { + let oldest: ServerWebSocket | undefined; + for (const ws of admittedSockets) { + const activity = ws.data.responsesActivity; + if (!activity?.idle || ws.data.nativeControl || ws.readyState !== 1 || ws.getBufferedAmount() !== 0) continue; + if (!oldest || activity.lastUsed < oldest.data.responsesActivity!.lastUsed) oldest = ws; + } + if (!oldest) return; + // Fence late frames before closing. Only the actual close callback releases admission. + oldest.data.responsesActivity!.closing = true; + try { + oldest.close(1001, "Idle Responses connection reclaimed"); + } catch { + // A failed/delayed close never manufactures a free slot or retries another victim. + if (oldest.readyState === 1) oldest.data.responsesActivity!.closing = false; + } +} export function tryReserveCodexWebSocket(): AdmissionReservation> | null { + if (websocketGate.metrics().active >= MAX_TRACKED_CODEX_WEBSOCKETS) reclaimOldestIdleSocket(); const gateLease = websocketGate.tryAcquire(); if (!gateLease) return null; let socket: ServerWebSocket | undefined; let active = true; return { bind(value) { - if (!active) return; + if (!active || socket) return; socket = value; + socket.data.responsesActivity = new CodexWebSocketActivity(); + admittedSockets.add(socket); }, release() { if (!active) return; active = false; + if (socket) admittedSockets.delete(socket); socket = undefined; gateLease.release(); }, diff --git a/src/server/index/live-sideband.ts b/src/server/index/live-sideband.ts index 8c4b639eb5f..d0840bdc626 100644 --- a/src/server/index/live-sideband.ts +++ b/src/server/index/live-sideband.ts @@ -40,8 +40,9 @@ export const MAX_WS_FRAME_BYTES = 50 * 1024 * 1024; * * What must never happen is neither: an immortal socket plus a destination that silently accepts * the orphaned delta. `tests/responses/ws-endpoint.test.ts` holds exactly that pair together. - * Raising the timeout off 0 is still worth doing for its own reasons (a dead peer holds a socket - * forever today), and Bun caps the value at 255 seconds, well inside the bound below. + * Responses admission separately reclaims its least recently used idle connection under + * capacity pressure. That is not a time limit, so expired-reference validation remains required. + * Keeping this shared timeout at 0 also preserves quiet live/voice sessions. */ export const WEBSOCKET_IDLE_TIMEOUT_SECONDS = 0; /** Ceiling a finite websocket idle timeout must stay under, in seconds. See above. */ diff --git a/src/server/index/websocket-handler.ts b/src/server/index/websocket-handler.ts index b0e4057517e..e39f2e71403 100644 --- a/src/server/index/websocket-handler.ts +++ b/src/server/index/websocket-handler.ts @@ -89,6 +89,7 @@ import { import type { ServeOptionsContext } from "./serve-options"; import type { RequestMetricsRecorder } from "../request-metrics"; import { resolveInboundBodyLimitBytes } from "../request-decompress"; +import { responseToolOwnership } from "../responses/response-tool-ownership"; /** * The WebSocket half of the Bun.serve options, split out of serve-options.ts to keep that file @@ -184,6 +185,8 @@ export function createWebsocketHandler( } return; } + if (ws.data.responsesActivity?.closing) return; + ws.data.responsesActivity?.touch(); const rawBytes = typeof raw === "string" ? Buffer.byteLength(raw) : raw.byteLength; if (rawBytes > MAX_WS_FRAME_BYTES) { sendJsonFrame(ws, buildWsErrorFrame(413, { @@ -223,6 +226,7 @@ export function createWebsocketHandler( if (frame.type === "response.inject") { if (!ws.data.nativeControl?.inject) throw new NativeSteeringError("injection_not_supported", "Native injection is disabled or unavailable on this route."); ws.data.nativeControl.inject(frame); + ws.data.responsesActivity?.acceptResults(frame.input); return; } if (frame.type === "response.steer") { @@ -230,7 +234,10 @@ export function createWebsocketHandler( ws.data.nativeControl.steer(frame); return; } - if (ws.data.nativeControl?.continue(frame)) return; + if (ws.data.nativeControl?.continue(frame)) { + ws.data.responsesActivity?.acceptResults(frame.input); + return; + } } catch (error) { sendJsonFrame(ws, buildWsErrorFrame(400, { type: "invalid_request_error", @@ -294,6 +301,7 @@ export function createWebsocketHandler( } // Only a genuinely admitted turn may receive steering or continuations. + const finishActivity = ws.data.responsesActivity?.beginTurn(frame.input); ws.data.nativeControl = nativeControl; const payload: Record = { ...frame }; delete payload.type; @@ -346,11 +354,16 @@ export function createWebsocketHandler( terminalRecorder = recorder; }, }); + const observeLog = nativeControl?.relayActive + ? createNativeSteeringLogObserver(logCtx, () => recordFirstOutput(logCtx, start)) + : (payload: string) => inspectResponseLogSsePayload(logCtx, payload); + const providerExecuted = responseToolOwnership(response); await sendResponseToWebSocket(ws, response, isCurrent, { untilEof: nativeControl?.relayActive === true, - onSsePayload: nativeControl?.relayActive - ? createNativeSteeringLogObserver(logCtx, () => recordFirstOutput(logCtx, start)) - : payload => inspectResponseLogSsePayload(logCtx, payload), + onSsePayload: payload => { + ws.data.responsesActivity?.observePayload(payload, providerExecuted); + observeLog(payload); + }, onTerminal: status => { terminalRecorder?.(status, logCtx.terminalHttpStatus); finalizeLog(httpStatusForRequestLogTerminal(status, logCtx), { @@ -392,6 +405,7 @@ export function createWebsocketHandler( /* socket already gone or send dropped */ } } finally { + finishActivity?.(); turnAdmissionLease.release(); if (ws.data.nativeControl === nativeControl) ws.data.nativeControl = undefined; if (!logged && turnAbort.signal.aborted) finalizeLog(499); diff --git a/src/server/lifecycle.ts b/src/server/lifecycle.ts index 97d1d5caa72..c9e6670ee14 100644 --- a/src/server/lifecycle.ts +++ b/src/server/lifecycle.ts @@ -38,6 +38,7 @@ export interface ActiveTurnLease extends AdmissionLease { attach(lease: AdmissionLease): void; bindAbortSignal(signal: AbortSignal): () => void; bindAbortController(ac: AbortController): void; + unbindAbortController(ac: AbortController): void; beginCodexAccountSelection(): CodexAccountSelectionAdmission; isTransferred(): boolean; } @@ -231,6 +232,17 @@ export function tryAdmitTurn(sessionLaneId?: string, clientSignal?: AbortSignal) activeTurns.set(ac, lease); lease.bindAbortSignal(ac.signal); }, + unbindAbortController(ac) { + if (!controllers.delete(ac)) return; + if (activeTurns.get(ac) === lease) activeTurns.delete(ac); + const listener = abortListeners.get(ac.signal); + if (listener) { + ac.signal.removeEventListener("abort", listener); + abortListeners.delete(ac.signal); + } + // A completed child relay cannot retire its still-running parent turn. + if (controllers.size === 0) lease.release(); + }, beginCodexAccountSelection() { const mainProfileDraining = nativeMainDrainOwners.size > 0; let selectionActive = active && !mainProfileDraining; @@ -327,7 +339,7 @@ export function unregisterTurn(ac: AbortController): void { turnReleaseMisses += 1; return; } - lease.release(); + lease.unbindAbortController(ac); } export function isDraining(): boolean { return shutdownDraining || temporaryDrainOwners.size > 0; } export function getActiveTurnCount(): number { return turnGate.metrics().active; } diff --git a/src/server/responses-undeclared-tool-guard.ts b/src/server/responses-undeclared-tool-guard.ts index e0a8b7029ad..d64a802b2ab 100644 --- a/src/server/responses-undeclared-tool-guard.ts +++ b/src/server/responses-undeclared-tool-guard.ts @@ -298,7 +298,7 @@ export function collectProviderExecutedCallTypes(body: unknown): Set, callTypes: ProviderExecutedCallTypes, ): boolean { diff --git a/src/server/responses/combo-stream-preflight.ts b/src/server/responses/combo-stream-preflight.ts index 390097b0762..5cdc1b4ddb6 100644 --- a/src/server/responses/combo-stream-preflight.ts +++ b/src/server/responses/combo-stream-preflight.ts @@ -4,6 +4,7 @@ import { comboFailureDecision } from "../../combos"; import { httpStatusFromTerminalError } from "../../lib/errors"; import type { RequestFailureStage } from "../../lib/request-failure-model"; import type { RequestLogContext } from "../request-log"; +import { carryResponseToolOwnership } from "./response-tool-ownership"; import { createSseInspector } from "../relay"; import { MAX_CLIENT_SSE_FRAME_BYTES } from "../sse-frame-buffer"; @@ -198,11 +199,11 @@ function replayBufferedResponse( reader.cancel(reason).catch(() => undefined); }, }); - return new Response(body, { + return carryResponseToolOwnership(response, new Response(body, { status: response.status, statusText: response.statusText, headers: response.headers, - }); + })); } function failedTerminalResponse( diff --git a/src/server/responses/compaction-recovery.ts b/src/server/responses/compaction-recovery.ts index 1b23aa6744a..f35c6a8d39f 100644 --- a/src/server/responses/compaction-recovery.ts +++ b/src/server/responses/compaction-recovery.ts @@ -1,4 +1,5 @@ import type { AdapterEvent, OcxConfig } from "../../types"; +import { carryResponseToolOwnership } from "./response-tool-ownership"; import { routeConcreteModel, type RouteResult } from "../../router"; import { copyPlainData } from "../../lib/plain-data"; import { jsonUtf8Bytes } from "../../lib/json-byte-size"; @@ -78,6 +79,7 @@ async function bufferedJson(response: Response, signal: AbortSignal): Promise<{ headers.delete("content-length"); headers.delete("content-encoding"); const replacement = new Response(bytes.bytes, { status: response.status, statusText: response.statusText, headers }); + carryResponseToolOwnership(response, replacement); if (isNonReplayableResponse(response)) markResponseNonReplayable(replacement); try { const json: unknown = JSON.parse(new TextDecoder("utf-8", { fatal: true }).decode(bytes.bytes)); diff --git a/src/server/responses/core-lifetime.ts b/src/server/responses/core-lifetime.ts index 2dd337ae534..4050f4be44f 100644 --- a/src/server/responses/core-lifetime.ts +++ b/src/server/responses/core-lifetime.ts @@ -1,4 +1,5 @@ import type { TranslatorBudget } from "../../lib/translator-budget"; +import { carryResponseToolOwnership } from "./response-tool-ownership"; import { finalizeTranslatorBudgetResponse } from "../../lib/translator-budget"; import { isNonReplayableResponse, markResponseNonReplayable } from "../../lib/upstream-retry"; import { @@ -50,7 +51,7 @@ export function finalizeOwnedTranslatorBudget(response: Response, budget: Transl if (isPreinspectedJsonResponse(response)) markPreinspectedJsonResponse(finalizedResponse); if (isNonReplayableResponse(response)) markResponseNonReplayable(finalizedResponse); if (runTurnAdapterSseResponses.has(response)) runTurnAdapterSseResponses.add(finalizedResponse); - return finalizedResponse; + return carryResponseToolOwnership(response, finalizedResponse); } /** Release a serving-account slot when the client body ends, errors or is cancelled. */ @@ -75,7 +76,7 @@ export function finalizeAccountLease(response: Response, release: () => void, on if (isEagerRelaySseResponse(response)) markEagerRelaySseResponse(wrapped); if (isPreinspectedJsonResponse(response)) markPreinspectedJsonResponse(wrapped); if (runTurnAdapterSseResponses.has(response)) runTurnAdapterSseResponses.add(wrapped); - return wrapped; + return carryResponseToolOwnership(response, wrapped); } diff --git a/src/server/responses/passthrough-execution.ts b/src/server/responses/passthrough-execution.ts index c41c1c3acc9..5d72e561493 100644 --- a/src/server/responses/passthrough-execution.ts +++ b/src/server/responses/passthrough-execution.ts @@ -14,6 +14,7 @@ import { guardDirectPassthroughBodyInactivity } from "../../lib/response-body-in import { resolveStallTimeoutMs } from "../../stall-timeout"; import { releaseUpstreamHostAdmission } from "../../codex/upstream-host-health"; import { releaseCodexAuthContextProbeLease } from "../../codex/auth-context"; +import { attachResponseToolOwnership } from "./response-tool-ownership"; /** Owns the native host lease across dispatch, recovery, and response construction. */ export async function executePassthroughResponse( @@ -50,11 +51,11 @@ export async function executePassthroughResponse( // Delivery has already classified SSE (including a missing upstream content type) // and consumed bounded JSON/errors. Guard only its remaining direct body, outside // the relay/lifetime wrappers so their prefetch cannot arm our inactivity clock. - return guardDirectPassthroughBodyInactivity( + return attachResponseToolOwnership(guardDirectPassthroughBodyInactivity( response, nativeExchange.upstream.signal, resolveStallTimeoutMs(requestContext.config.stallTimeoutSec, { localUpstream: nativeExchange.localUpstream }), - ); + ), nativeExchange.providerExecutedCallTypes); } finally { if (nativeHostState.lease) { releaseUpstreamHostAdmission(nativeHostState.lease); diff --git a/src/server/responses/response-tool-ownership.ts b/src/server/responses/response-tool-ownership.ts new file mode 100644 index 00000000000..323095a24da --- /dev/null +++ b/src/server/responses/response-tool-ownership.ts @@ -0,0 +1,21 @@ +import type { ProviderExecutedCallType } from "../responses-undeclared-tool-guard"; + +// Route-authorized execution ownership follows only the response selected for delivery. +// It never becomes a wire header or mutable state shared by competing route attempts. +const providerCalls = new WeakMap>(); + +export function attachResponseToolOwnership(response: Response, types: ReadonlySet): Response { + if (types.size > 0) providerCalls.set(response, new Set(types)); + return response; +} + +export function responseToolOwnership(response: Response): ReadonlySet | undefined { + return providerCalls.get(response); +} + +/** Carry ownership only when rebuilding the same selected response body. */ +export function carryResponseToolOwnership(source: Response, target: Response): Response { + const types = providerCalls.get(source); + if (types) providerCalls.set(target, types); + return target; +} diff --git a/src/server/ws-bridge.ts b/src/server/ws-bridge.ts index da2d2c50883..4009a48c1d7 100644 --- a/src/server/ws-bridge.ts +++ b/src/server/ws-bridge.ts @@ -11,6 +11,7 @@ import { BoundedSseFrameBuffer } from "./sse-frame-buffer"; import { sseDataPayload as parseSseBlock } from "./sse-payload-rewrite"; import { safeResponseHeaders } from "./safe-response-headers"; import type { AudioSocketTarget } from "./audio-dictation"; +import type { CodexWebSocketActivity } from "../codex/websocket-activity"; export { safeResponseHeaders } from "./safe-response-headers"; @@ -19,6 +20,7 @@ type ResponsesTerminalReporter = (status: ResponsesTerminalStatus) => void; type ResponsesPayloadObserver = (payload: string) => void; export interface WsData { + responsesActivity?: CodexWebSocketActivity; nativeControl?: NativeResponseControl; /** Content-free per-turn explanation; never a model capability assertion. */ nativeSteeringUnavailable?: string; diff --git a/structure/transports/responses.md b/structure/transports/responses.md index 771530b8d2f..97b4e0b747d 100644 --- a/structure/transports/responses.md +++ b/structure/transports/responses.md @@ -71,7 +71,7 @@ silence. It cannot observe the client's workspace or prove whether a successful repository state. Consequently, the active-turn and session-lane gates are concurrency admission limits, the translator budget follows [live retained-byte ownership](byte-accounting.md#stream-buffer-accounting), the response-state caps are cache retention limits, and the stall watchdog is a silence limit. None is a cumulative continuation or -semantic no-progress budget. The turn lease in `src/server/lifecycle.ts` subscribes to HTTP ingress cancellation, internal decision cancellation and every registered controller. Cancellation snapshots only that turn's controllers, releases its admission, session-lane reference, attached workflow leases and native-main ownership, then aborts the snapshot with the cancellation reason. It does not await upstream work or owner finalization, even when the owner remains suspended; other turns retain their controllers and admission. Late controllers abort immediately without registering or transferring ownership, and cancelled selections cannot claim a native-main account. Live/Realtime acquisition binds ingress cancellation until WebSocket upgrade transfers ownership. Normal completion detaches the listeners; late finalizers remain idempotent and do not release newly admitted turns. `tests/codex-integration/active-turn-disconnect.test.ts` covers these cancellation and ownership boundaries. +semantic no-progress budget. The turn lease in `src/server/lifecycle.ts` subscribes to HTTP ingress cancellation, internal decision cancellation and every registered controller. Cancellation snapshots only that turn's controllers, releases its admission, session-lane reference, attached workflow leases and native-main ownership, then aborts the snapshot with the cancellation reason. It does not await upstream work or owner finalization, even when the owner remains suspended; other turns retain their controllers and admission. Late controllers abort immediately without registering or transferring ownership, and cancelled selections cannot claim a native-main account. Live/Realtime acquisition binds ingress cancellation until WebSocket upgrade transfers ownership. Normal controller completion detaches only that controller and its listener; the last owner or an explicit parent release settles the lease. A completed combo relay cannot retire its parent WebSocket turn or cancel the next attempt, and a completed child's late abort cannot cancel former siblings. Late finalizers remain idempotent. `tests/codex-integration/active-registry-admission.test.ts` covers these shared-owner boundaries. `tests/codex-integration/active-turn-disconnect.test.ts` covers these cancellation and ownership boundaries. > Decision record: [ADR-0031](../decisions/ADR-0031-responses-http-sse.md) > Decision record: [ADR-0032](../decisions/ADR-0032-responses-http-sse.md) diff --git a/structure/transports/streaming-health.md b/structure/transports/streaming-health.md index a3a8becb95a..d93dad0b0bb 100644 --- a/structure/transports/streaming-health.md +++ b/structure/transports/streaming-health.md @@ -256,6 +256,38 @@ egress. HTTP/SSE fallback retains Bun fetch's own proxy rules, which do not cons The endpoint handles `response.create`, ignores `response.processed`, supports warmup `generate: false`, and feeds the same request pipeline as HTTP/SSE. +Incoming Responses connections share a 128-socket admission limit. Only when a new upgrade +finds that limit full, the handler reclaims the least recently used eligible idle connection. +Message handling and turn completion refresh a monotonic use order; connection age alone +does not choose the victim. Eligibility requires an open socket, no active turn, no pending +client tool or approval result, no retained native steering/injection owner, and no buffered +outbound bytes. A superseded turn remains active until its own cleanup finishes. + +Tool and approval identities are tracked across output items and terminal snapshots. Function argument +and custom-tool input completion also establish waits, correlated through `response.output_item.added`, +even if EOF precedes the completed item or terminal. An added item with unfinished arguments alone +establishes no wait. Only matching client or upstream results retire waits, suppressing completion echoes. +Tracking retains bounded identity hashes; overflow or ambiguous completion identities protect the +connection until close. A terminal does not make a connection idle while its client tool runs. + +The tracker reuses the final route's provider-executed tool classification, including declarations +injected into the outbound request. `src/server/responses/response-tool-ownership.ts` carries that +classification with the selected Response through body wrappers; failed combo attempts cannot +authorize the winning response's calls. Hosted xAI search therefore creates no client-result wait. +Each admitted ordinary turn starts a fresh returned-result deduplication scope after settling prior +results, even if its upstream omits `response.created`. Native control keeps same-response echo +suppression. `tests/responses/ws-capacity-tool-ownership.test.ts` covers these ownership boundaries. + +The victim is marked closing before sending close code 1001 with reason +`Idle Responses connection reclaimed`. Only its ordinary close callback releases the +admission lease. After requesting closure, the current upgrade immediately attempts reservation once. +If every socket is protected, the close fails, or the callback has not yet released the lease, +that upgrade receives the existing HTTP 503 `Codex WebSockets capacity reached` response. +A later upgrade can acquire the released slot; admission never exceeds the limit. Reclamation +does not delete saved continuation state; a client can reconnect through the existing upgrade +and replay path. This policy adds no idle timer or configuration and does not change the +separate Live/Realtime sockets or upstream connection reuse. Client disconnect releases socket admission independently of upstream cleanup: the eager relay may continue its existing bounded post-cancel drain for up to 15 seconds to observe a real terminal while a replacement socket already completes new work. + Registry-declared per-model compatibility hints (`modelResponsesUpstreamStreaming`) may ask the upstream Responses endpoint for bounded JSON on ANY client transport — WebSocket or ordinary HTTP/SSE. The bridge reframes that JSON into the same Responses event sequence diff --git a/tests/codex-integration/active-registry-admission.test.ts b/tests/codex-integration/active-registry-admission.test.ts index 18d79349cee..e531dcc718c 100644 --- a/tests/codex-integration/active-registry-admission.test.ts +++ b/tests/codex-integration/active-registry-admission.test.ts @@ -2,7 +2,7 @@ import { afterEach, beforeEach, describe, expect, test } from "bun:test"; import { mkdirSync, mkdtempSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; -import { MAX_ACTIVE_TURNS, abortAndReleaseAllTurns, activeRegistryMetrics, trackStreamLifetime, tryAdmitTurn, unregisterTurn } from "../../src/server/lifecycle"; +import { MAX_ACTIVE_TURNS, abortAndReleaseAllTurns, activeRegistryMetrics, registerTurn, trackStreamLifetime, tryAdmitTurn, unregisterTurn } from "../../src/server/lifecycle"; import { workflowBudgetSnapshot } from "../../src/lib/workflow-budget"; import { MAX_TRACKED_CODEX_WEBSOCKETS, @@ -243,6 +243,75 @@ describe("active registry admission", () => { expect(attachedReleases).toBe(1); }); + test("a completed child keeps its parent admitted and cannot cancel the next child", () => { + const before = activeRegistryMetrics().activeTurns; + const lease = tryAdmitTurn()!; + const parent = new AbortController(); + const first = new AbortController(); + const next = new AbortController(); + let released = 0; + lease.attach({ release() { released += 1; } }); + try { + registerTurn(parent, lease); + registerTurn(first, lease); + unregisterTurn(first); + expect(activeRegistryMetrics().activeTurns.active).toBe(before.active + 1); + expect(released).toBe(0); + registerTurn(next, lease); + expect(next.signal.aborted).toBe(false); + first.abort(new Error("late cleanup of completed child")); + unregisterTurn(first); + expect(parent.signal.aborted).toBe(false); + expect(next.signal.aborted).toBe(false); + expect(activeRegistryMetrics().activeTurns.active).toBe(before.active + 1); + unregisterTurn(next); + expect(released).toBe(0); + unregisterTurn(parent); + expect(released).toBe(1); + expect(activeRegistryMetrics().activeTurns.active).toBe(before.active); + expect(activeRegistryMetrics().activeTurns.releaseMisses).toBe(before.releaseMisses); + } finally { lease.release(); } + }); + + test("registering the same sole controller twice still releases on its one completion", () => { + const before = activeRegistryMetrics().activeTurns; + const lease = tryAdmitTurn()!; + const controller = new AbortController(); + try { + registerTurn(controller, lease); + registerTurn(controller, lease); + unregisterTurn(controller); + expect(activeRegistryMetrics().activeTurns.active).toBe(before.active); + unregisterTurn(controller); + expect(activeRegistryMetrics().activeTurns.releaseMisses).toBe(before.releaseMisses); + } finally { lease.release(); } + }); + + test("parent cancellation after child completion still releases and aborts the active child", () => { + const before = activeRegistryMetrics().activeTurns; + const lease = tryAdmitTurn()!; + const parent = new AbortController(); + const completed = new AbortController(); + const pending = new AbortController(); + const reason = new Error("parent turn superseded"); + let released = 0; + lease.attach({ release() { released += 1; } }); + try { + registerTurn(parent, lease); + registerTurn(completed, lease); + unregisterTurn(completed); + registerTurn(pending, lease); + parent.abort(reason); + expect(pending.signal.reason).toBe(reason); + expect(completed.signal.aborted).toBe(false); + expect(released).toBe(1); + expect(activeRegistryMetrics().activeTurns.active).toBe(before.active); + unregisterTurn(parent); + unregisterTurn(pending); + expect(activeRegistryMetrics().activeTurns.releaseMisses).toBe(before.releaseMisses); + } finally { lease.release(); } + }); + test("the 256-turn gate bounds concurrency, not sequential continuation count", () => { const before = activeRegistryMetrics().activeTurns.active; for (let index = 0; index <= MAX_ACTIVE_TURNS; index += 1) { diff --git a/tests/codex-integration/websocket-capacity-activity.test.ts b/tests/codex-integration/websocket-capacity-activity.test.ts new file mode 100644 index 00000000000..4744b38777d --- /dev/null +++ b/tests/codex-integration/websocket-capacity-activity.test.ts @@ -0,0 +1,454 @@ +import { afterEach, describe, expect, test } from "bun:test"; +import type { ServerWebSocket } from "bun"; +import { CodexWebSocketActivity } from "../../src/codex/websocket-activity"; +import { + MAX_TRACKED_CODEX_WEBSOCKETS, codexWebSocketAdmissionMetrics, + registerCodexWebSocket, tryReserveCodexWebSocket, unregisterCodexWebSocket, +} from "../../src/codex/websocket-registry"; +import type { AdmissionReservation } from "../../src/lib/admission"; +import { createWebsocketHandler } from "../../src/server/index/websocket-handler"; +import type { ServeOptionsContext } from "../../src/server/index/serve-options"; +import { MAX_ACTIVE_TURNS, tryAdmitTurn } from "../../src/server/lifecycle"; +import type { NativeResponseControl } from "../../src/server/responses/native-response-control"; +import { sendResponseToWebSocket, type WsData } from "../../src/server/ws-bridge"; + +type Frame = Record; +const leases: AdmissionReservation>[] = []; +const sockets: ServerWebSocket[] = []; +const handler = createWebsocketHandler({ config: {}, deps: {} } as ServeOptionsContext); + +afterEach(() => { + for (const ws of sockets.splice(0)) unregisterCodexWebSocket(ws); + for (const lease of leases.splice(0)) lease.release(); +}); + +function reserve() { + const lease = tryReserveCodexWebSocket(); + if (lease) leases.push(lease); + return lease; +} + +function admittedSocket(mode: "normal" | "delayed" | "throw" = "normal") { + const lease = reserve(); + expect(lease).not.toBeNull(); + const sent: Frame[] = []; + const closed: Array<{ code?: number; reason?: string }> = []; + let state = 1; + let buffered = 0; + const ws = { + data: { admissionLease: lease! } as WsData, + get readyState() { return state; }, + getBufferedAmount: () => buffered, + send(text: string) { sent.push(JSON.parse(text)); return 1; }, + close(code?: number, reason?: string) { + closed.push({ code, reason }); + if (mode === "throw") throw new Error("fixture close failure"); + if (mode === "normal") finishClose(); + }, + } as unknown as ServerWebSocket; + function finishClose() { + state = 3; + handler.close(ws, 1000, "fixture close"); + } + lease!.bind(ws); + registerCodexWebSocket(ws); + sockets.push(ws); + return { ws, sent, closed, finishClose, activity: ws.data.responsesActivity!, + setBuffered(bytes: number) { buffered = bytes; } }; +} + +function output(activity: CodexWebSocketActivity, item: Frame) { + activity.observePayload(JSON.stringify({ type: "response.output_item.done", item })); +} + +const argumentCompletions = [ + ["function_call", "response.function_call_arguments.done", "function_call_output"], + ["custom_tool_call", "response.custom_tool_call_input.done", "custom_tool_call_output"], +] as const; + +describe("Responses final tool arguments", () => { + test.each(argumentCompletions)("%s stays protected after EOF until the real handler accepts its result", async (type, done, result) => { + const client = admittedSocket(); + const finish = client.activity.beginTurn([]); + const argumentField = type === "function_call" ? "arguments" : "input"; + const frames = [ + { type: "response.output_item.added", item: { type, id: "item", call_id: "call", name: "fixture_tool", [argumentField]: "" } }, + { type: done, item_id: "item", [argumentField]: "{}" }, + ]; + await sendResponseToWebSocket(client.ws, new Response( + frames.map(frame => `data: ${JSON.stringify(frame)}\n\n`).join(""), + { headers: { "content-type": "text/event-stream" } }, + ), () => true, { onSsePayload: payload => client.activity.observePayload(payload) }); + finish(); + expect(client.sent.map(frame => frame.type)).toEqual(["response.output_item.added", done, "error"]); + expect(client.sent.at(-1)).toMatchObject({ error: { message: "Upstream stream ended before response terminal event" } }); + expect(client.ws.data.cancel).toBeUndefined(); + const others = Array.from({ length: MAX_TRACKED_CODEX_WEBSOCKETS - 1 }, () => admittedSocket()); + const finishOthers = others.map(other => other.activity.beginTurn([])); + try { + expect(reserve()).toBeNull(); + expect(client.closed).toEqual([]); + let accepted = false; + client.ws.data.nativeControl = { inject: () => { accepted = true; } } as unknown as NativeResponseControl; + handler.message(client.ws, JSON.stringify({ type: "response.inject", input: [{ type: result, call_id: "call", output: "done" }] })); + expect(accepted).toBe(true); + client.ws.data.nativeControl = undefined; + expect(reserve()).not.toBeNull(); + expect(client.closed).toEqual([{ code: 1001, reason: "Idle Responses connection reclaimed" }]); + } finally { + client.ws.data.nativeControl = undefined; + for (const end of finishOthers) end(); + } + }); + + test.each(argumentCompletions)("unfinished %s arguments do not invent pending client work", (type) => { + const activity = new CodexWebSocketActivity(); + const finish = activity.beginTurn([]); + activity.observePayload(JSON.stringify({ type: "response.output_item.added", item: { type, id: "item", call_id: "call" } })); + finish(); + expect(activity.idle).toBe(true); + }); + + test.each(argumentCompletions)("%s completion echoes stay settled and reused identities belong to the next turn", (type, done, resultType) => { + const activity = new CodexWebSocketActivity(); + const call = { type, id: "reused-item", call_id: "reused-call" }; + const result = { type: resultType, call_id: call.call_id, output: "done" }; + const added = JSON.stringify({ type: "response.output_item.added", item: call }); + const completed = JSON.stringify({ type: done, item_id: call.id }); + for (const created of [true, false]) { + const finish = activity.beginTurn([result]); + if (created) activity.observePayload(JSON.stringify({ type: "response.created", response: { id: "response" } })); + activity.observePayload(added); + activity.observePayload(added); // A repeated identity is not a new call. + activity.observePayload(completed); + finish(); + expect(activity.idle).toBe(false); + activity.acceptResults([result]); + activity.observePayload(completed); + output(activity, call); + activity.observePayload(JSON.stringify({ type: "response.completed", response: { output: [call] } })); + expect(activity.idle).toBe(true); + } + }); + + test("native response.created resets argument identities without suppressing a reused call", () => { + const activity = new CodexWebSocketActivity(); + const finish = activity.beginTurn([]); + for (const id of ["first", "second"]) { + activity.observePayload(JSON.stringify({ type: "response.created", response: { id } })); + activity.observePayload(JSON.stringify({ type: "response.output_item.added", item: { type: "function_call", id: "item", call_id: id } })); + activity.observePayload(JSON.stringify({ type: "response.function_call_arguments.done", item_id: "item" })); + if (id === "first") activity.acceptResults([{ type: "function_call_output", call_id: id }]); + } + finish(); + expect(activity.idle).toBe(false); + activity.acceptResults([{ type: "function_call_output", call_id: "second" }]); + expect(activity.idle).toBe(true); + }); + + test("hosted custom input completion stays idle only with this response's provider authorization", () => { + const activity = new CodexWebSocketActivity(); + const hosted = new Set([{ itemType: "custom_tool_call", callIdPrefix: "xs_call-" }]); + const added = JSON.stringify({ type: "response.output_item.added", item: { type: "custom_tool_call", id: "item", call_id: "xs_call-1" } }); + const done = JSON.stringify({ type: "response.custom_tool_call_input.done", item_id: "item" }); + const first = activity.beginTurn([]); + activity.observePayload(added, hosted); + activity.observePayload(done, hosted); + first(); + expect(activity.idle).toBe(true); + const second = activity.beginTurn([]); + activity.observePayload(added); + activity.observePayload(done); + second(); + expect(activity.idle).toBe(false); + activity.acceptResults([{ type: "custom_tool_call_output", call_id: "xs_call-1" }]); + expect(activity.idle).toBe(true); + }); + + test("unknown, oversized and mismatched final argument identities remain protected", () => { + for (const item of [ + undefined, + { type: "function_call", id: "", call_id: "call" }, + { type: "function_call", id: "item", call_id: "" }, + { type: "function_call", id: "item", call_id: "x".repeat(1025) }, + { type: "function_call", id: "x".repeat(1025), call_id: "call" }, + { type: "custom_tool_call", id: "item", call_id: "call" }, + ]) { + const activity = new CodexWebSocketActivity(); + const finish = activity.beginTurn([]); + activity.observePayload(JSON.stringify({ type: "response.output_item.added", item })); + activity.observePayload(JSON.stringify({ type: "response.function_call_arguments.done", item_id: item?.id ?? "unknown" })); + finish(); + activity.acceptResults([{ type: "function_call_output", call_id: "call" }]); + expect(activity.idle).toBe(false); + } + }); + + test("conflicting or over-budget argument identities conservatively protect the socket", () => { + for (const conflict of [true, false]) { + const activity = new CodexWebSocketActivity(); + const finish = activity.beginTurn([]); + for (let index = 0; index < (conflict ? 2 : 1025); index += 1) { + activity.observePayload(JSON.stringify({ type: "response.output_item.added", item: { + type: "function_call", id: conflict ? "conflict" : `item-${index}`, call_id: `call-${index}`, + } })); + } + finish(); + expect(activity.idle).toBe(false); + } + }); + + test.each(argumentCompletions)("invalid %s redefinitions cannot reuse a prior result identity", (type, done, resultType) => { + for (const returnedFirst of [false, true]) { + for (const call_id of [undefined, "", "x".repeat(1025)]) { + const activity = new CodexWebSocketActivity(); + const finish = activity.beginTurn([]); + const emit = (frame: Frame) => activity.observePayload(JSON.stringify(frame)); + const result = { type: resultType, call_id: "prior", output: "done" }; + emit({ type: "response.output_item.added", item: { type, id: "item", call_id: "prior" } }); + emit({ type: done, item_id: "item" }); + if (returnedFirst) activity.acceptResults([result]); + emit({ type: "response.output_item.added", item: { type, id: "item", call_id } }); + emit({ type: done, item_id: "item" }); + finish(); + expect(activity.idle).toBe(false); + if (!returnedFirst) activity.acceptResults([result]); + expect(activity.idle).toBe(false); + } + } + }); +}); + +describe("Responses socket activity", () => { + test("overlapping superseded turns each keep their own idempotent busy reference", () => { + const activity = new CodexWebSocketActivity(); + const first = activity.beginTurn([]); + const second = activity.beginTurn([]); + second(); + second(); + expect(activity.idle).toBe(false); + first(); + expect(activity.idle).toBe(true); + const third = activity.beginTurn([]); + first(); + expect(activity.idle).toBe(false); + third(); + expect(activity.idle).toBe(true); + }); + + test.each([ + ["function_call", "function_call_output"], + ["custom_tool_call", "custom_tool_call_output"], + ["local_shell_call", "function_call_output"], + ["shell_call", "shell_call_output"], + ["computer_call", "computer_call_output"], + ["apply_patch_call", "apply_patch_call_output"], + ["tool_search_call", "tool_search_output"], + ])("protects client %s work until its %s arrives", (callType, resultType) => { + const activity = new CodexWebSocketActivity(); + const finish = activity.beginTurn([]); + output(activity, { type: callType, call_id: "client-call", execution: "client" }); + activity.observePayload(JSON.stringify({ type: "response.completed", response: { output: [] } })); + finish(); + expect(activity.idle).toBe(false); + const continueTurn = activity.beginTurn([{ type: resultType, call_id: "client-call", output: "done" }]); + expect(activity.idle).toBe(false); + continueTurn(); + expect(activity.idle).toBe(true); + }); + + test("terminal-only tool calls and partial result batches retain all outstanding work", () => { + const activity = new CodexWebSocketActivity(); + const finish = activity.beginTurn([]); + output(activity, { type: "function_call", call_id: "first" }); + activity.observePayload(JSON.stringify({ type: "response.completed", response: { output: [ + { type: "function_call", call_id: "first" }, + { type: "custom_tool_call", call_id: "second" }, + { type: "mcp_approval_request", id: "approval" }, + ] } })); + finish(); + activity.beginTurn([{ type: "function_call_output", call_id: "first", output: "done" }])(); + expect(activity.idle).toBe(false); + activity.beginTurn([{ type: "function_call_output", call_id: "second", output: "wrong type" }, + { type: "mcp_approval_response", approval_request_id: "unrelated", approve: true }])(); + expect(activity.idle).toBe(false); + activity.beginTurn([{ type: "custom_tool_call_output", call_id: "second", output: "done" }])(); + expect(activity.idle).toBe(false); + activity.beginTurn([{ type: "mcp_approval_response", approval_request_id: "approval", approve: false }])(); + expect(activity.idle).toBe(true); + }); + + test("hosted tools and malformed observations do not invent client work", () => { + const activity = new CodexWebSocketActivity(); + const finish = activity.beginTurn([]); + for (const type of ["web_search_call", "file_search_call", "code_interpreter_call", "mcp_call", "tool_search_call", "message"]) { + output(activity, { type, id: "hosted-item" }); + } + output(activity, { type: "tool_search_call", execution: "server", call_id: "hosted-search" }); + for (const payload of ["bad JSON", "null", "[]", JSON.stringify({ type: "response.completed", response: null })]) { + activity.observePayload(payload); + } + finish(); + expect(activity.idle).toBe(true); + }); + + test("returned native results ignore terminal echoes but cannot mask a later call reusing its ID", () => { + const activity = new CodexWebSocketActivity(); + const call = { type: "function_call", call_id: "reused" }; + const result = { type: "function_call_output", call_id: "reused", output: "done" }; + const finish = activity.beginTurn([]); + activity.observePayload(JSON.stringify({ type: "response.created", response: { id: "first" } })); + output(activity, call); + activity.acceptResults([result]); + activity.observePayload(JSON.stringify({ type: "response.completed", response: { output: [call] } })); + finish(); + expect(activity.idle).toBe(true); + const finishNext = activity.beginTurn([result]); // Full replay repeats a historical result. + activity.observePayload(JSON.stringify({ type: "response.created", response: { id: "second" } })); + output(activity, call); + finishNext(); + expect(activity.idle).toBe(false); + activity.acceptResults([result]); + expect(activity.idle).toBe(true); + }); + + test("hosted shell output clears its call without waiting for a client result", () => { + const activity = new CodexWebSocketActivity(); + const finish = activity.beginTurn([]); + const call = { type: "shell_call", call_id: "hosted-shell", action: { commands: ["pwd"] } }; + const result = { type: "shell_call_output", call_id: "hosted-shell", output: [] }; + output(activity, call); + output(activity, result); + activity.observePayload(JSON.stringify({ type: "response.completed", response: { output: [call, result] } })); + finish(); + expect(activity.idle).toBe(true); + }); + + test.each(["response.inject", "response.create"])("real handler accounts for accepted native results on %s", type => { + const client = admittedSocket(); + const finish = client.activity.beginTurn([]); + const call = { type: "function_call", call_id: "native-call" }; + output(client.activity, call); + let accepted = 0; + const accept = () => { accepted += 1; return true; }; + client.ws.data.nativeControl = { relayActive: true, inject: accept, continue: accept } as unknown as NativeResponseControl; + handler.message(client.ws, JSON.stringify({ type, input: [ + { type: "function_call_output", call_id: "native-call", output: "done" }, + ] })); + expect(accepted).toBe(1); + expect(client.activity.idle).toBe(false); + client.activity.observePayload(JSON.stringify({ type: "response.completed", response: { output: [call] } })); + finish(); + expect(client.activity.idle).toBe(true); + expect(client.sent).toHaveLength(0); + }); + + test.each(["response.inject", "response.create"])("rejected native results on %s cannot release a client tool wait", type => { + const client = admittedSocket(); + output(client.activity, { type: "function_call", call_id: "native-call" }); + const reject = () => { throw new Error("fixture rejected result"); }; + client.ws.data.nativeControl = { relayActive: true, inject: reject, continue: reject } as unknown as NativeResponseControl; + handler.message(client.ws, JSON.stringify({ type, input: [ + { type: "function_call_output", call_id: "native-call", output: "done" }, + ] })); + expect(client.sent).toHaveLength(1); + expect(client.sent[0]).toMatchObject({ type: "error", status: 400 }); + expect(client.activity.idle).toBe(false); + }); + + test("unknown and oversized client result identities remain conservatively protected", () => { + for (const callId of [undefined, "", "x".repeat(1025)]) { + const activity = new CodexWebSocketActivity(); + output(activity, { type: "function_call", call_id: callId }); + activity.beginTurn([{ type: "function_call_output", call_id: callId, output: "done" }])(); + expect(activity.idle).toBe(false); + } + const activity = new CodexWebSocketActivity(); + const results = Array.from({ length: 1025 }, (_, index) => { + const call_id = `bounded-${index}`; + output(activity, { type: "function_call", call_id }); + return { type: "function_call_output", call_id, output: "done" }; + }); + activity.beginTurn(results)(); + expect(activity.idle).toBe(false); + }); + + test("real handler warmup and turn-capacity refusal do not consume pending client results", () => { + const client = admittedSocket(); + output(client.activity, { type: "function_call", call_id: "pending" }); + const input = [{ type: "function_call_output", call_id: "pending", output: "done" }]; + handler.message(client.ws, JSON.stringify({ type: "response.create", generate: false, input })); + expect(client.sent.map(frame => frame.type)).toEqual(["response.created", "response.completed"]); + expect(client.activity.idle).toBe(false); + const turns = Array.from({ length: MAX_ACTIVE_TURNS }, () => tryAdmitTurn()); + try { + expect(turns.every(Boolean)).toBe(true); + handler.message(client.ws, JSON.stringify({ type: "response.create", model: "fixture", input })); + expect(client.sent.at(-1)).toMatchObject({ type: "error", status: 503 }); + expect(client.activity.idle).toBe(false); + } finally { + for (const turn of turns) turn?.release(); + } + }); +}); + +describe("Responses WebSocket capacity reclamation", () => { + test("below capacity leaves idle connections open; at capacity reclaims least recently used", () => { + const clients = Array.from({ length: MAX_TRACKED_CODEX_WEBSOCKETS }, () => admittedSocket()); + expect(clients.every(client => client.closed.length === 0)).toBe(true); + clients[0].activity.beginTurn([])(); + const before = codexWebSocketAdmissionMetrics(); + expect(reserve()).not.toBeNull(); + expect(clients[0].closed).toHaveLength(0); + expect(clients[1].closed).toEqual([{ code: 1001, reason: "Idle Responses connection reclaimed" }]); + expect(clients.slice(2).every(client => client.closed.length === 0)).toBe(true); + expect(codexWebSocketAdmissionMetrics()).toMatchObject({ active: MAX_TRACKED_CODEX_WEBSOCKETS, peak: MAX_TRACKED_CODEX_WEBSOCKETS, + admitted: before.admitted + 1, rejected: before.rejected }); + }); + + test("skips active turns, client tools, buffered output and native control owners", () => { + const clients = Array.from({ length: MAX_TRACKED_CODEX_WEBSOCKETS }, () => admittedSocket()); + const finish = clients[0].activity.beginTurn([]); + output(clients[1].activity, { type: "function_call", call_id: "pending" }); + clients[2].setBuffered(1); + clients[3].ws.data.nativeControl = { relayActive: true } as NativeResponseControl; + expect(reserve()).not.toBeNull(); + expect(clients.slice(0, 4).every(client => client.closed.length === 0)).toBe(true); + expect(clients[4].closed).toHaveLength(1); + finish(); + }); + + test("all busy or unbound reservations leave the cap intact and reject new admission", () => { + const clients = Array.from({ length: MAX_TRACKED_CODEX_WEBSOCKETS - 1 }, () => admittedSocket()); + const finish = clients.map(client => client.activity.beginTurn([])); + expect(reserve()).not.toBeNull(); // An upgrade reserved before its open callback. + const before = codexWebSocketAdmissionMetrics(); + expect(reserve()).toBeNull(); + expect(clients.every(client => client.closed.length === 0)).toBe(true); + expect(codexWebSocketAdmissionMetrics()).toMatchObject({ active: MAX_TRACKED_CODEX_WEBSOCKETS, rejected: before.rejected + 1 }); + for (const end of finish) end(); + }); + + test.each(["delayed", "throw"] as const)("%s close cannot free admission before the close callback", mode => { + const victim = admittedSocket(mode); + const others = Array.from({ length: MAX_TRACKED_CODEX_WEBSOCKETS - 1 }, () => admittedSocket()); + const finish = others.map(client => client.activity.beginTurn([])); + expect(reserve()).toBeNull(); + expect(victim.closed).toHaveLength(1); + expect(victim.activity.closing).toBe(mode === "delayed"); + expect(codexWebSocketAdmissionMetrics().active).toBe(MAX_TRACKED_CODEX_WEBSOCKETS); + handler.message(victim.ws, JSON.stringify({ type: "response.create", generate: false })); + // A closer that throws while still OPEN must not silently disable the old client. + expect(victim.sent).toHaveLength(mode === "delayed" ? 0 : 2); + expect(victim.ws.data.turnId).toBe(mode === "delayed" ? undefined : 1); + expect(reserve()).toBeNull(); + expect(victim.closed).toHaveLength(mode === "delayed" ? 1 : 2); + expect(others.every(client => client.closed.length === 0)).toBe(true); + victim.finishClose(); + victim.finishClose(); + expect(codexWebSocketAdmissionMetrics().active).toBe(MAX_TRACKED_CODEX_WEBSOCKETS - 1); + expect(reserve()).not.toBeNull(); + expect(codexWebSocketAdmissionMetrics().active).toBe(MAX_TRACKED_CODEX_WEBSOCKETS); + for (const end of finish) end(); + }); +}); diff --git a/tests/fixtures/test-layout-expected.json b/tests/fixtures/test-layout-expected.json index 9ae6ad80eb6..f265c9ea497 100644 --- a/tests/fixtures/test-layout-expected.json +++ b/tests/fixtures/test-layout-expected.json @@ -557,7 +557,8 @@ "codex-user-identity.test.ts": "codex-integration", "codex-v2-gate.test.ts": "codex-integration", "codex-agent-thread-limit.test.ts": "codex-integration", "codex-warmup.test.ts": "codex-integration", "codex-web-search-switch.test.ts": "codex-integration", - "codex-websocket-registry.test.ts": "codex-integration", "codex-write-lock.test.ts": "codex-integration", + "codex-websocket-registry.test.ts": "codex-integration", + "websocket-capacity-activity.test.ts": "codex-integration", "codex-write-lock.test.ts": "codex-integration", "coding-agent-tool-result-images.test.ts": "adapters", "cold-spawn-warmup.test.ts": "ci-workflows", "combo-authoritative-reset.test.ts": "codex-integration", "combo-codex-exhaustion-cooldown.test.ts": "codex-integration", @@ -1277,6 +1278,8 @@ "winsw.test.ts": "service", "workflow-budget.test.ts": "lib", "ws-ambiguous-resend.test.ts": "responses", "ws-pre-output-refusal.test.ts": "responses", "responses-pool-overload-failover.test.ts": "responses", "compact-transient-alternate.test.ts": "responses", + "ws-capacity-reclamation.test.ts": "responses", + "ws-capacity-tool-ownership.test.ts": "responses", "ws-endpoint.test.ts": "responses", "ws-failure-stage.test.ts": "responses", "ws-native-injection.test.ts": "responses", diff --git a/tests/helpers/responses-core-source.ts b/tests/helpers/responses-core-source.ts index d487d4609b5..ad3bba4675a 100644 --- a/tests/helpers/responses-core-source.ts +++ b/tests/helpers/responses-core-source.ts @@ -47,6 +47,7 @@ export const RESPONSES_CORE_MODULES = [ "request-transport.ts", "request-sidecar-auth.ts", "response-effects.ts", + "response-tool-ownership.ts", "request-send-budget.ts", "passthrough-execution.ts", "passthrough-dispatch.ts", diff --git a/tests/responses/responses-core-modules.test.ts b/tests/responses/responses-core-modules.test.ts index ed1c3fb5a8f..b640d36d7a4 100644 --- a/tests/responses/responses-core-modules.test.ts +++ b/tests/responses/responses-core-modules.test.ts @@ -85,14 +85,15 @@ describe("Responses core module boundaries", () => { expect(ingress.indexOf("await executePassthroughResponse(")) .toBeLessThan(ingress.indexOf("const sidecarPlans = await executeResponsesSidecars(")); expect(native).toContain('if (nativeExchange instanceof Response || nativeExchange.kind === "adapter-handoff") return nativeExchange;'); - // Delivery is awaited inside the try, and its direct body is wrapped before - // the return. What matters is that both awaits stay inside the lease owner, - // not that the delivery call is itself the return expression. + // Delivery is awaited inside the try, and its direct body is wrapped with + // tool ownership before return. Both awaits must stay inside the lease owner; + // the delivery call itself need not be the return expression. expect(native).toContain("const response = await deliverPassthroughResponse("); - expect(native).toContain("return guardDirectPassthroughBodyInactivity("); + const guardedReturn = "return attachResponseToolOwnership(guardDirectPassthroughBodyInactivity("; + expect(native).toContain(guardedReturn); expect(native.indexOf("await deliverPassthroughResponse(")) - .toBeLessThan(native.indexOf("return guardDirectPassthroughBodyInactivity(")); - expect(native.indexOf("return guardDirectPassthroughBodyInactivity(")) + .toBeLessThan(native.indexOf(guardedReturn)); + expect(native.indexOf(guardedReturn)) .toBeLessThan(native.indexOf("} finally {")); expect(native.indexOf("admissionState.pendingHostAdmissionLease = null;")) .toBeLessThan(native.indexOf("await preparePassthroughExchange(")); diff --git a/tests/responses/ws-capacity-reclamation.test.ts b/tests/responses/ws-capacity-reclamation.test.ts new file mode 100644 index 00000000000..1a6da6f4abe --- /dev/null +++ b/tests/responses/ws-capacity-reclamation.test.ts @@ -0,0 +1,328 @@ +import { afterEach, beforeEach, expect, spyOn, test } from "bun:test"; +import { mkdtempSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { saveConfig } from "../../src/config"; +import { codexWebSocketAdmissionMetrics, MAX_TRACKED_CODEX_WEBSOCKETS } from "../../src/codex/websocket-registry"; +import { + clearResponseStateForTests, clearResponseStateMemoryForTests, RESPONSE_TTL_MS, + responseStateMetrics, sweepExpiredResponseStates, +} from "../../src/responses/state"; +import { startServer } from "../../src/server"; +import type { OcxConfig } from "../../src/types"; +import { installIsolatedCodexHome, type IsolatedCodexHome } from "../helpers/isolated-codex-home"; +import { removeTreeWithRetry } from "../helpers/remove-tree"; + +const originalFetch = globalThis.fetch; +const previousHome = process.env.OPENCODEX_HOME; +const MODEL = "fixture/capacity-model"; +type Frame = Record; +type Client = { socket: WebSocket; frames: Frame[]; closed?: CloseEvent; error?: Event }; +let home: string; +let codex: IsolatedCodexHome; +let proxy: ReturnType | undefined; +let upstream: ReturnType | undefined; +let finishStream: (() => void) | undefined; +let upstreamRequests = 0; +let upstreamCancellations = 0; +let heldUpstreamSignal: AbortSignal | undefined; +const upstreamBodies: Frame[] = []; +const clients: Client[] = []; +const refusedHandshakes = new Map(); +let closeReleaseDelayMs = 0; + +async function until(condition: () => boolean, description: string, timeoutMs = 5_000): Promise { + const deadline = Date.now() + timeoutMs; + while (!condition()) { + if (Date.now() >= deadline) throw new Error(`Timed out: ${description}`); + await Bun.sleep(2); + } +} + +function sse(event: Frame): Uint8Array { + return new TextEncoder().encode(`event: ${event.type}\ndata: ${JSON.stringify(event)}\n\n`); +} + +beforeEach(() => { + home = mkdtempSync(join(tmpdir(), "ocx-ws-capacity-")); + process.env.OPENCODEX_HOME = home; + clearResponseStateMemoryForTests(); + codex = installIsolatedCodexHome("ocx-ws-capacity-codex-"); + upstreamRequests = 0; + upstreamCancellations = 0; + heldUpstreamSignal = undefined; + upstreamBodies.length = 0; + finishStream = undefined; + refusedHandshakes.clear(); + closeReleaseDelayMs = 0; + globalThis.fetch = ((input: RequestInfo | URL, init?: RequestInit) => { + const url = new URL(input instanceof Request ? input.url : String(input)); + if (url.hostname !== "127.0.0.1" && url.hostname !== "localhost") { + throw new Error("Capacity fixture forbids external requests"); + } + if (url.port === String(upstream?.port) && String(init?.body).includes("hold-fixture-stream")) { + heldUpstreamSignal = init?.signal ?? (input instanceof Request ? input.signal : undefined); + } + return originalFetch(input, init); + }) as typeof fetch; + upstream = Bun.serve({ + // Only proxy cancellation may end a held response; Bun's default idle timeout must not. + hostname: "127.0.0.1", port: 0, idleTimeout: 0, + async fetch(request) { + const body = await request.json() as Frame; + upstreamRequests++; + upstreamBodies.push(body); + const id = `resp_capacity_${upstreamRequests}`; + const tool = JSON.stringify(body.input).includes("request-fixture-tool"); + const hold = JSON.stringify(body.input).includes("hold-fixture-stream"); + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(sse({ type: "response.created", response: { id, status: "in_progress", output: [] } })); + const output = tool + ? [{ type: "function_call", id: "fc_capacity", call_id: "call_capacity", name: "fixture_tool", arguments: "{}" }] + : [{ type: "message", id: "msg_capacity", role: "assistant", content: [{ type: "output_text", text: "capacity fixture completed" }] }]; + const finish = () => { + if (tool) controller.enqueue(sse({ type: "response.output_item.done", output_index: 0, item: output[0] })); + controller.enqueue(sse({ type: "response.completed", response: { id, status: "completed", output, + usage: { input_tokens: 1, output_tokens: 1, total_tokens: 2 } } })); + controller.close(); + }; + if (hold) finishStream = finish; else finish(); + }, + cancel() { upstreamCancellations++; }, + }); + return new Response(stream, { headers: { "content-type": "text/event-stream" } }); + }, + }); + saveConfig({ + port: 0, hostname: "127.0.0.1", codexAutoStart: false, websockets: true, + defaultProvider: "fixture", providers: { fixture: { + adapter: "openai-responses", baseUrl: `${upstream.url}v1`, apiKey: "sk-capacity-fixture", upstreamWebsocket: false, + allowPrivateNetwork: true, statelessResponses: true, + } }, + } as OcxConfig); + expect(codexWebSocketAdmissionMetrics().active).toBe(0); + // Observe actual upgrade refusals; delay only the fixture close callback in the delayed row. + const originalServe = Bun.serve; + const serveSpy = spyOn(Bun, "serve").mockImplementation(((options: any) => originalServe({ + ...options, + async fetch(request: Request, server: any) { + const response = await options.fetch(request, server); + const attempt = new URL(request.url).searchParams.get("capacity-attempt"); + if (attempt !== null && response) refusedHandshakes.set(attempt, + { status: response.status, retryAfter: response.headers.get("retry-after") }); + return response; + }, + websocket: { ...options.websocket, close(...args: any[]) { + if (closeReleaseDelayMs) setTimeout(() => options.websocket.close(...args), closeReleaseDelayMs); + else options.websocket.close(...args); + } }, + })) as typeof Bun.serve); + try { proxy = startServer(0); } finally { serveSpy.mockRestore(); } +}); + +afterEach(async () => { + try { + for (const client of clients) client.socket.close(); + await proxy?.stop(true); + await upstream?.stop(true); + await until(() => codexWebSocketAdmissionMetrics().active === 0, "all socket reservations released"); + } finally { + clients.length = 0; + proxy = undefined; + upstream = undefined; + globalThis.fetch = originalFetch; + clearResponseStateForTests(); + codex.restore(); + removeTreeWithRetry(home); + if (previousHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = previousHome; + } +}); + +async function connect(retriesRemaining = MAX_TRACKED_CODEX_WEBSOCKETS): Promise { + const url = new URL("/v1/responses", proxy!.url); + url.protocol = "ws:"; + const attempt = String(clients.length); + url.searchParams.set("capacity-attempt", attempt); + const client: Client = { socket: new WebSocket(url), frames: [] }; + clients.push(client); + client.socket.addEventListener("message", event => client.frames.push(JSON.parse(String(event.data)))); + client.socket.addEventListener("close", event => { client.closed = event; }); + client.socket.addEventListener("error", event => { client.error = event; }); + await until(() => client.socket.readyState === WebSocket.OPEN || Boolean(client.error || client.closed), "socket handshake"); + if (client.socket.readyState !== WebSocket.OPEN) { + expect(refusedHandshakes.get(attempt)).toEqual({ status: 503, retryAfter: "1" }); + await until(() => codexWebSocketAdmissionMetrics().active < MAX_TRACKED_CODEX_WEBSOCKETS, + "reclaimed socket releases admission before retry"); + expect(retriesRemaining).toBeGreaterThan(0); + return connect(retriesRemaining - 1); + } + return client; +} + +async function fill(): Promise { + const opened: Client[] = []; + for (let i = 0; i < MAX_TRACKED_CODEX_WEBSOCKETS; i++) opened.push(await connect()); + return opened; +} + +async function exchange(client: Client, fields: Frame, terminal = "response.completed"): Promise { + const start = client.frames.length; + client.socket.send(JSON.stringify({ type: "response.create", model: MODEL, input: "fixture input", ...fields })); + await until(() => client.frames.slice(start).some(frame => frame.type === terminal || frame.type === "error") + || Boolean(client.closed), `socket ${terminal}`); + const result = client.frames.slice(start).find(frame => frame.type === terminal || frame.type === "error"); + expect(result?.type).toBe(terminal); + return result!; +} + +function expectBounded(): void { + const metrics = codexWebSocketAdmissionMetrics(); + expect(metrics.active).toBe(MAX_TRACKED_CODEX_WEBSOCKETS); + expect(metrics.peak).toBe(MAX_TRACKED_CODEX_WEBSOCKETS); + expect(metrics.releaseMisses).toBe(0); +} + +test.each([0, 20])("the 129th handshake reclaims the least recently used idle socket and reconnects cleanly; close delay=%s", async delay => { + closeReleaseDelayMs = delay; + const opened = await fill(); + expect(opened.every(client => client.socket.readyState === WebSocket.OPEN)).toBe(true); + await exchange(opened[0]!, { generate: false }); + expect(upstreamRequests).toBe(0); + const replacement = await connect(); + await until(() => Boolean(opened[1]!.closed), "least recently used idle socket closed"); + expect(opened[1]!.closed?.code).toBe(1001); + if (delay) expect([...refusedHandshakes.values()].some(response => response.status === 503)).toBe(true); + expect(opened[0]!.socket.readyState).toBe(WebSocket.OPEN); + expectBounded(); + await exchange(replacement, { input: "fresh request after capacity recovery" }); + expect(upstreamRequests).toBe(1); + const reconnected = await connect(); + await exchange(reconnected, { input: "resubmitted full conversation" }); + expect(upstreamRequests).toBe(2); + expectBounded(); +}, 15_000); + +test("concurrent handshakes replace idle sockets without over-admitting or leaking leases", async () => { + const opened = await fill(); + const added = await Promise.all(Array.from({ length: 16 }, () => connect())); + await until(() => opened.filter(client => client.closed).length === 16, "all reclaimed clients closed"); + expect(opened.slice(0, 16).every(client => client.closed?.code === 1001)).toBe(true); + expect(opened.slice(16).every(client => client.socket.readyState === WebSocket.OPEN)).toBe(true); + expect(added.every(client => client.socket.readyState === WebSocket.OPEN)).toBe(true); + expectBounded(); +}, 15_000); + +test("capacity pressure preserves a streaming response until it completes", async () => { + const opened = await fill(); + await exchange(opened[0]!, { input: "hold-fixture-stream" }, "response.created"); + // Refresh every idle socket so the busy connection is the oldest by activity too. + for (const client of opened.slice(1)) await exchange(client, { generate: false }); + await connect(); + await until(() => Boolean(opened[1]!.closed), "oldest eligible idle socket closed"); + expect(opened[0]!.socket.readyState).toBe(WebSocket.OPEN); + finishStream!(); + await until(() => opened[0]!.frames.some(frame => frame.type === "response.completed"), "protected stream completes"); + expect(opened[0]!.frames.some(frame => frame.type === "error")).toBe(false); + expectBounded(); +}, 15_000); + +test("client tool execution stays protected, and all-busy saturation keeps the 503 refusal", async () => { + const opened = await fill(); + for (const client of opened) { + await exchange(client, { input: "request-fixture-tool", tools: [{ type: "function", name: "fixture_tool", + parameters: { type: "object", properties: {} } }] }); + } + const rejectedBefore = codexWebSocketAdmissionMetrics().rejected; + const response = await fetch(new URL("/v1/responses", proxy!.url), { headers: { + connection: "upgrade", upgrade: "websocket", "sec-websocket-version": "13", + "sec-websocket-key": Buffer.from("capacity-fixture").toString("base64"), + } }); + expect(response.status).toBe(503); + expect(response.headers.get("retry-after")).toBe("1"); + expect(await response.text()).toContain("Codex WebSockets capacity reached"); + expect(codexWebSocketAdmissionMetrics().rejected).toBe(rejectedBefore + 1); + expect(opened.every(client => client.socket.readyState === WebSocket.OPEN)).toBe(true); + await exchange(opened[0]!, { input: [{ type: "function_call_output", call_id: "call_capacity", output: "fixture tool result" }] }); + await connect(); + await until(() => Boolean(opened[0]!.closed), "completed tool workflow becomes reclaimable"); + expect(opened.slice(1).every(client => client.socket.readyState === WebSocket.OPEN)).toBe(true); + expectBounded(); +}, 20_000); + +test.each([false, true])("disconnect releases admission and allows a fresh request; abrupt=%s", async abrupt => { + const baseline = codexWebSocketAdmissionMetrics(); + const client = await connect(); + expect(codexWebSocketAdmissionMetrics().active).toBe(baseline.active + 1); + if (abrupt) { + await exchange(client, { input: "hold-fixture-stream" }, "response.created"); + expect(heldUpstreamSignal).toBeDefined(); + expect(heldUpstreamSignal!.aborted).toBe(false); + // Bun's client API terminates the real TCP connection without a closing handshake. + client.socket.terminate(); + } else client.socket.close(1000, "fixture normal disconnect"); + await until(() => codexWebSocketAdmissionMetrics().active === baseline.active, + "disconnected socket releases its reservation", 10_000); + await until(() => Boolean(client.closed), "client observes close", 10_000); + expect(client.closed?.code).toBe(abrupt ? 1006 : 1000); + const replacement = await connect(); + await exchange(replacement, { input: "fresh request after disconnect" }); + expect(upstreamRequests).toBe(abrupt ? 2 : 1); + replacement.socket.close(); + await until(() => codexWebSocketAdmissionMetrics().active === baseline.active, "replacement reservation released"); + expect(codexWebSocketAdmissionMetrics().releaseMisses).toBe(baseline.releaseMisses); + // Eager relay preserves a real terminal during its existing 15-second discard-drain. + // Socket admission and the replacement request must recover before that separate cleanup. + if (abrupt) { + await until(() => heldUpstreamSignal!.aborted, "abandoned upstream request aborted", 20_000); + await until(() => upstreamCancellations === 1, "upstream observed request cancellation"); + } +}, 25_000); + +test("failed HTTP upgrade releases its reservation before the next handshake", async () => { + const baseline = codexWebSocketAdmissionMetrics(); + const response = await fetch(new URL("/v1/responses", proxy!.url), { headers: { + connection: "upgrade", upgrade: "websocket", "sec-websocket-version": "13", + // Missing Sec-WebSocket-Key reaches Bun's unsuccessful upgrade after reservation. + } }); + expect(response.status).toBe(426); + expect(await response.text()).toContain("WebSocket upgrade failed"); + expect(codexWebSocketAdmissionMetrics()).toMatchObject({ active: baseline.active, + admitted: baseline.admitted + 1, rejected: baseline.rejected, releaseMisses: baseline.releaseMisses }); + const replacement = await connect(); + await exchange(replacement, { input: "fresh request after failed handshake" }); + expect(upstreamRequests).toBe(1); + expect(codexWebSocketAdmissionMetrics().active).toBe(baseline.active + 1); +}); + +test.each([false, true])("stateless continuation after capacity reclamation preserves full history; expired=%s", async expired => { + const original = await connect(); + const firstInput = { role: "user", content: [{ type: "input_text", text: "Remember the fixture number 47." }] }; + const first = await exchange(original, { input: [firstInput], store: false }); + const previousId = first.response.id; + expect(typeof previousId).toBe("string"); + expect(previousId.length).toBeGreaterThan(0); + await until(() => responseStateMetrics().count === 1, "first response retains real continuation state"); + for (let i = 1; i < MAX_TRACKED_CODEX_WEBSOCKETS; i++) await connect(); + const reconnected = await connect(); + await until(() => original.closed?.code === 1001, "socket holding the response ID reclaimed"); + if (expired) { + expect(sweepExpiredResponseStates(Date.now() + RESPONSE_TTL_MS + 60_000)).toBe(1); + expect(responseStateMetrics().count).toBe(0); + } + const delta = { role: "user", content: [{ type: "input_text", text: "Use that number in the next answer." }] }; + const outcome = await exchange(reconnected, { previous_response_id: previousId, input: [delta], store: false }, + expired ? "error" : "response.completed"); + if (expired) { + expect(outcome).toMatchObject({ type: "error", status: 400, error: { code: "previous_response_not_found" } }); + expect(upstreamRequests).toBe(1); // Never forward the delta without its missing prefix. + await exchange(reconnected, { input: [firstInput, ...first.response.output, delta], store: false }); + } + expect(upstreamRequests).toBe(2); // One initial response and one complete continuation, without a duplicate execution. + expect(upstreamBodies[1]!.previous_response_id).toBeUndefined(); + expect(upstreamBodies[1]!.input).toHaveLength(3); + expect(upstreamBodies[1]!.input).toMatchObject([firstInput, + { type: "message", role: "assistant", content: [{ type: "output_text", text: "capacity fixture completed" }] }, delta]); + expectBounded(); +}, 15_000); diff --git a/tests/responses/ws-capacity-tool-ownership.test.ts b/tests/responses/ws-capacity-tool-ownership.test.ts new file mode 100644 index 00000000000..468a6487381 --- /dev/null +++ b/tests/responses/ws-capacity-tool-ownership.test.ts @@ -0,0 +1,254 @@ +import { afterEach, beforeEach, expect, test } from "bun:test"; +import type { ServerWebSocket } from "bun"; +import { clearComboSelectionState, clearComboTargetCooldowns } from "../../src/combos"; +import { + codexWebSocketAdmissionMetrics, MAX_TRACKED_CODEX_WEBSOCKETS, tryReserveCodexWebSocket, +} from "../../src/codex/websocket-registry"; +import type { AdmissionReservation } from "../../src/lib/admission"; +import { runOptionalShutdownHooks } from "../../src/lib/optional-shutdown-hooks"; +import { clearKeyCooldowns } from "../../src/providers/key-failover"; +import { createWebsocketHandler } from "../../src/server/index/websocket-handler"; +import type { ServeOptionsContext } from "../../src/server/index/serve-options"; +import { clearRequestLogsForTests } from "../../src/server/request-log"; +import { setRelayPlatformForTests } from "../../src/server/responses/passthrough-delivery"; +import type { WsData } from "../../src/server/ws-bridge"; +import type { OcxConfig } from "../../src/types"; +import { acquireOwnedSpendHome } from "../helpers/owned-spend-home"; +import { createTempHome, type TempHome } from "../helpers/temp-home"; + +type Frame = Record; +type Handler = ReturnType; +type Client = { ws: ServerWebSocket; handler: Handler; sent: Frame[]; closed: number[] }; +const originalFetch = globalThis.fetch; +const proxyKeys = ["HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "NO_PROXY", "http_proxy", "https_proxy", "all_proxy", "no_proxy"]; +const clients: Client[] = []; +const extraLeases: AdmissionReservation>[] = []; +const finishBusy: Array<() => void> = []; +let home: TempHome; +let releaseSpendHome: () => void; +let savedProxy: Record; +let upstream: (body: Frame, url: URL) => Response; +let outbound: Frame[]; + +async function until(condition: () => boolean): Promise { + const deadline = Date.now() + 5_000; + while (!condition()) { + if (Date.now() >= deadline) throw new Error("WebSocket tool ownership fixture did not settle"); + await Bun.sleep(1); + } +} + +beforeEach(() => { + clearComboSelectionState(); + clearComboTargetCooldowns(); + clearKeyCooldowns(); + home = createTempHome("ocx-ws-tool-ownership-"); + releaseSpendHome = acquireOwnedSpendHome(); + savedProxy = Object.fromEntries(proxyKeys.map(key => [key, process.env[key]])); + for (const key of proxyKeys) delete process.env[key]; + outbound = []; + upstream = () => { throw new Error("Unexpected upstream request in tool ownership fixture"); }; + globalThis.fetch = (async (input, init) => { + const body = JSON.parse(String(init?.body)) as Frame; + outbound.push(body); + return upstream(body, new URL(input instanceof Request ? input.url : String(input))); + }) as typeof fetch; + clearRequestLogsForTests(); + expect(codexWebSocketAdmissionMetrics().active).toBe(0); +}); + +afterEach(async () => { + try { + for (const client of clients) client.ws.close(1000, "fixture teardown"); + await until(() => clients.every(({ ws }) => !ws.data.cancel && !ws.data.nativeControl)); + } finally { + for (const finish of finishBusy.splice(0)) finish(); + for (const lease of extraLeases.splice(0)) lease.release(); + clients.length = 0; + clearComboSelectionState(); + clearComboTargetCooldowns(); + clearKeyCooldowns(); + try { + runOptionalShutdownHooks(); + } finally { + try { releaseSpendHome(); } + finally { + setRelayPlatformForTests(undefined); + globalThis.fetch = originalFetch; + for (const key of proxyKeys) { + if (savedProxy[key] === undefined) delete process.env[key]; + else process.env[key] = savedProxy[key]; + } + home.remove(); + } + } + } + expect(codexWebSocketAdmissionMetrics().active).toBe(0); +}); + +function config(baseUrl = "https://api.x.ai/v1", injectXSearch = false): OcxConfig { + return { port: 0, defaultProvider: "fixture", websockets: true, providers: { fixture: { + adapter: "openai-responses", baseUrl, authMode: "key", apiKey: "fixture-tool-ownership-key", upstreamWebsocket: false, + ...(injectXSearch ? { xaiResponsesXSearch: true, supportsOpenAiWebSearchToolFields: false } : {}), + } } } as OcxConfig; +} + +function client(settings = config()): Client { + const lease = tryReserveCodexWebSocket(); + expect(lease).not.toBeNull(); + const handler = createWebsocketHandler({ config: settings, deps: {} } as ServeOptionsContext); + const sent: Frame[] = []; + const closed: number[] = []; + let state = 1; + const ws = { + data: { admissionLease: lease!, headers: new Headers({ "thread-id": `tool-ownership-${clients.length}` }) } as WsData, + get readyState() { return state; }, + getBufferedAmount: () => 0, + send(text: string) { sent.push(JSON.parse(text)); return 1; }, + close(code = 1000, reason = "") { + if (state === 3) return; + state = 3; + closed.push(code); + handler.close(ws, code, reason); + }, + } as unknown as ServerWebSocket; + handler.open(ws); + const result = { ws, handler, sent, closed }; + clients.push(result); + return result; +} + +async function turn(target: Client, fields: Frame): Promise { + const start = target.sent.length; + target.handler.message(target.ws, JSON.stringify({ type: "response.create", model: "fixture/grok-4.6", input: "fixture input", ...fields })); + await until(() => target.sent.slice(start).some(frame => frame.type === "error" || frame.type === "response.completed") + && !target.ws.data.cancel && !target.ws.data.nativeControl); + const frames = target.sent.slice(start); + expect(frames.filter(frame => frame.type === "error")).toEqual([]); + expect(frames.at(-1)).toMatchObject({ type: "response.completed", response: { status: "completed" } }); + return frames; +} + +function fillBusyAndReserve(): AdmissionReservation> | null { + while (codexWebSocketAdmissionMetrics().active < MAX_TRACKED_CODEX_WEBSOCKETS) { + finishBusy.push(client().ws.data.responsesActivity!.beginTurn([])); + } + const lease = tryReserveCodexWebSocket(); + if (lease) extraLeases.push(lease); + expect(codexWebSocketAdmissionMetrics().active).toBe(MAX_TRACKED_CODEX_WEBSOCKETS); + return lease; +} + +const hostedCall = { type: "custom_tool_call", id: "ctc_search", call_id: "xs_call-1", + name: "x_keyword_search", input: "{}", status: "completed" }; +const clientTool = { type: "function", name: "x_keyword_search", parameters: { type: "object", properties: {} } }; + +test.each([false, true])("completed xAI hosted search is reclaimable; route-injected declaration=%s", async injected => { + const target = client(config(injected ? "https://cli-chat-proxy.grok.com/v1" : "https://api.x.ai/v1", injected)); + upstream = () => Response.json({ id: "resp_hosted", status: "completed", output: [hostedCall] }); + const frames = await turn(target, { tools: injected + ? [{ type: "web_search", external_web_access: true }] + : [{ type: "x_search" }] }); + expect(outbound).toHaveLength(1); + expect(outbound[0].tools.some((tool: Frame) => tool.type === "x_search")).toBe(true); + expect(frames.find(frame => frame.type === "response.output_item.done")?.item).toMatchObject(hostedCall); + expect(fillBusyAndReserve()).not.toBeNull(); + expect(target.closed).toEqual([1001]); + expect(clients.slice(1).every(other => other.ws.readyState === 1)).toBe(true); +}); + +test.each([ + ["https://api.x.ai/v1", false], + ["https://provider.example/v1", true], +] as const)("declared client custom tool remains protected at %s; x_search declaration=%s", async (baseUrl, declareHosted) => { + const target = client(config(baseUrl)); + upstream = () => Response.json({ id: "resp_client", status: "completed", output: [hostedCall] }); + const frames = await turn(target, { tools: [clientTool, ...(declareHosted ? [{ type: "x_search" }] : [])] }); + expect(frames.find(frame => frame.type === "response.output_item.done")?.item).toMatchObject(hostedCall); + expect(fillBusyAndReserve()).toBeNull(); + expect(target.closed).toEqual([]); +}); + +function toolStream(id: string, call: Frame, includeCreated: boolean): Response { + const frames = [ + ...(includeCreated ? [{ type: "response.created", response: { id, status: "in_progress", output: [] } }] : []), + { type: "response.output_item.done", output_index: 0, item: call }, + { type: "response.completed", response: { id, status: "completed", output: [call] } }, + ]; + return new Response(frames.map(frame => `event: ${frame.type}\ndata: ${JSON.stringify(frame)}\n\n`).join(""), + { headers: { "content-type": "text/event-stream" } }); +} + +test.each([false, true])("a reused call ID on an ordinary next turn stays protected; response.created=%s", async includeCreated => { + const target = client(config("https://provider.example/v1")); + const tools = [{ type: "function", name: "fixture_tool", parameters: { type: "object", properties: {} } }]; + const firstCall = { type: "function_call", id: "fc_first", call_id: "call_0", name: "fixture_tool", arguments: '{"step":1}' }; + const nextCall = { ...firstCall, id: "fc_next", arguments: '{"step":2}' }; + upstream = () => toolStream("resp_first", firstCall, true); + await turn(target, { tools }); + expect(target.ws.data.responsesActivity!.idle).toBe(false); + upstream = () => toolStream("resp_next", nextCall, includeCreated); + const frames = await turn(target, { tools, input: [{ type: "function_call_output", call_id: "call_0", output: "first result" }] }); + expect(outbound).toHaveLength(2); + expect(frames.some(frame => frame.type === "response.created")).toBe(includeCreated); + expect(frames.find(frame => frame.type === "response.output_item.done")?.item).toMatchObject(nextCall); + expect(fillBusyAndReserve()).toBeNull(); + expect(target.closed).toEqual([]); +}); + +function twoProviderConfig(): OcxConfig { + const settings = config(); + settings.providers.ordinary = config("https://provider.example/v1").providers.fixture!; + return settings; +} + +test("switching provider on the same socket does not retain the prior turn's hosted authorization", async () => { + const target = client(twoProviderConfig()); + const hosts: string[] = []; + upstream = (_body, url) => { + hosts.push(url.host); + return Response.json({ id: `resp_switch_${hosts.length}`, status: "completed", output: [hostedCall] }); + }; + await turn(target, { tools: [{ type: "x_search" }] }); + expect(target.ws.data.responsesActivity!.idle).toBe(true); + const frames = await turn(target, { model: "ordinary/fixture-model", tools: [clientTool, { type: "x_search" }] }); + expect(hosts).toEqual(["api.x.ai", "provider.example"]); + expect(frames.find(frame => frame.type === "response.output_item.done")?.item).toMatchObject(hostedCall); + expect(fillBusyAndReserve()).toBeNull(); + expect(target.closed).toEqual([]); +}); + +test.each([ + ["darwin", false], ["darwin", true], ["win32", false], ["win32", true], +] as const)("combo fallback uses the final response's tool ownership; relay=%s hosted final=%s", async (platform, hostedFinal) => { + // Exercise the eager Windows relay without changing filesystem/platform identity. + setRelayPlatformForTests(platform); + const settings = twoProviderConfig(); + const targets = [{ provider: "fixture", model: "grok-4.6" }, { provider: "ordinary", model: "fixture-model" }]; + if (hostedFinal) targets.reverse(); + settings.combos = { ownership: { strategy: "failover", targets } }; + const target = client(settings); + const hosts: string[] = []; + upstream = (_body, url) => { + hosts.push(url.host); + if (hosts.length === 1) { + const response = { id: "resp_failed", status: "failed", output: [], + error: { type: "server_error", code: "upstream_server_error", message: "fixture first target unavailable" } }; + return new Response(`data: ${JSON.stringify({ type: "response.failed", response })}\n\n`, + { headers: { "content-type": "text/event-stream" } }); + } + return toolStream("resp_fallback", hostedCall, true); + }; + const frames = await turn(target, { model: "combo/ownership", tools: [clientTool, { type: "x_search" }] }); + expect(hosts).toEqual(hostedFinal ? ["provider.example", "api.x.ai"] : ["api.x.ai", "provider.example"]); + expect(outbound).toHaveLength(2); + expect(frames.find(frame => frame.type === "response.output_item.done")?.item).toMatchObject(hostedCall); + const replacement = fillBusyAndReserve(); + if (hostedFinal) { + expect(replacement).not.toBeNull(); + expect(target.closed).toEqual([1001]); + } else { + expect(replacement).toBeNull(); + expect(target.closed).toEqual([]); + } +});