fix(ai): corrected codex websocket stale-frame cleanup between turns

- Added `PI_CODEX_WEBSOCKET_MAX_IDLE_REUSE_MS` to gate `getOrCreateCodexWebSocketConnection()` reuse for stale sockets (#1450). - Fixed stale websocket responses by dropping non-error frames before each `streamRequest()`, preventing terminal frames from a prior turn from ending the next. - Changed heartbeat logic to fail sockets when no inbound data arrives after ping timeout, even without pong events.
2026-05-27 14:59:41 +02:00
parent 99bae2ce6c
commit 5d63af8211
3 changed files with 308 additions and 11 deletions
@@ -1,6 +1,15 @@
 # Changelog

 ## [Unreleased]
+### Added
+
+- Added `PI_CODEX_WEBSOCKET_MAX_IDLE_REUSE_MS` to control how long an idle Codex WebSocket stays eligible for reuse, with `0` disabling the check
+
+### Fixed
+
+- Fixed reused Codex WebSocket connections that had gone silent without activity to be dropped and replaced with a fresh handshake after the idle-reuse threshold, preventing stalled next requests
+- Fixed stale response frames left in the websocket queue from a completed turn so subsequent requests no longer process terminal frames from the previous response
+- Fixed websocket dead-socket detection to fail a stale connection when no inbound traffic or pong is observed after a ping timeout, improving recovery on runtimes that do not emit pong events

 ## [15.5.5] - 2026-05-27

@@ -96,6 +96,20 @@ const CODEX_WEBSOCKET_CONNECT_TIMEOUT_MS = 10000;
 const CODEX_WEBSOCKET_PING_INTERVAL_MS = 10_000;
 const CODEX_WEBSOCKET_PONG_TIMEOUT_MS = 60_000;
 const CODEX_WEBSOCKET_MESSAGE_QUEUE_CAPACITY = 4096;
+/**
+ * Maximum quiet period (no inbound frames AND no observed pong) we'll trust a
+ * reused WebSocket for before forcing a fresh handshake. Codex backends and
+ * intermediaries occasionally evict idle sockets server-side without sending a
+ * FIN, leaving the local `readyState` as OPEN while the next `send()` becomes a
+ * write into a half-open buffer. Reusing such a socket parks the next request
+ * at `#nextMessage` until the first-event/idle timeout fires (issue #1450). The
+ * heartbeat below also catches dead sockets, but only after `pongTimeoutMs`
+ * (default 60s) and only while a request is active — this gate closes the door
+ * earlier and even when the gap between requests is purely client-side (tool
+ * execution, user typing, etc.). Set `PI_CODEX_WEBSOCKET_MAX_IDLE_REUSE_MS=0`
+ * to disable.
+ */
+const CODEX_WEBSOCKET_MAX_IDLE_REUSE_MS = 30_000;
 /**
 * Steady-state liveness ceiling for the Codex WebSocket transport. Distinct from
 * the OMP-wide stream watchdog removed in #1392: a WebSocket can stay TCP-open
@@ -290,6 +304,10 @@ function getCodexWebSocketMessageQueueCapacity(): number {
 	);
 }

+function getCodexWebSocketMaxIdleReuseMs(): number {
+	return parseCodexNonNegativeInteger($env.PI_CODEX_WEBSOCKET_MAX_IDLE_REUSE_MS, CODEX_WEBSOCKET_MAX_IDLE_REUSE_MS);
+}
+
 function createCodexProviderSessionState(): CodexProviderSessionState {
 	const state: CodexProviderSessionState = {
 		webSocketSessions: new Map(),
@@ -2028,8 +2046,15 @@ class CodexWebSocketConnection {
 	#streamObserver?: (event: RawSseEvent) => void;
 	#heartbeatInterval: NodeJS.Timeout | undefined;
 	#removePongListener?: () => void;
-	#lastPongAt = 0;
-	#observedPong = false;
+	/**
+	 * Wall-clock of the most recent inbound activity on this socket — any
+	 * decoded message, any pong, or the moment the handshake completed. Used
+	 * by {@link isHealthyForReuse} so we don't write a continuation frame into
+	 * a TCP-open-but-server-evicted socket whose `readyState` still says OPEN.
+	 */
+	#lastInboundAt = 0;
+	/** Wall-clock of the last heartbeat ping we issued; 0 if none yet. */
+	#lastPingAt = 0;

 	constructor(url: string, headers: Record<string, string>, options: CodexWebSocketConnectionOptions) {
 		this.#url = url;
@@ -2041,6 +2066,29 @@ class CodexWebSocketConnection {
 		return this.#socket?.readyState === WebSocket.OPEN;
 	}

+	/**
+	 * Stricter variant of {@link isOpen} for the connection-pool reuse gate.
+	 * Refuses sockets that have been silent past {@link CODEX_WEBSOCKET_MAX_IDLE_REUSE_MS}.
+	 *
+	 * Bun's `WebSocket` does not always surface server-side eviction (no
+	 * `onclose`, no `onerror`), so a socket can sit in readyState OPEN long
+	 * after the upstream has dropped it. Reusing such a socket sends the next
+	 * `response.create` into a half-open write buffer and parks the reader
+	 * until the first-event / idle timeout fires (issue #1450). Forcing a
+	 * reconnect on any suspect socket trades a sub-second handshake for a
+	 * 60–300 s stall.
+	 */
+	isHealthyForReuse(): boolean {
+		if (!this.isOpen()) return false;
+		const maxIdleMs = getCodexWebSocketMaxIdleReuseMs();
+		if (maxIdleMs <= 0) return true;
+		// Initial connect sets #lastInboundAt; any later message or pong refreshes
+		// it. A zero value means the field was never initialized, which itself is
+		// a desync — treat as unhealthy.
+		if (this.#lastInboundAt === 0) return false;
+		return Date.now() - this.#lastInboundAt <= maxIdleMs;
+	}
+
 	matchesAuth(headers: Record<string, string>): boolean {
 		return this.#headers.authorization === headers.authorization;
 	}
@@ -2103,6 +2151,7 @@ class CodexWebSocketConnection {
 			if (!settled) {
 				settled = true;
 				clearPending();
+				this.#lastInboundAt = Date.now();
 				this.#captureHandshakeHeaders(socket, event);
 				this.#startHeartbeat(socket);
 				resolve();
@@ -2136,6 +2185,10 @@ class CodexWebSocketConnection {
 			this.#push(null);
 		};
 		socket.onmessage = event => {
+			// Stamp inbound activity before parsing so even malformed frames refresh
+			// the liveness clock — what matters for reuse health is that the upstream
+			// is still talking to us, not that every frame is well-formed.
+			this.#lastInboundAt = Date.now();
 			try {
 				const text = typeof event.data === "string" ? event.data : Buffer.from(event.data).toString("utf-8");
 				if (!text) return;
@@ -2179,6 +2232,17 @@ class CodexWebSocketConnection {
 		}
 		this.#activeRequest = true;
 		this.#streamObserver = onSseEvent;
+		// Drain any non-error frames left over from a prior request before sending.
+		// `processCodexResponseStream` breaks its `for-await` on the terminal event,
+		// which interrupts our generator at `yield next` (the post-yield `break`
+		// never runs). Any frame that landed between the consumer's break and the
+		// generator's `finally` lingers in `#queue` and would otherwise become the
+		// first frame of THIS request — a stale `response.completed` would end the
+		// turn immediately with empty output, and a stale non-progress frame would
+		// flip `sawFirstEvent` and silently downgrade the first-event timeout to
+		// the longer idle timeout. Transport errors are preserved so we surface
+		// the death signal instead of writing into a dead socket.
+		this.#dropStaleFrames();
 		const onAbort = () => {
 			this.close("aborted");
 			this.#push(createCodexWebSocketTransportError("request was aborted"));
@@ -2287,12 +2351,12 @@ class CodexWebSocketConnection {
 		const intervalMs = getCodexWebSocketPingIntervalMs();
 		if (intervalMs <= 0) return;

-		this.#lastPongAt = Date.now();
-		this.#observedPong = false;
+		this.#lastPingAt = 0;
 		const socketEventTarget = socket as EventTarget;
 		const onPong = () => {
-			this.#observedPong = true;
-			this.#lastPongAt = Date.now();
+			// Pongs are inbound activity — refresh the reuse-health clock so a quiet
+			// but ping-responsive socket stays trustworthy across requests.
+			this.#lastInboundAt = Date.now();
 		};
 		if (
 			typeof socketEventTarget.addEventListener === "function" &&
@@ -2307,8 +2371,21 @@ class CodexWebSocketConnection {
 				this.#stopHeartbeat();
 				return;
 			}
+			// Fail-closed on missing pongs even when no pong has ever been observed.
+			// The previous `#observedPong &&` guard disabled the timeout entirely on
+			// runtimes where Bun does not surface a `pong` event for our outgoing
+			// pings (issue #1450) — letting truly dead sockets sail through the
+			// pool until the per-request first-event / idle timeout (60–300 s)
+			// finally fired. Instead, trigger on inbound silence: if we sent a
+			// ping at least `pongTimeoutMs` ago and have received no traffic of
+			// any kind (data frame or pong) since, the socket is unhealthy.
 			const pongTimeoutMs = getCodexWebSocketPongTimeoutMs();
-			if (this.#observedPong && pongTimeoutMs > 0 && Date.now() - this.#lastPongAt > pongTimeoutMs) {
+			if (
+				pongTimeoutMs > 0 &&
+				this.#lastPingAt > 0 &&
+				this.#lastPingAt > this.#lastInboundAt &&
+				Date.now() - this.#lastPingAt > pongTimeoutMs
+			) {
 				this.#failQueue(createCodexWebSocketTransportError("websocket pong timeout"), "pong-timeout");
 				return;
 			}
@@ -2318,6 +2395,7 @@ class CodexWebSocketConnection {
 			}
 			try {
 				socket.ping();
+				this.#lastPingAt = Date.now();
 			} catch (error) {
 				this.#failQueue(
 					createCodexWebSocketTransportError(
@@ -2339,7 +2417,7 @@ class CodexWebSocketConnection {
 			this.#removePongListener();
 			this.#removePongListener = undefined;
 		}
-		this.#observedPong = false;
+		this.#lastPingAt = 0;
 	}

 	#failQueue(error: Error, closeReason: string): void {
@@ -2350,6 +2428,25 @@ class CodexWebSocketConnection {
 		this.#wakeWaiters();
 	}

+	/**
+	 * Discard data frames from a previous request that remained in `#queue`
+	 * after the consumer broke out on the terminal event. Preserves any queued
+	 * transport error (from `onerror` / `onclose` / `#failQueue`) so the next
+	 * `#nextMessage` surfaces the death signal instead of waiting it out.
+	 *
+	 * Returns the number of frames dropped (test/debug visibility only).
+	 */
+	#dropStaleFrames(): number {
+		if (this.#queue.length === 0) return 0;
+		const surviving = this.#queue.filter(item => item instanceof Error);
+		const dropped = this.#queue.length - surviving.length;
+		if (dropped === 0) return 0;
+		this.#queue.length = 0;
+		for (const item of surviving) this.#queue.push(item);
+		logCodexDebug("codex websocket dropped stale frames before request", { dropped });
+		return dropped;
+	}
+
 	#wakeWaiters(): void {
 		for (;;) {
 			const waiter = this.#waiters.shift();
@@ -2418,12 +2515,22 @@ async function getOrCreateCodexWebSocketConnection(
 ): Promise<CodexWebSocketConnection> {
 	const headerRecord = headersToRecord(headers);
 	if (state.connection?.isOpen()) {
-		if (state.connection.matchesAuth(headerRecord)) {
+		if (!state.connection.matchesAuth(headerRecord)) {
+			state.connection.close("token-refresh");
+			resetCodexWebSocketAppendState(state);
+		} else if (state.connection.isHealthyForReuse()) {
 			logger.time("codexWs:reuseOpenSocket");
 			return state.connection;
+		} else {
+			// Open in readyState but no inbound traffic recently — likely server-
+			// evicted (issue #1450). Force a fresh handshake instead of writing
+			// `response.create` into a half-open buffer and waiting out the
+			// first-event timeout. Drop append state because the new socket
+			// won't carry the prior `previous_response_id` context.
+			logCodexDebug("codex websocket reuse rejected by health check", {});
+			state.connection.close("stale-reuse");
+			resetCodexWebSocketAppendState(state);
 		}
-		state.connection.close("token-refresh");
-		resetCodexWebSocketAppendState(state);
 	}
 	state.connection?.close("reconnect");
 	resetCodexWebSocketAppendState(state);
@@ -20,6 +20,7 @@ const originalCodexWebSocketFirstEventTimeoutMs = Bun.env.PI_CODEX_WEBSOCKET_FIR
 const originalCodexWebSocketPingIntervalMs = Bun.env.PI_CODEX_WEBSOCKET_PING_INTERVAL_MS;
 const originalCodexWebSocketPongTimeoutMs = Bun.env.PI_CODEX_WEBSOCKET_PONG_TIMEOUT_MS;
 const originalCodexWebSocketMessageQueueCapacity = Bun.env.PI_CODEX_WEBSOCKET_MESSAGE_QUEUE_CAPACITY;
+const originalCodexWebSocketMaxIdleReuseMs = Bun.env.PI_CODEX_WEBSOCKET_MAX_IDLE_REUSE_MS;

 function restoreEnv(name: string, value: string | undefined): void {
 	if (value === undefined) {
@@ -41,6 +42,7 @@ afterEach(() => {
 	restoreEnv("PI_CODEX_WEBSOCKET_PING_INTERVAL_MS", originalCodexWebSocketPingIntervalMs);
 	restoreEnv("PI_CODEX_WEBSOCKET_PONG_TIMEOUT_MS", originalCodexWebSocketPongTimeoutMs);
 	restoreEnv("PI_CODEX_WEBSOCKET_MESSAGE_QUEUE_CAPACITY", originalCodexWebSocketMessageQueueCapacity);
+	restoreEnv("PI_CODEX_WEBSOCKET_MAX_IDLE_REUSE_MS", originalCodexWebSocketMaxIdleReuseMs);
 	vi.restoreAllMocks();
 });

@@ -2864,4 +2866,183 @@ describe("openai-codex streaming", () => {
 		expect(requestTurnStates[0]).toBeNull();
 		expect(requestTurnStates[1]).toBe("turn-state-1");
 	});
+
+	it("forces a fresh websocket when the prior connection has been idle past PI_CODEX_WEBSOCKET_MAX_IDLE_REUSE_MS", async () => {
+		const tempDir = TempDir.createSync("@pi-codex-stream-");
+		setAgentDir(tempDir.path());
+
+		// Tight reuse window so the test doesn't have to sleep for seconds. Disable
+		// the heartbeat so it doesn't independently kill the idle socket and mask
+		// the reuse-gate behaviour we're trying to verify.
+		Bun.env.PI_CODEX_WEBSOCKET_MAX_IDLE_REUSE_MS = "10";
+		Bun.env.PI_CODEX_WEBSOCKET_PING_INTERVAL_MS = "0";
+		Bun.env.PI_CODEX_WEBSOCKET_RETRY_BUDGET = "0";
+
+		const payload = Buffer.from(
+			JSON.stringify({ "https://api.openai.com/auth": { chatgpt_account_id: "acc_test" } }),
+			"utf8",
+		).toBase64();
+		const token = `aaa.${payload}.bbb`;
+
+		const fetchMock = vi.fn(async () => {
+			throw new Error("SSE fallback should not be called");
+		});
+		global.fetch = fetchMock as unknown as typeof fetch;
+
+		let constructorCount = 0;
+		let sendCount = 0;
+		class IdleReuseWebSocket extends MockWebSocket {
+			constructor(url: string, options?: { headers?: WsHeaders }) {
+				super(url, options);
+				constructorCount += 1;
+				this.scheduleOpen();
+			}
+
+			send(_data: string): void {
+				sendCount += 1;
+				this.emitCodexResponse({
+					messageId: `msg_${sendCount}`,
+					responseId: `resp_${sendCount}`,
+					text: `Hello ${sendCount}`,
+				});
+			}
+		}
+
+		global.WebSocket = IdleReuseWebSocket as unknown as typeof WebSocket;
+
+		const model = createCodexTestModel("https://chatgpt.com/backend-api");
+		const providerSessionState = new Map<string, ProviderSessionState>();
+		const firstContext: Context = {
+			systemPrompt: ["You are a helpful assistant."],
+			messages: [{ role: "user", content: "First", timestamp: Date.now() }],
+		};
+		const secondContext: Context = {
+			systemPrompt: ["You are a helpful assistant."],
+			messages: [
+				{ role: "user", content: "First", timestamp: Date.now() },
+				{ role: "user", content: "Second", timestamp: Date.now() },
+			],
+		};
+
+		await streamOpenAICodexResponses(model, firstContext, {
+			apiKey: token,
+			sessionId: "ws-idle-reuse-session",
+			providerSessionState,
+		}).result();
+
+		// Simulate the gap between a tool result and the continuation request: the
+		// socket sat quiet long enough that we shouldn't trust it without a fresh
+		// handshake. 30 ms > MAX_IDLE_REUSE_MS (10).
+		await new Promise(resolve => setTimeout(resolve, 30));
+
+		const second = await streamOpenAICodexResponses(model, secondContext, {
+			apiKey: token,
+			sessionId: "ws-idle-reuse-session",
+			providerSessionState,
+		}).result();
+
+		expect(second.stopReason).toBe("stop");
+		expect(constructorCount).toBe(2);
+		expect(sendCount).toBe(2);
+		expect(fetchMock).not.toHaveBeenCalled();
+	});
+
+	it("drops stale frames from a prior response before sending the next websocket request", async () => {
+		const tempDir = TempDir.createSync("@pi-codex-stream-");
+		setAgentDir(tempDir.path());
+
+		// Generous reuse window so the connection is happily reused across turns —
+		// the queue-drain behaviour is the only variable here.
+		Bun.env.PI_CODEX_WEBSOCKET_MAX_IDLE_REUSE_MS = "60000";
+		Bun.env.PI_CODEX_WEBSOCKET_PING_INTERVAL_MS = "0";
+		Bun.env.PI_CODEX_WEBSOCKET_RETRY_BUDGET = "0";
+
+		const payload = Buffer.from(
+			JSON.stringify({ "https://api.openai.com/auth": { chatgpt_account_id: "acc_test" } }),
+			"utf8",
+		).toBase64();
+		const token = `aaa.${payload}.bbb`;
+
+		const fetchMock = vi.fn(async () => {
+			throw new Error("SSE fallback should not be called");
+		});
+		global.fetch = fetchMock as unknown as typeof fetch;
+
+		let constructorCount = 0;
+		let sendCount = 0;
+		class LateFrameWebSocket extends MockWebSocket {
+			constructor(url: string, options?: { headers?: WsHeaders }) {
+				super(url, options);
+				constructorCount += 1;
+				this.scheduleOpen();
+			}
+
+			send(_data: string): void {
+				sendCount += 1;
+				if (sendCount === 1) {
+					this.emitCodexResponse({
+						messageId: "msg_1",
+						responseId: "resp_1",
+						text: "First",
+					});
+					// Stale frame that lands AFTER the consumer breaks on
+					// response.completed. Without the queue-drain at the top of
+					// streamRequest, this becomes the first frame of the next
+					// request: a stale terminal event would resolve the new turn
+					// with empty content, never reaching the model's real response.
+					this.sendJson({
+						type: "response.completed",
+						response: { id: "resp_stale", status: "completed", usage: DEFAULT_USAGE },
+					});
+					return;
+				}
+				this.emitCodexResponse({
+					messageId: "msg_2",
+					responseId: "resp_2",
+					text: "Second",
+				});
+			}
+		}
+
+		global.WebSocket = LateFrameWebSocket as unknown as typeof WebSocket;
+
+		const model = createCodexTestModel("https://chatgpt.com/backend-api");
+		const providerSessionState = new Map<string, ProviderSessionState>();
+		const firstContext: Context = {
+			systemPrompt: ["You are a helpful assistant."],
+			messages: [{ role: "user", content: "First", timestamp: Date.now() }],
+		};
+		const secondContext: Context = {
+			systemPrompt: ["You are a helpful assistant."],
+			messages: [
+				{ role: "user", content: "First", timestamp: Date.now() },
+				{ role: "user", content: "Second", timestamp: Date.now() },
+			],
+		};
+
+		const first = await streamOpenAICodexResponses(model, firstContext, {
+			apiKey: token,
+			sessionId: "ws-stale-frame-session",
+			providerSessionState,
+		}).result();
+		expect(first.stopReason).toBe("stop");
+
+		const second = await streamOpenAICodexResponses(model, secondContext, {
+			apiKey: token,
+			sessionId: "ws-stale-frame-session",
+			providerSessionState,
+		}).result();
+
+		expect(second.stopReason).toBe("stop");
+		expect(constructorCount).toBe(1);
+		expect(sendCount).toBe(2);
+		// Second turn must reflect the second response, not the stale terminal frame
+		// from the first turn's tail.
+		expect(second.responseId).toBe("resp_2");
+		const text = second.content
+			.filter((c): c is Extract<typeof c, { type: "text" }> => c.type === "text")
+			.map(c => c.text)
+			.join("");
+		expect(text).toBe("Second");
+	});
 });