From ad48b52de4a9b5ab3820b9539098fe5b1eef4676 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Sat, 14 Mar 2026 12:23:17 +0100 Subject: [PATCH] fix(ai): align codex websocket headers and terminate SSE closes #1961 --- .../src/providers/openai-codex-responses.ts | 130 +++++++++---- packages/ai/test/openai-codex-stream.test.ts | 173 ++++++++++++++++++ 2 files changed, 265 insertions(+), 38 deletions(-) diff --git a/packages/ai/src/providers/openai-codex-responses.ts b/packages/ai/src/providers/openai-codex-responses.ts index 8c9b8aae..a544e3d4 100644 --- a/packages/ai/src/providers/openai-codex-responses.ts +++ b/packages/ai/src/providers/openai-codex-responses.ts @@ -145,7 +145,15 @@ export const streamOpenAICodexResponses: StreamFunction<"openai-codex-responses" if (nextBody !== undefined) { body = nextBody as RequestBody; } - const headers = buildHeaders(model.headers, options?.headers, accountId, apiKey, options?.sessionId); + const websocketRequestId = options?.sessionId || createCodexRequestId(); + const sseHeaders = buildSSEHeaders(model.headers, options?.headers, accountId, apiKey, options?.sessionId); + const websocketHeaders = buildWebSocketHeaders( + model.headers, + options?.headers, + accountId, + apiKey, + websocketRequestId, + ); const bodyJson = JSON.stringify(body); const transport = options?.transport || "sse"; @@ -155,7 +163,7 @@ export const streamOpenAICodexResponses: StreamFunction<"openai-codex-responses" await processWebSocketStream( resolveCodexWebSocketUrl(model.baseUrl), body, - headers, + websocketHeaders, output, stream, model, @@ -194,7 +202,7 @@ export const streamOpenAICodexResponses: StreamFunction<"openai-codex-responses" try { response = await fetch(resolveCodexUrl(model.baseUrl), { method: "POST", - headers, + headers: sseHeaders, body: bodyJson, signal: options?.signal, }); @@ -378,13 +386,13 @@ async function* mapCodexEvents(events: AsyncIterable>): throw new Error(msg || "Codex response failed"); } - if (type === "response.done" || type === "response.completed") { + if (type === "response.done" || type === "response.completed" || type === "response.incomplete") { const response = (event as { response?: { status?: unknown } }).response; const normalizedResponse = response ? { ...response, status: normalizeCodexStatus(response.status) } : response; yield { ...event, type: "response.completed", response: normalizedResponse } as ResponseStreamEvent; - continue; + return; } yield event as unknown as ResponseStreamEvent; @@ -407,30 +415,39 @@ async function* parseSSE(response: Response): AsyncGenerator l.startsWith("data:")) - .map((l) => l.slice(5).trim()); - if (dataLines.length > 0) { - const data = dataLines.join("\n").trim(); - if (data && data !== "[DONE]") { - try { - yield JSON.parse(data); - } catch {} + const dataLines = chunk + .split("\n") + .filter((l) => l.startsWith("data:")) + .map((l) => l.slice(5).trim()); + if (dataLines.length > 0) { + const data = dataLines.join("\n").trim(); + if (data && data !== "[DONE]") { + try { + yield JSON.parse(data); + } catch {} + } } + idx = buffer.indexOf("\n\n"); } - idx = buffer.indexOf("\n\n"); } + } finally { + try { + await reader.cancel(); + } catch {} + try { + reader.releaseLock(); + } catch {} } } @@ -513,7 +530,7 @@ async function connectWebSocket(url: string, headers: Headers, signal?: AbortSig } const wsHeaders = headersToRecord(headers); - wsHeaders["OpenAI-Beta"] = OPENAI_BETA_RESPONSES_WEBSOCKETS; + delete wsHeaders["OpenAI-Beta"]; return new Promise((resolve, reject) => { let settled = false; @@ -533,16 +550,18 @@ async function connectWebSocket(url: string, headers: Headers, signal?: AbortSig resolve(socket); }; const onError: WebSocketListener = (event) => { + const error = extractWebSocketError(event); if (settled) return; settled = true; cleanup(); - reject(extractWebSocketError(event)); + reject(error); }; const onClose: WebSocketListener = (event) => { + const error = extractWebSocketCloseError(event); if (settled) return; settled = true; cleanup(); - reject(extractWebSocketCloseError(event)); + reject(error); }; const onAbort = () => { if (settled) return; @@ -702,7 +721,7 @@ async function* parseWebSocket(socket: WebSocketLike, signal?: AbortSignal): Asy try { const parsed = JSON.parse(text) as Record; const type = typeof parsed.type === "string" ? parsed.type : ""; - if (type === "response.completed" || type === "response.done") { + if (type === "response.completed" || type === "response.done" || type === "response.incomplete") { sawCompletion = true; done = true; } @@ -847,25 +866,42 @@ function extractAccountId(token: string): string { } } -function buildHeaders( +function createCodexRequestId(): string { + if (typeof globalThis.crypto?.randomUUID === "function") { + return globalThis.crypto.randomUUID(); + } + return `codex_${Date.now()}_${Math.random().toString(36).slice(2, 10)}`; +} + +function buildBaseCodexHeaders( + initHeaders: Record | undefined, + additionalHeaders: Record | undefined, + accountId: string, + token: string, +): Headers { + const headers = new Headers(initHeaders); + for (const [key, value] of Object.entries(additionalHeaders || {})) { + headers.set(key, value); + } + headers.set("Authorization", `Bearer ${token}`); + headers.set("chatgpt-account-id", accountId); + headers.set("originator", "pi"); + const userAgent = _os ? `pi (${_os.platform()} ${_os.release()}; ${_os.arch()})` : "pi (browser)"; + headers.set("User-Agent", userAgent); + return headers; +} + +function buildSSEHeaders( initHeaders: Record | undefined, additionalHeaders: Record | undefined, accountId: string, token: string, sessionId?: string, ): Headers { - const headers = new Headers(initHeaders); - headers.set("Authorization", `Bearer ${token}`); - headers.set("chatgpt-account-id", accountId); + const headers = buildBaseCodexHeaders(initHeaders, additionalHeaders, accountId, token); headers.set("OpenAI-Beta", "responses=experimental"); - headers.set("originator", "pi"); - const userAgent = _os ? `pi (${_os.platform()} ${_os.release()}; ${_os.arch()})` : "pi (browser)"; - headers.set("User-Agent", userAgent); headers.set("accept", "text/event-stream"); headers.set("content-type", "application/json"); - for (const [key, value] of Object.entries(additionalHeaders || {})) { - headers.set(key, value); - } if (sessionId) { headers.set("session_id", sessionId); @@ -873,3 +909,21 @@ function buildHeaders( return headers; } + +function buildWebSocketHeaders( + initHeaders: Record | undefined, + additionalHeaders: Record | undefined, + accountId: string, + token: string, + requestId: string, +): Headers { + const headers = buildBaseCodexHeaders(initHeaders, additionalHeaders, accountId, token); + headers.delete("accept"); + headers.delete("content-type"); + headers.delete("OpenAI-Beta"); + headers.delete("openai-beta"); + headers.set("OpenAI-Beta", OPENAI_BETA_RESPONSES_WEBSOCKETS); + headers.set("x-client-request-id", requestId); + headers.set("session_id", requestId); + return headers; +} diff --git a/packages/ai/test/openai-codex-stream.test.ts b/packages/ai/test/openai-codex-stream.test.ts index 2220c398..d3db8c29 100644 --- a/packages/ai/test/openai-codex-stream.test.ts +++ b/packages/ai/test/openai-codex-stream.test.ts @@ -18,6 +18,61 @@ afterEach(() => { vi.restoreAllMocks(); }); +function mockToken(): string { + const payload = Buffer.from( + JSON.stringify({ "https://api.openai.com/auth": { chatgpt_account_id: "acc_test" } }), + "utf8", + ).toString("base64"); + return `aaa.${payload}.bbb`; +} + +function buildSSEPayload({ + status, + includeDone = false, +}: { + status: "completed" | "incomplete"; + includeDone?: boolean; +}): string { + const terminalType = status === "incomplete" ? "response.incomplete" : "response.completed"; + const events = [ + `data: ${JSON.stringify({ + type: "response.output_item.added", + item: { type: "message", id: "msg_1", role: "assistant", status: "in_progress", content: [] }, + })}`, + `data: ${JSON.stringify({ type: "response.content_part.added", part: { type: "output_text", text: "" } })}`, + `data: ${JSON.stringify({ type: "response.output_text.delta", delta: "Hello" })}`, + `data: ${JSON.stringify({ + type: "response.output_item.done", + item: { + type: "message", + id: "msg_1", + role: "assistant", + status: "completed", + content: [{ type: "output_text", text: "Hello" }], + }, + })}`, + `data: ${JSON.stringify({ + type: terminalType, + response: { + status, + incomplete_details: status === "incomplete" ? { reason: "max_output_tokens" } : null, + usage: { + input_tokens: 5, + output_tokens: 3, + total_tokens: 8, + input_tokens_details: { cached_tokens: 0 }, + }, + }, + })}`, + ]; + + if (includeDone) { + events.push("data: [DONE]"); + } + + return `${events.join("\n\n")}\n\n`; +} + describe("openai-codex streaming", () => { it("streams SSE responses into AssistantMessageEventStream", async () => { const tempDir = mkdtempSync(join(tmpdir(), "pi-codex-stream-")); @@ -130,6 +185,124 @@ describe("openai-codex streaming", () => { expect(sawDone).toBe(true); }); + it("completes after response.completed even when the SSE body stays open", async () => { + const tempDir = mkdtempSync(join(tmpdir(), "pi-codex-stream-")); + process.env.PI_CODING_AGENT_DIR = tempDir; + const token = mockToken(); + const encoder = new TextEncoder(); + const sse = buildSSEPayload({ status: "completed", includeDone: true }); + + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode(sse)); + }, + }); + + global.fetch = vi.fn(async (input: string | URL) => { + const url = typeof input === "string" ? input : input.toString(); + if (url === "https://api.github.com/repos/openai/codex/releases/latest") { + return new Response(JSON.stringify({ tag_name: "rust-v0.0.0" }), { status: 200 }); + } + if (url.startsWith("https://raw.githubusercontent.com/openai/codex/")) { + return new Response("PROMPT", { status: 200, headers: { etag: '"etag"' } }); + } + if (url === "https://chatgpt.com/backend-api/codex/responses") { + return new Response(stream, { + status: 200, + headers: { "content-type": "text/event-stream" }, + }); + } + return new Response("not found", { status: 404 }); + }) as typeof fetch; + + 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 context: Context = { + systemPrompt: "You are a helpful assistant.", + messages: [{ role: "user", content: "Say hello", timestamp: Date.now() }], + }; + + const result = await Promise.race([ + streamOpenAICodexResponses(model, context, { apiKey: token }).result(), + new Promise((_, reject) => { + setTimeout(() => reject(new Error("Timed out waiting for completed SSE stream")), 1000); + }), + ]); + + expect(result.content.find((c) => c.type === "text")?.text).toBe("Hello"); + expect(result.stopReason).toBe("stop"); + }); + + it("maps response.incomplete to stopReason length even when the SSE body stays open", async () => { + const tempDir = mkdtempSync(join(tmpdir(), "pi-codex-stream-")); + process.env.PI_CODING_AGENT_DIR = tempDir; + const token = mockToken(); + const encoder = new TextEncoder(); + const sse = buildSSEPayload({ status: "incomplete" }); + + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode(sse)); + }, + }); + + global.fetch = vi.fn(async (input: string | URL) => { + const url = typeof input === "string" ? input : input.toString(); + if (url === "https://api.github.com/repos/openai/codex/releases/latest") { + return new Response(JSON.stringify({ tag_name: "rust-v0.0.0" }), { status: 200 }); + } + if (url.startsWith("https://raw.githubusercontent.com/openai/codex/")) { + return new Response("PROMPT", { status: 200, headers: { etag: '"etag"' } }); + } + if (url === "https://chatgpt.com/backend-api/codex/responses") { + return new Response(stream, { + status: 200, + headers: { "content-type": "text/event-stream" }, + }); + } + return new Response("not found", { status: 404 }); + }) as typeof fetch; + + 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 context: Context = { + systemPrompt: "You are a helpful assistant.", + messages: [{ role: "user", content: "Say hello", timestamp: Date.now() }], + }; + + const result = await Promise.race([ + streamOpenAICodexResponses(model, context, { apiKey: token }).result(), + new Promise((_, reject) => { + setTimeout(() => reject(new Error("Timed out waiting for incomplete SSE stream")), 1000); + }), + ]); + + expect(result.content.find((c) => c.type === "text")?.text).toBe("Hello"); + expect(result.stopReason).toBe("length"); + }); + it("sets conversation_id/session_id headers and prompt_cache_key when sessionId is provided", async () => { const tempDir = mkdtempSync(join(tmpdir(), "pi-codex-stream-")); process.env.PI_CODING_AGENT_DIR = tempDir;