From bd751f0515aa60cce16670071a907e0ef6dcf4ab Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Sat, 7 Mar 2026 14:34:12 +0100 Subject: [PATCH] chore(coding-agent): add codex tool-loop cache probe --- .../test/sdk-codex-cache-probe-tool-loop.ts | 365 ++++++++++++++++++ 1 file changed, 365 insertions(+) create mode 100644 packages/coding-agent/test/sdk-codex-cache-probe-tool-loop.ts diff --git a/packages/coding-agent/test/sdk-codex-cache-probe-tool-loop.ts b/packages/coding-agent/test/sdk-codex-cache-probe-tool-loop.ts new file mode 100644 index 00000000..731bf54c --- /dev/null +++ b/packages/coding-agent/test/sdk-codex-cache-probe-tool-loop.ts @@ -0,0 +1,365 @@ +#!/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 AssistantMessage, getModel, Type } from "@mariozechner/pi-ai"; +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" | "auto"; + +interface Args { + turns: number; + sessionPath: string; + transport: Transport; + maxTokens: 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 "--session": { + const value = argv[++i]; + if (!value) throw new Error("Missing value for --session"); + sessionPath = resolve(value); + break; + } + case "--transport": { + const value = argv[++i]; + if (value !== "sse" && value !== "websocket" && 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 Number of turns to run. Must be between ${MIN_TURNS} and ${MAX_TURNS}. Default: ${DEFAULT_TURNS} + --session Specific session jsonl file to write + --transport sse | websocket | auto. Default: sse + --max-tokens Max output tokens per subrequest. Default: ${DEFAULT_MAX_TOKENS} + --help Show this message + +Notes: + - Uses createAgentSession() from the coding-agent SDK + - Provider/model fixed to openai-codex/gpt-5.4 + - Thinking level fixed to medium + - 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: () => [], + getPathMetadata: () => new Map(), + extendResources: () => {}, + reload: async () => {}, + }; +} + +function getAssistantText(message: AssistantMessage): string { + return message.content + .filter((block): block is Extract => 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 { + 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 { + const args = parseArgs(process.argv.slice(2)); + mkdirSync(dirname(args.sessionPath), { recursive: true }); + + const authStorage = AuthStorage.create(); + const modelRegistry = new ModelRegistry(authStorage); + + const model = getModel("openai-codex", "gpt-5.4"); + if (!model) { + throw new Error("Model openai-codex/gpt-5.4 not found"); + } + const baseModel = { ...model, maxTokens: args.maxTokens }; + + 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: "medium", + 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[] = []; + let previousCacheRead: number | null = null; + + console.log(`provider openai-codex, model gpt-5.4`); + console.log(`session ${session.sessionFile}`); + console.log(`turns ${args.turns}, transport ${args.transport}, reasoning medium, 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 startedAt = Date.now(); + await session.prompt(prompt); + const elapsedMs = Date.now() - startedAt; + + 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; + } + + 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(" | "), + ); + } + + 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 => value !== null); + + console.log(""); + 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; +});