fix(ai): roll up RUN_FINISHED usage across iterations onto root span

.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`.

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
 

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

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<string, unknown> | undefined
   private eventToolNames?: Array<string>
   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(

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<TProviderDetails = ProviderUsageDetails>(
     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<T>(
+  a: T | undefined,
+  b: T | undefined,
+): T | undefined {
+  if (!a && !b) return undefined
+  const out: Record<string, number> = {}
+  const keys = new Set<string>([
+    ...(a ? Object.keys(a as Record<string, unknown>) : []),
+    ...(b ? Object.keys(b as Record<string, unknown>) : []),
+  ])
+  let hasAny = false
+  for (const k of keys) {
+    const va = (a as Record<string, number | undefined> | undefined)?.[k]
+    const vb = (b as Record<string, number | undefined> | 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<PromptTokensDetails>(
+    acc.promptTokensDetails,
+    delta.promptTokensDetails,
+  )
+  if (promptTokensDetails) {
+    merged.promptTokensDetails = promptTokensDetails
+  }
+
+  const completionTokensDetails = mergeNumericBag<CompletionTokensDetails>(
+    acc.completionTokensDetails,
+    delta.completionTokensDetails,
+  )
+  if (completionTokensDetails) {
+    merged.completionTokensDetails = completionTokensDetails
+  }
+
+  const costDetails = mergeNumericBag<UsageCostBreakdown>(
+    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
+}