From ab1460e9a68a8368f7c858603ebdd09306089046 Mon Sep 17 00:00:00 2001 From: Mahmoud Mabrouk Date: Mon, 3 Aug 2026 22:32:06 +0200 Subject: [PATCH] fix(runner): stop reporting context-window size as a run's token usage When the harness's PromptResponse carried no usage, the runner fell back to the ACP stream's `usage_update.used` value and reported it as the run's token total. That value is not a token count of the run. It is how full the agent's context window is at that point in the turn. Measured on one stack: 6,041 root spans shaped `input 0 / output 0 / total 63369`, with no cost. A wrong number here is worse than no number, because a reader cannot tell a real 63,369-token run from a context-size artifact, and any token aggregate built on that path is silently poisoned. Fixing the usage merge alone would not have removed those rows. The tracer keeps its own usage record built straight from the stream, and handing it nothing is a no-op rather than a clear, so it kept stamping the context size on the root span anyway. So the stream handler now drops `used` and keeps only the reported cost, and an update carrying neither is not recorded at all. The stamping function now also skips a usage record with no tokens and no cost. All zeros is the absence of a measurement, not a measured zero, and asserting that a run spent nothing is its own wrong answer. Nothing consumes the context size. The only readers of the merged usage object are the span attributes and the Python SDK's `record_usage`, and both treat every number in it as a token count. Tests: the full runner suite passes, 1,486 tests across 97 files, plus nine new cases covering a context-size-only stream, a cost-only stream, a half-reported split, and a normal split delivered through the harness. Claude-Session: https://claude.ai/code/session_01RkWWQUNNzRbaB5jnCAdjYA --- .../runner/src/engines/sandbox_agent/usage.ts | 29 ++- services/runner/src/protocol.ts | 8 +- services/runner/src/tracing/otel.ts | 19 +- .../unit/otel-usage-context-size.test.ts | 174 ++++++++++++++++++ .../tests/unit/sandbox-agent-usage.test.ts | 82 +++++++++ 5 files changed, 299 insertions(+), 13 deletions(-) create mode 100644 services/runner/tests/unit/otel-usage-context-size.test.ts diff --git a/services/runner/src/engines/sandbox_agent/usage.ts b/services/runner/src/engines/sandbox_agent/usage.ts index 1d08b3447b..afb17d8e01 100644 --- a/services/runner/src/engines/sandbox_agent/usage.ts +++ b/services/runner/src/engines/sandbox_agent/usage.ts @@ -25,7 +25,20 @@ export async function readRunUsage( } } -/** Combine prompt token counts with stream cost when no Pi usage writeback exists. */ +/** + * Combine prompt token counts with stream cost when no Pi usage writeback exists. + * + * The token total is ONLY ever the harness-reported split. There is no fallback token + * source: the ACP stream's `usage_update.used` is the agent's context-window occupancy, + * not a count of the tokens this run spent, so it must never become a token total. When + * the harness reports no split, this returns no tokens at all (cost alone still counts as + * usage) — absent data has to read as absent, because a plausible-looking wrong total + * silently poisons every aggregate built on it. + * + * Cost follows the same rule via omission: an unreported cost leaves the key OFF, because a + * substituted `0` would claim the run was measured and free. A reported cost is passed through + * as-is, including a genuine `0`. + */ export function mergePromptAndStreamUsage( promptResult: any, streamUsage: AgentUsage | undefined, @@ -33,10 +46,16 @@ export function mergePromptAndStreamUsage( const promptUsage = promptResult?.usage; const inputTokens = promptUsage?.inputTokens ?? streamUsage?.input ?? 0; const outputTokens = promptUsage?.outputTokens ?? streamUsage?.output ?? 0; - const total = inputTokens + outputTokens || streamUsage?.total || 0; - const cost = streamUsage?.cost ?? 0; - return total > 0 || cost > 0 - ? { input: inputTokens, output: outputTokens, total, cost } + const total = inputTokens + outputTokens; + const cost = streamUsage?.cost; + const hasCost = cost != null; + return total > 0 || hasCost + ? { + input: inputTokens, + output: outputTokens, + total, + ...(hasCost ? { cost } : {}), + } : undefined; } diff --git a/services/runner/src/protocol.ts b/services/runner/src/protocol.ts index f06b25b2c2..6b20457812 100644 --- a/services/runner/src/protocol.ts +++ b/services/runner/src/protocol.ts @@ -427,7 +427,13 @@ export interface AgentUsage { input: number; output: number; total: number; - cost: number; + /** + * INVARIANT: absent means the cost is UNKNOWN (the harness reported none); a present `0` is a + * measured zero — a free model or a fully cached turn. Consumers read presence as evidence of + * a measurement, so a producer must never substitute a zero for an absence: doing so records + * an unpriced run as a free one, which every downstream aggregate then believes. + */ + cost?: number; } export interface AgentRunRequest { diff --git a/services/runner/src/tracing/otel.ts b/services/runner/src/tracing/otel.ts index ea5f3fcccd..2b173b8cb2 100644 --- a/services/runner/src/tracing/otel.ts +++ b/services/runner/src/tracing/otel.ts @@ -1163,7 +1163,9 @@ export function createSandboxAgentOtel( } function stampUsage(span: Span, u: AgentUsage | undefined): void { - if (!u) return; + // No tokens and no cost is not a measured zero, it is the absence of a measurement: + // stamping zeros would assert a run cost nothing and spent nothing. + if (!u || (u.total <= 0 && u.cost <= 0)) return; span.setAttribute("gen_ai.usage.input_tokens", u.input); span.setAttribute("gen_ai.usage.output_tokens", u.output); span.setAttribute("gen_ai.usage.prompt_tokens", u.input); @@ -1446,16 +1448,19 @@ export function createSandboxAgentOtel( } if (kind === "usage_update") { - // ACP usage_update carries only `used` (context tokens) and `cost.amount`. The - // per-call input/output split is NOT on the stream; it rides on the PromptResponse, - // which the sandbox-agent engine reads. Keep total + cost here and leave the split to the caller. + // ACP usage_update carries `used` and `cost.amount`. `used` is the agent's CONTEXT-WINDOW + // occupancy at this point in the turn, NOT the tokens this run spent, so it is dropped: + // reporting it as a total produced runs shaped `input 0 / output 0 / total `. + // The token split is not on the stream at all; it rides on the PromptResponse, which the + // sandbox-agent engine reads and hands back through `setUsage`. Cost is the only run total + // here, so an update without one carries nothing worth recording. const cost = update.cost?.amount; - const total = update.used; + if (typeof cost !== "number") return; usage = { input: usage?.input ?? 0, output: usage?.output ?? 0, - total: typeof total === "number" ? total : usage?.total ?? 0, - cost: typeof cost === "number" ? cost : usage?.cost ?? 0, + total: usage?.total ?? 0, + cost, }; record({ type: "usage", ...usage }); } diff --git a/services/runner/tests/unit/otel-usage-context-size.test.ts b/services/runner/tests/unit/otel-usage-context-size.test.ts new file mode 100644 index 0000000000..b2818a502e --- /dev/null +++ b/services/runner/tests/unit/otel-usage-context-size.test.ts @@ -0,0 +1,174 @@ +/** + * Unit tests for how the sandbox-agent ACP tracer treats `usage_update`. + * + * ACP's `usage_update.used` is the agent's context-window occupancy, not the tokens a run + * spent. Reporting it as a token total produced runs shaped `input 0 / output 0 / total + * ` with no cost — a number a reader cannot tell apart from a real total. These + * tests pin that the context size never reaches a `usage` event, `usage()`, or a span, while + * the harness-reported split (delivered through `setUsage`) still does. + * + * Spans export over OTLP from a module-level provider, so we spy on the OTel API tracer and + * capture what each span records (same approach as otel-skills-error.test.ts). + * + * Run: pnpm test (or: pnpm exec vitest run tests/unit/otel-usage-context-size.test.ts) + */ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { trace, type Span } from "@opentelemetry/api"; + +import { createSandboxAgentOtel } from "../../src/tracing/otel.ts"; +import type { AgentEvent } from "../../src/protocol.ts"; + +interface FakeSpan { + name: string; + attributes: Record; +} + +/** Replace the OTel tracer so every span built records into a captured array. */ +function spyTracer(): FakeSpan[] { + const spans: FakeSpan[] = []; + const makeSpan = (name: string): Span => { + const span: FakeSpan = { name, attributes: {} }; + spans.push(span); + const api = { + setAttribute(key: string, value: unknown) { + span.attributes[key] = value; + return api; + }, + setAttributes(attrs: Record) { + Object.assign(span.attributes, attrs); + return api; + }, + recordException() {}, + setStatus() { + return api; + }, + end() {}, + spanContext() { + return { + traceId: "0".repeat(32), + spanId: "0".repeat(16), + traceFlags: 1, + }; + }, + isRecording: () => true, + addEvent: () => api, + updateName: () => api, + }; + return api as unknown as Span; + }; + vi.spyOn(trace, "getTracer").mockReturnValue({ + startSpan: (name: string) => makeSpan(name), + startActiveSpan: ((name: string, fn: (s: Span) => unknown) => + fn(makeSpan(name))) as any, + } as any); + return spans; +} + +const contextSizeUpdate = (used: number, cost?: number) => ({ + sessionUpdate: "usage_update", + used, + ...(cost === undefined ? {} : { cost: { amount: cost } }), +}); + +const usageKeys = (span: FakeSpan | undefined) => + Object.keys(span?.attributes ?? {}).filter((k) => + k.startsWith("gen_ai.usage."), + ); + +afterEach(() => { + vi.restoreAllMocks(); +}); + +describe("usage_update carries context size, not tokens", () => { + it("reports nothing when the stream knows only the context size", () => { + const spans = spyTracer(); + const emitted: AgentEvent[] = []; + const otel = createSandboxAgentOtel({ + harness: "claude", + model: "anthropic/claude-haiku", + emit: (e) => emitted.push(e), + }); + otel.start({ prompt: "hi" }); + otel.handleUpdate(contextSizeUpdate(63369)); + + expect(otel.usage()).toBeUndefined(); + expect(emitted.filter((e) => e.type === "usage")).toEqual([]); + + // The engine finds no harness split either, so nothing overrides the stream before finish. + otel.setUsage(undefined); + otel.finish(); + expect(usageKeys(spans.find((s) => s.name === "invoke_agent"))).toEqual([]); + expect( + emitted.some((e) => e.type === "usage"), + "no usage event anywhere in the run", + ).toBe(false); + }); + + it("keeps the stream cost, without inventing a token total from the context size", () => { + const spans = spyTracer(); + const otel = createSandboxAgentOtel({ + harness: "claude", + model: "anthropic/claude-haiku", + }); + otel.start({ prompt: "hi" }); + otel.handleUpdate(contextSizeUpdate(63369, 0.04)); + + expect(otel.usage()).toEqual({ input: 0, output: 0, total: 0, cost: 0.04 }); + otel.finish(); + const agentSpan = spans.find((s) => s.name === "invoke_agent"); + expect(agentSpan?.attributes["gen_ai.usage.cost"]).toBe(0.04); + expect(agentSpan?.attributes["gen_ai.usage.total_tokens"]).toBe(0); + }); + + it("emits a token-only usage event with no cost key when the harness priced nothing", () => { + const spans = spyTracer(); + const emitted: AgentEvent[] = []; + const otel = createSandboxAgentOtel({ + harness: "codex", + model: "openai-codex/gpt-5.5", + emit: (e) => emitted.push(e), + emitSpans: true, + }); + otel.start({ prompt: "hi" }); + // No `usage_update` at all — codex reports a token split but never a cost. + otel.setUsage({ input: 12, output: 3, total: 15 }); + otel.finish(); + + const usageEvent = emitted.find((e) => e.type === "usage") as Record< + string, + unknown + >; + expect(usageEvent).toEqual({ + type: "usage", + input: 12, + output: 3, + total: 15, + }); + expect("cost" in usageEvent).toBe(false); + + // Nothing downstream may see a zero: the agent span carries no cost either. + const agentSpan = spans.find((s) => s.name === "invoke_agent"); + expect(agentSpan?.attributes["gen_ai.usage.cost"]).toBeUndefined(); + const chatSpan = spans.find((s) => s.name.startsWith("chat")); + expect(chatSpan?.attributes["gen_ai.usage.total_tokens"]).toBe(15); + expect(chatSpan?.attributes["gen_ai.usage.cost"]).toBeUndefined(); + }); + + it("stamps the harness-reported split, which arrives through setUsage", () => { + const spans = spyTracer(); + const otel = createSandboxAgentOtel({ + harness: "claude", + model: "anthropic/claude-haiku", + }); + otel.start({ prompt: "hi" }); + otel.handleUpdate(contextSizeUpdate(63369, 0.04)); + otel.setUsage({ input: 12, output: 3, total: 15, cost: 0.04 }); + otel.finish(); + + const agentSpan = spans.find((s) => s.name === "invoke_agent"); + expect(agentSpan?.attributes["gen_ai.usage.input_tokens"]).toBe(12); + expect(agentSpan?.attributes["gen_ai.usage.output_tokens"]).toBe(3); + expect(agentSpan?.attributes["gen_ai.usage.total_tokens"]).toBe(15); + expect(agentSpan?.attributes["gen_ai.usage.cost"]).toBe(0.04); + }); +}); diff --git a/services/runner/tests/unit/sandbox-agent-usage.test.ts b/services/runner/tests/unit/sandbox-agent-usage.test.ts index a43b7c2bb9..4080c1ccfb 100644 --- a/services/runner/tests/unit/sandbox-agent-usage.test.ts +++ b/services/runner/tests/unit/sandbox-agent-usage.test.ts @@ -59,6 +59,75 @@ describe("mergePromptAndStreamUsage", () => { it("returns undefined when no usage was reported", () => { assert.equal(mergePromptAndStreamUsage({}, undefined), undefined); }); + + it("reports no usage when only the ACP context size is known", () => { + // `usage_update.used` is context-window occupancy, not tokens spent: with no harness + // split there is nothing honest to report, so the caller writes no usage attributes. + assert.equal( + mergePromptAndStreamUsage({ stopReason: "end_turn" }, { input: 0, output: 0, total: 63369 }), + undefined, + ); + }); + + it("keeps a stream cost even when the harness reported no token split", () => { + assert.deepEqual( + mergePromptAndStreamUsage({}, { input: 0, output: 0, total: 63369, cost: 0.04 }), + { input: 0, output: 0, total: 0, cost: 0.04 }, + ); + }); + + it("uses the harness split unchanged, never the stream total", () => { + assert.deepEqual( + mergePromptAndStreamUsage( + { usage: { inputTokens: 12, outputTokens: 3 } }, + { input: 0, output: 0, total: 63369 }, + ), + { input: 12, output: 3, total: 15 }, + ); + }); + + it("keeps a half-reported split (output only) as the total", () => { + assert.deepEqual( + mergePromptAndStreamUsage({ usage: { outputTokens: 5 } }, undefined), + { input: 0, output: 5, total: 5 }, + ); + }); + + it("omits cost entirely when the harness reported none", () => { + // A substituted `0` would say the run was measured and free. Consumers (the SDK's + // `record_usage`, the Vercel projection) read the key's PRESENCE as the measurement, so + // an unpriced harness — codex among them — has to leave it off. + const merged = mergePromptAndStreamUsage( + { usage: { inputTokens: 12, outputTokens: 3 } }, + undefined, + ); + assert.deepEqual(merged, { input: 12, output: 3, total: 15 }); + assert.equal("cost" in merged!, false); + }); + + it("keeps a reported zero cost, which is a measurement", () => { + // A free model or a fully cached turn really does cost 0. Now that absence is expressible, + // a reported zero must survive as the number it is. + const merged = mergePromptAndStreamUsage( + { usage: { inputTokens: 12, outputTokens: 3 } }, + { input: 0, output: 0, total: 0, cost: 0 }, + ); + assert.deepEqual(merged, { input: 12, output: 3, total: 15, cost: 0 }); + }); + + it("keeps a reported zero cost even with no token split at all", () => { + assert.deepEqual( + mergePromptAndStreamUsage({}, { input: 0, output: 0, total: 0, cost: 0 }), + { input: 0, output: 0, total: 0, cost: 0 }, + ); + }); + + it("falls back to already-known stream tokens, which only a prior setUsage can set", () => { + assert.deepEqual( + mergePromptAndStreamUsage({}, { input: 8, output: 2, total: 10, cost: 0.01 }), + { input: 8, output: 2, total: 10, cost: 0.01 }, + ); + }); }); describe("resolveRunUsage", () => { @@ -79,4 +148,17 @@ describe("resolveRunUsage", () => { { input: 3, output: 4, total: 7, cost: 0.03 }, ); }); + + it("resolves to no usage when neither the writeback nor the harness reported a split", async () => { + assert.equal( + await resolveRunUsage({ + sandbox: {}, + usageOutPath: undefined, + isDaytona: false, + promptResult: { stopReason: "end_turn" }, + streamUsage: { input: 0, output: 0, total: 63369 }, + }), + undefined, + ); + }); });