feat(ai): add cached codex websocket transport
This commit is contained in:
Mario Zechner
@@ -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:");
|
||||
|
||||
Reference in New Issue
Block a user