@tanstack/ai 0.47.3 → 0.49.1

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/interrupt-resume.ts

@@ -445,6 +445,20 @@ export async function validateInterruptResumeBatch(
       continue
     }
 
+    if (entry.status === 'cancelled') {
+      if (entry.payload !== undefined) {
+        errors.push(
+          interruptItemError(
+            input,
+            record.interruptId,
+            'invalid-payload',
+            `Cancelled interrupt ${record.interruptId} must not include a payload.`,
+          ),
+        )
+      }
+      continue
+    }
+
     const tool = runtimeTool(input.tools, binding.toolName)
     if (!tool) {
       errors.push(
@@ -511,19 +525,6 @@ export async function validateInterruptResumeBatch(
       }
     }
 
-    if (entry.status === 'cancelled') {
-      if (entry.payload !== undefined) {
-        errors.push(
-          interruptItemError(
-            input,
-            record.interruptId,
-            'invalid-payload',
-            `Cancelled interrupt ${record.interruptId} must not include a payload.`,
-          ),
-        )
-      }
-      continue
-    }
     if (schemaDrifted) continue
 
     if (binding.kind === 'client-tool-execution') {

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)
 }