fix(ai): forward token usage on the structured-output fallback path

season179 · season179 · commit d69f3a929020 · 2026-06-19T14:16:46.000+08:00
`fallbackStructuredOutputStream` — used by `chat({ outputSchema, stream: true })`
whenever an adapter resolves the schema through the non-streaming
`structuredOutput()` rather than a native streaming or combined path (Ollama,
plus Anthropic and Gemini models that predate combined tools+schema support) —
wrapped `structuredOutput()` but dropped the `usage` from its result. Consumers
reading `RUN_FINISHED.usage` saw `undefined`, and the engine's `runOnUsage`
middleware hook (gated on `chunk.usage`) never fired, so cost-tracking and
observability layers reported zero token counts on that path.

The synthesized `RUN_FINISHED` now carries the adapter-reported `usage`,
matching the native streaming path. Adapters that don't report usage are
unaffected — the conditional spread omits the key entirely.

Adds unit coverage in chat-structured-output-stream.test.ts (usage forwarded;
omitted when absent) and an e2e regression (anthropic-structured-usage) that
drives the Anthropic adapter through the fallback against an aimock mount and
asserts usage reaches RUN_FINISHED.usage.
diff --git a/.changeset/fallback-structured-output-usage.md b/.changeset/fallback-structured-output-usage.md
@@ -0,0 +1,18 @@
+---
+'@tanstack/ai': patch
+---
+
+fix: forward token usage on the structured-output fallback path
+
+`fallbackStructuredOutputStream` — used by `chat({ outputSchema, stream: true })`
+whenever an adapter resolves the schema through the non-streaming
+`structuredOutput()` rather than a native streaming or combined path (in practice
+Ollama, plus Anthropic and Gemini models that predate combined tools+schema
+support) — wrapped `structuredOutput()` but dropped the `usage` from its result.
+Consumers reading `RUN_FINISHED.usage` saw `undefined`, and the engine's
+`runOnUsage` middleware hook (gated on `chunk.usage`) never fired, so
+cost-tracking and observability layers reported zero token counts on that path.
+
+The synthesized `RUN_FINISHED` now carries the adapter-reported `usage`, matching
+the native streaming path. Adapters that don't report usage are unaffected (no
+`usage` key is emitted).
diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts
@@ -33,7 +33,11 @@ import type {
   ClientToolRequest,
   ToolResult,
 } from './tools/tool-calls'
