ton/pi_harness
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(ai): use zstd compression for codex sse transport

Browse Source
This commit is contained in:
Vegard Stikbakke
2026-06-30 14:29:51 +02:00
parent a3cc169d97
commit 0ac3cfe09b
3 changed files with 138 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ai/CHANGELOG.md
+4
View File
@@ -6,6 +6,10 @@
- Fixed OpenAI Codex user-agent construction to synchronously load Node OS metadata, avoiding a startup race that could report `pi (browser)` in Node/Bun.
### Added
- Added zstd request-body compression for the OpenAI Codex Responses SSE transport. Requests are sent with `Content-Encoding: zstd` when Node/Bun zstd support is available; the WebSocket transport is unchanged.
## [0.80.3] - 2026-06-30
### Added
packages/ai/src/api/openai-codex-responses.ts
+47 -1
View File
@@ -1,4 +1,5 @@
import type * as NodeOs from "node:os";
import type * as NodeZlib from "node:zlib";
import type {
Tool as OpenAITool,
ResponseCreateParamsStreaming,
@@ -58,6 +59,9 @@ const DEFAULT_MAX_RETRIES = 0;
const BASE_DELAY_MS = 1000;
const DEFAULT_MAX_RETRY_DELAY_MS = 60_000;
const DEFAULT_WEBSOCKET_CONNECT_TIMEOUT_MS = 15_000;
// The Codex backend accepts zstd-compressed request bodies on the SSE responses
// endpoint (the same endpoint the official Codex client compresses against).
const REQUEST_COMPRESSION_ZSTD_LEVEL = 3;
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";
@@ -177,6 +181,39 @@ function normalizeTimeoutMs(value: number | undefined): number | undefined {
return Math.floor(value);
}
// ============================================================================
// Request Compression
// ============================================================================
type ProcessWithBuiltinModule = typeof process & {
getBuiltinModule?: (id: "node:zlib") => typeof NodeZlib;
};
function loadNodeZlib(): typeof NodeZlib | null {
if (typeof process === "undefined" || !(process.versions?.node || process.versions?.bun)) {
return null;
}
return (process as ProcessWithBuiltinModule).getBuiltinModule?.("node:zlib") ?? null;
}
// Returns the zstd-compressed body bytes, or null when compression is
// unavailable (browser/Vite builds). Callers fall back to sending the
// uncompressed JSON when this returns null.
function compressRequestBodyZstd(bodyJson: string): Uint8Array | null {
const zlib = loadNodeZlib();
if (!zlib || typeof zlib.zstdCompressSync !== "function") {
return null;
}
try {
const compressed = zlib.zstdCompressSync(bodyJson, {
params: { [zlib.constants.ZSTD_c_compressionLevel]: REQUEST_COMPRESSION_ZSTD_LEVEL },
});
return new Uint8Array(compressed.buffer, compressed.byteOffset, compressed.byteLength);
} catch {
return null;
}
}
// ============================================================================
// Main Stream Function
// ============================================================================
@@ -298,6 +335,15 @@ export const stream: StreamFunction<"openai-codex-responses", OpenAICodexRespons
}
}
// Compress the request body once for the SSE path. The Codex backend
// decodes Content-Encoding: zstd; the WebSocket transport above sends the
// uncompressed JSON frame, matching the official Codex client.
const compressedBody = compressRequestBodyZstd(bodyJson);
if (compressedBody) {
sseHeaders.set("content-encoding", "zstd");
}
const sseBody: Uint8Array | string = compressedBody ?? bodyJson;
// Fetch with retry logic for rate limits and transient errors
let response: Response | undefined;
let lastError: Error | undefined;
@@ -316,7 +362,7 @@ export const stream: StreamFunction<"openai-codex-responses", OpenAICodexRespons
response = await fetch(resolveCodexUrl(model.baseUrl), {
method: "POST",
headers: sseHeaders,
body: bodyJson,
body: sseBody,
signal: combinedSignal.signal,
});
} catch (error) {
packages/ai/test/openai-codex-stream.test.ts
+87 -3
View File
@@ -1,6 +1,7 @@
import { mkdtempSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { zstdDecompressSync } from "node:zlib";
import { afterEach, describe, expect, it, vi } from "vitest";
import {
getOpenAICodexWebSocketDebugStats,
@@ -32,6 +33,16 @@ function mockToken(): string {
return `aaa.${payload}.bbb`;
}
function decodeCodexRequestBody(body: RequestInit["body"] | undefined): Record<string, unknown> | null {
if (typeof body === "string") {
return JSON.parse(body) as Record<string, unknown>;
}
if (body instanceof Uint8Array) {
return JSON.parse(Buffer.from(zstdDecompressSync(body)).toString("utf8")) as Record<string, unknown>;
}
return null;
}
function buildSSEPayload({
status,
includeDone = false,
@@ -541,7 +552,7 @@ describe("openai-codex streaming", () => {
expect(headers?.get("x-client-request-id")).toBe(sessionId);
// Verify sessionId is set in request body as prompt_cache_key
const body = typeof init?.body === "string" ? (JSON.parse(init.body) as Record<string, unknown>) : null;
const body = decodeCodexRequestBody(init?.body);
expect(body?.prompt_cache_key).toBe(sessionId);
return new Response(stream, {
@@ -649,7 +660,7 @@ describe("openai-codex streaming", () => {
return new Response("PROMPT", { status: 200, headers: { etag: '"etag"' } });
}
if (url === "https://chatgpt.com/backend-api/codex/responses") {
const body = typeof init?.body === "string" ? (JSON.parse(init.body) as Record<string, unknown>) : null;
const body = decodeCodexRequestBody(init?.body);
requestedReasoning = body?.reasoning;
return new Response(stream, {
status: 200,
@@ -746,7 +757,7 @@ describe("openai-codex streaming", () => {
return new Response("PROMPT", { status: 200, headers: { etag: '"etag"' } });
}
if (url === "https://chatgpt.com/backend-api/codex/responses") {
const body = typeof init?.body === "string" ? (JSON.parse(init.body) as Record<string, unknown>) : null;
const body = decodeCodexRequestBody(init?.body);
requestedReasoning = body?.reasoning;
return new Response(stream, {
@@ -1654,6 +1665,79 @@ describe("openai-codex streaming", () => {
expect(codexRequests).toBe(2);
});
it("zstd-compresses SSE request bodies", async () => {
const token = mockToken();
const encoder = new TextEncoder();
const sse = buildSSEPayload({ status: "completed" });
let capturedEncoding: string | null = null;
let capturedBody: Uint8Array | string | undefined;
const fetchMock = vi.fn(async (input: string | URL, init?: RequestInit) => {
const url = typeof input === "string" ? input : input.toString();
if (url !== "https://chatgpt.com/backend-api/codex/responses") {
throw new Error(`Unexpected URL: ${url}`);
}
const headers = init?.headers instanceof Headers ? init.headers : undefined;
capturedEncoding = headers?.get("content-encoding") ?? null;
capturedBody = init?.body as Uint8Array | string | undefined;
return new Response(
new ReadableStream<Uint8Array>({
start(controller) {
controller.enqueue(encoder.encode(sse));
controller.close();
},
}),
{ status: 200, headers: { "content-type": "text/event-stream" } },
);
});
vi.stubGlobal("fetch", fetchMock);
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 largeText = "compress me ".repeat(400);
await streamOpenAICodexResponses(
model,
{
systemPrompt: "You are a helpful assistant.",
messages: [{ role: "user", content: largeText, timestamp: 1 }],
},
{ apiKey: token, transport: "sse" },
).result();
expect(capturedEncoding).toBe("zstd");
expect(capturedBody).toBeInstanceOf(Uint8Array);
const decoded = JSON.parse(Buffer.from(zstdDecompressSync(capturedBody as Uint8Array)).toString("utf8")) as {
input: Array<{ content: Array<{ text: string }> }>;
};
expect(decoded.input[0].content[0].text).toBe(largeText);
capturedEncoding = null;
capturedBody = undefined;
await streamOpenAICodexResponses(
model,
{
systemPrompt: "You are a helpful assistant.",
messages: [{ role: "user", content: "hi", timestamp: 1 }],
},
{ apiKey: token, transport: "sse" },
).result();
expect(capturedEncoding).toBe("zstd");
expect(capturedBody).toBeInstanceOf(Uint8Array);
});
it("uses exponential backoff across repeated SSE retries without retry headers", async () => {
vi.useFakeTimers();
vi.setSystemTime(new Date("2026-05-13T00:00:00Z"));