diff --git a/.changeset/ai-otel-root-span-usage-rollup.md b/.changeset/ai-otel-root-span-usage-rollup.md new file mode 100644 index 000000000..7e765d6df --- /dev/null +++ b/.changeset/ai-otel-root-span-usage-rollup.md @@ -0,0 +1,29 @@ +--- +"@tanstack/ai": patch +--- + +Fix root `chat` span carrying only last iteration's usage (#916) + +`otelMiddleware`'s `onFinish` stamps `FinishInfo.usage` onto the root +`chat` span, but the chat engine previously passed only the final +iteration's `RUN_FINISHED.usage` to that hook. `handleRunFinishedEvent` +overwrote `finishedEvent` each iteration, and `beginIteration()` reset +it to `null`, so any multi-iteration run (e.g. tool call → final +answer) under-reported the root span's `gen_ai.usage.input_tokens` / +`gen_ai.usage.output_tokens` — and any other middleware reading +`FinishInfo.usage` saw the same truncated total. + +The fix introduces a separate `accumulatedUsage` field on `TextEngine` +that survives `beginIteration()`'s per-iteration reset, accumulates +every usage-bearing `RUN_FINISHED` chunk (the same summation `onUsage` +consumers already do), and is what `onFinish` receives. A new shared +helper `accumulateTokenUsage` sums core token counts plus every +optional numeric field (cost, cache/reasoning breakdowns, upstream cost +split, duration-based billing, units billed); provider-shaped +`providerUsageDetails` is preserved as latest-wins, matching how +`onUsage` consumers see the most recent per-iteration bag. + +Per-iteration span attributes and the `gen_ai.client.token.usage` +histogram are unchanged — they still carry each iteration's +incremental values. The roll-up lives on the root span only, matching +the documented contract in `docs/advanced/otel.md`. diff --git a/docs/advanced/middleware.md b/docs/advanced/middleware.md index a416e8f70..9fa802ce0 100644 --- a/docs/advanced/middleware.md +++ b/docs/advanced/middleware.md @@ -433,11 +433,9 @@ Exactly **one** terminal hook fires per `chat()` invocation. They are mutually e > **`onFinish` info fields and structured-output runs:** the `info` object reflects the **agent loop's** terminal state — finalization state is intentionally segregated to keep agent-loop semantics clean. > > - `info.content` — the agent loop's accumulated text. Finalization JSON deltas are **not** included here. The structured-output result is delivered via the `structured-output.complete` CUSTOM event, which middleware observes via `onChunk` (with `ctx.phase === 'structuredOutput'`). -> - `info.usage` — the agent loop's last `RUN_FINISHED.usage`. For a tools-less structured-output run (no agent-loop iteration produces `RUN_FINISHED`), this is `undefined`. To capture finalization tokens, use `onUsage` — that hook fires for **every** `RUN_FINISHED` carrying usage, including the finalization call. +> - `info.usage` — usage rolled up across all agent-loop iterations' `RUN_FINISHED` chunks. Every numeric field on `TokenUsage` (input/output tokens, cost, cache/reasoning breakdowns, upstream cost split, etc.) is summed across iterations, so this carries the full-run total — the documented contract for the root `chat` span's `gen_ai.usage.*` attributes. For a tools-less structured-output run (no agent-loop iteration produces `RUN_FINISHED`), this is `undefined`. To capture finalization tokens separately, use `onUsage` — that hook fires for **every** `RUN_FINISHED` carrying usage, including the finalization call, with each iteration's incremental values. > - `info.finishReason` — the agent loop's last `finishReason`. `null` when no agent-loop iteration produced `RUN_FINISHED` (e.g. a tools-less structured-output run). > - `info.duration` — wall-clock duration of the entire `chat()` invocation, including finalization. -> -> To aggregate usage across the whole run, accumulate from `onUsage` callbacks rather than relying on `info.usage`. ```typescript import { type ChatMiddleware } from "@tanstack/ai"; @@ -467,7 +465,7 @@ The `info` object for `onFinish` (`FinishInfo`): | `finishReason` | `string \| null` | The agent loop's last `finishReason`. `null` when no agent-loop iteration produced `RUN_FINISHED` (e.g. a tools-less `chat({ outputSchema })` run). | | `duration` | `number` | Total run duration in milliseconds, including any structured-output finalization. | | `content` | `string` | The agent loop's accumulated text content. Does **not** include finalization JSON deltas — for that, observe the `structured-output.complete` CUSTOM event via `onChunk`. | -| `usage` | `{ promptTokens; completionTokens; totalTokens } \| undefined` | **Optional.** The agent loop's last `RUN_FINISHED.usage`. **Does not include finalization tokens** — use `onUsage` to observe those. Always guard with `if (info.usage)` or `info.usage?.`. | +| `usage` | `TokenUsage \| undefined` | **Optional.** Usage rolled up across all agent-loop iterations' `RUN_FINISHED` chunks. **Does not include finalization tokens** — use `onUsage` to observe those. Always guard with `if (info.usage)` or `info.usage?.`. | ## Context Object diff --git a/docs/advanced/otel.md b/docs/advanced/otel.md index 9f6b78ee4..41c830272 100644 --- a/docs/advanced/otel.md +++ b/docs/advanced/otel.md @@ -90,6 +90,8 @@ Iteration spans are numbered (`#0`, `#1`, ...) so distinct iterations of the sam | tool | `gen_ai.tool.type` | `function` | | tool | `tanstack.ai.tool.outcome` | `success` / `error` | +**Usage accumulation.** The root span rolls up `gen_ai.usage.input_tokens` and `gen_ai.usage.output_tokens` across all iterations — every numeric field on `TokenUsage` (cost, total tokens, cache/reasoning breakdowns, upstream cost split, duration-based billing, units billed) is summed the same way. Per-iteration span attributes and the `gen_ai.client.token.usage` histogram are unaffected; they still carry each iteration's incremental values. The roll-up is what `onFinish`'s `info.usage` carries, so any middleware reading `FinishInfo.usage` sees the full-run total. + Usage attributes beyond input/output tokens are emitted only when the provider reports them, so spans stay clean otherwise. Cache and reasoning breakdowns use the official GenAI semconv names; `gen_ai.usage.cost` and `gen_ai.usage.total_tokens` are de-facto extensions consumed directly by backends like PostHog — without them, backends re-derive cost from their own price tables and lose cache discounts and gateway markup. Fields with no established convention (duration-based billing, the upstream cost split) are TanStack-namespaced. ### Metrics diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index 93e7a6481..95f5f14e2 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -11,6 +11,7 @@ import { stripToSpecMiddleware } from '../../strip-to-spec-middleware' import { streamToText } from '../../stream-to-response.js' import { resolveDebugOption } from '../../logger/resolve' import { EventType } from '../../types' +import { accumulateTokenUsage } from '../../utilities/usage' import { normalizeToolResult } from '../../utilities/tool-result' import { isProviderExecutedToolCall } from '../../utilities/provider-executed' import { LazyToolManager } from './tools/lazy-tool-manager' @@ -59,6 +60,7 @@ import type { StructuredOutputStream, TextMessageContentEvent, TextOptions, + TokenUsage, ToolCall, ToolCallArgsEvent, ToolCallEndEvent, @@ -531,6 +533,16 @@ class TextEngine< private eventOptions?: Record | undefined private eventToolNames?: Array private finishedEvent: RunFinishedEvent | null = null + /** + * Cross-iteration roll-up of `RUN_FINISHED.usage` from every iteration of + * the agent loop. `finishedEvent` (above) is overwritten each iteration to + * carry the latest `finishReason` / lazy-tool chunk context, so it can't + * hold the roll-up. `accumulatedUsage` is the running total passed to + * `onFinish` — the documented contract for the root `chat` span's + * `gen_ai.usage.*` attributes (docs/advanced/otel.md). Not reset by + * `beginIteration()`; reset only at run start. + */ + private accumulatedUsage: TokenUsage | null = null private earlyTermination = false private toolPhase: ToolPhaseResult = 'continue' private cyclePhase: CyclePhase = 'processText' @@ -909,7 +921,11 @@ class TextEngine< finishReason: this.lastFinishReason, duration: Date.now() - this.streamStartTime, content: this.accumulatedContent, - usage: this.finishedEvent?.usage, + // Cross-iteration roll-up — see `handleRunFinishedEvent`. Falls + // back to `finishedEvent?.usage` only when no RUN_FINISHED chunk + // carried usage at all (preserves the prior undefined behavior + // for adapters that never report usage). + usage: this.accumulatedUsage ?? this.finishedEvent?.usage, }) } } @@ -1237,6 +1253,18 @@ class TextEngine< private handleRunFinishedEvent(chunk: RunFinishedEvent): void { this.finishedEvent = chunk this.lastFinishReason = chunk.finishReason ?? null + // Accumulate per-iteration usage into the cross-iteration roll-up that + // `onFinish` will hand to middleware. `finishedEvent` is overwritten per + // iteration, so the usage must live in a separate field that survives + // `beginIteration()`'s reset. Per-iteration `onUsage` is still fired + // separately by `streamModelResponse` (with that iteration's incremental + // usage), so histograms and per-iteration span attributes are unaffected. + if (chunk.usage) { + this.accumulatedUsage = accumulateTokenUsage( + this.accumulatedUsage, + chunk.usage, + ) + } } private handleRunErrorEvent( diff --git a/packages/ai/src/utilities/usage.ts b/packages/ai/src/utilities/usage.ts index 06b08f108..e5e9f678c 100644 --- a/packages/ai/src/utilities/usage.ts +++ b/packages/ai/src/utilities/usage.ts @@ -1,4 +1,10 @@ -import type { ProviderUsageDetails, TokenUsage } from '../types' +import type { + CompletionTokensDetails, + PromptTokensDetails, + ProviderUsageDetails, + TokenUsage, + UsageCostBreakdown, +} from '../types' /** * Input parameters for building base TokenUsage. @@ -39,3 +45,128 @@ export function buildBaseUsage( totalTokens: input.totalTokens, } } + +/** + * Sum two optional finite numbers. Returns `undefined` only when neither side + * reports a finite value — when only one side is present, that value wins + * (so a single iteration's cost still surfaces on the roll-up). + */ +function sumOptional( + a: number | undefined, + b: number | undefined, +): number | undefined { + const hasA = typeof a === 'number' && Number.isFinite(a) + const hasB = typeof b === 'number' && Number.isFinite(b) + if (hasA && hasB) return a + b + if (hasA) return a + if (hasB) return b + return undefined +} + +/** + * Merge two optional numeric-detail bags field-by-field. Fields present on + * only one side are preserved; fields present on both are summed; the result + * is `undefined` only when both inputs are absent (no empty object littering + * the merged `TokenUsage`). + */ +function mergeNumericBag( + a: T | undefined, + b: T | undefined, +): T | undefined { + if (!a && !b) return undefined + const out: Record = {} + const keys = new Set([ + ...(a ? Object.keys(a as Record) : []), + ...(b ? Object.keys(b as Record) : []), + ]) + let hasAny = false + for (const k of keys) { + const va = (a as Record | undefined)?.[k] + const vb = (b as Record | undefined)?.[k] + const v = sumOptional(va, vb) + if (v !== undefined) { + out[k] = v + hasAny = true + } + } + return hasAny ? (out as unknown as T) : undefined +} + +/** + * Accumulate `delta` usage into a running `acc` total. Used by the chat + * engine to roll up per-iteration `RUN_FINISHED.usage` across the agent + * loop so `FinishInfo.usage` carries cross-iteration totals — the + * documented contract for the root `chat` span's `gen_ai.usage.*` + * attributes (see `docs/advanced/otel.md`). + * + * Core counts (`promptTokens`, `completionTokens`, `totalTokens`) are + * always summed. Optional numeric fields (cost, cache/reasoning + * breakdowns, duration-based billing, upstream cost split) sum when both + * sides report them; when only one side reports a field, that value is + * preserved so a single iteration's cost still surfaces on the roll-up. + * `providerUsageDetails` is provider-specific and not generally summable, + * so the latest value (from `delta`) wins — matching how `onUsage` + * consumers see the most recent per-iteration value. + * + * Returns a fresh object so callers can store the result without aliasing + * either input. + */ +export function accumulateTokenUsage( + acc: TokenUsage | null | undefined, + delta: TokenUsage, +): TokenUsage { + if (!acc) { + return { ...delta } + } + + const merged: TokenUsage = { + promptTokens: acc.promptTokens + delta.promptTokens, + completionTokens: acc.completionTokens + delta.completionTokens, + totalTokens: acc.totalTokens + delta.totalTokens, + } + + const cost = sumOptional(acc.cost, delta.cost) + if (cost !== undefined) merged.cost = cost + const durationSeconds = sumOptional( + acc.durationSeconds, + delta.durationSeconds, + ) + if (durationSeconds !== undefined) merged.durationSeconds = durationSeconds + const unitsBilled = sumOptional(acc.unitsBilled, delta.unitsBilled) + if (unitsBilled !== undefined) merged.unitsBilled = unitsBilled + + const promptTokensDetails = mergeNumericBag( + acc.promptTokensDetails, + delta.promptTokensDetails, + ) + if (promptTokensDetails) { + merged.promptTokensDetails = promptTokensDetails + } + + const completionTokensDetails = mergeNumericBag( + acc.completionTokensDetails, + delta.completionTokensDetails, + ) + if (completionTokensDetails) { + merged.completionTokensDetails = completionTokensDetails + } + + const costDetails = mergeNumericBag( + acc.costDetails, + delta.costDetails, + ) + if (costDetails) { + merged.costDetails = costDetails + } + + // providerUsageDetails is provider-shaped (not generally summable). + // Latest value wins — matches how `onUsage` consumers observe the most + // recent per-iteration bag. + if (delta.providerUsageDetails !== undefined) { + merged.providerUsageDetails = delta.providerUsageDetails + } else if (acc.providerUsageDetails !== undefined) { + merged.providerUsageDetails = acc.providerUsageDetails + } + + return merged +} diff --git a/packages/ai/tests/middleware.test.ts b/packages/ai/tests/middleware.test.ts index 4630468d1..e1c968de1 100644 --- a/packages/ai/tests/middleware.test.ts +++ b/packages/ai/tests/middleware.test.ts @@ -1723,6 +1723,161 @@ describe('chat() middleware', () => { expect(info.usage).toEqual(usage) }) + it('rolls up usage across iterations in FinishInfo (root span contract)', async () => { + // Regression guard for issue #916: the root `chat` span's + // `gen_ai.usage.input_tokens` / `gen_ai.usage.output_tokens` are + // documented to roll up across all iterations (docs/advanced/otel.md). + // The chat engine previously passed only the final iteration's usage + // to `onFinish` (overwriting `finishedEvent` per iteration), so the + // root span under-reported any multi-iteration run. + const onFinish = vi.fn() + const onUsage = vi.fn() + + // Iteration 1: 500/50 (tool call → tool_calls finishReason) + // Iteration 2: 700/80 (final answer → stop finishReason) + // Expected roll-up on root: 1200/130 + const { adapter } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.toolStart('tc-1', 'status'), + ev.toolArgs('tc-1', '{}'), + ev.toolEnd('tc-1', 'status', { input: {} }), + ev.runFinished('tool_calls', 'run-1', { + promptTokens: 500, + completionTokens: 50, + totalTokens: 550, + }), + ], + [ + ev.runStarted(), + ev.textContent('done'), + ev.runFinished('stop', 'run-1', { + promptTokens: 700, + completionTokens: 80, + totalTokens: 780, + }), + ], + ], + }) + + const tool = serverTool('status', () => ({ ok: true })) + + const stream = chat({ + adapter, + messages: [{ role: 'user', content: 'go' }], + tools: [tool], + middleware: [{ name: 'capture', onFinish, onUsage }], + }) + await collectChunks(stream as AsyncIterable) + + expect(onFinish).toHaveBeenCalledOnce() + const info = onFinish.mock.calls[0]![1] + expect(info.usage).toBeDefined() + expect(info.usage!.promptTokens).toBe(1200) + expect(info.usage!.completionTokens).toBe(130) + expect(info.usage!.totalTokens).toBe(1330) + + // Per-iteration onUsage is unchanged — each call carries that + // iteration's incremental usage (not the roll-up). + expect(onUsage).toHaveBeenCalledTimes(2) + expect(onUsage.mock.calls[0]![1]).toEqual({ + promptTokens: 500, + completionTokens: 50, + totalTokens: 550, + }) + expect(onUsage.mock.calls[1]![1]).toEqual({ + promptTokens: 700, + completionTokens: 80, + totalTokens: 780, + }) + }) + + it('rolls up cost and detail breakdowns across iterations', async () => { + // Companion to the input/output roll-up test: the optional numeric + // fields (cost, cache/reasoning breakdowns, upstream cost split) must + // also accumulate so backends relying on `FinishInfo.usage` see full + // run totals — not just the final iteration's slice. + const onFinish = vi.fn() + + const { adapter } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.toolStart('tc-1', 'status'), + ev.toolArgs('tc-1', '{}'), + ev.toolEnd('tc-1', 'status', { input: {} }), + { + ...ev.runFinished('tool_calls', 'run-1', { + promptTokens: 500, + completionTokens: 50, + totalTokens: 550, + }), + usage: { + promptTokens: 500, + completionTokens: 50, + totalTokens: 550, + promptTokensDetails: { + cachedTokens: 80, + cacheWriteTokens: 10, + }, + completionTokensDetails: { reasoningTokens: 15 }, + cost: 0.01, + costDetails: { + upstreamCost: 0.008, + upstreamInputCost: 0.006, + upstreamOutputCost: 0.002, + }, + }, + }, + ], + [ + ev.runStarted(), + ev.textContent('done'), + { + ...ev.runFinished('stop', 'run-1', { + promptTokens: 700, + completionTokens: 80, + totalTokens: 780, + }), + usage: { + promptTokens: 700, + completionTokens: 80, + totalTokens: 780, + promptTokensDetails: { cachedTokens: 200 }, + completionTokensDetails: { reasoningTokens: 30 }, + cost: 0.02, + costDetails: { upstreamCost: 0.015 }, + }, + }, + ], + ], + }) + + const tool = serverTool('status', () => ({ ok: true })) + + const stream = chat({ + adapter, + messages: [{ role: 'user', content: 'go' }], + tools: [tool], + middleware: [{ name: 'capture', onFinish }], + }) + await collectChunks(stream as AsyncIterable) + + const usage = onFinish.mock.calls[0]![1].usage! + expect(usage.promptTokens).toBe(1200) + expect(usage.completionTokens).toBe(130) + expect(usage.totalTokens).toBe(1330) + // Detail fields sum when both iterations report them. + expect(usage.promptTokensDetails?.cachedTokens).toBe(280) + expect(usage.promptTokensDetails?.cacheWriteTokens).toBe(10) + expect(usage.completionTokensDetails?.reasoningTokens).toBe(45) + expect(usage.cost).toBeCloseTo(0.03, 10) + expect(usage.costDetails?.upstreamCost).toBeCloseTo(0.023, 10) + expect(usage.costDetails?.upstreamInputCost).toBe(0.006) + expect(usage.costDetails?.upstreamOutputCost).toBe(0.002) + }) + it('should accumulate content across multiple text chunks', async () => { const onFinish = vi.fn() diff --git a/packages/ai/tests/middlewares/otel.test.ts b/packages/ai/tests/middlewares/otel.test.ts index f7fa553ac..ef741c286 100644 --- a/packages/ai/tests/middlewares/otel.test.ts +++ b/packages/ai/tests/middlewares/otel.test.ts @@ -331,6 +331,85 @@ describe('otelMiddleware — duration histogram and rollup', () => { expect(root.attributes['tanstack.ai.iterations']).toBe(1) expect(root.ended).toBe(true) }) + + it('rolls up usage across iterations onto the root span (issue #916)', async () => { + // Regression for the documented contract in docs/advanced/otel.md: + // "The root span rolls up `gen_ai.usage.input_tokens` and + // `gen_ai.usage.output_tokens` across all iterations." + // + // Two-iteration scripted flow (tool_calls → stop), iteration 1 reports + // 500/50 and iteration 2 reports 700/80. The root `chat` span must + // carry the roll-up: 1200/130 — not just the final iteration's 700/80. + const { tracer, spans } = createFakeTracer() + const { meter, records } = createFakeMeter() + const mw = otelMiddleware({ tracer, meter }) + const ctx = makeCtx() + + // --- Iteration 1 (tool_calls) --- + await runToIterationStart(mw, ctx) + await mw.onChunk?.(ctx, { + ...ev.runFinished('tool_calls'), + model: 'scripted-model', + usage: { promptTokens: 500, completionTokens: 50, totalTokens: 550 }, + }) + await mw.onUsage?.(ctx, { + promptTokens: 500, + completionTokens: 50, + totalTokens: 550, + }) + + // --- Iteration 2 (stop) --- + ctx.iteration = 1 + await mw.onConfig?.(ctx, { messages: [], systemPrompts: [], tools: [] }) + await mw.onChunk?.(ctx, { + ...ev.runFinished('stop'), + model: 'scripted-model', + usage: { promptTokens: 700, completionTokens: 80, totalTokens: 780 }, + }) + await mw.onUsage?.(ctx, { + promptTokens: 700, + completionTokens: 80, + totalTokens: 780, + }) + + await mw.onFinish?.(ctx, { + finishReason: 'stop', + duration: 1250, + content: 'done', + // The engine now hands the cross-iteration roll-up here (1200/130). + usage: { promptTokens: 1200, completionTokens: 130, totalTokens: 1330 }, + }) + + const root = spans[0]! + expect(root.attributes['gen_ai.usage.input_tokens']).toBe(1200) + expect(root.attributes['gen_ai.usage.output_tokens']).toBe(130) + expect(root.attributes['gen_ai.usage.total_tokens']).toBe(1330) + expect(root.attributes['tanstack.ai.iterations']).toBe(2) + + // Per-iteration span attributes remain the per-iteration (incremental) + // values — the roll-up lives on the root span only. + const iter1 = spans[1]! + expect(iter1.attributes['gen_ai.usage.input_tokens']).toBe(500) + expect(iter1.attributes['gen_ai.usage.output_tokens']).toBe(50) + const iter2 = spans[2]! + expect(iter2.attributes['gen_ai.usage.input_tokens']).toBe(700) + expect(iter2.attributes['gen_ai.usage.output_tokens']).toBe(80) + + // The token histogram records each iteration's tokens (4 records: 2 + // iterations × input+output). Sum equals the roll-up. + const tokenRecords = records.filter( + (r) => r.name === 'gen_ai.client.token.usage', + ) + expect(tokenRecords).toHaveLength(4) + const totalInput = tokenRecords + .filter((r) => r.attributes!['gen_ai.token.type'] === 'input') + .reduce((sum, r) => sum + r.value, 0) + const totalOutput = tokenRecords + .filter((r) => r.attributes!['gen_ai.token.type'] === 'output') + .reduce((sum, r) => sum + r.value, 0) + expect(totalInput).toBe(1200) + expect(totalOutput).toBe(130) + }) }) describe('otelMiddleware — full usage emission', () => {