-import type { AnyTextAdapter, StructuredOutputOptions } from './adapter'
+import type {
+  AnyTextAdapter,
+  StructuredOutputOptions,
+  StructuredOutputResult,
+} from './adapter'
 import type {
   AgentLoopStrategy,
   AnyTool,
@@ -2861,7 +2865,7 @@ async function* fallbackStructuredOutputStream(
     timestamp,
   }
 
-  let result: { data: unknown; rawText: string }
+  let result: StructuredOutputResult<unknown>
   try {
     result = await adapter.structuredOutput(options)
   } catch (error) {
@@ -2917,6 +2921,12 @@ async function* fallbackStructuredOutputStream(
     model,
     timestamp,
     finishReason: 'stop',
+    // Forward adapter-reported token usage so consumers reading
+    // `RUN_FINISHED.usage` (and the engine's `runOnUsage` middleware hook) see
+    // it on the fallback path, mirroring the native streaming path. The
+    // conditional spread avoids emitting `usage: undefined` for adapters that
+    // don't report it. See #758.
+    ...(result.usage ? { usage: result.usage } : {}),
   }
 }
 
diff --git a/packages/ai/tests/chat-structured-output-stream.test.ts b/packages/ai/tests/chat-structured-output-stream.test.ts
@@ -18,7 +18,7 @@ import { describe, expect, it } from 'vitest'
 import { z } from 'zod'
 import { chat } from '../src/activities/chat/index'
 import { EventType } from '../src/types'
-import type { StreamChunk } from '../src/types'
+import type { RunFinishedEvent, StreamChunk, TokenUsage } from '../src/types'
 import type { AnyTextAdapter } from '../src/activities/chat/adapter'
 import { collectChunks } from './test-utils'
 
@@ -42,7 +42,9 @@ const validPerson: Person = {
  */
 function makeAdapter(opts: {
   structuredOutputStream?: (o: unknown) => AsyncIterable<StreamChunk>
-  structuredOutput?: (o: unknown) => Promise<{ data: unknown; rawText: string }>
+  structuredOutput?: (
+    o: unknown,
+  ) => Promise<{ data: unknown; rawText: string; usage?: TokenUsage }>
 }): AnyTextAdapter {
   return {
     kind: 'text' as const,
@@ -375,6 +377,72 @@ describe('chat({ outputSchema, stream: true })', () => {
       expect(complete).toBeDefined()
       expect(complete!.value.object).toEqual(invalidObject)
     })
+
+    it('forwards adapter-reported usage onto RUN_FINISHED (#758)', async () => {
+      // Regression for #758: the fallback wraps the non-streaming
+      // `structuredOutput`, whose `{ data, rawText, usage }` result includes
+      // token usage. Before the fix the synthesized RUN_FINISHED dropped
+      // `usage`, so consumers (and the `runOnUsage` middleware hook) saw
+      // `undefined` on every fallback-path provider (Anthropic, Gemini, Ollama).
+      const usage: TokenUsage = {
+        promptTokens: 125,
+        completionTokens: 1346,
+        totalTokens: 1471,
+        promptTokensDetails: { cachedTokens: 5760 },
+      }
+      const adapter = makeAdapter({
+        structuredOutput: async () => ({
+          data: validPerson,
+          rawText: JSON.stringify(validPerson),
+          usage,
+        }),
+      })
+
+      const stream = chat({
+        adapter,
+        messages: [{ role: 'user', content: 'extract' }],
+        outputSchema: PersonSchema,
+        stream: true,
+      })
+
+      const chunks = await collectChunks(
+        stream as unknown as AsyncIterable<StreamChunk>,
+      )
+
+      const finished = chunks.find(
+        (c) => c.type === EventType.RUN_FINISHED,
+      ) as RunFinishedEvent | undefined
+      expect(finished).toBeDefined()
+      expect(finished!.usage).toEqual(usage)
+    })
+
+    it('omits usage on RUN_FINISHED when the adapter does not report it', async () => {
+      // The conditional spread must not synthesize `usage: undefined` for
+      // adapters whose `structuredOutput` returns no usage.
+      const adapter = makeAdapter({
+        structuredOutput: async () => ({
+          data: validPerson,
+          rawText: JSON.stringify(validPerson),
+        }),
+      })
+
+      const stream = chat({
+        adapter,
+        messages: [{ role: 'user', content: 'extract' }],
+        outputSchema: PersonSchema,
+        stream: true,
+      })
+
+      const chunks = await collectChunks(
+        stream as unknown as AsyncIterable<StreamChunk>,
+      )
+
+      const finished = chunks.find(
+        (c) => c.type === EventType.RUN_FINISHED,
+      ) as RunFinishedEvent | undefined
+      expect(finished).toBeDefined()
+      expect('usage' in finished!).toBe(false)
+    })
   })
 
   describe('lifecycle ordering', () => {
diff --git a/testing/e2e/global-setup.ts b/testing/e2e/global-setup.ts
@@ -73,6 +73,16 @@ export default async function globalSetup() {
   // `promptTokensDetails.cachedTokens` / `completionTokensDetails.reasoningTokens`.
   mock.mount('/openai-usage-details', openaiUsageDetailsMount())
 
+  // Anthropic structured-output fallback usage (#758). The Anthropic text
+  // adapter has no native `structuredOutputStream`, so streaming structured
+  // output runs through the activity layer's `fallbackStructuredOutputStream`,
+  // which wraps the non-streaming `structuredOutput()`. aimock's native
+  // Anthropic helper doesn't synthesize a tool-forced `structured_output`
+  // response with usage, so this mount hand-crafts the non-streaming
+  // `/v1/messages` JSON the adapter expects. The companion spec asserts the
+  // `usage` survives onto `RUN_FINISHED.usage` on the fallback path.
+  mock.mount('/anthropic-structured-usage', anthropicStructuredUsageMount())
+
   await mock.start()
   console.log(`[aimock] started on port 4010`)
   ;(globalThis as any).__aimock = mock
@@ -547,6 +557,60 @@ function openaiUsageDetailsMount(): Mountable {
   }
 }
 
+/**
+ * Mounts the non-streaming Anthropic `/v1/messages` response the text adapter's
+ * `structuredOutput()` expects: a tool-forced `structured_output` `tool_use`
+ * block plus a `usage` object carrying `input_tokens` / `output_tokens` /
+ * `cache_read_input_tokens`. `buildAnthropicUsage` normalizes those into
+ * `promptTokens` / `completionTokens` / `promptTokensDetails.cachedTokens`.
+ * Drives the #758 fallback-path usage regression.
+ */
+function anthropicStructuredUsageMount(): Mountable {
+  return {
+    async handleRequest(
+      req: http.IncomingMessage,
+      res: http.ServerResponse,
+      // The mount prefix (/anthropic-structured-usage) is stripped before
+      // dispatch; the Anthropic SDK posts to <baseURL>/v1/messages and aimock
+      // strips the ?beta=... query string from `pathname`.
+      pathname: string,
+    ): Promise<boolean> {
+      if (req.method !== 'POST' || !pathname.startsWith('/v1/messages')) {
+        return false
+      }
+      // structuredOutput() makes a non-streaming request (stream: false), so
+      // respond with a single JSON message rather than an SSE stream.
+      await drainBody(req)
+      res.statusCode = 200
+      res.setHeader('Content-Type', 'application/json')
+      res.end(
+        JSON.stringify({
+          id: 'msg_structured_usage_e2e',
+          type: 'message',
+          role: 'assistant',
+          model: 'claude-opus-4-1',
+          content: [
+            {
+              type: 'tool_use',
+              id: 'toolu_structured_output',
+              name: 'structured_output',
+              input: { recommendation: 'Fender Stratocaster', price: 1299 },
+            },
+          ],
+          stop_reason: 'tool_use',
+          stop_sequence: null,
+          usage: {
+            input_tokens: 125,
+            output_tokens: 1346,
+            cache_read_input_tokens: 5760,
+          },
+        }),
+      )
+      return true
+    },
+  }
+}
+
 function buildToolPlusServerToolEvents(): Array<Record<string, unknown>> {
   const messageId = 'msg_bug_604'
   const model = 'claude-sonnet-4-5'
diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts
@@ -43,6 +43,7 @@ import { Route as ApiImageRouteImport } from './routes/api.image'
 import { Route as ApiChatRouteImport } from './routes/api.chat'
 import { Route as ApiAudioRouteImport } from './routes/api.audio'
 import { Route as ApiArktypeToolWireRouteImport } from './routes/api.arktype-tool-wire'
+import { Route as ApiAnthropicStructuredUsageRouteImport } from './routes/api.anthropic-structured-usage'
 import { Route as ApiAnthropicSkillsWireRouteImport } from './routes/api.anthropic-skills-wire'
 import { Route as ApiAnthropicBugTestRouteImport } from './routes/api.anthropic-bug-test'
 import { Route as ProviderFeatureRouteImport } from './routes/$provider/$feature'
@@ -226,6 +227,12 @@ const ApiArktypeToolWireRoute = ApiArktypeToolWireRouteImport.update({
   path: '/api/arktype-tool-wire',
   getParentRoute: () => rootRouteImport,
 } as any)
+const ApiAnthropicStructuredUsageRoute =
+  ApiAnthropicStructuredUsageRouteImport.update({
+    id: '/api/anthropic-structured-usage',
+    path: '/api/anthropic-structured-usage',
+    getParentRoute: () => rootRouteImport,
+  } as any)
 const ApiAnthropicSkillsWireRoute = ApiAnthropicSkillsWireRouteImport.update({
   id: '/api/anthropic-skills-wire',
   path: '/api/anthropic-skills-wire',
@@ -282,6 +289,7 @@ export interface FileRoutesByFullPath {
   '/$provider/$feature': typeof ProviderFeatureRoute
   '/api/anthropic-bug-test': typeof ApiAnthropicBugTestRoute
   '/api/anthropic-skills-wire': typeof ApiAnthropicSkillsWireRoute
+  '/api/anthropic-structured-usage': typeof ApiAnthropicStructuredUsageRoute
   '/api/arktype-tool-wire': typeof ApiArktypeToolWireRoute
   '/api/audio': typeof ApiAudioRouteWithChildren
   '/api/chat': typeof ApiChatRoute
@@ -326,6 +334,7 @@ export interface FileRoutesByTo {
   '/$provider/$feature': typeof ProviderFeatureRoute
   '/api/anthropic-bug-test': typeof ApiAnthropicBugTestRoute
   '/api/anthropic-skills-wire': typeof ApiAnthropicSkillsWireRoute
+  '/api/anthropic-structured-usage': typeof ApiAnthropicStructuredUsageRoute
   '/api/arktype-tool-wire': typeof ApiArktypeToolWireRoute
   '/api/audio': typeof ApiAudioRouteWithChildren
   '/api/chat': typeof ApiChatRoute
@@ -371,6 +380,7 @@ export interface FileRoutesById {
   '/$provider/$feature': typeof ProviderFeatureRoute
   '/api/anthropic-bug-test': typeof ApiAnthropicBugTestRoute
   '/api/anthropic-skills-wire': typeof ApiAnthropicSkillsWireRoute
+  '/api/anthropic-structured-usage': typeof ApiAnthropicStructuredUsageRoute
   '/api/arktype-tool-wire': typeof ApiArktypeToolWireRoute
   '/api/audio': typeof ApiAudioRouteWithChildren
   '/api/chat': typeof ApiChatRoute
@@ -417,6 +427,7 @@ export interface FileRouteTypes {
     | '/$provider/$feature'
     | '/api/anthropic-bug-test'
     | '/api/anthropic-skills-wire'
+    | '/api/anthropic-structured-usage'
     | '/api/arktype-tool-wire'
     | '/api/audio'
     | '/api/chat'
@@ -461,6 +472,7 @@ export interface FileRouteTypes {
     | '/$provider/$feature'
     | '/api/anthropic-bug-test'
     | '/api/anthropic-skills-wire'
+    | '/api/anthropic-structured-usage'
     | '/api/arktype-tool-wire'
     | '/api/audio'
     | '/api/chat'
@@ -505,6 +517,7 @@ export interface FileRouteTypes {
     | '/$provider/$feature'
     | '/api/anthropic-bug-test'
     | '/api/anthropic-skills-wire'
+    | '/api/anthropic-structured-usage'
     | '/api/arktype-tool-wire'
     | '/api/audio'
     | '/api/chat'
@@ -550,6 +563,7 @@ export interface RootRouteChildren {
   ProviderFeatureRoute: typeof ProviderFeatureRoute
   ApiAnthropicBugTestRoute: typeof ApiAnthropicBugTestRoute
   ApiAnthropicSkillsWireRoute: typeof ApiAnthropicSkillsWireRoute
+  ApiAnthropicStructuredUsageRoute: typeof ApiAnthropicStructuredUsageRoute
   ApiArktypeToolWireRoute: typeof ApiArktypeToolWireRoute
   ApiAudioRoute: typeof ApiAudioRouteWithChildren
   ApiChatRoute: typeof ApiChatRoute
@@ -815,6 +829,13 @@ declare module '@tanstack/react-router' {
       preLoaderRoute: typeof ApiArktypeToolWireRouteImport
       parentRoute: typeof rootRouteImport
     }
+    '/api/anthropic-structured-usage': {
+      id: '/api/anthropic-structured-usage'
+      path: '/api/anthropic-structured-usage'
+      fullPath: '/api/anthropic-structured-usage'
+      preLoaderRoute: typeof ApiAnthropicStructuredUsageRouteImport
+      parentRoute: typeof rootRouteImport
+    }
     '/api/anthropic-skills-wire': {
       id: '/api/anthropic-skills-wire'
       path: '/api/anthropic-skills-wire'
@@ -947,6 +968,7 @@ const rootRouteChildren: RootRouteChildren = {
   ProviderFeatureRoute: ProviderFeatureRoute,
   ApiAnthropicBugTestRoute: ApiAnthropicBugTestRoute,
   ApiAnthropicSkillsWireRoute: ApiAnthropicSkillsWireRoute,
+  ApiAnthropicStructuredUsageRoute: ApiAnthropicStructuredUsageRoute,
   ApiArktypeToolWireRoute: ApiArktypeToolWireRoute,
   ApiAudioRoute: ApiAudioRouteWithChildren,
   ApiChatRoute: ApiChatRoute,
diff --git a/testing/e2e/src/routes/api.anthropic-structured-usage.ts b/testing/e2e/src/routes/api.anthropic-structured-usage.ts
@@ -0,0 +1,73 @@
+import { createFileRoute } from '@tanstack/react-router'
+import { chat, createChatOptions } from '@tanstack/ai'
+import { createAnthropicChat } from '@tanstack/ai-anthropic'
+import { z } from 'zod'
+
+const LLMOCK_DEFAULT_BASE = process.env.LLMOCK_URL || 'http://127.0.0.1:4010'
+const DUMMY_KEY = 'sk-e2e-test-dummy-key'
+
+/**
+ * Drives the Anthropic text adapter (`AnthropicTextAdapter`) through the
+ * streaming structured-output path. That adapter has no native
+ * `structuredOutputStream`, so `chat({ outputSchema, stream: true })` routes
+ * through the activity layer's `fallbackStructuredOutputStream`, which wraps the
+ * non-streaming `structuredOutput()`. The mounted `/anthropic-structured-usage`
+ * aimock path returns a tool-forced `structured_output` response whose `usage`
+ * carries `input_tokens` / `output_tokens` / `cache_read_input_tokens`.
+ *
+ * Regression for #758: before the fix the fallback dropped `result.usage`, so
+ * `RUN_FINISHED.usage` was `undefined` on every fallback-path provider. The
+ * companion spec asserts the usage now reaches `RUN_FINISHED.usage`.
+ */
+export const Route = createFileRoute('/api/anthropic-structured-usage')({
+  server: {
+    handlers: {
+      POST: async () => {
+        // `claude-opus-4-1` is intentionally a *pre-4.5* model: it is NOT in
+        // `ANTHROPIC_COMBINED_TOOLS_AND_SCHEMA_MODELS`, so
+        // `supportsCombinedToolsAndSchema()` is false and the engine routes
+        // `chat({ outputSchema, stream: true })` through the non-streaming
+        // `structuredOutput()` wrapped by `fallbackStructuredOutputStream` —
+        // exactly the path #758 fixes. A 4.5+ model would use the native
+        // combined path and never touch the fallback.
+        const adapter = createAnthropicChat('claude-opus-4-1', DUMMY_KEY, {
+          baseURL: `${LLMOCK_DEFAULT_BASE}/anthropic-structured-usage`,
+        })
+
+        const options = createChatOptions({
+          adapter,
+          outputSchema: z.object({
+            recommendation: z.string(),
+            price: z.number(),
+          }),
+          stream: true,
+        })
+
+        let usage: Record<string, unknown> | undefined
+        try {
+          for await (const chunk of chat({
+            ...options,
+            messages: [{ role: 'user', content: 'recommend a guitar as json' }],
+          })) {
+            if (chunk.type === 'RUN_FINISHED') {
+              usage = chunk.usage as Record<string, unknown> | undefined
+            }
+          }
+        } catch (error) {
+          return new Response(
+            JSON.stringify({
+              ok: false,
+              error: error instanceof Error ? error.message : String(error),
+            }),
+            { status: 200, headers: { 'Content-Type': 'application/json' } },
+          )
+        }
+
+        return new Response(JSON.stringify({ ok: true, usage }), {
+          status: 200,
+          headers: { 'Content-Type': 'application/json' },
+        })
+      },
+    },
+  },
+})
diff --git a/testing/e2e/tests/anthropic-structured-usage.spec.ts b/testing/e2e/tests/anthropic-structured-usage.spec.ts
