From 6477d28360497d830b107f9842c9db50259de8f7 Mon Sep 17 00:00:00 2001 From: Mahmoud Mabrouk Date: Mon, 3 Aug 2026 23:23:44 +0200 Subject: [PATCH] fix(runner): let one span own each token observation The runner published the same token count twice inside one OTLP batch: on the per-call `chat` leaf spans, and again as the run total on their `invoke_agent` ancestor. `gen_ai.usage.*_tokens` is the incremental bucket by contract, so ingest read both as separate observations. Nothing showed it, because the roll-up never ran on the runner's batch at all. The span tree was seeded only from spans with no parent, and the runner's batch has none. Once the roll-up runs, it adds the parent's own incremental value on top of the children it just summed, and the agent span reports exactly twice the real count. Measured live before this change, on a four-turn Pi run: invoke_agent incremental 7776 cumulative 15552 turn 0 -> chat 1850 turn 1 -> chat 1926 turn 2 -> chat 1983 turn 3 -> chat 2017 The four leaves sum to 7776, and the parent reported 15552. The producer owns this, not the roll-up. Adding a span's own incremental value to its children's cumulative is the correct general definition, and making the API skip a span by name would trade one bug for a wrong rule that breaks any producer whose parent span legitimately makes its own model call. So the token attributes come off `invoke_agent` in both tracers. `gen_ai.usage.cost` stays, because ingest maps it to the cumulative bucket, which makes it the explicit summary a parent should carry rather than a repeated incremental one. Both harnesses had the same shape. Pi gave each turn's `chat` span that turn's usage and the agent span their sum, so the error scaled with turn count. Claude ACP creates one `chat` leaf per prompt in `start()` and stamped the same run total on both it and the agent span. Moving ownership to the leaf loses nothing in either case. Turn spans were already clean; they never carried usage. Live after this change: a one-turn run reports 1739 against a leaf sum of 1739, a three-turn run reports 5413 against 1777 plus 1809 plus 1827, and a four-turn run reports 7414 against a leaf sum of 7414. Cost survives on every run. Tests: the full runner suite passes, 1,490 tests across 98 files, and `tsc --noEmit` is clean. Four new tests assert, for both tracers, that no span with children carries token attributes and that the leaf sum equals the run total. Claude-Session: https://claude.ai/code/session_01RkWWQUNNzRbaB5jnCAdjYA --- services/runner/src/tracing/otel.ts | 74 +++-- .../unit/otel-usage-context-size.test.ts | 15 +- .../tests/unit/otel-usage-ownership.test.ts | 267 ++++++++++++++++++ 3 files changed, 331 insertions(+), 25 deletions(-) create mode 100644 services/runner/tests/unit/otel-usage-ownership.test.ts diff --git a/services/runner/src/tracing/otel.ts b/services/runner/src/tracing/otel.ts index 2b173b8cb2..5e85eec3d0 100644 --- a/services/runner/src/tracing/otel.ts +++ b/services/runner/src/tracing/otel.ts @@ -560,6 +560,19 @@ function lastAssistantText(messages: any): string { return ""; } +/** + * Stamp a run's cost, and only its cost, on a PARENT span (`invoke_agent`). + * + * A parent must not repeat its children's `gen_ai.usage.*_tokens`: ingest maps those to the + * incremental bucket, which exactly one span may own, and the parent's own token total is + * the roll-up of the leaves below it. `gen_ai.usage.cost` is different — ingest maps it to + * the cumulative bucket, so it states the subtree total it actually is. + */ +function stampRunCost(span: Span, cost: number | undefined): void { + if (cost == null || cost <= 0) return; + span.setAttribute("gen_ai.usage.cost", cost); +} + /** Fill an LLM span from a finished assistant message (model, tokens, finish, output). */ /** Returns the error message when the assistant turn failed (stopReason/errorMessage), else * undefined — so the caller can emit a matching `error` event, not just stamp the span. */ @@ -624,7 +637,7 @@ export interface AgentaOtel { /** Flush this run's trace to Agenta. Await before the process/response ends. */ flush: () => Promise; /** Run totals (tokens + cost) summed across turns, for roll-up onto the parent. */ - usage: () => { input: number; output: number; total: number; cost: number }; + usage: () => AgentUsage; } /** @@ -665,6 +678,9 @@ export function createAgentaOtel( // (the agent and workflow spans are exported in separate OTLP batches, so Agenta's // per-batch cumulative roll-up cannot bridge them on its own). const runUsage = { input: 0, output: 0, total: 0, cost: 0 }; + // Whether ANY turn reported a cost. Without it a run the harness never priced is + // indistinguishable from one it priced at zero, and the sum below would report the second. + let costReported = false; function accumulateUsage(msg: any): void { const u = msg?.usage; @@ -674,7 +690,10 @@ export function createAgentaOtel( runUsage.input += input; runUsage.output += output; runUsage.total += u.totalTokens ?? input + output; - if (u.cost?.total != null) runUsage.cost += u.cost.total; + if (u.cost?.total != null) { + runUsage.cost += u.cost.total; + costReported = true; + } } const register = (pi: ExtensionAPI): void => { @@ -828,20 +847,15 @@ export function createAgentaOtel( lastAssistantText(event?.messages), config.captureContent, ); - // Stamp the run total on the agent span so it shows the agent's tokens/cost even - // though Agenta cannot roll the per-turn LLM spans up across batches. - if (runUsage.total > 0) { - agentSpan.setAttribute("gen_ai.usage.input_tokens", runUsage.input); - agentSpan.setAttribute("gen_ai.usage.output_tokens", runUsage.output); - agentSpan.setAttribute("gen_ai.usage.prompt_tokens", runUsage.input); - agentSpan.setAttribute( - "gen_ai.usage.completion_tokens", - runUsage.output, - ); - agentSpan.setAttribute("gen_ai.usage.total_tokens", runUsage.total); - if (runUsage.cost > 0) - agentSpan.setAttribute("gen_ai.usage.cost", runUsage.cost); - } + // OWNERSHIP: exactly one span owns each incremental observation. The per-turn `chat` + // spans already carry this run's tokens, and Agenta ingests `gen_ai.usage.*_tokens` + // as INCREMENTAL — repeating the run total here would make the parent's roll-up add + // it on top of the very tokens it summed, reporting twice the real count. The agent + // span's token total is the roll-up of its turns (same batch, so it is complete). + // Cost is the exception by contract, not by accident: `gen_ai.usage.cost` ingests as + // an explicitly CUMULATIVE subtree total, which is what a run total is, and it is the + // harness's billed figure rather than a recompute — so the parent keeps carrying it. + stampRunCost(agentSpan, costReported ? runUsage.cost : undefined); agentSpan.end(); agentSpan = undefined; agentCtx = undefined; @@ -853,10 +867,21 @@ export function createAgentaOtel( register, config, flush: () => flushTrace(config.traceId, config.redactor, runId), - usage: () => ({ ...runUsage }), + // The usage writeback (`extensions/agenta.ts`) serializes this straight onto the wire, so + // an unreported cost has to leave the key off rather than ship the running sum's 0. + usage: () => (costReported ? { ...runUsage } : stripCost(runUsage)), }; } +/** Drop the cost key, so an unpriced run reads as unknown rather than measured-free. */ +function stripCost(usage: { + input: number; + output: number; + total: number; +}): AgentUsage { + return { input: usage.input, output: usage.output, total: usage.total }; +} + // --------------------------------------------------------------------------- // sandbox-agent / ACP tracer (one per run; state is closure-scoped) // --------------------------------------------------------------------------- @@ -1162,16 +1187,19 @@ export function createSandboxAgentOtel( } } + /** Stamp the full usage split on the LEAF model span, the one span that owns it. */ function stampUsage(span: Span, u: AgentUsage | undefined): void { // 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; + // stamping zeros would assert a run cost nothing and spent nothing. An ABSENT cost counts + // as no cost here, exactly like a reported 0, so it cannot carry a token-less record. + if (!u || (u.total <= 0 && !(u.cost != null && 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); span.setAttribute("gen_ai.usage.completion_tokens", u.output); span.setAttribute("gen_ai.usage.total_tokens", u.total); - if (u.cost > 0) span.setAttribute("gen_ai.usage.cost", u.cost); + if (u.cost != null && u.cost > 0) + span.setAttribute("gen_ai.usage.cost", u.cost); } function setUsage(finalUsage: AgentUsage | undefined): void { @@ -1624,7 +1652,11 @@ export function createSandboxAgentOtel( } if (agentSpan) { setOutput(agentSpan, text, capture); - stampUsage(agentSpan, usage); + // Tokens belong to the `chat` leaf stamped just above — this tracer's `usage` IS the + // run total and there is exactly one model span to own it. Stamping it here as well + // would double the agent span's rolled-up total; only the cumulative cost is a + // parent's to report. See stampRunCost. + stampRunCost(agentSpan, usage?.cost); agentSpan.end(); agentSpan = undefined; } diff --git a/services/runner/tests/unit/otel-usage-context-size.test.ts b/services/runner/tests/unit/otel-usage-context-size.test.ts index b2818a502e..4ad69a8844 100644 --- a/services/runner/tests/unit/otel-usage-context-size.test.ts +++ b/services/runner/tests/unit/otel-usage-context-size.test.ts @@ -117,7 +117,10 @@ describe("usage_update carries context size, not tokens", () => { 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); + // Cost is the only usage a parent reports; tokens belong to the `chat` leaf. + expect(usageKeys(agentSpan)).toEqual(["gen_ai.usage.cost"]); + const chatSpan = spans.find((s) => s.name.startsWith("chat")); + expect(chatSpan?.attributes["gen_ai.usage.total_tokens"]).toBe(0); }); it("emits a token-only usage event with no cost key when the harness priced nothing", () => { @@ -165,10 +168,14 @@ describe("usage_update carries context size, not tokens", () => { otel.setUsage({ input: 12, output: 3, total: 15, cost: 0.04 }); otel.finish(); + // The split lands on the model span that owns it (see otel-usage-ownership.test.ts). + const chatSpan = spans.find((s) => s.name.startsWith("chat")); + expect(chatSpan?.attributes["gen_ai.usage.input_tokens"]).toBe(12); + expect(chatSpan?.attributes["gen_ai.usage.output_tokens"]).toBe(3); + expect(chatSpan?.attributes["gen_ai.usage.total_tokens"]).toBe(15); + expect(chatSpan?.attributes["gen_ai.usage.cost"]).toBe(0.04); + 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/otel-usage-ownership.test.ts b/services/runner/tests/unit/otel-usage-ownership.test.ts new file mode 100644 index 0000000000..17180b8d7d --- /dev/null +++ b/services/runner/tests/unit/otel-usage-ownership.test.ts @@ -0,0 +1,267 @@ +/** + * Who owns a token observation in the span trees `tracing/otel.ts` builds. + * + * Agenta ingests `gen_ai.usage.*_tokens` as an INCREMENTAL measurement and rolls each span's + * cumulative total up from its children. A whole run ships in ONE OTLP batch, so a parent that + * repeats its children's totals gets them counted twice: `invoke_agent` used to report the run + * total that its own `chat` spans had already reported, and every agent run showed exactly + * twice its real token count (measured live: 3,843 real read as 7,686). + * + * INVARIANT these tests pin, for both tracers: a span that has children never carries a + * `gen_ai.usage.*_tokens` attribute. Only the leaf model spans do, and their sum is the run + * total. Cost is the deliberate exception — `gen_ai.usage.cost` ingests as an explicitly + * CUMULATIVE subtree total, so a parent may carry it and `invoke_agent` still does. + * + * Run: pnpm exec vitest run tests/unit/otel-usage-ownership.test.ts + */ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { trace, type Span } from "@opentelemetry/api"; + +import { + createAgentaOtel, + createSandboxAgentOtel, +} from "../../src/tracing/otel.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 TOKEN_KEYS = [ + "gen_ai.usage.input_tokens", + "gen_ai.usage.output_tokens", + "gen_ai.usage.prompt_tokens", + "gen_ai.usage.completion_tokens", + "gen_ai.usage.total_tokens", +]; + +const isLeafModelSpan = (span: FakeSpan) => span.name.startsWith("chat"); + +const tokenTotal = (span: FakeSpan) => + Number(span.attributes["gen_ai.usage.total_tokens"] ?? 0); + +/** + * The producer's half of the roll-up invariant: no span that has children reports a token + * measurement, so a parent's cumulative can only ever be the sum of its leaves. Returns that + * sum so a caller can compare it against the run total the harness reported. + */ +function assertOnlyLeavesOwnTokens(spans: FakeSpan[]): number { + for (const span of spans) { + if (isLeafModelSpan(span)) continue; + for (const key of TOKEN_KEYS) { + expect( + span.attributes[key], + `${span.name} must not report ${key}: its children already do`, + ).toBeUndefined(); + } + } + return spans + .filter(isLeafModelSpan) + .reduce((sum, s) => sum + tokenTotal(s), 0); +} + +/** Drive the Pi extension lifecycle by capturing the handlers it registers. */ +function piHandlers(otel: ReturnType) { + const handlers: Record Promise> = {}; + otel.register({ + on: (name: string, fn: (e: any, ctx?: any) => Promise) => { + handlers[name] = fn; + }, + } as any); + return handlers; +} + +const assistantMessage = (input: number, output: number, cost: number) => ({ + role: "assistant", + model: "gpt-5.6-luna", + provider: "openai", + content: "ok", + usage: { + input, + output, + totalTokens: input + output, + cost: { total: cost }, + }, +}); + +afterEach(() => { + vi.restoreAllMocks(); +}); + +describe("the ACP tracer stamps the run's tokens on one span only", () => { + it("gives the chat leaf the token split and the agent span only the cost", () => { + const spans = spyTracer(); + const otel = createSandboxAgentOtel({ + harness: "claude", + model: "anthropic/claude-haiku", + emitSpans: true, + }); + otel.start({ prompt: "hi" }); + otel.setUsage({ input: 3595, output: 248, total: 3843, cost: 0.0219 }); + otel.finish(); + + const agentSpan = spans.find((s) => s.name === "invoke_agent"); + const chatSpan = spans.find(isLeafModelSpan); + + expect(chatSpan?.attributes["gen_ai.usage.total_tokens"]).toBe(3843); + expect(agentSpan?.attributes["gen_ai.usage.cost"]).toBe(0.0219); + + // The roll-up over this batch yields 3,843 — the real total, not 7,686. + expect(assertOnlyLeavesOwnTokens(spans)).toBe(3843); + }); + + it("keeps the turn span free of usage so an intermediate level cannot repeat it", () => { + const spans = spyTracer(); + const otel = createSandboxAgentOtel({ + harness: "claude", + model: "anthropic/claude-haiku", + emitSpans: true, + }); + otel.start({ prompt: "hi" }); + otel.setUsage({ input: 100, output: 20, total: 120, cost: 0.01 }); + otel.finish(); + + const turnSpan = spans.find((s) => s.name === "turn 0"); + expect(turnSpan).toBeDefined(); + expect( + Object.keys(turnSpan?.attributes ?? {}).filter((k) => + k.startsWith("gen_ai.usage."), + ), + ).toEqual([]); + }); +}); + +describe("the Pi tracer stamps each turn's tokens on that turn's chat span", () => { + it("rolls a multi-turn run up to the real total exactly once", async () => { + const spans = spyTracer(); + const otel = createAgentaOtel({ + captureContent: false, + requestModel: "gpt-5.6-luna", + }); + const handlers = piHandlers(otel); + const turns = [ + [1600, 250], + [1700, 226], + [1750, 233], + [1800, 217], + ]; + + await handlers["before_agent_start"]?.({ prompt: "hi" }); + await handlers["agent_start"]?.({}); + for (const [index, [input, output]] of turns.entries()) { + await handlers["turn_start"]?.({ turnIndex: index }); + await handlers["before_provider_request"]?.({}, {}); + await handlers["message_end"]?.({ + message: assistantMessage(input, output, 0.001), + }); + await handlers["turn_end"]?.({}); + } + await handlers["agent_end"]?.({ messages: [] }); + + const expectedTotal = turns.reduce((sum, [i, o]) => sum + i + o, 0); + expect(otel.usage().total).toBe(expectedTotal); + + const chatSpans = spans.filter(isLeafModelSpan); + expect(chatSpans).toHaveLength(turns.length); + expect(assertOnlyLeavesOwnTokens(spans)).toBe(expectedTotal); + + const agentSpan = spans.find((s) => s.name === "invoke_agent"); + expect(agentSpan?.attributes["gen_ai.usage.cost"]).toBeCloseTo(0.004, 10); + }); + + it("still reports the run's cost on the agent span when a turn reported none", async () => { + const spans = spyTracer(); + const otel = createAgentaOtel({ + captureContent: false, + requestModel: "gpt-5.6-luna", + }); + const handlers = piHandlers(otel); + + await handlers["before_agent_start"]?.({ prompt: "hi" }); + await handlers["agent_start"]?.({}); + await handlers["turn_start"]?.({ turnIndex: 0 }); + await handlers["before_provider_request"]?.({}, {}); + await handlers["message_end"]?.({ + message: { + role: "assistant", + model: "gpt-5.6-luna", + usage: { input: 10, output: 5, totalTokens: 15 }, + }, + }); + await handlers["turn_end"]?.({}); + await handlers["agent_end"]?.({ messages: [] }); + + const agentSpan = spans.find((s) => s.name === "invoke_agent"); + expect(agentSpan?.attributes["gen_ai.usage.cost"]).toBeUndefined(); + expect(assertOnlyLeavesOwnTokens(spans)).toBe(15); + + // The same record is serialized to the usage writeback the engine reads back and returns + // on the wire. An unpriced run must omit the key: a `0` there reads as "measured, free". + const usage = otel.usage(); + expect(usage).toEqual({ input: 10, output: 5, total: 15 }); + expect("cost" in usage).toBe(false); + }); + + it("reports a cost of zero when a turn actually priced the run at zero", async () => { + spyTracer(); + const otel = createAgentaOtel({ + captureContent: false, + requestModel: "gpt-5.6-luna", + }); + const handlers = piHandlers(otel); + + await handlers["before_agent_start"]?.({ prompt: "hi" }); + await handlers["agent_start"]?.({}); + await handlers["turn_start"]?.({ turnIndex: 0 }); + await handlers["before_provider_request"]?.({}, {}); + await handlers["message_end"]?.({ + message: assistantMessage(10, 5, 0), + }); + await handlers["turn_end"]?.({}); + await handlers["agent_end"]?.({ messages: [] }); + + // A free model is a measurement, not an absence — it survives as the 0 it is. + expect(otel.usage()).toEqual({ input: 10, output: 5, total: 15, cost: 0 }); + }); +});