fix(ai): roll up OpenTelemetry usage across iterations

.changeset/otel-root-usage-rollup.md

@@ -0,0 +1,5 @@
+---
+'@tanstack/ai': patch
+---
+
+Fix OpenTelemetry root spans to report usage across all chat iterations, including error and abort exits.

packages/ai/src/middlewares/otel.ts

@@ -28,6 +28,7 @@ import type {
   GenerationMiddleware,
   GenerationMiddlewareContext,
 } from '../activities/middleware/types'
+import type { TokenUsage } from '../types'
 
 /**
  * Scope (role) of an OTel span emitted by this middleware.
@@ -138,13 +139,38 @@ interface RequestState {
    * from the (base-shaped) finish info, which doesn't carry it.
    */
   lastFinishReason: string | null
+  rootUsageAttributes: Record<string, number> | null
+  rootUsageApplied: boolean
 }
 
 const stateByCtx = new WeakMap<ChatMiddlewareContext, RequestState>()
 
 const DEFAULT_MAX_CONTENT_LENGTH = 100_000
 const REDACTION_FAILED_SENTINEL = '[redaction_failed]'
 
+function accumulateUsageAttributes(
+  current: Record<string, number> | null,
+  usage: TokenUsage,
+): Record<string, number> {
+  const accumulated = current ?? {}
+  for (const [key, value] of Object.entries(usageAttributes(usage))) {
+    if (typeof value === 'number') {
+      accumulated[key] = (accumulated[key] ?? 0) + value
+    }
+  }
+  return accumulated
+}
+
+function applyRootUsage(state: RequestState, fallbackUsage?: TokenUsage): void {
+  if (state.rootUsageApplied) return
+
+  const attributes =
+    state.rootUsageAttributes ??
+    (fallbackUsage ? usageAttributes(fallbackUsage) : null)
+  if (attributes) state.rootSpan.setAttributes(attributes)
+  state.rootUsageApplied = true
+}
+
 function serializeContent(content: unknown): string {
   if (typeof content === 'string') return content
   if (!Array.isArray(content)) return ''
@@ -395,6 +421,8 @@ export function otelMiddleware(
           assistantTextBufferTruncated: false,
           startTime: Date.now(),
           lastFinishReason: null,
+          rootUsageAttributes: null,
+          rootUsageApplied: false,
         })
       })
     },
@@ -650,6 +678,11 @@ export function otelMiddleware(
         const state = stateByCtx.get(chatCtx)
         if (!state) return
 
+        state.rootUsageAttributes = accumulateUsageAttributes(
+          state.rootUsageAttributes,
+          usage,
+        )
+
         // Always record the token histogram — metrics don't depend on having
         // an iteration span, and skipping here would drop metric data if an
         // adapter emits `onUsage` outside the iteration window.
@@ -889,6 +922,7 @@ export function otelMiddleware(
           })
         }
 
+        applyRootUsage(state)
         safeCall('otel.onSpanEnd', () =>
           onSpanEnd?.({ kind: 'chat', ctx: chatCtx }, state.rootSpan),
         )
@@ -970,6 +1004,7 @@ export function otelMiddleware(
           })
         }
 
+        applyRootUsage(state)
         safeCall('otel.onSpanEnd', () =>
           onSpanEnd?.({ kind: 'chat', ctx: chatCtx }, state.rootSpan),
         )
@@ -1027,9 +1062,6 @@ export function otelMiddleware(
           })
         }
 
