diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 98c607877..0b56d0805 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -9,6 +9,7 @@ - Fixed Fireworks GLM 5.2 Fast to use the OpenAI-compatible endpoint and `thinkingLevelMap`, aligning it with GLM 5.2 ([#6195](https://github.com/earendil-works/pi/issues/6195)). - Fixed Amazon Bedrock prompt-cache points for Claude Fable 5 and Claude Sonnet 5 ([#6235](https://github.com/earendil-works/pi/issues/6235)). - Fixed DS4 server context overflow detection for `Prompt has ... tokens, but the configured context size is ... tokens` errors ([#6262](https://github.com/earendil-works/pi/issues/6262)). +- Fixed OpenAI Codex WebSocket sessions to rotate cached connections before the backend's 60-minute limit, avoiding connection-limit failures on long sessions ([#6268](https://github.com/earendil-works/pi/issues/6268)). ### Added diff --git a/packages/ai/src/api/openai-codex-responses.ts b/packages/ai/src/api/openai-codex-responses.ts index d09a8d674..16734a756 100644 --- a/packages/ai/src/api/openai-codex-responses.ts +++ b/packages/ai/src/api/openai-codex-responses.ts @@ -750,6 +750,7 @@ async function* parseSSE(response: Response, signal?: AbortSignal): AsyncGenerat const OPENAI_BETA_RESPONSES_WEBSOCKETS = "responses_websockets=2026-02-06"; const SESSION_WEBSOCKET_CACHE_TTL_MS = 5 * 60 * 1000; +const SESSION_WEBSOCKET_MAX_AGE_MS = 55 * 60 * 1000; type WebSocketEventType = "open" | "message" | "error" | "close"; type WebSocketListener = (event: unknown) => void; @@ -770,6 +771,7 @@ interface CachedWebSocketContinuationState { interface CachedWebSocketConnection { socket: WebSocketLike; busy: boolean; + createdAt: number; idleTimer?: ReturnType; continuation?: CachedWebSocketContinuationState; } @@ -934,6 +936,10 @@ function isWebSocketReusable(socket: WebSocketLike): boolean { return readyState === undefined || readyState === 1; } +function isWebSocketSessionExpired(entry: CachedWebSocketConnection): boolean { + return Date.now() - entry.createdAt >= SESSION_WEBSOCKET_MAX_AGE_MS; +} + function closeWebSocketSilently(socket: WebSocketLike, code = 1000, reason = "done"): void { try { socket.close(code, reason); @@ -1057,7 +1063,10 @@ async function acquireWebSocket( clearTimeout(cached.idleTimer); cached.idleTimer = undefined; } - if (!cached.busy && isWebSocketReusable(cached.socket)) { + if (!cached.busy && isWebSocketSessionExpired(cached)) { + closeWebSocketSilently(cached.socket, 1000, "connection_age_limit"); + websocketSessionCache.delete(sessionId); + } else if (!cached.busy && isWebSocketReusable(cached.socket)) { cached.busy = true; return { socket: cached.socket, @@ -1091,7 +1100,7 @@ async function acquireWebSocket( } const socket = await connectWebSocket(url, headers, signal, connectTimeoutMs, env); - const entry: CachedWebSocketConnection = { socket, busy: true }; + const entry: CachedWebSocketConnection = { socket, busy: true, createdAt: Date.now() }; websocketSessionCache.set(sessionId, entry); return { socket, diff --git a/packages/ai/test/openai-codex-stream.test.ts b/packages/ai/test/openai-codex-stream.test.ts index 3a2c06177..016acc822 100644 --- a/packages/ai/test/openai-codex-stream.test.ts +++ b/packages/ai/test/openai-codex-stream.test.ts @@ -4,6 +4,7 @@ import { join } from "node:path"; import { zstdDecompressSync } from "node:zlib"; import { afterEach, describe, expect, it, vi } from "vitest"; import { + closeOpenAICodexWebSocketSessions, getOpenAICodexWebSocketDebugStats, resetOpenAICodexWebSocketDebugStats, stream as streamOpenAICodexResponses, @@ -20,6 +21,7 @@ afterEach(() => { } else { process.env.PI_CODING_AGENT_DIR = originalAgentDir; } + closeOpenAICodexWebSocketSessions(); resetOpenAICodexWebSocketDebugStats(); vi.useRealTimers(); vi.restoreAllMocks(); @@ -1444,6 +1446,111 @@ describe("openai-codex streaming", () => { expect(fetchMock).not.toHaveBeenCalled(); }); + it("opens a fresh cached websocket before the backend connection age limit", async () => { + vi.useFakeTimers(); + const startedAt = new Date("2026-07-03T00:00:00Z"); + vi.setSystemTime(startedAt); + const token = mockToken(); + const sentConnectionIds: number[] = []; + let connections = 0; + + class MockWebSocket { + static OPEN = 1; + static CLOSED = 3; + readyState = MockWebSocket.OPEN; + private readonly connectionId = ++connections; + private listeners = new Map void>>(); + + constructor(_url: string, _protocols?: string | string[] | { headers?: Record }) { + queueMicrotask(() => this.dispatch("open", {})); + } + + addEventListener(type: string, listener: (event: unknown) => void): void { + let listeners = this.listeners.get(type); + if (!listeners) { + listeners = new Set(); + this.listeners.set(type, listeners); + } + listeners.add(listener); + } + + removeEventListener(type: string, listener: (event: unknown) => void): void { + this.listeners.get(type)?.delete(listener); + } + + send(): void { + sentConnectionIds.push(this.connectionId); + const responseId = `resp_${this.connectionId}`; + queueMicrotask(() => { + this.dispatch("message", { + data: JSON.stringify({ + type: "response.completed", + response: { + id: responseId, + status: "completed", + usage: { input_tokens: 5, output_tokens: 3, total_tokens: 8 }, + }, + }), + }); + }); + } + + close(): void { + this.readyState = MockWebSocket.CLOSED; + } + + private dispatch(type: string, event: unknown): void { + for (const listener of this.listeners.get(type) ?? []) { + listener(event); + } + } + } + + vi.stubGlobal("WebSocket", MockWebSocket); + + const model: Model<"openai-codex-responses"> = { + id: "gpt-5.1-codex", + name: "GPT-5.1 Codex", + api: "openai-codex-responses", + provider: "openai-codex", + baseUrl: "https://chatgpt.com/backend-api", + reasoning: true, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 400000, + maxTokens: 128000, + }; + const sessionId = "aged-ws-session"; + const firstContext: Context = { + systemPrompt: "You are a helpful assistant.", + messages: [{ role: "user", content: "Say hello", timestamp: 1 }], + }; + + const first = await streamOpenAICodexResponses(model, firstContext, { + apiKey: token, + sessionId, + transport: "websocket-cached", + }).result(); + vi.setSystemTime(new Date(startedAt.getTime() + 56 * 60 * 1000)); + const secondContext: Context = { + systemPrompt: "You are a helpful assistant.", + messages: [...firstContext.messages, first, { role: "user", content: "Now finish", timestamp: 2 }], + }; + + await streamOpenAICodexResponses(model, secondContext, { + apiKey: token, + sessionId, + transport: "websocket-cached", + }).result(); + + expect(connections).toBe(2); + expect(sentConnectionIds).toEqual([1, 2]); + expect(getOpenAICodexWebSocketDebugStats(sessionId)).toMatchObject({ + connectionsCreated: 2, + connectionsReused: 0, + }); + }); + it("sends only response input deltas in websocket-cached mode", async () => { const token = mockToken(); const sentBodies: unknown[] = [];