diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index e6b3deb8..3bb43438 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Added + +- Added `websocket-cached` transport support for OpenAI Codex Responses used with ChatGPT subscription auth. This keeps the same WebSocket open for a session and, after the first request, sends only new conversation items instead of resending the full chat history when possible. + ## [0.71.0] - 2026-04-30 ### Breaking Changes diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index b1e7f66e..8798b024 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -12,7 +12,15 @@ export type { GoogleOptions } from "./providers/google.js"; export type { GoogleThinkingLevel } from "./providers/google-shared.js"; export type { GoogleVertexOptions } from "./providers/google-vertex.js"; export type { MistralOptions } from "./providers/mistral.js"; -export type { OpenAICodexResponsesOptions } from "./providers/openai-codex-responses.js"; +export type { + OpenAICodexResponsesOptions, + OpenAICodexWebSocketDebugStats, +} from "./providers/openai-codex-responses.js"; +export { + closeOpenAICodexWebSocketSessions, + getOpenAICodexWebSocketDebugStats, + resetOpenAICodexWebSocketDebugStats, +} from "./providers/openai-codex-responses.js"; export type { OpenAICompletionsOptions } from "./providers/openai-completions.js"; export type { OpenAIResponsesOptions } from "./providers/openai-responses.js"; export * from "./providers/register-builtins.js"; diff --git a/packages/ai/src/providers/openai-codex-responses.ts b/packages/ai/src/providers/openai-codex-responses.ts index 541cb98a..6469e741 100644 --- a/packages/ai/src/providers/openai-codex-responses.ts +++ b/packages/ai/src/providers/openai-codex-responses.ts @@ -74,6 +74,7 @@ interface RequestBody { store?: boolean; stream?: boolean; instructions?: string; + previous_response_id?: string; input?: ResponseInput; tools?: OpenAITool[]; tool_choice?: "auto"; @@ -193,7 +194,7 @@ export const streamOpenAICodexResponses: StreamFunction<"openai-codex-responses" stream.end(); return; } catch (error) { - if (transport === "websocket" || websocketStarted) { + if (transport === "websocket" || transport === "websocket-cached" || websocketStarted) { throw error; } } @@ -536,13 +537,82 @@ interface WebSocketLike { removeEventListener(type: WebSocketEventType, listener: WebSocketListener): void; } +interface CachedWebSocketContinuationState { + lastRequestBody: RequestBody; + lastResponseId: string; + lastResponseItems: ResponseInput; +} + interface CachedWebSocketConnection { socket: WebSocketLike; busy: boolean; idleTimer?: ReturnType; + continuation?: CachedWebSocketContinuationState; +} + +export interface OpenAICodexWebSocketDebugStats { + requests: number; + connectionsCreated: number; + connectionsReused: number; + cachedContextRequests: number; + storeTrueRequests: number; + fullContextRequests: number; + deltaRequests: number; + lastInputItems: number; + lastDeltaInputItems?: number; + lastPreviousResponseId?: string; } const websocketSessionCache = new Map(); +const websocketDebugStats = new Map(); + +function getOrCreateWebSocketDebugStats(sessionId: string): OpenAICodexWebSocketDebugStats { + let stats = websocketDebugStats.get(sessionId); + if (!stats) { + stats = { + requests: 0, + connectionsCreated: 0, + connectionsReused: 0, + cachedContextRequests: 0, + storeTrueRequests: 0, + fullContextRequests: 0, + deltaRequests: 0, + lastInputItems: 0, + }; + websocketDebugStats.set(sessionId, stats); + } + return stats; +} + +export function getOpenAICodexWebSocketDebugStats(sessionId: string): OpenAICodexWebSocketDebugStats | undefined { + const stats = websocketDebugStats.get(sessionId); + return stats ? { ...stats } : undefined; +} + +export function resetOpenAICodexWebSocketDebugStats(sessionId?: string): void { + if (sessionId) { + websocketDebugStats.delete(sessionId); + return; + } + websocketDebugStats.clear(); +} + +export function closeOpenAICodexWebSocketSessions(sessionId?: string): void { + const closeEntry = (entry: CachedWebSocketConnection) => { + if (entry.idleTimer) clearTimeout(entry.idleTimer); + closeWebSocketSilently(entry.socket, 1000, "debug_close"); + }; + if (sessionId) { + const entry = websocketSessionCache.get(sessionId); + if (entry) closeEntry(entry); + websocketSessionCache.delete(sessionId); + return; + } + for (const entry of websocketSessionCache.values()) { + closeEntry(entry); + } + websocketSessionCache.clear(); +} type WebSocketConstructor = new ( url: string, @@ -650,11 +720,17 @@ async function acquireWebSocket( headers: Headers, sessionId: string | undefined, signal?: AbortSignal, -): Promise<{ socket: WebSocketLike; release: (options?: { keep?: boolean }) => void }> { +): Promise<{ + socket: WebSocketLike; + entry?: CachedWebSocketConnection; + reused: boolean; + release: (options?: { keep?: boolean }) => void; +}> { if (!sessionId) { const socket = await connectWebSocket(url, headers, signal); return { socket, + reused: false, release: ({ keep } = {}) => { if (keep === false) { closeWebSocketSilently(socket); @@ -675,6 +751,8 @@ async function acquireWebSocket( cached.busy = true; return { socket: cached.socket, + entry: cached, + reused: true, release: ({ keep } = {}) => { if (!keep || !isWebSocketReusable(cached.socket)) { closeWebSocketSilently(cached.socket); @@ -690,6 +768,7 @@ async function acquireWebSocket( const socket = await connectWebSocket(url, headers, signal); return { socket, + reused: false, release: () => { closeWebSocketSilently(socket); }, @@ -706,6 +785,8 @@ async function acquireWebSocket( websocketSessionCache.set(sessionId, entry); return { socket, + entry, + reused: false, release: ({ keep } = {}) => { if (!keep || !isWebSocketReusable(entry.socket)) { closeWebSocketSilently(entry.socket); @@ -850,6 +931,60 @@ async function* parseWebSocket(socket: WebSocketLike, signal?: AbortSignal): Asy } } +function requestBodyWithoutInput(body: RequestBody): RequestBody { + const { input: _input, previous_response_id: _previousResponseId, ...rest } = body; + return rest; +} + +function responseInputsEqual(a: ResponseInput | undefined, b: ResponseInput | undefined): boolean { + return JSON.stringify(a ?? []) === JSON.stringify(b ?? []); +} + +function requestBodiesMatchExceptInput(a: RequestBody, b: RequestBody): boolean { + return JSON.stringify(requestBodyWithoutInput(a)) === JSON.stringify(requestBodyWithoutInput(b)); +} + +function getCachedWebSocketInputDelta( + body: RequestBody, + continuation: CachedWebSocketContinuationState, +): ResponseInput | undefined { + if (!requestBodiesMatchExceptInput(body, continuation.lastRequestBody)) { + return undefined; + } + + const currentInput = body.input ?? []; + const baseline = [...(continuation.lastRequestBody.input ?? []), ...continuation.lastResponseItems]; + if (currentInput.length < baseline.length) { + return undefined; + } + + const prefix = currentInput.slice(0, baseline.length); + if (!responseInputsEqual(prefix, baseline)) { + return undefined; + } + + return currentInput.slice(baseline.length); +} + +function buildCachedWebSocketRequestBody(entry: CachedWebSocketConnection, body: RequestBody): RequestBody { + const continuation = entry.continuation; + if (!continuation) { + return body; + } + + const delta = getCachedWebSocketInputDelta(body, continuation); + if (!delta || !continuation.lastResponseId) { + entry.continuation = undefined; + return body; + } + + return { + ...body, + previous_response_id: continuation.lastResponseId, + input: delta, + }; +} + async function processWebSocketStream( url: string, body: RequestBody, @@ -860,10 +995,33 @@ async function processWebSocketStream( onStart: () => void, options?: OpenAICodexResponsesOptions, ): Promise { - const { socket, release } = await acquireWebSocket(url, headers, options?.sessionId, options?.signal); + const { socket, entry, reused, release } = await acquireWebSocket(url, headers, options?.sessionId, options?.signal); let keepConnection = true; + const useCachedContext = options?.transport === "websocket-cached"; + // ChatGPT Codex Responses rejects `store: true` ("Store must be set to false"). + // WebSocket continuation still works via connection-scoped previous_response_id state. + const fullBody = body; + const requestBody = useCachedContext && entry ? buildCachedWebSocketRequestBody(entry, fullBody) : fullBody; + const stats = options?.sessionId ? getOrCreateWebSocketDebugStats(options.sessionId) : undefined; + if (stats) { + stats.requests++; + if (reused) stats.connectionsReused++; + else stats.connectionsCreated++; + if (useCachedContext) stats.cachedContextRequests++; + if (requestBody.store === true) stats.storeTrueRequests++; + stats.lastInputItems = requestBody.input?.length ?? 0; + if (requestBody.previous_response_id) { + stats.deltaRequests++; + stats.lastDeltaInputItems = requestBody.input?.length ?? 0; + stats.lastPreviousResponseId = requestBody.previous_response_id; + } else { + stats.fullContextRequests++; + stats.lastDeltaInputItems = undefined; + stats.lastPreviousResponseId = undefined; + } + } try { - socket.send(JSON.stringify({ type: "response.create", ...body })); + socket.send(JSON.stringify({ type: "response.create", ...requestBody })); onStart(); stream.push({ type: "start", partial: output }); await processResponsesStream(mapCodexEvents(parseWebSocket(socket, options?.signal)), output, stream, model, { @@ -873,8 +1031,20 @@ async function processWebSocketStream( }); if (options?.signal?.aborted) { keepConnection = false; + } else if (useCachedContext && entry && output.responseId) { + const responseItems = convertResponsesMessages(model, { messages: [output] }, CODEX_TOOL_CALL_PROVIDERS, { + includeSystemPrompt: false, + }).filter((item) => item.type !== "function_call_output"); + entry.continuation = { + lastRequestBody: fullBody, + lastResponseId: output.responseId, + lastResponseItems: responseItems, + }; } } catch (error) { + if (entry) { + entry.continuation = undefined; + } keepConnection = false; throw error; } finally { diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index bdca2a84..ae89baf7 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -58,7 +58,7 @@ export interface ThinkingBudgets { // Base options all providers share export type CacheRetention = "none" | "short" | "long"; -export type Transport = "sse" | "websocket" | "auto"; +export type Transport = "sse" | "websocket" | "websocket-cached" | "auto"; export interface ProviderResponse { status: number; diff --git a/packages/ai/test/codex-websocket-cached-probe.ts b/packages/ai/test/codex-websocket-cached-probe.ts new file mode 100644 index 00000000..ccc875f8 --- /dev/null +++ b/packages/ai/test/codex-websocket-cached-probe.ts @@ -0,0 +1,291 @@ +#!/usr/bin/env tsx +/** + * Live probe for OpenAI Codex Responses websocket-cached mode. + * + * Runs a simple tool loop directly against the pi-ai provider source so it does not + * depend on built dist packages or coding-agent SDK wiring. + */ + +import { tmpdir } from "node:os"; +import { join, resolve } from "node:path"; +import { Type } from "typebox"; +import { AuthStorage } from "../../coding-agent/src/core/auth-storage.js"; +import { getModel } from "../src/models.js"; +import { + closeOpenAICodexWebSocketSessions, + getOpenAICodexWebSocketDebugStats, + resetOpenAICodexWebSocketDebugStats, + streamOpenAICodexResponses, +} from "../src/providers/openai-codex-responses.js"; +import type { AssistantMessage, Context, Message, Model, Tool, ToolResultMessage, Transport } from "../src/types.js"; + +type ThinkingLevel = "minimal" | "low" | "medium" | "high" | "xhigh"; + +interface Args { + turns: number; + transport: Transport; + maxTokens: number; + reasoning: ThinkingLevel; + sessionId: string; +} + +const DEFAULT_TURNS = 20; +const DEFAULT_MAX_TOKENS = 64; + +function parseArgs(argv: string[]): Args { + let turns = DEFAULT_TURNS; + let transport: Transport = "websocket-cached"; + let maxTokens = DEFAULT_MAX_TOKENS; + let reasoning: ThinkingLevel = "low"; + let sessionId = `pi-ai-codex-ws-cached-probe-${Date.now()}`; + + for (let i = 0; i < argv.length; i++) { + const arg = argv[i]; + switch (arg) { + case "--turns": + turns = Number.parseInt(required(argv[++i], arg), 10); + break; + case "--transport": { + const value = required(argv[++i], arg); + if (value !== "sse" && value !== "websocket" && value !== "websocket-cached" && value !== "auto") { + throw new Error(`Invalid --transport: ${value}`); + } + transport = value; + break; + } + case "--max-tokens": + maxTokens = Number.parseInt(required(argv[++i], arg), 10); + break; + case "--reasoning": { + const value = required(argv[++i], arg); + if (value !== "minimal" && value !== "low" && value !== "medium" && value !== "high" && value !== "xhigh") { + throw new Error(`Invalid --reasoning: ${value}`); + } + reasoning = value; + break; + } + case "--session-id": + sessionId = required(argv[++i], arg); + break; + case "--help": + printHelp(); + process.exit(0); + break; + default: + throw new Error(`Unknown argument: ${arg}`); + } + } + + return { turns, transport, maxTokens, reasoning, sessionId }; +} + +function required(value: string | undefined, flag: string): string { + if (!value) throw new Error(`Missing value for ${flag}`); + return value; +} + +function printHelp(): void { + console.log(`Usage: npx tsx test/codex-websocket-cached-probe.ts [options] + +Options: + --turns Number of user turns. Default: ${DEFAULT_TURNS} + --transport sse | websocket | websocket-cached | auto. Default: websocket-cached + --reasoning minimal | low | medium | high | xhigh. Default: low + --max-tokens Max output tokens per model request. Default: ${DEFAULT_MAX_TOKENS} + --session-id Session id for websocket/cache state +`); +} + +function buildPrompt(turn: number): string { + const marker = `TURN-${String(turn).padStart(2, "0")}-MARKER-${(turn * 17 + 13) % 97}`; + const lines = [ + "This is an automated OpenAI Codex Responses websocket cache probe.", + `Task for turn ${turn}: call deterministic_probe exactly once before your final answer.`, + `Use tool arguments: turn=${turn}, marker=${marker}`, + `After the tool result arrives, reply exactly: TURN ${turn} OK ${marker}`, + "The following repeated block is intentional benchmark padding.", + ]; + for (let i = 1; i <= 180; i++) { + lines.push( + `Turn ${turn} synthetic record ${String(i).padStart(3, "0")}: alpha beta gamma delta epsilon zeta eta theta iota kappa lambda mu nu xi omicron pi rho sigma tau upsilon phi chi psi omega.`, + ); + } + return lines.join("\n"); +} + +function deterministicProbeTool(): Tool { + return { + name: "deterministic_probe", + description: "Mandatory benchmark tool. Call exactly once with the turn and marker from the user prompt.", + parameters: Type.Object({ + turn: Type.Number(), + marker: Type.String(), + }), + }; +} + +function executeTool(call: Extract): ToolResultMessage { + return { + role: "toolResult", + toolCallId: call.id, + toolName: call.name, + content: [{ type: "text", text: `deterministic_probe_result ${JSON.stringify(call.arguments)} fixed=OK` }], + details: { fixed: "OK" }, + isError: false, + timestamp: Date.now(), + }; +} + +function textOf(message: AssistantMessage): string { + return message.content + .filter((block): block is Extract => block.type === "text") + .map((block) => block.text) + .join("\n") + .trim(); +} + +function average(values: number[]): number { + return values.reduce((sum, value) => sum + value, 0) / Math.max(1, values.length); +} + +function percentile(values: number[], p: number): number { + if (values.length === 0) return 0; + const sorted = [...values].sort((a, b) => a - b); + return sorted[Math.min(sorted.length - 1, Math.max(0, Math.ceil((p / 100) * sorted.length) - 1))]; +} + +async function main(): Promise { + const args = parseArgs(process.argv.slice(2)); + const model = getModel("openai-codex", "gpt-5.5") as Model<"openai-codex-responses"> | undefined; + if (!model) throw new Error("Model openai-codex/gpt-5.5 not found"); + const modelWithMaxTokens = { ...model, maxTokens: args.maxTokens }; + const authStorage = AuthStorage.create(); + const apiKey = (await authStorage.getApiKey("openai-codex")) ?? (await authStorage.getApiKey("openai")); + if (!apiKey) { + throw new Error("No OpenAI Codex API key found in coding-agent auth storage."); + } + const context: Context = { + systemPrompt: + "You are participating in a benchmark. For each benchmark turn, call deterministic_probe exactly once before the final answer. Keep final answers minimal.", + messages: [], + tools: [deterministicProbeTool()], + }; + const elapsed: number[] = []; + resetOpenAICodexWebSocketDebugStats(args.sessionId); + + console.log(`provider openai-codex, model gpt-5.5`); + console.log(`sessionId ${args.sessionId}`); + console.log( + `turns ${args.turns}, transport ${args.transport}, reasoning ${args.reasoning}, maxTokens ${args.maxTokens}`, + ); + console.log(`scratch ${resolve(join(tmpdir(), args.sessionId))}`); + console.log(""); + + for (let turn = 1; turn <= args.turns; turn++) { + context.messages.push({ role: "user", content: buildPrompt(turn), timestamp: Date.now() }); + const beforeStats = getOpenAICodexWebSocketDebugStats(args.sessionId); + const started = Date.now(); + let requests = 0; + let assistantCount = 0; + let toolResults = 0; + let finalText = ""; + let turnInput = 0; + let turnOutput = 0; + let turnCacheRead = 0; + let turnCacheWrite = 0; + + while (true) { + requests++; + const message = await streamOpenAICodexResponses(modelWithMaxTokens, context, { + apiKey, + sessionId: args.sessionId, + transport: args.transport, + reasoningEffort: args.reasoning, + maxTokens: args.maxTokens, + }).result(); + assistantCount++; + context.messages.push(message); + turnInput += message.usage.input; + turnOutput += message.usage.output; + turnCacheRead += message.usage.cacheRead; + turnCacheWrite += message.usage.cacheWrite; + const toolCalls = message.content.filter( + (block): block is Extract => + block.type === "toolCall", + ); + console.log( + [ + `turn ${String(turn).padStart(2, "0")}.${requests}`, + `stop ${message.stopReason}`, + `in ${message.usage.input}`, + `out ${message.usage.output}`, + `cache ${message.usage.cacheRead}/${message.usage.cacheWrite}`, + `tools ${toolCalls.length}`, + ].join(" | "), + ); + if (message.stopReason === "error" || message.stopReason === "aborted") { + throw new Error(message.errorMessage ?? `request failed on turn ${turn}.${requests}`); + } + if (toolCalls.length === 0) { + finalText = textOf(message); + break; + } + for (const call of toolCalls) { + context.messages.push(executeTool(call) as Message); + toolResults++; + } + if (requests > 4) throw new Error(`Too many requests for turn ${turn}`); + } + + const elapsedMs = Date.now() - started; + elapsed.push(elapsedMs); + const afterStats = getOpenAICodexWebSocketDebugStats(args.sessionId); + const statLine = afterStats + ? `ws requests ${afterStats.requests - (beforeStats?.requests ?? 0)} | new/reused ${afterStats.connectionsCreated - (beforeStats?.connectionsCreated ?? 0)}/${afterStats.connectionsReused - (beforeStats?.connectionsReused ?? 0)} | cached ${afterStats.cachedContextRequests - (beforeStats?.cachedContextRequests ?? 0)} | store ${afterStats.storeTrueRequests - (beforeStats?.storeTrueRequests ?? 0)} | full/delta ${afterStats.fullContextRequests - (beforeStats?.fullContextRequests ?? 0)}/${afterStats.deltaRequests - (beforeStats?.deltaRequests ?? 0)}` + : "ws none"; + console.log( + [ + `turn ${String(turn).padStart(2, "0")} agg`, + `elapsed ${(elapsedMs / 1000).toFixed(1)}s`, + `assistant ${assistantCount}`, + `toolResults ${toolResults}`, + `in ${turnInput}`, + `out ${turnOutput}`, + `cache ${turnCacheRead}/${turnCacheWrite}`, + statLine, + `final ${JSON.stringify(finalText).slice(0, 80)}`, + ].join(" | "), + ); + } + + const stats = getOpenAICodexWebSocketDebugStats(args.sessionId); + console.log(""); + console.log( + [ + "timing", + `turns ${elapsed.length}`, + `total ${(elapsed.reduce((sum, value) => sum + value, 0) / 1000).toFixed(1)}s`, + `avg ${(average(elapsed) / 1000).toFixed(2)}s`, + `p50 ${(percentile(elapsed, 50) / 1000).toFixed(2)}s`, + `p95 ${(percentile(elapsed, 95) / 1000).toFixed(2)}s`, + `max ${(Math.max(...elapsed) / 1000).toFixed(2)}s`, + ].join(" | "), + ); + console.log( + [ + "transport summary", + `requested ${args.transport}`, + `observed ${stats && stats.requests > 0 ? "websocket" : "sse/no-websocket"}`, + `storeTrue ${stats ? `${stats.storeTrueRequests}/${stats.requests}` : "0/0"}`, + `full/delta ${stats ? `${stats.fullContextRequests}/${stats.deltaRequests}` : "0/0"}`, + `connections created/reused ${stats ? `${stats.connectionsCreated}/${stats.connectionsReused}` : "0/0"}`, + `lastPreviousResponseId ${stats?.lastPreviousResponseId ?? "n/a"}`, + ].join(" | "), + ); + closeOpenAICodexWebSocketSessions(args.sessionId); +} + +main().catch((error: unknown) => { + console.error(error instanceof Error ? error.message : String(error)); + process.exitCode = 1; +}); diff --git a/packages/ai/test/openai-codex-stream.test.ts b/packages/ai/test/openai-codex-stream.test.ts index 90be500b..ab51333a 100644 --- a/packages/ai/test/openai-codex-stream.test.ts +++ b/packages/ai/test/openai-codex-stream.test.ts @@ -3,21 +3,26 @@ import { tmpdir } from "node:os"; import { join } from "node:path"; import { afterEach, describe, expect, it, vi } from "vitest"; import { + getOpenAICodexWebSocketDebugStats, + resetOpenAICodexWebSocketDebugStats, streamOpenAICodexResponses, streamSimpleOpenAICodexResponses, } from "../src/providers/openai-codex-responses.js"; import type { Context, Model } from "../src/types.js"; const originalFetch = global.fetch; +const originalWebSocket = globalThis.WebSocket; const originalAgentDir = process.env.PI_CODING_AGENT_DIR; afterEach(() => { global.fetch = originalFetch; + globalThis.WebSocket = originalWebSocket; if (originalAgentDir === undefined) { delete process.env.PI_CODING_AGENT_DIR; } else { process.env.PI_CODING_AGENT_DIR = originalAgentDir; } + resetOpenAICodexWebSocketDebugStats(); vi.restoreAllMocks(); }); @@ -746,4 +751,150 @@ describe("openai-codex streaming", () => { const streamResult = streamOpenAICodexResponses(model, context, { apiKey: token }); await streamResult.result(); }); + it("sends only response input deltas in websocket-cached mode", async () => { + const token = mockToken(); + const sentBodies: unknown[] = []; + const responses = [ + { responseId: "resp_1", messageId: "msg_1", text: "Hello" }, + { responseId: "resp_2", messageId: "msg_2", text: "Done" }, + ]; + + class MockWebSocket { + static OPEN = 1; + readyState = MockWebSocket.OPEN; + 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(data: string): void { + sentBodies.push(JSON.parse(data)); + const response = responses.shift(); + if (!response) throw new Error("unexpected websocket request"); + const events = [ + { type: "response.created", response: { id: response.responseId } }, + { + type: "response.output_item.added", + item: { + type: "message", + id: response.messageId, + role: "assistant", + status: "in_progress", + content: [], + }, + }, + { type: "response.content_part.added", part: { type: "output_text", text: "" } }, + { type: "response.output_text.delta", delta: response.text }, + { + type: "response.output_item.done", + item: { + type: "message", + id: response.messageId, + role: "assistant", + status: "completed", + content: [{ type: "output_text", text: response.text }], + }, + }, + { + type: "response.completed", + response: { + id: response.responseId, + status: "completed", + usage: { + input_tokens: 5, + output_tokens: 3, + total_tokens: 8, + input_tokens_details: { cached_tokens: 0 }, + }, + }, + }, + ]; + queueMicrotask(() => { + for (const event of events) { + this.dispatch("message", { data: JSON.stringify(event) }); + } + }); + } + + close(): void { + this.readyState = 3; + } + + private dispatch(type: string, event: unknown): void { + for (const listener of this.listeners.get(type) ?? []) { + listener(event); + } + } + } + + globalThis.WebSocket = MockWebSocket as unknown as typeof WebSocket; + + 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 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: "session-1", + transport: "websocket-cached", + }).result(); + + 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: "session-1", + transport: "websocket-cached", + }).result(); + + expect(sentBodies).toHaveLength(2); + const firstBody = sentBodies[0] as { input: unknown[]; previous_response_id?: string; store?: boolean }; + const secondBody = sentBodies[1] as { input: unknown[]; previous_response_id?: string; store?: boolean }; + expect(firstBody.store).toBe(false); + expect(firstBody.previous_response_id).toBeUndefined(); + expect(firstBody.input).toEqual([{ role: "user", content: [{ type: "input_text", text: "Say hello" }] }]); + expect(secondBody.store).toBe(false); + expect(secondBody.previous_response_id).toBe("resp_1"); + expect(secondBody.input).toEqual([{ role: "user", content: [{ type: "input_text", text: "Now finish" }] }]); + expect(getOpenAICodexWebSocketDebugStats("session-1")).toMatchObject({ + requests: 2, + connectionsCreated: 1, + connectionsReused: 1, + cachedContextRequests: 2, + storeTrueRequests: 0, + fullContextRequests: 1, + deltaRequests: 1, + lastDeltaInputItems: 1, + lastPreviousResponseId: "resp_1", + }); + }); }); diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index ff91ea43..987224c1 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Added + +- Added `websocket-cached` to the transport setting options for the OpenAI Codex provider used with ChatGPT subscription auth. This keeps the same WebSocket open for a session and, after the first request, sends only the new conversation items instead of resending the full chat history when possible. + ## [0.71.0] - 2026-04-30 ### Breaking Changes diff --git a/packages/coding-agent/src/modes/interactive/components/settings-selector.ts b/packages/coding-agent/src/modes/interactive/components/settings-selector.ts index 0bc92052..dff07822 100644 --- a/packages/coding-agent/src/modes/interactive/components/settings-selector.ts +++ b/packages/coding-agent/src/modes/interactive/components/settings-selector.ts @@ -235,7 +235,7 @@ export class SettingsSelectorComponent extends Container { label: "Transport", description: "Preferred transport for providers that support multiple transports", currentValue: config.transport, - values: ["sse", "websocket", "auto"], + values: ["sse", "websocket", "websocket-cached", "auto"], }, { id: "hide-thinking", diff --git a/packages/coding-agent/test/sdk-codex-cache-probe-tool-loop.ts b/packages/coding-agent/test/sdk-codex-cache-probe-tool-loop.ts index 96a9b77d..19aa1bac 100644 --- a/packages/coding-agent/test/sdk-codex-cache-probe-tool-loop.ts +++ b/packages/coding-agent/test/sdk-codex-cache-probe-tool-loop.ts @@ -11,7 +11,20 @@ import { mkdirSync } from "node:fs"; import { tmpdir } from "node:os"; import { dirname, join, resolve } from "node:path"; import process from "node:process"; -import { type AssistantMessage, getModel, Type } from "@mariozechner/pi-ai"; +import { + type Api, + type AssistantMessage, + type AssistantMessageEventStream, + type Context, + getModel, + type Model, + type SimpleStreamOptions, + Type, +} from "@mariozechner/pi-ai"; +import { + getOpenAICodexWebSocketDebugStats, + streamSimpleOpenAICodexResponses, +} from "../../ai/src/providers/openai-codex-responses.js"; import { AuthStorage } from "../src/core/auth-storage.js"; import { createExtensionRuntime } from "../src/core/extensions/loader.js"; import type { ToolDefinition } from "../src/core/extensions/types.js"; @@ -21,7 +34,7 @@ import { createAgentSession } from "../src/core/sdk.js"; import { SessionManager } from "../src/core/session-manager.js"; import { SettingsManager } from "../src/core/settings-manager.js"; -type Transport = "sse" | "websocket" | "auto"; +type Transport = "sse" | "websocket" | "websocket-cached" | "auto"; interface Args { turns: number; @@ -30,6 +43,16 @@ interface Args { maxTokens: number; } +interface WebSocketStatsSnapshot { + requests: number; + connectionsCreated: number; + connectionsReused: number; + cachedContextRequests: number; + storeTrueRequests: number; + fullContextRequests: number; + deltaRequests: number; +} + interface SubrequestRecord { turn: number; subrequest: number; @@ -67,7 +90,7 @@ function parseArgs(argv: string[]): Args { } case "--transport": { const value = argv[++i]; - if (value !== "sse" && value !== "websocket" && value !== "auto") { + if (value !== "sse" && value !== "websocket" && value !== "websocket-cached" && value !== "auto") { throw new Error(`Invalid --transport value: ${value}`); } transport = value; @@ -105,14 +128,14 @@ function printHelp(): void { Options: --turns Number of turns to run. Must be between ${MIN_TURNS} and ${MAX_TURNS}. Default: ${DEFAULT_TURNS} --session Specific session jsonl file to write - --transport sse | websocket | auto. Default: sse + --transport sse | websocket | websocket-cached | auto. Default: sse --max-tokens Max output tokens per subrequest. Default: ${DEFAULT_MAX_TOKENS} --help Show this message Notes: - Uses createAgentSession() from the coding-agent SDK - - Provider/model fixed to openai-codex/gpt-5.4 - - Thinking level fixed to medium + - Provider/model fixed to openai-codex/gpt-5.5 + - Thinking level fixed to low - Activates exactly one deterministic custom tool - Prompts are intentionally > 1024 tokens and explicitly describe the test `); @@ -161,6 +184,54 @@ function createMinimalResourceLoader(systemPrompt: string): ResourceLoader { }; } +function average(values: number[]): number { + return values.length === 0 ? 0 : values.reduce((sum, value) => sum + value, 0) / values.length; +} + +function percentile(values: number[], percentileValue: number): number { + if (values.length === 0) return 0; + const sorted = [...values].sort((a, b) => a - b); + const index = Math.min(sorted.length - 1, Math.max(0, Math.ceil((percentileValue / 100) * sorted.length) - 1)); + return sorted[index]; +} + +function getWebSocketStatsSnapshot(sessionId: string): WebSocketStatsSnapshot { + const stats = getOpenAICodexWebSocketDebugStats(sessionId); + return { + requests: stats?.requests ?? 0, + connectionsCreated: stats?.connectionsCreated ?? 0, + connectionsReused: stats?.connectionsReused ?? 0, + cachedContextRequests: stats?.cachedContextRequests ?? 0, + storeTrueRequests: stats?.storeTrueRequests ?? 0, + fullContextRequests: stats?.fullContextRequests ?? 0, + deltaRequests: stats?.deltaRequests ?? 0, + }; +} + +function diffWebSocketStats(after: WebSocketStatsSnapshot, before: WebSocketStatsSnapshot): WebSocketStatsSnapshot { + return { + requests: after.requests - before.requests, + connectionsCreated: after.connectionsCreated - before.connectionsCreated, + connectionsReused: after.connectionsReused - before.connectionsReused, + cachedContextRequests: after.cachedContextRequests - before.cachedContextRequests, + storeTrueRequests: after.storeTrueRequests - before.storeTrueRequests, + fullContextRequests: after.fullContextRequests - before.fullContextRequests, + deltaRequests: after.deltaRequests - before.deltaRequests, + }; +} + +function formatWebSocketStats(label: string, stats: WebSocketStatsSnapshot): string { + if (stats.requests === 0) return `${label} websocket none`; + return [ + `${label} websocket`, + `requests ${stats.requests}`, + `new/reused ${stats.connectionsCreated}/${stats.connectionsReused}`, + `cached ${stats.cachedContextRequests}`, + `store ${stats.storeTrueRequests}`, + `full/delta ${stats.fullContextRequests}/${stats.deltaRequests}`, + ].join(" | "); +} + function getAssistantText(message: AssistantMessage): string { return message.content .filter((block): block is Extract => block.type === "text") @@ -206,11 +277,24 @@ async function main(): Promise { const authStorage = AuthStorage.create(); const modelRegistry = ModelRegistry.create(authStorage); - const model = getModel("openai-codex", "gpt-5.4"); + const model = getModel("openai-codex", "gpt-5.5"); if (!model) { - throw new Error("Model openai-codex/gpt-5.4 not found"); + throw new Error("Model openai-codex/gpt-5.5 not found"); } const baseModel = { ...model, maxTokens: args.maxTokens }; + const streamSimpleOpenAICodexResponsesForRegistry = ( + registryModel: Model, + context: Context, + options?: SimpleStreamOptions, + ): AssistantMessageEventStream => + streamSimpleOpenAICodexResponses(registryModel as Model<"openai-codex-responses">, context, options); + modelRegistry.registerProvider("openai-codex", { + api: "openai-codex-responses", + baseUrl: baseModel.baseUrl, + apiKey: "!echo source-provider-override-uses-auth-storage", + streamSimple: streamSimpleOpenAICodexResponsesForRegistry, + models: [baseModel], + }); const settingsManager = SettingsManager.inMemory({ compaction: { enabled: false }, @@ -226,7 +310,7 @@ async function main(): Promise { cwd: process.cwd(), agentDir: dirname(args.sessionPath), model: baseModel, - thinkingLevel: "medium", + thinkingLevel: "low", customTools: [deterministicProbeTool() as unknown as ToolDefinition], resourceLoader, sessionManager: SessionManager.open(args.sessionPath), @@ -239,20 +323,23 @@ async function main(): Promise { const unsubscribe = session.subscribe(() => {}); const records: SubrequestRecord[] = []; + const turnElapsedMs: number[] = []; let previousCacheRead: number | null = null; - console.log(`provider openai-codex, model gpt-5.4`); + console.log(`provider openai-codex, model gpt-5.5`); console.log(`session ${session.sessionFile}`); - console.log(`turns ${args.turns}, transport ${args.transport}, reasoning medium, maxTokens ${args.maxTokens}`); + console.log(`turns ${args.turns}, transport ${args.transport}, reasoning low, maxTokens ${args.maxTokens}`); console.log(""); for (let turn = 1; turn <= args.turns; turn++) { const prompt = buildPrompt(turn); const promptTokens = estimateTokens(prompt); const previousMessagesLength = session.messages.length; + const websocketStatsBefore = getWebSocketStatsSnapshot(session.sessionId); const startedAt = Date.now(); await session.prompt(prompt); const elapsedMs = Date.now() - startedAt; + turnElapsedMs.push(elapsedMs); const newMessages = session.messages.slice(previousMessagesLength); const assistantMessages = newMessages.filter((message): message is AssistantMessage => @@ -316,6 +403,8 @@ async function main(): Promise { previousCacheRead = assistant.usage.cacheRead; } + const websocketStatsAfter = getWebSocketStatsSnapshot(session.sessionId); + const websocketStatsForTurn = diffWebSocketStats(websocketStatsAfter, websocketStatsBefore); console.log( [ `turn ${String(turn).padStart(2, "0")} agg`, @@ -327,6 +416,7 @@ async function main(): Promise { `total ${turnTotal}`, ].join(" | "), ); + console.log(formatWebSocketStats(`turn ${String(turn).padStart(2, "0")}`, websocketStatsForTurn)); } const violations = records @@ -343,7 +433,50 @@ async function main(): Promise { }) .filter((value): value is NonNullable => value !== null); + const totalElapsedMs = turnElapsedMs.reduce((sum, value) => sum + value, 0); console.log(""); + console.log( + [ + "timing", + `turns ${turnElapsedMs.length}`, + `total ${(totalElapsedMs / 1000).toFixed(1)}s`, + `avg ${(average(turnElapsedMs) / 1000).toFixed(2)}s`, + `p50 ${(percentile(turnElapsedMs, 50) / 1000).toFixed(2)}s`, + `p95 ${(percentile(turnElapsedMs, 95) / 1000).toFixed(2)}s`, + `max ${(Math.max(...turnElapsedMs) / 1000).toFixed(2)}s`, + ].join(" | "), + ); + const websocketStats = getOpenAICodexWebSocketDebugStats(session.sessionId); + const requestedWebsocket = + args.transport === "websocket" || args.transport === "websocket-cached" || args.transport === "auto"; + const observedWebsocket = Boolean(websocketStats && websocketStats.requests > 0); + console.log( + [ + "transport summary", + `requested ${args.transport}`, + `observed ${observedWebsocket ? "websocket" : "sse/no-websocket"}`, + `sseFallbackSuspected ${requestedWebsocket && !observedWebsocket ? "yes" : "no"}`, + `cachedContext ${websocketStats?.cachedContextRequests ? "yes" : "no"}`, + `storeTrue ${websocketStats ? `${websocketStats.storeTrueRequests}/${websocketStats.requests}` : "0/0"}`, + `delta ${websocketStats ? `${websocketStats.deltaRequests}/${websocketStats.requests}` : "0/0"}`, + `full ${websocketStats ? `${websocketStats.fullContextRequests}/${websocketStats.requests}` : "0/0"}`, + ].join(" | "), + ); + if (websocketStats) { + console.log( + [ + "websocket details", + `requests ${websocketStats.requests}`, + `connections created/reused ${websocketStats.connectionsCreated}/${websocketStats.connectionsReused}`, + `cachedContext ${websocketStats.cachedContextRequests}`, + `storeTrue ${websocketStats.storeTrueRequests}`, + `full/delta ${websocketStats.fullContextRequests}/${websocketStats.deltaRequests}`, + `lastInputItems ${websocketStats.lastInputItems}`, + `lastDeltaItems ${websocketStats.lastDeltaInputItems ?? "n/a"}`, + `lastPreviousResponseId ${websocketStats.lastPreviousResponseId ?? "n/a"}`, + ].join(" | "), + ); + } console.log(`subrequest cache read monotonic: ${violations.length === 0 ? "yes" : "NO"}`); if (violations.length > 0) { console.log("violations:");