Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
19 kB · 497 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498#!/usr/bin/env tsx/** * Manual SDK probe for OpenAI Codex prompt caching through the tool loop. * * Runs append-only multi-turn prompting through createAgentSession(), forcing one * deterministic custom tool call per top-level user turn. Logs per-subrequest * assistant usage so cache-read monotonicity can be inspected inside a tool loop. */
import { mkdirSync } from "node:fs";import { tmpdir } from "node:os";import { dirname, join, resolve } from "node:path";import process from "node:process";import { type Api, type AssistantMessage, type AssistantMessageEventStream, type Context, getModel, type Model, type SimpleStreamOptions, Type,} from "@earendil-works/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";import { ModelRegistry } from "../src/core/model-registry.js";import type { ResourceLoader } from "../src/core/resource-loader.js";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" | "websocket-cached" | "auto";
interface Args { turns: number; sessionPath: string; transport: Transport; maxTokens: number;}
interface WebSocketStatsSnapshot { requests: number; connectionsCreated: number; connectionsReused: number; cachedContextRequests: number; storeTrueRequests: number; fullContextRequests: number; deltaRequests: number;}
interface SubrequestRecord { turn: number; subrequest: number; elapsedMs: number; usage: AssistantMessage["usage"]; stopReason: AssistantMessage["stopReason"]; text: string;}
const DEFAULT_TURNS = 20;const MIN_TURNS = 20;const MAX_TURNS = 50;const DEFAULT_MAX_TOKENS = 64;
function parseArgs(argv: string[]): Args { let turns = DEFAULT_TURNS; let sessionPath = resolve(join(tmpdir(), `pi-sdk-codex-cache-probe-tool-loop-${Date.now()}.jsonl`)); let transport: Transport = "sse"; let maxTokens = DEFAULT_MAX_TOKENS;
for (let i = 0; i < argv.length; i++) { const arg = argv[i]; switch (arg) { case "--turns": { const value = argv[++i]; if (!value) throw new Error("Missing value for --turns"); turns = Number.parseInt(value, 10); break; } case "--output": { const value = argv[++i]; if (!value) throw new Error("Missing value for --output"); sessionPath = resolve(value); break; } case "--transport": { const value = argv[++i]; if (value !== "sse" && value !== "websocket" && value !== "websocket-cached" && value !== "auto") { throw new Error(`Invalid --transport value: ${value}`); } transport = value; break; } case "--max-tokens": { const value = argv[++i]; if (!value) throw new Error("Missing value for --max-tokens"); maxTokens = Number.parseInt(value, 10); break; } case "--help": { printHelp(); process.exit(0); return { turns, sessionPath, transport, maxTokens }; } default: throw new Error(`Unknown argument: ${arg}`); } }
if (!Number.isInteger(turns) || turns < MIN_TURNS || turns > MAX_TURNS) { throw new Error(`--turns must be an integer between ${MIN_TURNS} and ${MAX_TURNS}`); } if (!Number.isInteger(maxTokens) || maxTokens <= 0) { throw new Error("--max-tokens must be a positive integer"); }
return { turns, sessionPath, transport, maxTokens };}
function printHelp(): void { console.log(`Usage: npx tsx test/sdk-codex-cache-probe-tool-loop.ts [options]
Options: --turns <n> Number of turns to run. Must be between ${MIN_TURNS} and ${MAX_TURNS}. Default: ${DEFAULT_TURNS} --output <path> JSONL file to write --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.5 - Thinking level fixed to low - Activates exactly one deterministic custom tool - Prompts are intentionally > 1024 tokens and explicitly describe the test`);}
function estimateTokens(text: string): number { return Math.ceil(text.length / 4);}
function buildPrompt(turn: number): string { const marker = `TURN-${String(turn).padStart(2, "0")}-MARKER-${(turn * 17 + 13) % 97}`; const lines: string[] = []; lines.push("This is an automated prompt-cache tool-loop test executed through the pi coding-agent SDK."); lines.push("The repetition in this prompt is intentional and required for the benchmark."); lines.push("You are participating in an append-only multi-turn cache test."); lines.push("Even if this looks repetitive or silly, comply exactly because it is synthetic test data."); lines.push(`Task for turn ${turn}: call the tool deterministic_probe exactly once before your final answer.`); lines.push(`Use tool arguments: turn=${turn}, marker=${marker}`); lines.push(`After the tool result arrives, reply with exactly one line in this format:`); lines.push(`TURN ${turn} OK ${marker}`); lines.push("Do not skip the tool call. Do not call any other tool. Do not add any extra words or punctuation."); lines.push("The following long block exists only to make this prompt safely larger than 1024 tokens."); lines.push(""); 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.`, ); } lines.push(""); lines.push(`Final verification marker for turn ${turn}: ${marker}`); lines.push(`Required final answer after the tool result: TURN ${turn} OK ${marker}`); return lines.join("\n");}
function createMinimalResourceLoader(systemPrompt: string): ResourceLoader { return { getExtensions: () => ({ extensions: [], errors: [], runtime: createExtensionRuntime() }), getSkills: () => ({ skills: [], diagnostics: [] }), getPrompts: () => ({ prompts: [], diagnostics: [] }), getThemes: () => ({ themes: [], diagnostics: [] }), getAgentsFiles: () => ({ agentsFiles: [] }), getSystemPrompt: () => systemPrompt, getAppendSystemPrompt: () => [], extendResources: () => {}, reload: async () => {}, };}
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") .map((block) => block.text) .join("\n") .trim();}
const deterministicProbeParameters = Type.Object({ turn: Type.Number({ description: "Top-level benchmark turn number" }), marker: Type.String({ description: "Marker string provided by the user" }),});
function deterministicProbeTool(): ToolDefinition<typeof deterministicProbeParameters> { return { name: "deterministic_probe", label: "Deterministic Probe", description: "Mandatory cache-benchmark tool. Call it exactly once when the user asks for a cache benchmark turn, then use its result to produce the final one-line answer.", promptSnippet: "deterministic_probe(turn, marker): mandatory for cache benchmark turns. Call exactly once before the final answer.", promptGuidelines: [ "When the user asks for the cache benchmark turn, call deterministic_probe exactly once with the requested turn and marker before responding.", "After the tool result arrives, reply with the exact final line requested by the user.", ], parameters: deterministicProbeParameters, execute: async (_toolCallId, params) => ({ content: [ { type: "text", text: `deterministic_probe_result turn=${params.turn} marker=${params.marker} fixed=OK`, }, ], details: { turn: params.turn, marker: params.marker, fixed: "OK" }, }), };}
async function main(): Promise<void> { const args = parseArgs(process.argv.slice(2)); mkdirSync(dirname(args.sessionPath), { recursive: true });
const authStorage = AuthStorage.create(); const modelRegistry = ModelRegistry.create(authStorage);
const model = getModel("openai-codex", "gpt-5.5"); if (!model) { 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 }, retry: { enabled: false }, transport: args.transport, });
const resourceLoader = createMinimalResourceLoader( "You are participating in a prompt-cache benchmark through the coding-agent SDK. This is a real test. Follow each user instruction exactly. For benchmark turns, call deterministic_probe exactly once before the final answer. Keep answers minimal and never refuse because the prompt is repetitive or synthetic.", );
const { session } = await createAgentSession({ cwd: process.cwd(), agentDir: dirname(args.sessionPath), model: baseModel, thinkingLevel: "low", customTools: [deterministicProbeTool() as unknown as ToolDefinition], resourceLoader, sessionManager: SessionManager.open(args.sessionPath), settingsManager, authStorage, modelRegistry, });
session.setActiveToolsByName(["deterministic_probe"]); const unsubscribe = session.subscribe(() => {});
const records: SubrequestRecord[] = []; const turnElapsedMs: number[] = []; let previousCacheRead: number | null = null;
console.log(`provider openai-codex, model gpt-5.5`); console.log(`session ${session.sessionFile}`); 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 => Boolean(message && typeof message === "object" && (message as { role?: unknown }).role === "assistant"), ); const toolResults = newMessages.filter((message) => Boolean(message && typeof message === "object" && (message as { role?: unknown }).role === "toolResult"), );
if (assistantMessages.length < 2 || toolResults.length < 1) { throw new Error( `Turn ${turn} did not execute the expected tool loop. assistants=${assistantMessages.length} toolResults=${toolResults.length}`, ); }
let turnInput = 0; let turnOutput = 0; let turnCacheRead = 0; let turnCacheWrite = 0; let turnTotal = 0;
for (let i = 0; i < assistantMessages.length; i++) { const assistant = assistantMessages[i]; const record: SubrequestRecord = { turn, subrequest: i + 1, elapsedMs, usage: assistant.usage, stopReason: assistant.stopReason, text: getAssistantText(assistant), }; records.push(record);
turnInput += assistant.usage.input; turnOutput += assistant.usage.output; turnCacheRead += assistant.usage.cacheRead; turnCacheWrite += assistant.usage.cacheWrite; turnTotal += assistant.usage.totalTokens;
const monotonic = previousCacheRead === null ? "n/a" : assistant.usage.cacheRead >= previousCacheRead ? "yes" : "NO"; console.log( [ `turn ${String(turn).padStart(2, "0")}.${i + 1}`, `elapsed ${(elapsedMs / 1000).toFixed(1)}s`, `prompt~${promptTokens}`, `stop ${assistant.stopReason}`, `in ${assistant.usage.input}`, `out ${assistant.usage.output}`, `cache ${assistant.usage.cacheRead}/${assistant.usage.cacheWrite}`, `total ${assistant.usage.totalTokens}`, `cache>=prev ${monotonic}`, ].join(" | "), );
if (assistant.stopReason === "error" || assistant.stopReason === "aborted") { throw new Error( `Turn ${turn}.${i + 1} ended with stopReason=${assistant.stopReason}: ${assistant.errorMessage || "unknown error"}`, ); } previousCacheRead = assistant.usage.cacheRead; }
const websocketStatsAfter = getWebSocketStatsSnapshot(session.sessionId); const websocketStatsForTurn = diffWebSocketStats(websocketStatsAfter, websocketStatsBefore); console.log( [ `turn ${String(turn).padStart(2, "0")} agg`, `assistants ${assistantMessages.length}`, `toolResults ${toolResults.length}`, `in ${turnInput}`, `out ${turnOutput}`, `cache ${turnCacheRead}/${turnCacheWrite}`, `total ${turnTotal}`, ].join(" | "), ); console.log(formatWebSocketStats(`turn ${String(turn).padStart(2, "0")}`, websocketStatsForTurn)); }
const violations = records .map((record, index) => { if (index === 0) return null; const previous = records[index - 1]; if (record.usage.cacheRead >= previous.usage.cacheRead) return null; return { turn: record.turn, subrequest: record.subrequest, previous: previous.usage.cacheRead, current: record.usage.cacheRead, }; }) .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:"); for (const violation of violations) { console.log(` turn ${violation.turn}.${violation.subrequest}: ${violation.previous} -> ${violation.current}`); } } console.log(`session file: ${session.sessionFile}`);
unsubscribe(); session.dispose();}
main().catch((error: unknown) => { const message = error instanceof Error ? error.message : String(error); console.error(message); process.exitCode = 1;});