feat(ai): add cached codex websocket transport
This commit is contained in:
Mario Zechner
@@ -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
|
||||
|
||||
@@ -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";
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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;
|
||||
|
||||
291
packages/ai/test/codex-websocket-cached-probe.ts
Normal file
291
packages/ai/test/codex-websocket-cached-probe.ts
Normal 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;
|
||||
});
|
||||
@@ -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",
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user