diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 8a43f32f..694e25a7 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -20,6 +20,7 @@ ### Fixed +- Fixed OpenAI Codex Responses WebSocket sessions to reconnect once when OpenAI's connection limit is reached before output starts ([#5973](https://github.com/earendil-works/pi/issues/5973)). - Fixed OpenCode Go GLM-5.2 metadata to expose `xhigh` reasoning and send `reasoning_effort: "max"` ([#5967](https://github.com/earendil-works/pi/issues/5967)). - Fixed Claude Fable 5 thinking-off requests to omit Anthropic's unsupported `thinking.type: "disabled"` payload ([#5567](https://github.com/earendil-works/pi/pull/5567) by [@tmustier](https://github.com/tmustier)). diff --git a/packages/ai/src/api/openai-codex-responses.ts b/packages/ai/src/api/openai-codex-responses.ts index ee2503d5..6107bb22 100644 --- a/packages/ai/src/api/openai-codex-responses.ts +++ b/packages/ai/src/api/openai-codex-responses.ts @@ -61,6 +61,7 @@ const DEFAULT_SSE_HEADER_TIMEOUT_MS = 20_000; const DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS = 15_000; const CODEX_TOOL_CALL_PROVIDERS = new Set(["openai", "openai-codex", "opencode"]); const WEBSOCKET_MESSAGE_TOO_BIG_CLOSE_CODE = 1009; +const WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE = "websocket_connection_limit_reached"; const CODEX_RESPONSE_STATUSES = new Set([ "completed", @@ -253,52 +254,62 @@ export const stream: StreamFunction<"openai-codex-responses", OpenAICodexRespons if (transport !== "sse" && !websocketDisabledForSession) { let websocketStarted = false; - try { - await processWebSocketStream( - resolveCodexWebSocketUrl(model.baseUrl), - body, - websocketHeaders, - output, - stream, - model, - () => { - websocketStarted = true; - }, - idleTimeoutMs, - websocketConnectTimeoutMs, - options, - ); + let retriedWebSocketConnectionLimit = false; + while (true) { + websocketStarted = false; + try { + await processWebSocketStream( + resolveCodexWebSocketUrl(model.baseUrl), + body, + websocketHeaders, + output, + stream, + model, + () => { + websocketStarted = true; + }, + idleTimeoutMs, + websocketConnectTimeoutMs, + options, + ); - if (options?.signal?.aborted) { - throw new Error("Request was aborted"); + if (options?.signal?.aborted) { + throw new Error("Request was aborted"); + } + stream.push({ + type: "done", + reason: output.stopReason as "stop" | "length" | "toolUse", + message: output, + }); + stream.end(); + return; + } catch (error) { + const aborted = options?.signal?.aborted; + const connectionLimitBeforeStart = !websocketStarted && isWebSocketConnectionLimitReachedError(error); + if (!aborted && connectionLimitBeforeStart && !retriedWebSocketConnectionLimit) { + retriedWebSocketConnectionLimit = true; + continue; + } + if (aborted || (isCodexNonTransportError(error) && !connectionLimitBeforeStart)) { + throw error; + } + appendAssistantMessageDiagnostic( + output, + createAssistantMessageDiagnostic("provider_transport_failure", error, { + configuredTransport: transport, + fallbackTransport: websocketStarted ? undefined : "sse", + eventsEmitted: websocketStarted, + phase: websocketStarted ? "after_message_stream_start" : "before_message_stream_start", + requestBytes: new TextEncoder().encode(bodyJson).byteLength, + }), + ); + recordWebSocketFailure(options?.sessionId, error); + if (websocketStarted) { + throw error; + } + recordWebSocketSseFallback(options?.sessionId); + break; } - stream.push({ - type: "done", - reason: output.stopReason as "stop" | "length" | "toolUse", - message: output, - }); - stream.end(); - return; - } catch (error) { - const aborted = options?.signal?.aborted; - if (aborted || isCodexNonTransportError(error)) { - throw error; - } - appendAssistantMessageDiagnostic( - output, - createAssistantMessageDiagnostic("provider_transport_failure", error, { - configuredTransport: transport, - fallbackTransport: websocketStarted ? undefined : "sse", - eventsEmitted: websocketStarted, - phase: websocketStarted ? "after_message_stream_start" : "before_message_stream_start", - requestBytes: new TextEncoder().encode(bodyJson).byteLength, - }), - ); - recordWebSocketFailure(options?.sessionId, error); - if (websocketStarted) { - throw error; - } - recordWebSocketSseFallback(options?.sessionId); } } @@ -582,16 +593,32 @@ function isCodexNonTransportError(error: unknown): boolean { return error instanceof CodexApiError || error instanceof CodexProtocolError; } +function isWebSocketConnectionLimitReachedError(error: unknown): boolean { + return error instanceof CodexApiError && error.code === WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE; +} + +function extractCodexEventError(event: Record): { code?: string; message?: string } { + const nested = event.error && typeof event.error === "object" ? (event.error as Record) : undefined; + return { + code: typeof event.code === "string" ? event.code : typeof nested?.code === "string" ? nested.code : undefined, + message: + typeof event.message === "string" + ? event.message + : typeof nested?.message === "string" + ? nested.message + : undefined, + }; +} + async function* mapCodexEvents(events: AsyncIterable>): AsyncGenerator { for await (const event of events) { const type = typeof event.type === "string" ? event.type : undefined; if (!type) continue; if (type === "error") { - const code = (event as { code?: string }).code || ""; - const message = (event as { message?: string }).message || ""; + const { code, message } = extractCodexEventError(event); throw new CodexApiError(`Codex error: ${message || code || JSON.stringify(event)}`, { - code: code || undefined, + code, payload: event, }); } diff --git a/packages/ai/test/openai-codex-stream.test.ts b/packages/ai/test/openai-codex-stream.test.ts index 9c526151..b68d5b73 100644 --- a/packages/ai/test/openai-codex-stream.test.ts +++ b/packages/ai/test/openai-codex-stream.test.ts @@ -1195,6 +1195,68 @@ describe("openai-codex streaming", () => { }); }); + it("reconnects once when the websocket connection limit is reached before output starts", async () => { + const token = mockToken(); + let connections = 0; + + const fetchMock = vi.fn(); + vi.stubGlobal("fetch", fetchMock); + + class MockWebSocket extends EventTarget { + private readonly limitReached = connections++ === 0; + + constructor() { + super(); + queueMicrotask(() => this.dispatchEvent(new Event("open"))); + } + + send(): void { + const event = this.limitReached + ? { type: "error", error: { code: "websocket_connection_limit_reached" } } + : { + type: "response.completed", + response: { + id: "resp_1", + status: "completed", + usage: { input_tokens: 5, output_tokens: 3, total_tokens: 8 }, + }, + }; + queueMicrotask(() => { + this.dispatchEvent(Object.assign(new Event("message"), { data: JSON.stringify(event) })); + }); + } + + close(): void {} + } + + 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 result = await streamOpenAICodexResponses( + model, + { systemPrompt: "", messages: [] }, + { + apiKey: token, + }, + ).result(); + + expect(result.stopReason).toBe("stop"); + expect(connections).toBe(2); + expect(fetchMock).not.toHaveBeenCalled(); + }); + it("falls back to SSE when a websocket is idle before the first event", async () => { vi.useFakeTimers(); const token = mockToken();