-        if (info.usage) {
-          state.rootSpan.setAttributes(usageAttributes(info.usage))
-        }
         if (state.lastFinishReason) {
           state.rootSpan.setAttribute('gen_ai.response.finish_reasons', [
             state.lastFinishReason,
@@ -1040,6 +1072,7 @@ export function otelMiddleware(
           state.iterationCount,
         )
 
+        applyRootUsage(state, info.usage)
         safeCall('otel.onSpanEnd', () =>
           onSpanEnd?.({ kind: 'chat', ctx: chatCtx }, state.rootSpan),
         )

packages/ai/tests/middlewares/otel.test.ts

@@ -56,6 +56,30 @@ async function runToIterationStart(
   })
 }
 
+async function runUsageIterations(
+  mw: ChatMiddleware,
+  ctx: ChatMiddlewareContext,
+  usages: Array<TokenUsage | undefined>,
+) {
+  await mw.onStart?.(ctx)
+  for (const [iteration, usage] of usages.entries()) {
+    ctx.iteration = iteration
+    ctx.phase = 'beforeModel'
+    await mw.onConfig?.(ctx, {
+      messages: [],
+      systemPrompts: [],
+      tools: [],
+    })
+    await mw.onChunk?.(ctx, {
+      ...ev.runFinished(
+        iteration === usages.length - 1 ? 'stop' : 'tool_calls',
+      ),
+      ...(usage !== undefined ? { usage } : {}),
+    })
+    if (usage !== undefined) await mw.onUsage?.(ctx, usage)
+  }
+}
+
 class RateLimitError extends Error {
   override name = 'RateLimitError'
 }
@@ -286,6 +310,63 @@ describe('otelMiddleware — token histogram', () => {
   })
 })
 
