fix: capture per-model token usage in bench SDK driver #33
89
bench/sdk-driver.test.ts
Normal file
89
bench/sdk-driver.test.ts
Normal file
@@ -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,
|
||||||
|
});
|
||||||
|
});
|
||||||
|
});
|
||||||
@@ -38,7 +38,20 @@ export interface SdkDriverConfig {
|
|||||||
|
|
||||||
// --- The slice of the Claude Agent SDK this adapter consumes ------------------
|
// --- 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 {
|
interface SdkUsage {
|
||||||
input_tokens?: number;
|
input_tokens?: number;
|
||||||
output_tokens?: number;
|
output_tokens?: number;
|
||||||
@@ -46,13 +59,13 @@ interface SdkUsage {
|
|||||||
cache_read_input_tokens?: number;
|
cache_read_input_tokens?: number;
|
||||||
}
|
}
|
||||||
|
|
||||||
interface SdkResultMessage {
|
export interface SdkResultMessage {
|
||||||
type: "result";
|
type: "result";
|
||||||
/** `error_max_turns` when the run hit the turn cap. */
|
/** `error_max_turns` when the run hit the turn cap. */
|
||||||
subtype: string;
|
subtype: string;
|
||||||
usage?: SdkUsage;
|
usage?: SdkUsage;
|
||||||
/** Per-model usage, including the auxiliary small model the runtime invokes. */
|
/** Per-model usage, including the auxiliary small model the runtime invokes. */
|
||||||
modelUsage?: Record<string, SdkUsage>;
|
modelUsage?: Record<string, SdkModelUsage>;
|
||||||
total_cost_usd?: number;
|
total_cost_usd?: number;
|
||||||
num_turns?: number;
|
num_turns?: number;
|
||||||
result?: string;
|
result?: string;
|
||||||
@@ -105,21 +118,30 @@ interface SdkModule {
|
|||||||
/**
|
/**
|
||||||
* Sum the four token components across every model the run touched, so the
|
* 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
|
* 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
|
* requires. The per-model `modelUsage` (camelCase) is the primary source; the
|
||||||
* present.
|
* 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 {
|
export function sumTokens(result: SdkResultMessage): TokenComponents {
|
||||||
const usages = result.modelUsage
|
|
||||||
? Object.values(result.modelUsage)
|
|
||||||
: result.usage
|
|
||||||
? [result.usage]
|
|
||||||
: [];
|
|
||||||
const total: TokenComponents = { freshInput: 0, cacheCreation: 0, cacheRead: 0, output: 0 };
|
const total: TokenComponents = { freshInput: 0, cacheCreation: 0, cacheRead: 0, output: 0 };
|
||||||
for (const usage of usages) {
|
|
||||||
total.freshInput += usage.input_tokens ?? 0;
|
const perModel = result.modelUsage ? Object.values(result.modelUsage) : [];
|
||||||
total.cacheCreation += usage.cache_creation_input_tokens ?? 0;
|
if (perModel.length > 0) {
|
||||||
total.cacheRead += usage.cache_read_input_tokens ?? 0;
|
for (const usage of perModel) {
|
||||||
total.output += usage.output_tokens ?? 0;
|
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;
|
return total;
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user