From 4cb8f37202d1be833ce9eefa3b79da340ec3bc51 Mon Sep 17 00:00:00 2001 From: Ross Date: Mon, 26 Jan 2026 12:59:28 +0000 Subject: [PATCH] feat(cli): add configurable usageFields for backend token parsing Adds usageFields config option to CliBackendConfig allowing per-backend customization of token usage field names. This enables correct parsing of cache_creation_input_tokens from Anthropic's API (previously missed) while maintaining backwards compatibility through fallback defaults. - Add usageFields type to CliBackendConfig - Add Zod schema validation for usageFields - Configure default fields for Claude CLI (with Anthropic's actual field names) - Configure default fields for Codex CLI (OpenAI field names) - Update toUsage() to use backend config with fallback to hardcoded defaults --- src/agents/cli-backends.ts | 12 ++++++++++ src/agents/cli-runner/helpers.ts | 37 ++++++++++++++++++++++-------- src/config/types.agent-defaults.ts | 8 +++++++ src/config/zod-schema.core.ts | 9 ++++++++ 4 files changed, 57 insertions(+), 9 deletions(-) diff --git a/src/agents/cli-backends.ts b/src/agents/cli-backends.ts index f21c04f52..713d01fc5 100644 --- a/src/agents/cli-backends.ts +++ b/src/agents/cli-backends.ts @@ -48,6 +48,13 @@ const DEFAULT_CLAUDE_BACKEND: CliBackendConfig = { systemPromptWhen: "first", clearEnv: ["ANTHROPIC_API_KEY", "ANTHROPIC_API_KEY_OLD"], serialize: true, + usageFields: { + input: ["input_tokens", "inputTokens"], + output: ["output_tokens", "outputTokens"], + cacheRead: ["cache_read_input_tokens", "cached_input_tokens", "cacheRead"], + cacheWrite: ["cache_creation_input_tokens", "cache_write_input_tokens", "cacheWrite"], + total: ["total_tokens", "total"], + }, }; const DEFAULT_CODEX_BACKEND: CliBackendConfig = { @@ -72,6 +79,11 @@ const DEFAULT_CODEX_BACKEND: CliBackendConfig = { imageArg: "--image", imageMode: "repeat", serialize: true, + usageFields: { + input: ["prompt_tokens", "input_tokens"], + output: ["completion_tokens", "output_tokens"], + total: ["total_tokens"], + }, }; function normalizeBackendKey(key: string): string { diff --git a/src/agents/cli-runner/helpers.ts b/src/agents/cli-runner/helpers.ts index 6103125b0..50891959b 100644 --- a/src/agents/cli-runner/helpers.ts +++ b/src/agents/cli-runner/helpers.ts @@ -226,15 +226,34 @@ export function normalizeCliModel(modelId: string, backend: CliBackendConfig): s return trimmed; } -function toUsage(raw: Record): CliUsage | undefined { +function toUsage(raw: Record, backend?: CliBackendConfig): CliUsage | undefined { const pick = (key: string) => typeof raw[key] === "number" && raw[key] > 0 ? (raw[key] as number) : undefined; - const input = pick("input_tokens") ?? pick("inputTokens"); - const output = pick("output_tokens") ?? pick("outputTokens"); - const cacheRead = - pick("cache_read_input_tokens") ?? pick("cached_input_tokens") ?? pick("cacheRead"); - const cacheWrite = pick("cache_write_input_tokens") ?? pick("cacheWrite"); - const total = pick("total_tokens") ?? pick("total"); + + const pickFirst = (keys: string[] | undefined, fallback: string[]): number | undefined => { + const ordered = keys && keys.length > 0 ? keys : fallback; + for (const key of ordered) { + const value = pick(key); + if (value !== undefined) return value; + } + return undefined; + }; + + const fields = backend?.usageFields; + const input = pickFirst(fields?.input, ["input_tokens", "inputTokens"]); + const output = pickFirst(fields?.output, ["output_tokens", "outputTokens"]); + const cacheRead = pickFirst(fields?.cacheRead, [ + "cache_read_input_tokens", + "cached_input_tokens", + "cacheRead", + ]); + const cacheWrite = pickFirst(fields?.cacheWrite, [ + "cache_creation_input_tokens", + "cache_write_input_tokens", + "cacheWrite", + ]); + const total = pickFirst(fields?.total, ["total_tokens", "total"]); + if (!input && !output && !cacheRead && !cacheWrite && !total) return undefined; return { input, output, cacheRead, cacheWrite, total }; } @@ -284,7 +303,7 @@ export function parseCliJson(raw: string, backend: CliBackendConfig): CliOutput } if (!isRecord(parsed)) return null; const sessionId = pickSessionId(parsed, backend); - const usage = isRecord(parsed.usage) ? toUsage(parsed.usage) : undefined; + const usage = isRecord(parsed.usage) ? toUsage(parsed.usage, backend) : undefined; const text = collectText(parsed.message) || collectText(parsed.content) || @@ -315,7 +334,7 @@ export function parseCliJsonl(raw: string, backend: CliBackendConfig): CliOutput sessionId = parsed.thread_id.trim(); } if (isRecord(parsed.usage)) { - usage = toUsage(parsed.usage) ?? usage; + usage = toUsage(parsed.usage, backend) ?? usage; } const item = isRecord(parsed.item) ? parsed.item : null; if (item && typeof item.text === "string") { diff --git a/src/config/types.agent-defaults.ts b/src/config/types.agent-defaults.ts index 2a42d3623..e55ea2461 100644 --- a/src/config/types.agent-defaults.ts +++ b/src/config/types.agent-defaults.ts @@ -89,6 +89,14 @@ export type CliBackendConfig = { imageMode?: "repeat" | "list"; /** Serialize runs for this CLI. */ serialize?: boolean; + /** JSON field names for token usage parsing (in order of preference). */ + usageFields?: { + input?: string[]; + output?: string[]; + cacheRead?: string[]; + cacheWrite?: string[]; + total?: string[]; + }; }; export type AgentDefaultsConfig = { diff --git a/src/config/zod-schema.core.ts b/src/config/zod-schema.core.ts index 4a8c80bcc..956e44ff3 100644 --- a/src/config/zod-schema.core.ts +++ b/src/config/zod-schema.core.ts @@ -266,6 +266,15 @@ export const CliBackendSchema = z imageArg: z.string().optional(), imageMode: z.union([z.literal("repeat"), z.literal("list")]).optional(), serialize: z.boolean().optional(), + usageFields: z + .object({ + input: z.array(z.string()).optional(), + output: z.array(z.string()).optional(), + cacheRead: z.array(z.string()).optional(), + cacheWrite: z.array(z.string()).optional(), + total: z.array(z.string()).optional(), + }) + .optional(), }) .strict();