+describe('otelMiddleware — root usage rollup', () => {
+  it('rolls up mixed multi-iteration usage without changing iteration data or metrics', async () => {
+    const usages: Array<TokenUsage | undefined> = [
+      {
+        promptTokens: 0,
+        completionTokens: 0,
+        totalTokens: 0,
+        promptTokensDetails: { cachedTokens: 0, cacheWriteTokens: 0 },
+        completionTokensDetails: { reasoningTokens: 0 },
+      },
+      undefined,
+      {
+        promptTokens: 100,
+        completionTokens: 10,
+        totalTokens: 110,
+        promptTokensDetails: { cachedTokens: 20, cacheWriteTokens: 5 },
+        completionTokensDetails: { reasoningTokens: 1 },
+      },
+      { promptTokens: 200, completionTokens: 20, totalTokens: 220 },
+    ]
+    const expected: TokenUsage = {
+      promptTokens: 300,
+      completionTokens: 30,
+      totalTokens: 330,
+      promptTokensDetails: { cachedTokens: 20, cacheWriteTokens: 5 },
+      completionTokensDetails: { reasoningTokens: 1 },
+    }
+    const { tracer, spans } = createFakeTracer()
+    const { meter, records } = createFakeMeter()
+    const mw = otelMiddleware({ tracer, meter })
+    const ctx = makeCtx()
+
+    await runUsageIterations(mw, ctx, usages)
+    await mw.onFinish?.(ctx, {
+      finishReason: 'stop',
+      duration: 10,
+      content: '',
+      usage: usages.at(-1)!,
+    })
+
+    const [root, ...iterations] = spans
+    expect(root!.attributes).toMatchObject(usageAttributes(expected))
+    expect(root!.attributes['tanstack.ai.iterations']).toBe(usages.length)
+    expect(iterations).toHaveLength(usages.length)
+    expect(
+      iterations.map((span) => span.attributes['gen_ai.usage.input_tokens']),
+    ).toEqual([0, undefined, 100, 200])
+
+    const tokenRecords = records.filter(
+      (record) => record.name === 'gen_ai.client.token.usage',
+    )
+    expect(tokenRecords.map((record) => record.value)).toEqual([
+      0, 0, 100, 10, 200, 20,
+    ])
+  })
+})
+
 describe('otelMiddleware — duration histogram and rollup', () => {
   it('records duration histogram on onFinish and rolls up tokens onto root', async () => {
     const { tracer, spans } = createFakeTracer()
@@ -833,9 +914,15 @@ describe('otelMiddleware — error and abort paths', () => {
       toolCallId?: string | undefined
       ended: boolean
     }> = []
+    let rootInputAtEnd: unknown
     const mw = otelMiddleware({
       tracer,
       onSpanEnd: (info, span) => {
+        if (info.kind === 'chat') {
+          rootInputAtEnd = (span as FakeSpan).attributes[
+            'gen_ai.usage.input_tokens'
+          ]
+        }
         seen.push({
           kind: info.kind,
           toolName: info.kind === 'tool' ? info.toolName : undefined,
@@ -847,6 +934,11 @@ describe('otelMiddleware — error and abort paths', () => {
     const ctx = makeCtx({ hasTools: true })
 
     await runToIterationStart(mw, ctx)
+    await mw.onUsage?.(ctx, {
+      promptTokens: 12,
+      completionTokens: 3,
+      totalTokens: 15,
+    })
     await mw.onBeforeToolCall?.(ctx, {
       toolCall: makeToolCall({ id: 'tc-err', function: { name: 'my_tool' } }),
       tool: undefined,
@@ -862,6 +954,7 @@ describe('otelMiddleware — error and abort paths', () => {
     const toolCall = seen.find((s) => s.kind === 'tool')!
     expect(toolCall.toolName).toBe('my_tool')
     expect(toolCall.toolCallId).toBe('tc-err')
+    expect(rootInputAtEnd).toBe(12)
   })
 
   it('onAbort fires onSpanEnd for iteration, open tool spans, then root — in depth-first order', async () => {
@@ -872,9 +965,15 @@ describe('otelMiddleware — error and abort paths', () => {
       toolCallId?: string | undefined
       ended: boolean
     }> = []
+    let rootInputAtEnd: unknown
     const mw = otelMiddleware({
       tracer,
       onSpanEnd: (info, span) => {
+        if (info.kind === 'chat') {
+          rootInputAtEnd = (span as FakeSpan).attributes[
+            'gen_ai.usage.input_tokens'
+          ]
+        }
         seen.push({
           kind: info.kind,
           toolName: info.kind === 'tool' ? info.toolName : undefined,
@@ -886,6 +985,11 @@ describe('otelMiddleware — error and abort paths', () => {
     const ctx = makeCtx({ hasTools: true })
 
     await runToIterationStart(mw, ctx)
+    await mw.onUsage?.(ctx, {
+      promptTokens: 8,
+      completionTokens: 2,
+      totalTokens: 10,
+    })
     await mw.onBeforeToolCall?.(ctx, {
       toolCall: makeToolCall({
         id: 'tc-abort',
@@ -901,6 +1005,7 @@ describe('otelMiddleware — error and abort paths', () => {
 
     expect(seen.map((s) => s.kind)).toEqual(['iteration', 'tool', 'chat'])
     expect(seen.every((s) => s.ended === false)).toBe(true)
+    expect(rootInputAtEnd).toBe(8)
   })
 
   it('onFinish sweeps dangling tool spans with outcome=unknown before closing the iteration span', async () => {

testing/e2e/src/routes/api.otel-usage.ts

@@ -1,8 +1,10 @@
 import { createFileRoute } from '@tanstack/react-router'
-import { chat, createChatOptions } from '@tanstack/ai'
+import { chat, createChatOptions, toolDefinition } from '@tanstack/ai'
 import { otelMiddleware } from '@tanstack/ai/middlewares/otel'
 import { createOpenaiChatCompletions } from '@tanstack/ai-openai'
 import { createOpenRouterText } from '@tanstack/ai-openrouter'
+import { z } from 'zod'
+import { createTextAdapter } from '@/lib/providers'
 import type {
   AttributeValue,
   Context,
@@ -13,6 +15,11 @@ import type {
 
 const LLMOCK_DEFAULT_BASE = process.env.LLMOCK_URL || 'http://127.0.0.1:4010'
 const DUMMY_KEY = 'sk-e2e-test-dummy-key'
+const weatherTool = toolDefinition({
+  name: 'get_weather',
+  description: 'Get weather',
+  inputSchema: z.object({ city: z.string() }),
+}).server(async ({ city }) => ({ city, temperature: 72, condition: 'sunny' }))
 
 interface CapturedSpan {
   name: string
@@ -121,28 +128,42 @@ export const Route = createFileRoute('/api/otel-usage')({
     handlers: {
       POST: async ({ request }) => {
         let provider = 'openai'
+        let testId: string | undefined
         try {
-          const body = (await request.json()) as { provider?: string }
+          const body = (await request.json()) as {
+            provider?: string
+            testId?: string
+          }
           if (typeof body.provider === 'string') provider = body.provider
+          if (typeof body.testId === 'string') testId = body.testId
         } catch {
           // No/invalid body — default provider.
         }
 
         const adapter =
-          provider === 'openrouter'
-            ? createOpenRouterText('openai/gpt-4o', DUMMY_KEY, {
-                serverURL: `${LLMOCK_DEFAULT_BASE}/openrouter-cost/v1`,
-              })
-            : createOpenaiChatCompletions('gpt-4o', DUMMY_KEY, {
-                baseURL: `${LLMOCK_DEFAULT_BASE}/openai-usage-details/v1`,
-              })
+          provider === 'tool-loop'
+            ? createTextAdapter('openai', undefined, undefined, testId).adapter
+            : provider === 'openrouter'
+              ? createOpenRouterText('openai/gpt-4o', DUMMY_KEY, {
+                  serverURL: `${LLMOCK_DEFAULT_BASE}/openrouter-cost/v1`,
+                })
+              : createOpenaiChatCompletions('gpt-4o', DUMMY_KEY, {
+                  baseURL: `${LLMOCK_DEFAULT_BASE}/openai-usage-details/v1`,
+                })
 
         const { tracer, spans } = createLocalCaptureTracer()
 
         try {
           for await (const _chunk of chat({
             ...createChatOptions({ adapter }),
-            messages: [{ role: 'user', content: 'hi' }],
+            messages: [
+              {
+                role: 'user',
+                content:
+                  provider === 'tool-loop' ? '[with-tool] run test' : 'hi',
+              },
+            ],
+            ...(provider === 'tool-loop' ? { tools: [weatherTool] } : {}),
             middleware: [otelMiddleware({ tracer })],
           })) {
             // Drain — the assertions live on the captured spans.

testing/e2e/tests/middleware.spec.ts

@@ -228,6 +228,46 @@ test.describe('Middleware Lifecycle', () => {
     }
   })
 
+  test('otel middleware rolls up usage across a tool loop', async ({
+    request,
+    testId,
+  }) => {
+    const res = await request.post('/api/otel-usage', {
+      data: { provider: 'tool-loop', testId },
+    })
+    expect(res.ok()).toBe(true)
+    const { ok, error, spans } = await res.json()
+    expect(error ?? null).toBeNull()
+    expect(ok).toBe(true)
+
+    const iterationSpans = spans.filter(
+      (span: any) => span.kind === SpanKind.CLIENT,
+    )
+    expect(iterationSpans).toHaveLength(2)
+    expect(iterationSpans.every((span: any) => span.ended)).toBe(true)
+
+    const rootSpans = spans.filter(
+      (span: any) =>
+        span.kind === SpanKind.INTERNAL &&
+        span.attributes['tanstack.ai.iterations'] === 2,
+    )
+    expect(rootSpans).toHaveLength(1)
+    expect(rootSpans[0].ended).toBe(true)
+
+    for (const key of [
+      'gen_ai.usage.input_tokens',
+      'gen_ai.usage.output_tokens',
+      'gen_ai.usage.total_tokens',
+    ]) {
+      expect(rootSpans[0].attributes[key]).toBe(
+        iterationSpans.reduce(
+          (total: number, span: any) => total + span.attributes[key],
+          0,
+        ),
+      )
+    }
+  })
+
   test('otel middleware emits total/cache/reasoning usage details on spans', async ({
     request,
   }) => {