feat(ai): add cached codex websocket transport

This commit is contained in:
Mario Zechner
2026-05-01 13:07:40 +02:00
parent 8040dd6ded
commit 4745a95898
9 changed files with 779 additions and 18 deletions

View File

@@ -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

View File

@@ -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";

View File

@@ -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<typeof setTimeout>;
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<string, CachedWebSocketConnection>();
const websocketDebugStats = new Map<string, OpenAICodexWebSocketDebugStats>();
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<void> {
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 {

View File

@@ -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;

View File

@@ -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 <n> Number of user turns. Default: ${DEFAULT_TURNS}
--transport <mode> sse | websocket | websocket-cached | auto. Default: websocket-cached
--reasoning <level> minimal | low | medium | high | xhigh. Default: low
--max-tokens <n> Max output tokens per model request. Default: ${DEFAULT_MAX_TOKENS}
--session-id <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<AssistantMessage["content"][number], { type: "toolCall" }>): 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<AssistantMessage["content"][number], { type: "text" }> => 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<void> {
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<AssistantMessage["content"][number], { type: "toolCall" }> =>
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;
});

View File

@@ -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<string, Set<(event: unknown) => void>>();
constructor(_url: string, _protocols?: string | string[] | { headers?: Record<string, string> }) {
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",
});
});
});

View File

@@ -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

View File

@@ -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",

View File

@@ -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 <n> Number of turns to run. Must be between ${MIN_TURNS} and ${MAX_TURNS}. Default: ${DEFAULT_TURNS}
--session <path> Specific session jsonl file to write
--transport <mode> sse | websocket | auto. Default: sse
--transport <mode> sse | websocket | websocket-cached | auto. Default: sse
--max-tokens <n> 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<AssistantMessage["content"][number], { type: "text" }> => block.type === "text")
@@ -206,11 +277,24 @@ async function main(): Promise<void> {
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<Api>,
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<void> {
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<void> {
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<void> {
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<void> {
`total ${turnTotal}`,
].join(" | "),
);
console.log(formatWebSocketStats(`turn ${String(turn).padStart(2, "0")}`, websocketStatsForTurn));
}
const violations = records
@@ -343,7 +433,50 @@ async function main(): Promise<void> {
})
.filter((value): value is NonNullable<typeof value> => 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:");