shumengya/sproutclaw
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(ai): align codex websocket headers and terminate SSE closes #1961

Browse Source
This commit is contained in:
Mario Zechner
2026-03-14 12:23:17 +01:00
parent 1feccfedcb
commit ad48b52de4
2 changed files with 265 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

130
packages/ai/src/providers/openai-codex-responses.ts
View File

@@ -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<Record<string, unknown>>):
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<Record<string, unkn
const decoder = new TextDecoder();
let buffer = "";
while (true) {
const { done, value } = await reader.read();
if (done) break;
buffer += decoder.decode(value, { stream: true });
try {
while (true) {
const { done, value } = await reader.read();
if (done) break;
buffer += decoder.decode(value, { stream: true });
let idx = buffer.indexOf("\n\n");
while (idx !== -1) {
const chunk = buffer.slice(0, idx);
buffer = buffer.slice(idx + 2);
let idx = buffer.indexOf("\n\n");
while (idx !== -1) {
const chunk = buffer.slice(0, idx);
buffer = buffer.slice(idx + 2);
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 {}
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<WebSocketLike>((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<string, unknown>;
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<string, string> | undefined,
additionalHeaders: Record<string, string> | 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<string, string> | undefined,
additionalHeaders: Record<string, string> | 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<string, string> | undefined,
additionalHeaders: Record<string, string> | 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;
}

173
packages/ai/test/openai-codex-stream.test.ts
View File

@@ -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<Uint8Array>({
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<never>((_, 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<Uint8Array>({
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<never>((_, 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;