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

[feiiiiii5]

Closes #916.

Summary

The root chat span was carrying only the last iteration's gen_ai.usage.input_tokens / gen_ai.usage.output_tokens instead of the documented cross-iteration roll-up. Per docs/advanced/otel.md:

The root span rolls up gen_ai.usage.input_tokens and gen_ai.usage.output_tokens across all iterations.

Root cause

TextEngine.handleRunFinishedEvent(chunk) overwrote this.finishedEvent = chunk each iteration, and beginIteration() reset this.finishedEvent = null. The terminal hook then passed usage: this.finishedEvent?.usage to runOnFinish — i.e. only the final iteration's RUN_FINISHED.usage. Any multi-iteration run (e.g. tool call → final answer) under-reported the root span and any middleware reading FinishInfo.usage saw the same truncated total.

Adapters can't compensate: each iteration is a fresh chatStream() call on a stateless adapter, so a RUN_FINISHED chunk can only ever carry that one API call's usage.

Per-iteration data was correct — onUsage fired with each iteration's own (incremental) usage, iteration spans got the right values, and the gen_ai.client.token.usage histogram summed correctly. Only the root span's roll-up (and FinishInfo.usage for any middleware relying on it) under-reported multi-iteration runs.

Fix

Add a separate accumulatedUsage field on TextEngine that:

• Survives beginIteration()'s per-iteration reset (reset only at run start, via field initialization — TextEngine is instantiated per chat() call)
• Accumulates every usage-bearing RUN_FINISHED chunk via the new shared accumulateTokenUsage helper (the same summation onUsage consumers already do)
• Is what onFinish receives, with finishedEvent?.usage as a fallback only when no RUN_FINISHED chunk carried usage at all (preserves prior undefined behavior for adapters that never report usage)

accumulateTokenUsage helper

Sums core counts (promptTokens, completionTokens, totalTokens) plus every optional numeric field:

• cost, durationSeconds, unitsBilled
• promptTokensDetails (cachedTokens, cacheWriteTokens, audioTokens, videoTokens, imageTokens, textTokens, documentTokens)
• completionTokensDetails (reasoningTokens, audioTokens, videoTokens, imageTokens, textTokens, documentTokens)
• costDetails (upstreamCost, upstreamInputCost, upstreamOutputCost)

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-shaped and not generally summable, so the latest value wins — matching how onUsage consumers see the most recent per-iteration bag.

What's unchanged

• Per-iteration span attributes — otelMiddleware.onChunk's RUN_FINISHED handler still sets usageAttributes(chunk.usage) on the iteration span (each iteration's incremental values)
• gen_ai.client.token.usage histogram — otelMiddleware.onUsage still records per-iteration tokens; the roll-up lives on the root span only
• onUsage callback shape — still fires per-iteration with that iteration's incremental usage, not the roll-up
• info.finishReason — still the last iteration's finishReason

Test plan

Red→green verification (both tests failed with 700 !== 1200 before the fix, pass after):

• tests/middleware.test.ts — "rolls up usage across iterations in FinishInfo (root span contract)": two-iteration scripted adapter (500/50 → 700/80), asserts info.usage.promptTokens === 1200, completionTokens === 130, totalTokens === 1330, and that onUsage still fires per-iteration with incremental values
• tests/middleware.test.ts — "rolls up cost and detail breakdowns across iterations": asserts promptTokensDetails.cachedTokens, completionTokensDetails.reasoningTokens, cost, and costDetails.upstreamCost/upstreamInputCost/upstreamOutputCost all sum correctly
• tests/middlewares/otel.test.ts — "rolls up usage across iterations onto the root span (issue otelMiddleware root chat span carries last-iteration usage, not the documented cross-iteration roll-up #916)": asserts root span attributes (gen_ai.usage.input_tokens === 1200), per-iteration span attributes (500 and 700), and histogram records (4 records summing to 1200/130)

Full test suite

Test Files  63 passed (63)
     Tests  1139 passed (1139)
  Duration  2.17s

Plus tsc --noEmit clean.

Repro verified

The reporter's repro (two-iteration scripted adapter, 500/50 → 700/80) produced 700 80 on the root span before this fix. With this fix, the root span now carries 1200 130, matching the documented contract.

Summary by CodeRabbit

• Bug Fixes

  • Corrected token-usage reporting for multi-step agent runs by rolling up onFinish usage across all iterations.
  • Ensured OpenTelemetry root chat spans report combined token usage totals, while iteration-level metrics remain incremental.
  • Updated rollup behavior to accurately sum token counts and associated cost/breakdown fields across agent-loop finishes.

