From 690456d21ece7cf26282642c1948f2cebd0c6611 Mon Sep 17 00:00:00 2001 From: alexion Date: Thu, 16 Jul 2026 19:30:58 -0400 Subject: [PATCH] fix: capture per-model token usage in bench SDK driver MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit sumTokens read the Agent SDK's per-model `modelUsage` entries with snake_case field names, but the SDK reports those per-model entries in camelCase (`inputTokens`, `cacheReadInputTokens`, ...). Every token component therefore fell through to zero, silently zeroing the cost-equivalent-token headline metric — while `total_cost_usd` and `num_turns` (top-level snake_case) kept working and masked it. Read `modelUsage` with the correct camelCase fields, keeping the snake_case aggregate `usage` as the fallback. Export `sumTokens` and add a regression test covering both the per-model camelCase sum (folding in the auxiliary model) and the snake_case fallback, so a future SDK field-casing drift fails a test instead of producing zero-token samples. --- bench/sdk-driver.test.ts | 89 ++++++++++++++++++++++++++++++++++++++++ bench/sdk-driver.ts | 54 ++++++++++++++++-------- 2 files changed, 127 insertions(+), 16 deletions(-) create mode 100644 bench/sdk-driver.test.ts diff --git a/bench/sdk-driver.test.ts b/bench/sdk-driver.test.ts new file mode 100644 index 0000000..32170ed --- /dev/null +++ b/bench/sdk-driver.test.ts @@ -0,0 +1,89 @@ +import { describe, expect, it } from "vitest"; +import { sumTokens } from "./sdk-driver.js"; +import type { SdkResultMessage } from "./sdk-driver.js"; + +describe("sumTokens", () => { + // Behavior: sumTokens reads per-model token usage from the SDK result's + // `modelUsage` map and sums the four token components across EVERY model, + // folding in the auxiliary small model the runtime invokes for internal + // chores — because that is real consumption against the same allowance + // (see the token-components note in result.ts). Crucially it reads the SDK's + // camelCase field names (inputTokens / outputTokens / cacheCreationInputTokens + // / cacheReadInputTokens); this is a regression guard against a bug where the + // driver read snake_case keys, so every component silently summed to zero. + // + // The result carries two models — a main model and the aux small model — with + // DISTINCT numbers on every field, so a wrong field mapping cannot be masked + // by another and both models must be folded in to reach the totals. The + // expected sums are derived BY HAND from the two models, independent of how + // sumTokens computes them, per the metric mapping + // (inputTokens -> freshInput, outputTokens -> output, + // cacheCreationInputTokens -> cacheCreation, cacheReadInputTokens -> cacheRead): + // freshInput = 500 + 30 = 530 + // output = 200 + 8 = 208 + // cacheCreation = 3000 + 100 = 3100 + // cacheRead = 10000 + 400 = 10400 + it("sums the four camelCase token components across every model, folding in the auxiliary model", () => { + const result: SdkResultMessage = { + type: "result", + subtype: "success", + modelUsage: { + "claude-opus-4-8": { + inputTokens: 500, + outputTokens: 200, + cacheCreationInputTokens: 3000, + cacheReadInputTokens: 10000, + }, + "claude-haiku-aux": { + inputTokens: 30, + outputTokens: 8, + cacheCreationInputTokens: 100, + cacheReadInputTokens: 400, + }, + }, + }; + + expect(sumTokens(result)).toEqual({ + freshInput: 530, + output: 208, + cacheCreation: 3100, + cacheRead: 10400, + }); + }); + + // Behavior: when the result carries NO per-model `modelUsage` breakdown, + // sumTokens falls back to the aggregate `usage` block. Unlike the per-model + // map, the SDK reports this aggregate in snake_case (input_tokens / + // output_tokens / cache_creation_input_tokens / cache_read_input_tokens), so + // this pins that the fallback path reads the OTHER casing correctly and that + // the two shapes are not confused. This is a regression guard against reading + // the wrong casing on the fallback path. + // + // There is a single aggregate source, so each expected component equals its + // own field's value; the four numbers are DISTINCT so a wrong field mapping + // cannot be masked by another. Values are hand-worked from the aggregate, + // independent of how sumTokens computes them, per the mapping + // (input_tokens -> freshInput, output_tokens -> output, + // cache_creation_input_tokens -> cacheCreation, + // cache_read_input_tokens -> cacheRead): + // freshInput = 700, output = 90, cacheCreation = 4000, cacheRead = 20000. + it("falls back to the snake_case aggregate usage when no per-model modelUsage is present", () => { + const result: SdkResultMessage = { + type: "result", + subtype: "success", + usage: { + input_tokens: 700, + output_tokens: 90, + cache_creation_input_tokens: 4000, + cache_read_input_tokens: 20000, + }, + }; + + expect(sumTokens(result)).toEqual({ + freshInput: 700, + cacheCreation: 4000, + cacheRead: 20000, + output: 90, + }); + }); +}); diff --git a/bench/sdk-driver.ts b/bench/sdk-driver.ts index d820df4..5734e26 100644 --- a/bench/sdk-driver.ts +++ b/bench/sdk-driver.ts @@ -38,7 +38,20 @@ export interface SdkDriverConfig { // --- The slice of the Claude Agent SDK this adapter consumes ------------------ -/** The per-request token usage the SDK reports, per model. */ +/** + * Per-model token usage as the SDK's `modelUsage` reports it, one entry per model + * the run touched. The SDK reports these in camelCase — distinct from the aggregate + * {@link SdkUsage} below, which it reports in the Anthropic API's snake_case shape. + * Reading the wrong casing silently yields zeros, so the two are kept separate. + */ +interface SdkModelUsage { + inputTokens?: number; + outputTokens?: number; + cacheCreationInputTokens?: number; + cacheReadInputTokens?: number; +} + +/** The aggregate per-request usage (snake_case), the fallback when no per-model breakdown is present. */ interface SdkUsage { input_tokens?: number; output_tokens?: number; @@ -46,13 +59,13 @@ interface SdkUsage { cache_read_input_tokens?: number; } -interface SdkResultMessage { +export interface SdkResultMessage { type: "result"; /** `error_max_turns` when the run hit the turn cap. */ subtype: string; usage?: SdkUsage; /** Per-model usage, including the auxiliary small model the runtime invokes. */ - modelUsage?: Record; + modelUsage?: Record; total_cost_usd?: number; num_turns?: number; result?: string; @@ -105,21 +118,30 @@ interface SdkModule { /** * Sum the four token components across every model the run touched, so the * auxiliary small model the runtime invokes is folded in as the metric spec - * requires. Falls back to the aggregate `usage` when no per-model breakdown is - * present. + * requires. The per-model `modelUsage` (camelCase) is the primary source; the + * aggregate `usage` (snake_case) is the fallback when no per-model breakdown is + * present. The two shapes use different field casing, so each is read with its + * own names — reading the wrong casing is what silently produced zero tokens. */ -function sumTokens(result: SdkResultMessage): TokenComponents { - const usages = result.modelUsage - ? Object.values(result.modelUsage) - : result.usage - ? [result.usage] - : []; +export function sumTokens(result: SdkResultMessage): TokenComponents { const total: TokenComponents = { freshInput: 0, cacheCreation: 0, cacheRead: 0, output: 0 }; - for (const usage of usages) { - total.freshInput += usage.input_tokens ?? 0; - total.cacheCreation += usage.cache_creation_input_tokens ?? 0; - total.cacheRead += usage.cache_read_input_tokens ?? 0; - total.output += usage.output_tokens ?? 0; + + const perModel = result.modelUsage ? Object.values(result.modelUsage) : []; + if (perModel.length > 0) { + for (const usage of perModel) { + total.freshInput += usage.inputTokens ?? 0; + total.cacheCreation += usage.cacheCreationInputTokens ?? 0; + total.cacheRead += usage.cacheReadInputTokens ?? 0; + total.output += usage.outputTokens ?? 0; + } + return total; + } + + if (result.usage) { + total.freshInput += result.usage.input_tokens ?? 0; + total.cacheCreation += result.usage.cache_creation_input_tokens ?? 0; + total.cacheRead += result.usage.cache_read_input_tokens ?? 0; + total.output += result.usage.output_tokens ?? 0; } return total; } -- 2.47.3