otelMiddleware root chat span carries last-iteration usage, not the documented cross-iteration roll-up

[phas02]

Which project does this relate to?

AI

Describe the bug

The otel docs (docs/advanced/otel.md, "Usage Accumulation") say:

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

In 0.40.0 the root chat span instead carries only the last iteration's usage. otelMiddleware's onFinish stamps FinishInfo.usage onto the root span, and the chat engine never accumulates usage across iterations:

• beginIteration() resets this.finishedEvent = null each iteration (activities/chat/index.js)
• handleRunFinishedEvent(chunk) overwrites this.finishedEvent = chunk
• the terminal hook passes usage: this.finishedEvent?.usage — i.e. the final iteration's RUN_FINISHED usage only

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 (e.g. @tanstack/ai-anthropic builds it from a single response's message_delta).

The per-iteration data is correct — onUsage fires with each iteration's own (incremental) usage, iteration spans get the right values, and the gen_ai.client.token.usage histogram sums correctly. Only the root span's roll-up (and FinishInfo.usage for any middleware relying on it) under-reports multi-iteration runs.

A fix in the engine would be to accumulate usage-bearing RUN_FINISHED chunks across iterations (the same summation onUsage consumers do today) and pass the totals in FinishInfo.

Your Example Website or App

repro script below (Bun or Node, @tanstack/ai@0.40.0)

Steps to Reproduce the Bug or Issue

Run a two-iteration chat (tool call → finish) where iteration 1 reports 500/50 and iteration 2 reports 700/80, with otelMiddleware and an in-memory exporter:

import { chat, toolDefinition } from "@tanstack/ai";
import { otelMiddleware } from "@tanstack/ai/middlewares/otel";
import {
  BasicTracerProvider, InMemorySpanExporter, SimpleSpanProcessor,
} from "@opentelemetry/sdk-trace-base";
import { z } from "zod";

const exporter = new InMemorySpanExporter();
const provider = new BasicTracerProvider({
  spanProcessors: [new SimpleSpanProcessor(exporter)],
});

const ITER1 = { promptTokens: 500, completionTokens: 50, totalTokens: 550 };
const ITER2 = { promptTokens: 700, completionTokens: 80, totalTokens: 780 };

function scriptedAdapter() {
  let call = 0;
  return {
    kind: "text", name: "scripted", provider: "scripted", model: "scripted-model",
    async *chatStream(options) {
      call += 1;
      const base = { model: "scripted-model", timestamp: Date.now() };
      const runId = `run-${call}`, threadId = "t1";
      yield { type: "RUN_STARTED", runId, threadId, ...base };
      if (!options.messages.some((m) => m.role === "tool")) {
        yield { type: "TOOL_CALL_START", toolCallId: "tc1", toolCallName: "status", toolName: "status", index: 0, ...base };
        yield { type: "TOOL_CALL_ARGS", toolCallId: "tc1", delta: "{}", args: "{}", ...base };
        yield { type: "TOOL_CALL_END", toolCallId: "tc1", toolCallName: "status", toolName: "status", input: {}, ...base };
        yield { type: "RUN_FINISHED", runId, threadId, finishReason: "tool_calls", usage: ITER1, ...base };
      } else {
        yield { type: "TEXT_MESSAGE_START", messageId: "m1", role: "assistant", ...base };
        yield { type: "TEXT_MESSAGE_CONTENT", messageId: "m1", delta: "done", ...base };
        yield { type: "TEXT_MESSAGE_END", messageId: "m1", ...base };
        yield { type: "RUN_FINISHED", runId, threadId, finishReason: "stop", usage: ITER2, ...base };
      }
    },
  };
}

const statusTool = toolDefinition({
  name: "status", description: "status", inputSchema: z.object({}),
}).server(async () => ({ ok: true }));

const stream = chat({
  adapter: scriptedAdapter(),
  messages: [{ id: "u1", role: "user", content: "go" }],
  tools: [statusTool],
  middleware: [otelMiddleware({ tracer: provider.getTracer("repro") })],
});
for await (const _ of stream) {}

await provider.forceFlush();
for (const s of exporter.getFinishedSpans()) {
  console.log(s.name, s.attributes["gen_ai.usage.input_tokens"], s.attributes["gen_ai.usage.output_tokens"]);
}

Output:

execute_tool status undefined undefined
chat scripted-model #0 500 50
chat scripted-model #1 700 80
chat scripted-model 700 80        <-- root span: last iteration only

Expected behavior

Per the docs, the root chat scripted-model span should carry the roll-up: gen_ai.usage.input_tokens = 1200, gen_ai.usage.output_tokens = 130.

Platform

• OS: macOS (also reproduced on Linux)
• Runtime: Bun 1.3.14 / Node 22
• @tanstack/ai: 0.40.0

Additional context

Workaround we're shipping meanwhile: a small companion middleware that sums onUsage increments and overwrites the root span's gen_ai.usage.* from otelMiddleware's onSpanEnd hook (which conveniently fires before rootSpan.end()).