• Documentation

  • Clarified the onFinish FinishInfo.usage contract (rolled-up agent-loop usage, undefined when absent) and how it relates to onUsage.
  • Documented OpenTelemetry usage accumulation semantics in advanced middleware/OTel guidance.

• Tests

  • Added regression coverage for onFinish usage rollups and OTel duration/usage histogram alignment across iterations.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 641f2279-3e72-43d4-bf0c-160fc0219f0a

📥 Commits

Reviewing files that changed from the base of the PR and between a48a2bf and d783764.

📒 Files selected for processing (1)

• docs/advanced/middleware.md

🚧 Files skipped from review as they are similar to previous changes (1)

• docs/advanced/middleware.md

---

📝 Walkthrough

Walkthrough

This change accumulates RUN_FINISHED.usage across agent-loop iterations, passes the total to onFinish, and records the total on the root OpenTelemetry span. Optional usage details are merged, while per-iteration callbacks and telemetry remain incremental.

Changes

Usage rollup

Layer / File(s)	Summary
Token usage accumulator packages/ai/src/utilities/usage.ts	Adds accumulateTokenUsage to sum core and optional numeric usage fields, merge detail bags, and retain the latest provider usage details.
Engine and middleware integration packages/ai/src/activities/chat/index.ts, packages/ai/tests/middleware.test.ts, packages/ai/tests/middlewares/otel.test.ts, docs/advanced/middleware.md, docs/advanced/otel.md, .changeset/ai-otel-root-span-usage-rollup.md	TextEngine accumulates usage across iterations and supplies it to onFinish; regression tests and documentation cover root-span totals and unchanged per-iteration reporting.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Sequence Diagram(s)

sequenceDiagram
  participant TextEngine
  participant accumulateTokenUsage
  participant MiddlewareOnFinish
  participant RootChatSpan
  TextEngine->>accumulateTokenUsage: accumulate each RUN_FINISHED usage
  accumulateTokenUsage-->>TextEngine: return rolled-up TokenUsage
  TextEngine->>MiddlewareOnFinish: invoke onFinish with FinishInfo.usage
  MiddlewareOnFinish->>RootChatSpan: set rolled-up usage attributes

Loading

Possibly related PRs

• TanStack/ai#747: Updates OpenTelemetry usage attributes derived from onFinish usage.
• TanStack/ai#789: Forwards adapter usage into RUN_FINISHED events that this change accumulates.

Suggested reviewers: alemtuzlak

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The PR has strong technical detail, but it does not follow the required template sections for Changes, Checklist, and Release Impact.	Rewrite the description using the repository template and include the required Changes, Checklist, and Release Impact sections.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is concise and accurately describes the main fix to roll up RUN_FINISHED usage onto the root span.
Linked Issues check	✅ Passed	The code and tests implement #916 by accumulating usage across iterations and passing the totals to FinishInfo.usage/root span.
Out of Scope Changes check	✅ Passed	The docs, tests, and changeset all support the usage-rollup fix; no unrelated code changes stand out.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

docs/advanced/middleware.md (1)

461-469: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win

Stale FinishInfo.usage table entry contradicts the updated bullet above.

Line 468 still describes usage as "the agent loop's last RUN_FINISHED.usage" — the exact pre-fix behavior this PR corrects (per the new bullet at Line 436, which says usage is now "rolled up across all agent-loop iterations"). This table row wasn't updated and will mislead readers about the current contract.

📝 Proposed fix

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

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@docs/advanced/middleware.md` around lines 461 - 469, Update the
FinishInfo.usage row in the onFinish documentation table to describe usage as
rolled up across all agent-loop iterations, matching the updated contract and
nearby bullet. Preserve the optionality and clarification that finalization
tokens are excluded and should be observed via onUsage.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Outside diff comments:
In `@docs/advanced/middleware.md`:
- Around line 461-469: Update the FinishInfo.usage row in the onFinish
documentation table to describe usage as rolled up across all agent-loop
iterations, matching the updated contract and nearby bullet. Preserve the
optionality and clarification that finalization tokens are excluded and should
be observed via onUsage.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 9986b9e8-1247-4a2f-8153-b1bd75f375b9

📥 Commits

Reviewing files that changed from the base of the PR and between bf870fa and a48a2bf.

📒 Files selected for processing (7)

• .changeset/ai-otel-root-span-usage-rollup.md
• docs/advanced/middleware.md
• docs/advanced/otel.md
• packages/ai/src/activities/chat/index.ts
• packages/ai/src/utilities/usage.ts
• packages/ai/tests/middleware.test.ts
• packages/ai/tests/middlewares/otel.test.ts