From 027dbf8aea94323c05d86e678f56461e4f0a7c4e Mon Sep 17 00:00:00 2001 From: Dex Date: Sun, 12 Jul 2026 11:42:13 +0000 Subject: [PATCH] fix(ai): roll up OpenTelemetry usage across iterations Fixes #916 --- .changeset/otel-root-usage-rollup.md | 5 + packages/ai/src/middlewares/otel.ts | 39 +++++++- packages/ai/tests/middlewares/otel.test.ts | 105 +++++++++++++++++++++ testing/e2e/src/routes/api.otel-usage.ts | 41 ++++++-- testing/e2e/tests/middleware.spec.ts | 40 ++++++++ 5 files changed, 217 insertions(+), 13 deletions(-) create mode 100644 .changeset/otel-root-usage-rollup.md diff --git a/.changeset/otel-root-usage-rollup.md b/.changeset/otel-root-usage-rollup.md new file mode 100644 index 000000000..c533364bc --- /dev/null +++ b/.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. diff --git a/packages/ai/src/middlewares/otel.ts b/packages/ai/src/middlewares/otel.ts index f4ddba35e..43f6ff983 100644 --- a/packages/ai/src/middlewares/otel.ts +++ b/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,6 +139,8 @@ interface RequestState { * from the (base-shaped) finish info, which doesn't carry it. */ lastFinishReason: string | null + rootUsageAttributes: Record | null + rootUsageApplied: boolean } const stateByCtx = new WeakMap() @@ -145,6 +148,29 @@ const stateByCtx = new WeakMap() const DEFAULT_MAX_CONTENT_LENGTH = 100_000 const REDACTION_FAILED_SENTINEL = '[redaction_failed]' +function accumulateUsageAttributes( + current: Record | null, + usage: TokenUsage, +): Record { + 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), ) diff --git a/packages/ai/tests/middlewares/otel.test.ts b/packages/ai/tests/middlewares/otel.test.ts index f7fa553ac..ad500a1b0 100644 --- a/packages/ai/tests/middlewares/otel.test.ts +++ b/packages/ai/tests/middlewares/otel.test.ts @@ -56,6 +56,30 @@ async function runToIterationStart( }) } +async function runUsageIterations( + mw: ChatMiddleware, + ctx: ChatMiddlewareContext, + usages: Array, +) { + 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 = [ + { + 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 () => { diff --git a/testing/e2e/src/routes/api.otel-usage.ts b/testing/e2e/src/routes/api.otel-usage.ts index 367d83715..8769a6af8 100644 --- a/testing/e2e/src/routes/api.otel-usage.ts +++ b/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. diff --git a/testing/e2e/tests/middleware.spec.ts b/testing/e2e/tests/middleware.spec.ts index 6b15d6759..a9c7f540d 100644 --- a/testing/e2e/tests/middleware.spec.ts +++ b/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, }) => {