@tanstack/ai 0.47.3 → 0.48.0

package/src/activities/chat/tools/tool-calls.ts

@@ -1,4 +1,6 @@
 import { normalizeToolResult } from '../../../utilities/tool-result'
+import { tanstackMetadata } from '../../../utilities/merge-metadata'
+import type { AdapterYieldChunk } from '../../../utilities/adapter-yield-chunk'
 import { isStandardSchema, parseWithStandardSchema } from './schema-converter'
 import type { ToolApprovalResolution } from '../../../interrupts'
 import type {
@@ -231,10 +233,8 @@ export class ToolCallManager<
    * Add a TOOL_CALL_START event to begin tracking a tool call (AG-UI)
    */
   addToolCallStartEvent(event: ToolCallStartEvent): void {
-    const index = event.index ?? this.toolCallsMap.size
-    const runtimeEvent = event as Partial<ToolCallStartEvent> &
-      Pick<ToolCallStartEvent, 'toolName'>
-    const name = runtimeEvent.toolCallName ?? runtimeEvent.toolName
+    const index = (event as AdapterYieldChunk).index ?? this.toolCallsMap.size
+    const name = event.toolCallName ?? event.toolName
     this.toolCallsMap.set(index, {
       id: event.toolCallId,
       type: 'function',
@@ -250,10 +250,14 @@ export class ToolCallManager<
    * Add a TOOL_CALL_ARGS event to accumulate arguments (AG-UI)
    */
   addToolCallArgsEvent(event: ToolCallArgsEvent): void {
-    // Find the tool call by ID
+    const extra = event as AdapterYieldChunk
     for (const [, toolCall] of this.toolCallsMap.entries()) {
       if (toolCall.id === event.toolCallId) {
-        toolCall.function.arguments += event.delta
+        if (typeof extra.args === 'string' && extra.args !== '') {
+          toolCall.function.arguments = extra.args
+        } else {
+          toolCall.function.arguments += event.delta
+        }
         break
       }
     }
@@ -264,16 +268,13 @@ export class ToolCallManager<
    * Called when TOOL_CALL_END is received
    */
   completeToolCall(event: ToolCallEndEvent): void {
-    for (const [, toolCall] of this.toolCallsMap.entries()) {
-      if (toolCall.id === event.toolCallId) {
-        if (event.input !== undefined) {
-          // Normalize null/non-object to {} (e.g. Anthropic empty tool_use blocks)
-          const normalized =
-            event.input && typeof event.input === 'object' ? event.input : {}
-          toolCall.function.arguments = JSON.stringify(normalized)
-        }
-        break
-      }
+    for (const toolCall of this.toolCallsMap.values()) {
+      if (toolCall.id !== event.toolCallId) continue
+      if (event.input === undefined) return
+      const normalized =
+        event.input && typeof event.input === 'object' ? event.input : {}
+      toolCall.function.arguments = JSON.stringify(normalized)
+      return
     }
   }
 
@@ -301,7 +302,7 @@ export class ToolCallManager<
   async *executeTools(
     finishEvent: RunFinishedEvent,
     ...contextArgs: ExecuteToolsContextArgs<TContext>
-  ): AsyncGenerator<ToolCallEndEvent, Array<ModelMessage>, void> {
+  ): AsyncGenerator<AdapterYieldChunk, Array<ModelMessage>, void> {
     const toolCallsArray = this.getToolCalls()
     const toolResults: Array<ModelMessage> = []
     const hasRuntimeContext = contextArgs.length > 0
@@ -313,9 +314,6 @@ export class ToolCallManager<
       let toolResultContent: string | Array<ContentPart>
       let toolResultState: ToolOutputState | undefined
       // Holds the parsed/validated execution output before serialization.
-      // Surfaced on the emitted `TOOL_CALL_END` event as `output` so
-      // consumers can read it typed (via `TypedStreamChunk` distribution
-      // over the tools array) without re-parsing `result`.
       // Stays `undefined` when the tool has no `execute` (client-only
       // tools) or when execution throws.
       let toolOutput: unknown
@@ -397,7 +395,10 @@ export class ToolCallManager<
         toolCallId: toolCall.id,
         toolCallName: toolCall.function.name,
         toolName: toolCall.function.name,
-        model: finishEvent.model,
+        model: (() => {
+          const model = tanstackMetadata(finishEvent)?.model
+          return typeof model === 'string' ? model : undefined
+        })(),
         timestamp: Date.now(),
         // Typed parsed output (undefined for failed exec / client-only tools).
         ...(toolOutput !== undefined ? { output: toolOutput } : {}),
@@ -433,7 +434,7 @@ export interface ToolResult {
   duration?: number
   /**
    * Parsed tool input (after JSON parse + optional Standard Schema validation).
-   * Surfaced on engine-emitted `TOOL_CALL_END` events for TypedStreamChunk consumers.
+   * Parsed tool input after JSON parse + optional Standard Schema validation.
    */
   input?: unknown
   /**

package/src/activities/generateVideo/index.ts

@@ -33,6 +33,8 @@ import type {
   GenerationMiddlewareContext,
 } from '../middleware/types'
 import type { VideoAdapter } from './adapter'
+import { normalizeStreamChunk } from '../../utilities/normalize-stream-chunk'
+import type { AdapterYieldChunk } from '../../utilities/adapter-yield-chunk'
 import type {
   MediaPrompt,
   MediaPromptFor,
@@ -728,13 +730,13 @@ async function* runStreamingVideoGeneration<
           timestamp: Date.now(),
         }
 
-        yield {
+        yield* normalizeStreamChunk({
           type: 'RUN_FINISHED',
           runId,
           threadId: wireThreadId,
           finishReason: 'stop',
           timestamp: Date.now(),
-        } as StreamChunk
+        } as AdapterYieldChunk)
         return
       }
 
@@ -768,15 +770,14 @@ async function* runStreamingVideoGeneration<
       code: payload.code,
       source: 'generateVideo',
     })
-    yield {
+    yield* normalizeStreamChunk({
       type: 'RUN_ERROR',
       runId,
       threadId: wireThreadId,
       message: payload.message,
-      code: payload.code,
-      error: payload,
+      ...(payload.code !== undefined ? { code: payload.code } : {}),
       timestamp: Date.now(),
-    } as StreamChunk
+    } as AdapterYieldChunk)
   } finally {
     abortControls.clear()
     if (!settled) {

package/src/activities/stream-generation-result.ts

@@ -7,6 +7,7 @@
 import { EventType } from '@ag-ui/core'
 import { toRunErrorPayload } from './error-payload'
 import type { StreamChunk } from '../types'
+import { normalizeStreamChunk } from '../utilities/normalize-stream-chunk'
 
 function createId(prefix: string): string {
   return `${prefix}-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`
@@ -74,31 +75,26 @@ export async function* streamGenerationResult<TResult>(
       timestamp: Date.now(),
     }
 
-    yield {
+    yield* normalizeStreamChunk({
       type: EventType.RUN_FINISHED,
       runId,
       threadId,
       finishReason: 'stop',
       timestamp: Date.now(),
-    }
+    })
   } catch (error: unknown) {
     const payload = toRunErrorPayload(error, 'Generation failed')
     // `code` is omitted entirely when undefined so the event matches the
-    // AG-UI `code?: string` shape under `exactOptionalPropertyTypes`. The
-    // deprecated nested `error` form preserves the same conditional
-    // structure for backward compatibility.
+    // AG-UI `code?: string` shape under `exactOptionalPropertyTypes`.
     const codeFields =
       payload.code !== undefined ? { code: payload.code } : undefined
-    yield {
+    yield* normalizeStreamChunk({
       type: EventType.RUN_ERROR,
+      runId,
+      threadId,
       message: payload.message,
       ...codeFields,
-      // Deprecated nested form for backward compatibility
-      error: {
-        message: payload.message,
-        ...codeFields,
-      },
       timestamp: Date.now(),
-    }
+    })
   }
 }

package/src/activities/summarize/chat-stream-summarize.ts

@@ -1,14 +1,58 @@
 import { EventType } from '@ag-ui/core'
 import { toRunErrorPayload } from '../error-payload'
 import { MAX_TOKENS_KEYS } from '../../utilities/sampling-keys'
+import { rebuildTokenUsage } from '../../utilities/ag-ui-usage'
+import type { AdapterYieldChunk } from '../../utilities/adapter-yield-chunk'
+import { tanstackMetadata } from '../../utilities/merge-metadata'
+import { normalizeStreamChunk } from '../../utilities/normalize-stream-chunk'
 import { BaseSummarizeAdapter } from './adapter'
 import type {
   StreamChunk,
   SummarizationOptions,
   SummarizationResult,
   TextOptions,
+  TokenUsage,
 } from '../../types'
 
+function consumeSpecSummarizeChunk(
+  chunk: StreamChunk,
+  state: { summary: string; model: string; usage: TokenUsage },
+): void {
+  if (chunk.type === EventType.TEXT_MESSAGE_CONTENT) {
+    if (chunk.delta) state.summary += chunk.delta
+    return
+  }
+
+  const tanstack = tanstackMetadata(chunk)
+  if (
+    (chunk.type === EventType.RUN_STARTED ||
+      chunk.type === EventType.RUN_FINISHED ||
+      chunk.type === EventType.TEXT_MESSAGE_START) &&
+    typeof tanstack?.model === 'string'
+  ) {
+    state.model = tanstack.model
+  }
+
+  if (chunk.type === EventType.RUN_FINISHED) {
+    const rebuilt = rebuildTokenUsage(chunk.usage, tanstack?.usage)
+    if (rebuilt) state.usage = rebuilt
+  }
+}
+
+function throwRunError(
+  chunk: Extract<StreamChunk, { type: 'RUN_ERROR' }>,
+): never {
+  const message =
+    typeof chunk.message === 'string' && chunk.message.length > 0
+      ? chunk.message
+      : 'Summarization failed'
+  const err = new Error(message)
+  if (typeof chunk.code === 'string') {
+    ;(err as Error & { code?: string }).code = chunk.code
+  }
+  throw err
+}
+
 /**
  * Minimal contract for a text adapter that supports `chatStream`. Lets
  * `ChatStreamSummarizeAdapter` work with any text adapter without coupling
@@ -21,7 +65,7 @@ import type {
  * `SummarizationOptions<TProviderOptions>` on the wrapper itself.
  */
 export interface ChatStreamCapable {
-  chatStream: (options: TextOptions<any>) => AsyncIterable<StreamChunk>
+  chatStream: (options: TextOptions<any>) => AsyncIterable<AdapterYieldChunk>
 }
 
 /**
@@ -201,10 +245,12 @@ export class ChatStreamSummarizeAdapter<
   ): Promise<SummarizationResult> {
     const systemPrompt = this.buildSummarizationPrompt(options)
 
-    let summary = ''
     const id = this.generateId()
-    let model = options.model
-    let usage = { promptTokens: 0, completionTokens: 0, totalTokens: 0 }
+    const state = {
+      summary: '',
+      model: options.model,
+      usage: { promptTokens: 0, completionTokens: 0, totalTokens: 0 },
+    }
 
     options.logger.request(
       `activity=summarize provider=${this.name} model=${options.model} text-length=${options.text.length} maxLength=${options.maxLength ?? 'unset'}`,
@@ -212,41 +258,15 @@ export class ChatStreamSummarizeAdapter<
     )
 
     try {
-      for await (const chunk of this.textAdapter.chatStream(
+      for await (const raw of this.textAdapter.chatStream(
         this.buildTextOptions(options, systemPrompt),
       )) {
-        if (chunk.type === 'TEXT_MESSAGE_CONTENT') {
-          if (chunk.content) {
-            summary = chunk.content
-          } else if (chunk.delta) {
-            // Append delta only when present — a content-less chunk with no
-            // delta would otherwise concat literal `'undefined'`.
-            summary += chunk.delta
-          }
-          model = chunk.model || model
-        }
-        if (chunk.type === 'RUN_FINISHED') {
-          if (chunk.usage) {
-            usage = chunk.usage
-          }
-        }
-        // Surface failures: the underlying chatStream emits RUN_ERROR instead
-        // of throwing, so without this branch summarize() would return an
-        // empty summary and pretend a failed run succeeded.
-        if (chunk.type === 'RUN_ERROR') {
-          const message =
-            (chunk.error && typeof chunk.error.message === 'string'
-              ? chunk.error.message
-              : null) ?? 'Summarization failed'
-          const code =
-            chunk.error && typeof chunk.error.code === 'string'
-              ? chunk.error.code
-              : undefined
-          const err = new Error(message)
-          if (code) {
-            ;(err as Error & { code?: string }).code = code
-          }
-          throw err
+        for (const chunk of normalizeStreamChunk(raw as AdapterYieldChunk)) {
+          // Surface failures: the underlying chatStream emits RUN_ERROR instead
+          // of throwing, so without this branch summarize() would return an
+          // empty summary and pretend a failed run succeeded.
+          if (chunk.type === EventType.RUN_ERROR) throwRunError(chunk)
+          consumeSpecSummarizeChunk(chunk, state)
         }
       }
     } catch (error: unknown) {
@@ -259,7 +279,12 @@ export class ChatStreamSummarizeAdapter<
       throw error
     }
 
-    return { id, model, summary, usage }
+    return {
+      id,
+      model: state.model,
+      summary: state.summary,
+      usage: state.usage,
+    }
   }
 
   override async *summarizeStream(
@@ -273,46 +298,45 @@ export class ChatStreamSummarizeAdapter<
     )
 
     const id = this.generateId()
-    let summary = ''
-    let model = options.model
-    let usage: SummarizationResult['usage'] = {
-      promptTokens: 0,
-      completionTokens: 0,
-      totalTokens: 0,
+    const state = {
+      summary: '',
+      model: options.model,
+      usage: {
+        promptTokens: 0,
+        completionTokens: 0,
+        totalTokens: 0,
+      } satisfies SummarizationResult['usage'],
     }
 
     try {
-      for await (const chunk of this.textAdapter.chatStream(
+      for await (const raw of this.textAdapter.chatStream(
         this.buildTextOptions(options, systemPrompt),
       )) {
-        // Accumulate the same way `summarize()` does so consumers see deltas
-        // AND the terminal `generation:result` event below carries the same
-        // final summary that non-streaming returns.
-        if (chunk.type === 'TEXT_MESSAGE_CONTENT') {
-          if (chunk.content) {
-            summary = chunk.content
-          } else if (chunk.delta) {
-            summary += chunk.delta
+        for (const chunk of normalizeStreamChunk(raw as AdapterYieldChunk)) {
+          // Accumulate the same way `summarize()` does so consumers see deltas
+          // AND the terminal `generation:result` event below carries the same
+          // final summary that non-streaming returns.
+          consumeSpecSummarizeChunk(chunk, state)
+
+          // Emit the GenerationClient-shaped result event just before the
+          // terminal RUN_FINISHED so subscribers (useSummarize) populate
+          // `result` before flipping `status` to success.
+          if (chunk.type === EventType.RUN_FINISHED) {
+            yield {
+              type: EventType.CUSTOM,
+              name: 'generation:result',
+              value: {
+                id,
+                model: state.model,
+                summary: state.summary,
+                usage: state.usage,
+              } satisfies SummarizationResult,
+              timestamp: Date.now(),
+            }
           }
-          if (chunk.model) model = chunk.model
-        }
 
-        // Emit the GenerationClient-shaped result event just before the
-        // terminal RUN_FINISHED so subscribers (useSummarize) populate
-        // `result` before flipping `status` to success.
-        if (chunk.type === 'RUN_FINISHED') {
-          if (chunk.usage) usage = chunk.usage
-          if (chunk.model) model = chunk.model
-          yield {
-            type: EventType.CUSTOM,
-            name: 'generation:result',
-            value: { id, model, summary, usage } satisfies SummarizationResult,
-            model,
-            timestamp: Date.now(),
-          }
+          yield chunk
         }
-
-        yield chunk
       }
     } catch (error: unknown) {
       options.logger.errors(`${this.name}.summarizeStream fatal`, {

package/src/adapter-internals.ts

@@ -57,3 +57,5 @@ export {
   structuredOutputCompleteChunk,
   structuredOutputStartChunk,
 } from './utilities/structured-output-events'
+export { tanstackMetadata } from './utilities/merge-metadata'
+export { isSpecTopLevelKey } from './utilities/spec-event-keys'

package/src/client.ts

@@ -297,6 +297,17 @@ export type {
 } from './activities/chat/stream/index'
 
 export { uiMessagesToWire } from './utilities/ag-ui-wire'
+export {
+  mergeMetadata,
+  tanstackMetadata,
+  withTanstackMetadata,
+} from './utilities/merge-metadata'
+export { fromSpecTokenUsage, toSpecTokenUsage } from './utilities/ag-ui-usage'
+export type { SpecTokenUsage } from './utilities/ag-ui-usage'
+export { normalizeStreamChunk } from './utilities/normalize-stream-chunk'
+export { restoreInboundChunk } from './utilities/restore-inbound-chunk'
+export type { AdapterYieldChunk } from './utilities/adapter-yield-chunk'
+export { getChunkRunId, getChunkThreadId } from './utilities/chunk-ids'
 export type { WireMessage } from './utilities/ag-ui-wire'
 
 export type {
@@ -326,6 +337,8 @@ export type {
   StreamChunk,
   StructuredOutputPart,
   TextPart,
+  TanStackMessageMetadata,
+  TanStackRunMetadata,
   ThinkingPart,
   ToolCall,
   ToolCallPart,

package/src/index.ts

@@ -459,6 +459,12 @@ export { generationParamsFromBody, generationParamsFromRequest } from './client'
 
 // AG-UI wire serialization (used internally by @tanstack/ai-client)
 export { uiMessagesToWire } from './utilities/ag-ui-wire'
+export { mergeMetadata, withTanstackMetadata } from './utilities/merge-metadata'
+export { fromSpecTokenUsage, toSpecTokenUsage } from './utilities/ag-ui-usage'
+export type { SpecTokenUsage } from './utilities/ag-ui-usage'
+export { normalizeStreamChunk } from './utilities/normalize-stream-chunk'
+export type { AdapterYieldChunk } from './utilities/adapter-yield-chunk'
+export { getChunkRunId, getChunkThreadId } from './utilities/chunk-ids'
 export type { WireMessage } from './utilities/ag-ui-wire'
 export {
   isContentPart,

package/src/middlewares/otel.ts

@@ -10,6 +10,8 @@ import {
 } from '../utilities/sampling-keys'
 import { firstNumber } from '../utilities/numbers'
 import { errorMessage, errorTypeName } from '../utilities/errors'
+import { rebuildTokenUsage } from '../utilities/ag-ui-usage'
+import { tanstackMetadata } from '../utilities/merge-metadata'
 import { usageAttributes } from './usage-attributes'
 import type {
   AttributeValue,
@@ -631,26 +633,32 @@ export function otelMiddleware(
         }
 
         if (chunk.type !== 'RUN_FINISHED') return
+        const tanstack = tanstackMetadata(chunk)
+        const extra = chunk as {
+          finishReason?: string | null
+          model?: string
+        }
+        const finishReason = extra.finishReason ?? tanstack?.finishReason
+        const model = extra.model ?? tanstack?.model
         // Capture for the root-span finish_reasons attribute set in onFinish,
         // which receives base-shaped info without a finishReason field.
-        if (chunk.finishReason) state.lastFinishReason = chunk.finishReason
+        if (finishReason) state.lastFinishReason = finishReason
         const span = state.currentIterationSpan
         if (!span) return
 
-        if (chunk.finishReason) {
-          span.setAttribute('gen_ai.response.finish_reasons', [
-            chunk.finishReason,
-          ])
+        if (finishReason) {
+          span.setAttribute('gen_ai.response.finish_reasons', [finishReason])
         }
-        if (chunk.model) span.setAttribute('gen_ai.response.model', chunk.model)
+        if (model) span.setAttribute('gen_ai.response.model', model)
 
         // Set usage attributes on the iteration span directly from the chunk
         // so they're available before `onUsage` fires. Histogram recording is
         // deliberately NOT done here — the chat runner always invokes
         // `runOnUsage` when `chunk.usage` is present, and `onUsage` is the
         // canonical place for the metric. Recording in both would double-count.
-        if (chunk.usage) {
-          span.setAttributes(usageAttributes(chunk.usage))
+        const tokenUsage = rebuildTokenUsage(chunk.usage, tanstack?.usage)
+        if (tokenUsage) {
+          span.setAttributes(usageAttributes(tokenUsage))
         }
 
         if (captureContent && state.assistantTextBuffer.length > 0) {

package/src/stream-to-response.ts

@@ -8,6 +8,7 @@ import { wasRunDetached } from './delivery-detach'
 import { notifyRunDisconnected } from './delivery-disconnect'
 import { resolveResumeRunId } from './stream-durability'
 import { EventType } from './types'
+import { toWireChunk } from './strip-to-spec-middleware'
 import { resolveDebugOption } from './logger/resolve'
 import { runErrorEventToError } from './utilities/errors'
 import type { LockStore } from './activities/chat/middleware/locks'
@@ -287,10 +288,13 @@ function sseEncoders(
     encodeChunk: (chunk, index) => {
       const id = getId?.(chunk, index)
       const idLine = id === undefined ? '' : `id: ${id}\n`
-      return encoder.encode(`${idLine}data: ${JSON.stringify(chunk)}\n\n`)
+      const wire = toWireChunk(chunk)
+      return encoder.encode(`${idLine}data: ${JSON.stringify(wire)}\n\n`)
     },
     encodeError: (error) =>
-      encoder.encode(`data: ${JSON.stringify(runErrorChunk(error))}\n\n`),
+      encoder.encode(
+        `data: ${JSON.stringify(toWireChunk(runErrorChunk(error)))}\n\n`,
+      ),
   }
 }
 
@@ -1069,12 +1073,15 @@ function ndjsonEncoders(
   return {
     encodeChunk: (chunk, index) => {
       const id = getId?.(chunk, index)
+      const wire = toWireChunk(chunk)
       const line =
-        id === undefined ? JSON.stringify(chunk) : JSON.stringify({ id, chunk })
+        id === undefined
+          ? JSON.stringify(wire)
+          : JSON.stringify({ id, chunk: wire })
       return encoder.encode(`${line}\n`)
     },
     encodeError: (error) =>
-      encoder.encode(`${JSON.stringify(runErrorChunk(error))}\n`),
+      encoder.encode(`${JSON.stringify(toWireChunk(runErrorChunk(error)))}\n`),
   }
 }
 

package/src/stream-to-websocket.ts

@@ -1,5 +1,6 @@
 import { chatParamsFromRequestBody } from './utilities/chat-params'
 import { durableStreamSource, runErrorChunk } from './stream-to-response'
+import { toWireChunk } from './strip-to-spec-middleware'
 import { resolveDebugOption } from './logger/resolve'
 import type { StreamDurability } from './stream-durability'
 import type { DebugOption } from './logger/types'
@@ -35,7 +36,8 @@ export function encodeWsFrame(
   chunk: StreamChunk,
   id: string | undefined,
 ): string {
-  return JSON.stringify(id === undefined ? chunk : { id, chunk })
+  const wire = toWireChunk(chunk)
+  return JSON.stringify(id === undefined ? wire : { id, chunk: wire })
 }
 
 /**

package/src/strip-to-spec-middleware.ts

@@ -1,34 +1,58 @@
-import type { ChatMiddleware } from './activities/chat/middleware/types'
 import type { StreamChunk } from './types'
+import { EventType } from './types'
+import type { AdapterYieldChunk } from './utilities/adapter-yield-chunk'
+import { isTanstackUsage, toSpecTokenUsage } from './utilities/ag-ui-usage'
+import {
+  tanstackMetadata,
+  withTanstackMetadata,
+} from './utilities/merge-metadata'
+import { normalizeStreamChunk } from './utilities/normalize-stream-chunk'
+import { isSpecTopLevelKey } from './utilities/spec-event-keys'
 
 /**
- * Strip only the deprecated nested `error` object from RUN_ERROR events.
- * The flat `message`/`code` fields are the spec-compliant form.
- *
- * All other fields pass through unchanged. @ag-ui/core's BaseEventSchema
- * uses `.passthrough()`, so extra fields (model, content, usage,
- * finishReason, toolName, stepId, etc.) are allowed and won't break
- * spec validation or verifyEvents.
+ * Delete unknown top-level keys from a stream chunk.
+ * Keep only AG-UI spec keys for this event type.
+ * Convert TanStack TokenUsage objects to the spec `usage[]` array.
  */
-export function stripToSpec(chunk: StreamChunk): StreamChunk {
-  // Only strip the deprecated nested error object from RUN_ERROR.
-  if (chunk.type === 'RUN_ERROR' && 'error' in chunk) {
-    const { error: _deprecated, ...rest } = chunk
-    return rest
+export function stripToSpec(
+  chunk: StreamChunk | AdapterYieldChunk,
+): StreamChunk {
+  const out: Record<string, unknown> = {}
+  for (const [key, value] of Object.entries(chunk)) {
+    if (isSpecTopLevelKey(chunk.type, key) && value !== undefined) {
+      out[key] = value
+    }
   }
-  return chunk
+
+  if (
+    (chunk.type === EventType.RUN_FINISHED ||
+      chunk.type === EventType.RUN_ERROR) &&
+    isTanstackUsage(out.usage)
+  ) {
+    const model = tanstackMetadata(chunk)?.model
+    const { usage, leftover } = toSpecTokenUsage(out.usage, {
+      model: typeof model === 'string' ? model : undefined,
+    })
+    out.usage = usage
+    if (leftover !== undefined) {
+      return withTanstackMetadata(out as StreamChunk, {
+        usage: leftover,
+      }) as StreamChunk
+    }
+  }
+
+  return out as StreamChunk
 }
 
 /**
- * Middleware that ensures events are AG-UI spec compliant.
- * Currently only strips the deprecated nested `error` object from RUN_ERROR.
- * All other fields pass through unchanged (passthrough allowed by spec).
+ * Move TanStack extras into `metadata.tanstack`, then keep only spec keys.
+ * Custom servers that skip `chat()` still round-trip `finishReason` on SSE/HTTP/WS.
+ * Fan-out extras (encrypted-value, TOOL_CALL_RESULT) stay on the `chat()` path;
+ * this encoder is 1:1 with the durability log offset.
  */
-export function stripToSpecMiddleware(): ChatMiddleware {
-  return {
-    name: 'strip-to-spec',
-    onChunk(_ctx, chunk) {
-      return stripToSpec(chunk)
-    },
-  }
+export function toWireChunk(
+  chunk: StreamChunk | AdapterYieldChunk,
+): StreamChunk {
+  const [normalized] = normalizeStreamChunk(chunk)
+  return stripToSpec(normalized ?? chunk)
 }