diff --git a/src/server/responses/codex-ws-session.ts b/src/server/responses/codex-ws-session.ts index 87d8b81905d..36f7065792b 100644 --- a/src/server/responses/codex-ws-session.ts +++ b/src/server/responses/codex-ws-session.ts @@ -1,6 +1,43 @@ import { readCodexWsUpgradeHeaders } from "./codex-ws-metadata"; export const MAX_CODEX_WS_SESSION_EXCHANGES = 32; +export const CODEX_WS_SESSION_CLOSE_GRACE_MS = 1_000; + +/** A disposed lease must not leave a private socket waiting on an unresponsive peer. */ +function closeSessionSocket(socket: WebSocket & { terminate?: () => void }, immediate: boolean): void { + if (immediate) { + // Preserve the retained pool's immediate physical-capacity release. + try { socket.close(); } catch { /* already closing */ } + try { socket.terminate?.(); } catch { /* already closed */ } + return; + } + if (socket.readyState === 3) return; + if (typeof socket.terminate !== "function") { + // A close observer cannot force teardown without the optional native method. + try { socket.close(); } catch { /* best-effort graceful teardown */ } + return; + } + let timer: ReturnType | undefined; + const cleanup = (): void => { + if (timer !== undefined) clearTimeout(timer); + timer = undefined; + socket.removeEventListener("close", cleanup); + }; + const terminate = (): void => { + cleanup(); + if (socket.readyState !== 3) { + try { socket.terminate?.(); } catch { /* best-effort physical teardown */ } + } + }; + // Install before close(): older runtimes and test peers may close synchronously. + // Capture only the socket, not the disposed session or its former exchange owner. + socket.addEventListener("close", cleanup, { once: true }); + timer = setTimeout(terminate, CODEX_WS_SESSION_CLOSE_GRACE_MS); + timer.unref?.(); + try { socket.close(); } catch { terminate(); } + // close() may synchronously change the state narrowed by the earlier guard. + if ((socket.readyState as number) === 3) cleanup(); +} /** Owns one physical socket; request listeners belong to the exchange, not this object. */ export class CodexWsSession { @@ -69,10 +106,7 @@ export class CodexWsSession { try { owner?.(reason); } finally { this.busy = false; this.completedIds.clear(); - try { this.socket.close(); } catch { /* already closing */ } - if (this.retainable) { - try { (this.socket as WebSocket & { terminate?: () => void }).terminate?.(); } catch { /* already closed */ } - } + closeSessionSocket(this.socket, this.retainable); this.changed(); } } diff --git a/structure/transports/responses-wire-shapes.md b/structure/transports/responses-wire-shapes.md index 4bbe49be3e4..476f2267128 100644 --- a/structure/transports/responses-wire-shapes.md +++ b/structure/transports/responses-wire-shapes.md @@ -485,14 +485,14 @@ These are transport-fidelity guarantees, not a provider-billing guarantee. Every exchange also leaves a content-free stage record (`CodexWsStageRecord`, #4191): create-frame bytes (measured on failure only — the committed-success record keeps it null so the happy path never byte-counts a megabyte replay frame), send completion, numeric close code, elapsed, first-frame and first-response-event durations, frame counters, liveness ping/pong counts, pool reuse, the OCX/Bun versions, and (#6740) `terminalEvent` — the normalized `response.completed` / `response.failed` / `response.incomplete` / `error` frame that ended a plain exchange, observed before any synthetic tail — with `terminalErrorStatus`, the wrapper status of a terminal `error` frame. Both are null when no terminal frame arrived, on native-control exchanges (which can end on an acknowledgement), and on rows written before they existed; malformed persisted values read as null without dropping the stage. The exchange pins the record on the resolved Response (`markCodexWsStage`, the same marker seam as `markCodexWsResponse`); `handleResponses` adopts it onto the serving attempt, and usage.jsonl persists it per attempt behind a drop-guard normalizer, so hand-edited rows cannot inject strings into the DTO. Later snapshots update the same response-local record in place, so an attempt holding the committed reference observes final success or failure counters. Each exchange supplies a complete fresh snapshot; separate responses keep distinct records. On eager-relay cancel-drain expiry, upstream cancellation finalizes the transport snapshot before the cancellation hook writes the usage row; an actual terminal observed within the drain still wins over cancellation. The record never carries conversation text, headers, close-reason text, or account identifiers, and it is not a fallback-eligibility signal: nothing it says permits a resend. The one replacement an operator can grant after a socket dies is the resend gate's decision (see [ambiguous-resend gate](responses-failover.md#ambiguous-resend-gate)). -Eligible complete-input creates can retain a canonical upstream socket within -one selected account, credential, thread and turn. Model/tier and immutable -handshake headers and the selected outbound proxy must also match. Turn-state and turn-metadata headers are -projected into their same-name per-frame metadata slots; explicit body values win. -The pool retains at most 32 sockets, expires idle sockets after 30 seconds, and -retires a socket after five minutes or 32 successful exchanges (after active work -finishes). Cancellation, errors, idle unsolicited frames and shutdown dispose it. -A busy key uses a separate one-shot connection rather than interleaving requests. +Eligible complete-input creates can retain a canonical upstream socket within one selected account, credential, thread and turn. +Model/tier, immutable handshake headers and the selected outbound proxy must also match. Turn-state and turn-metadata headers are projected into their same-name per-frame metadata slots; explicit body values win. +The pool retains at most 32 sockets, expires idle sockets after 30 seconds, and retires a socket after five minutes or 32 successful exchanges (after active work finishes). +Cancellation, errors, idle unsolicited frames and shutdown dispose it. A busy key uses a separate one-shot connection rather than interleaving requests. +Private disposal requests graceful close, then schedules optional native `terminate()` after a one-second grace; successful retained release stays reusable. +That unrefed fallback holds only the socket and cleanup state, removes its close observer on timeout, and cancels both on early close. +Without native `terminate()`, disposal attempts `close()` without a fallback timer or observer; retained disposal requests close and optional terminate immediately. +`tests/responses/codex-ws-handshake-validation.test.ts` covers native stalled-peer teardown and optional-method cleanup. This is connection reuse, not native incremental-input synthesis: complete HTTP inputs are never trimmed and no previous response id is invented. Explicit diff --git a/tests/responses/codex-ws-handshake-validation.test.ts b/tests/responses/codex-ws-handshake-validation.test.ts index 5af7c3ef83b..31290af35e9 100644 --- a/tests/responses/codex-ws-handshake-validation.test.ts +++ b/tests/responses/codex-ws-handshake-validation.test.ts @@ -1,6 +1,7 @@ import { afterEach, expect, test } from "bun:test"; +import { watchdogMs } from "../helpers/ci-watchdog"; import { codexWsExchange } from "../../src/server/responses/codex-ws-exchange"; -import { CodexWsSession } from "../../src/server/responses/codex-ws-session"; +import { CODEX_WS_SESSION_CLOSE_GRACE_MS, CodexWsSession } from "../../src/server/responses/codex-ws-session"; import { CODEX_RESPONSES_HTTP_URL, prepareCodexWsRequest } from "../../src/server/responses/codex-ws-request"; import { CODEX_WS_METADATA_MAX_BYTES, CODEX_WS_METADATA_MAX_HEADERS, CODEX_WS_METADATA_MAX_VALUE_BYTES, @@ -123,3 +124,257 @@ test("non-canonical Responses transports keep their existing header policy", asy const response = await exchange(upgrade([flag, "true"]), "https://api.openai.com/v1/responses"); expect(response.headers.has(flag)).toBe(false); }); + +interface CloseTimer { callback: () => void; delay: number; unref: () => void } +class ClosingSocket extends EventTarget { + readyState = 0; + closeCalls = 0; + terminateCalls = 0; + closeMode: "stall" | "sync" | "throw" = "stall"; + terminateThrows = false; + readonly closeListeners = new Set(); + override addEventListener(type: string, listener: EventListenerOrEventListenerObject | null, options?: AddEventListenerOptions | boolean) { + super.addEventListener(type, listener, options); + if (type === "close" && listener) this.closeListeners.add(listener); + } + override removeEventListener(type: string, listener: EventListenerOrEventListenerObject | null, options?: EventListenerOptions | boolean) { + super.removeEventListener(type, listener, options); + if (type === "close" && listener) this.closeListeners.delete(listener); + } + open() { this.readyState = 1; this.dispatchEvent(new Event("open")); } + peerClose() { this.readyState = 3; this.dispatchEvent(new Event("close")); } + close() { + this.closeCalls++; + if (this.closeMode === "throw") throw new Error("fixture close failure"); + this.readyState = 2; + if (this.closeMode === "sync") this.peerClose(); + } + terminate?: () => void = () => { + this.terminateCalls++; + if (this.terminateThrows) throw new Error("fixture terminate failure"); + this.peerClose(); + }; +} + +/** Synchronous clock scope: no awaited work runs while global timers are replaced. */ +function withClosingSession(run: (fixture: { + session: CodexWsSession; socket: ClosingSocket; timers: Set; + expire(): void; unrefs(): number; changes(): number; +}) => void, retainable = false) { + const savedSocket = globalThis.WebSocket; + const savedSet = globalThis.setTimeout; + const savedClear = globalThis.clearTimeout; + const timers = new Set(); + let unrefs = 0, changes = 0; + globalThis.WebSocket = ClosingSocket as unknown as typeof WebSocket; + globalThis.setTimeout = ((callback: () => void, delay: number) => { + const timer = { callback, delay, unref: () => { unrefs++; } }; + timers.add(timer); + return timer; + }) as unknown as typeof setTimeout; + globalThis.clearTimeout = ((timer: CloseTimer) => { timers.delete(timer); }) as unknown as typeof clearTimeout; + try { + const session = new CodexWsSession("ws://127.0.0.1/fixture", {}, retainable, () => { changes++; }); + run({ session, socket: session.socket as unknown as ClosingSocket, timers, + unrefs: () => unrefs, changes: () => changes, + expire() { + for (const timer of [...timers]) { + timers.delete(timer); + timer.callback(); + } + }, + }); + } finally { + globalThis.setTimeout = savedSet; + globalThis.clearTimeout = savedClear; + globalThis.WebSocket = savedSocket; + } +} + +test("private closing handshake has one unrefed grace timer and one forced teardown", () => { + withClosingSession(({ session, socket, timers, expire, unrefs, changes }) => { + socket.open(); + expect(session.reserve()).toBe(true); + const reason = new Error("fixture cancellation"); + let observed: Error | undefined; + session.bindOwner(error => { observed = error; }); + session.dispose(reason); + expect(observed).toBe(reason); + expect(session.closed).toBe(true); + expect(session.busy).toBe(false); + expect(socket.closeCalls).toBe(1); + expect(socket.terminateCalls).toBe(0); + expect(timers.size).toBe(1); + expect([...timers][0]!.delay).toBe(CODEX_WS_SESSION_CLOSE_GRACE_MS); + expect(unrefs()).toBe(1); + session.dispose(); + expect(timers.size).toBe(1); + expire(); + expect(socket.terminateCalls).toBe(1); + expect(socket.readyState).toBe(3); + expect(timers.size).toBe(0); + expect(socket.closeListeners.size).toBe(0); + expect(changes()).toBe(1); + }); +}); + +test("private closing handshake cancels fallback on asynchronous and synchronous close", () => { + for (const mode of ["stall", "sync"] as const) withClosingSession(({ session, socket, timers, expire }) => { + socket.open(); socket.closeMode = mode; + session.dispose(); + if (mode === "stall") socket.peerClose(); + expect(timers.size).toBe(0); + expect(socket.closeListeners.size).toBe(0); + expire(); + expect(socket.terminateCalls).toBe(0); + }); +}); + +test("retained closing handshake keeps immediate disposal and successful reuse", () => { + withClosingSession(({ session, socket, timers }) => { + socket.open(); session.reserve(); session.release("fixture-completed"); + expect(session.hasCompleted("fixture-completed")).toBe(true); + expect(session.reserve()).toBe(true); + expect(socket.closeCalls).toBe(0); + expect(timers.size).toBe(0); + session.dispose(); + expect(socket.closeCalls).toBe(1); + expect(socket.terminateCalls).toBe(1); + expect(timers.size).toBe(0); + }, true); +}); + +test("private closing handshake contains close and terminate exceptions", () => { + withClosingSession(({ session, socket, timers, changes }) => { + socket.closeMode = "throw"; + session.dispose(); + expect(socket.terminateCalls).toBe(1); + expect(timers.size).toBe(0); + expect(socket.closeListeners.size).toBe(0); + expect(changes()).toBe(1); + }); + withClosingSession(({ session, socket, timers, expire }) => { + socket.open(); socket.terminateThrows = true; + session.dispose(); expire(); + expect(socket.terminateCalls).toBe(1); + expect(timers.size).toBe(0); + expect(socket.closeListeners.size).toBe(0); + }); +}); + +test("private closing handshake remains bounded when the exchange owner throws", () => { + withClosingSession(({ session, socket, expire, changes }) => { + socket.open(); session.reserve(); + session.bindOwner(() => { throw new Error("fixture owner failure"); }); + expect(() => session.dispose()).toThrow("fixture owner failure"); + expect(changes()).toBe(1); + expire(); + expect(socket.terminateCalls).toBe(1); + }); +}); + +test("private closing handshake covers completed releases, failed releases and connecting sockets", () => { + for (const completedId of ["fixture-completed", null]) withClosingSession(({ session, socket, expire }) => { + socket.open(); session.reserve(); session.release(completedId); expire(); + expect(socket.terminateCalls).toBe(1); + }); + withClosingSession(({ session, socket, expire }) => { + session.dispose(); expire(); + expect(socket.terminateCalls).toBe(1); + }); +}); + +test("private closing handshake creates no timer after a remote close", () => { + withClosingSession(({ session, socket, timers }) => { + socket.open(); socket.peerClose(); session.dispose(); + expect(socket.closeCalls).toBe(0); + expect(socket.terminateCalls).toBe(0); + expect(timers.size).toBe(0); + }); +}); + +test("private closing handshake adds no observer or timer without optional terminate", () => { + for (const mode of ["stall", "throw"] as const) withClosingSession(({ session, socket, timers, unrefs, changes }) => { + socket.open(); socket.terminate = undefined; socket.closeMode = mode; + session.dispose(); + expect(socket.closeCalls).toBe(1); + expect(timers.size).toBe(0); + expect(socket.closeListeners.size).toBe(0); + expect(unrefs()).toBe(0); + expect(changes()).toBe(1); + }); +}); + +test.each(["dispose", "release"] as const)("private %s closes a real peer that withholds its close reply", async mode => { + // A raw peer accepts the upgrade and reads Close, but deliberately never answers it. + // Keep the native client's close and terminate operations intact on both Bun versions. + const keys = ["HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY", "NO_PROXY", "http_proxy", "https_proxy", "all_proxy", "no_proxy"]; + const saved = Object.fromEntries(keys.map(key => [key, process.env[key]])); + for (const key of keys) delete process.env[key]; + let listener: ReturnType | undefined; + let session: CodexWsSession | undefined; + let timer: ReturnType | undefined; + let resolvePeerClosed!: () => void; + const peerClosed = new Promise(resolve => { resolvePeerClosed = resolve; }); + let headers = Buffer.alloc(0), closeFrame = Buffer.alloc(0), upgraded = false; + try { + listener = Bun.listen({ + hostname: "127.0.0.1", port: 0, + socket: { + data(socket, data) { + if (!upgraded) { + headers = Buffer.concat([headers, data]); + const boundary = headers.indexOf("\r\n\r\n"); + if (boundary < 0) return; + const key = /^Sec-WebSocket-Key:\s*(.+)\r?$/im.exec(headers.subarray(0, boundary).toString())?.[1]?.trim(); + if (!key) { socket.end("HTTP/1.1 400 Bad Request\r\nContent-Length: 0\r\n\r\n"); return; } + const accept = new Bun.CryptoHasher("sha1").update(key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").digest("base64"); + upgraded = true; + socket.write("HTTP/1.1 101 Switching Protocols\r\nConnection: Upgrade\r\nUpgrade: websocket\r\n" + + `Sec-WebSocket-Accept: ${accept}\r\n\r\n`); + data = headers.subarray(boundary + 4); + } + closeFrame = Buffer.concat([closeFrame, data]); + }, + end(socket) { socket.end(); resolvePeerClosed(); }, + close() { resolvePeerClosed(); }, + }, + }); + session = new CodexWsSession(`ws://127.0.0.1:${listener.port}`, {}); + const socket = session.socket; + await new Promise((resolve, reject) => { + const cleanup = () => { + clearTimeout(timer); timer = undefined; + socket.removeEventListener("open", onOpen); + socket.removeEventListener("error", onError); + }; + const onOpen = () => { cleanup(); resolve(); }; + const onError = () => { cleanup(); reject(new Error("fixture WebSocket failed to open")); }; + socket.addEventListener("open", onOpen, { once: true }); + socket.addEventListener("error", onError, { once: true }); + timer = setTimeout(onError, watchdogMs(3000)); + }); + expect(session.reserve()).toBe(true); + const deadline = new Promise((_, reject) => { + timer = setTimeout(() => reject(new Error("disposed private socket left its peer open")), watchdogMs(3000)); + }); + if (mode === "release") session.release("fixture-completed"); + else session.dispose(); + expect(session.closed).toBe(true); + expect(session.busy).toBe(false); + expect(session.reserve()).toBe(false); + await Promise.race([peerClosed, deadline]); + expect(closeFrame.length).toBeGreaterThanOrEqual(2); + expect(closeFrame[0]! & 0x0f).toBe(0x08); + expect(socket.readyState).toBe(3); + } finally { + clearTimeout(timer); + session?.dispose(); + try { (session?.socket as (WebSocket & { terminate?: () => void }) | undefined)?.terminate?.(); } catch { /* fixture teardown */ } + listener?.stop(true); + for (const key of keys) { + if (saved[key] === undefined) delete process.env[key]; + else process.env[key] = saved[key]; + } + } +}, watchdogMs(5000));