@tanstack/ai 0.47.3 → 0.48.0

package/src/utilities/adapter-yield-chunk.ts

@@ -0,0 +1,30 @@
+import type { InterruptSubmissionError } from '../interrupts'
+import type { ContentPart, StreamChunk, ToolOutputState } from '../types'
+
+/**
+ * Adapter / engine yield before normalize. Public StreamChunk is spec-only.
+ * This type still allows the old extra fields.
+ */
+export type AdapterYieldChunk = StreamChunk & {
+  model?: string
+  finishReason?: 'stop' | 'length' | 'content_filter' | 'tool_calls' | null
+  // `any` so TokenUsage adapter yields stay assignable after public usage[]
+  usage?: any
+  content?: string
+  args?: string
+  toolName?: string
+  toolCallName?: string
+  index?: number
+  input?: unknown
+  output?: unknown
+  result?: string | Array<ContentPart>
+  state?: ToolOutputState | Record<string, unknown>
+  stepId?: string
+  stepType?: string
+  delta?: string | ReadonlyArray<unknown>
+  signature?: string
+  error?: { message: string; code?: string }
+  'tanstack:interruptErrors'?: ReadonlyArray<InterruptSubmissionError>
+  threadId?: string
+  runId?: string
+}

package/src/utilities/ag-ui-usage.test.ts

@@ -0,0 +1,194 @@
+import { describe, expect, it } from 'vitest'
+import {
+  fromSpecTokenUsage,
+  rebuildTokenUsage,
+  toSpecTokenUsage,
+} from './ag-ui-usage'
+import type { TokenUsage } from '../types'
+
+const rich: TokenUsage = {
+  promptTokens: 100,
+  completionTokens: 50,
+  totalTokens: 150,
+  promptTokensDetails: { cachedTokens: 10 },
+  completionTokensDetails: { reasoningTokens: 7 },
+  billed: { quantity: 3, unit: 'units' },
+  cost: 0.02,
+  costDetails: { upstreamInputCost: 0.01, upstreamOutputCost: 0.01 },
+  providerUsageDetails: { foo: 1 },
+  durationSeconds: 2,
+  unitsBilled: 3,
+}
+
+describe('toSpecTokenUsage', () => {
+  it('maps core fields to spec usage[0] and keeps leftovers under tanstack.usage', () => {
+    const { usage, leftover } = toSpecTokenUsage(rich, {
+      provider: 'openai',
+      model: 'gpt-5.5',
+    })
+    expect(usage).toEqual([
+      {
+        provider: 'openai',
+        model: 'gpt-5.5',
+        inputTokens: 100,
+        outputTokens: 50,
+        totalTokens: 150,
+        cachedInputTokens: 10,
+        reasoningTokens: 7,
+      },
+    ])
+    expect(leftover).toEqual({
+      billed: { quantity: 3, unit: 'units' },
+      cost: 0.02,
+      costDetails: { upstreamInputCost: 0.01, upstreamOutputCost: 0.01 },
+      providerUsageDetails: { foo: 1 },
+      durationSeconds: 2,
+      unitsBilled: 3,
+    })
+  })
+
+  it('returns undefined leftover when only mapped token fields exist', () => {
+    const { leftover } = toSpecTokenUsage({
+      promptTokens: 1,
+      completionTokens: 2,
+      totalTokens: 3,
+      promptTokensDetails: { cachedTokens: 10 },
+      completionTokensDetails: { reasoningTokens: 7 },
+    })
+    expect(leftover).toBeUndefined()
+  })
+
+  it('keeps unmapped detail keys in leftover', () => {
+    const { usage, leftover } = toSpecTokenUsage({
+      promptTokens: 1,
+      completionTokens: 2,
+      totalTokens: 3,
+      promptTokensDetails: { cachedTokens: 10, audioTokens: 4 },
+      completionTokensDetails: { reasoningTokens: 7, audioTokens: 1 },
+    })
+    expect(usage[0]?.cachedInputTokens).toBe(10)
+    expect(usage[0]?.reasoningTokens).toBe(7)
+    expect(leftover).toEqual({
+      promptTokensDetails: { audioTokens: 4 },
+      completionTokensDetails: { audioTokens: 1 },
+    })
+  })
+})
+
+describe('fromSpecTokenUsage', () => {
+  it('rebuilds TokenUsage from spec array plus leftover', () => {
+    const { usage, leftover } = toSpecTokenUsage(rich, {
+      provider: 'openai',
+      model: 'gpt-5.5',
+    })
+    expect(fromSpecTokenUsage(usage, leftover)).toEqual(rich)
+  })
+
+  it('rebuilds from spec array alone', () => {
+    expect(
+      fromSpecTokenUsage([{ inputTokens: 4, outputTokens: 5, totalTokens: 9 }]),
+    ).toEqual({
+      promptTokens: 4,
+      completionTokens: 5,
+      totalTokens: 9,
+    })
+  })
+
+  it('rebuilds spec-only cachedInputTokens and reasoningTokens into details', () => {
+    expect(
+      fromSpecTokenUsage([
+        {
+          inputTokens: 1,
+          outputTokens: 2,
+          totalTokens: 3,
+          cachedInputTokens: 10,
+          reasoningTokens: 7,
+        },
+      ]),
+    ).toEqual({
+      promptTokens: 1,
+      completionTokens: 2,
+      totalTokens: 3,
+      promptTokensDetails: { cachedTokens: 10 },
+      completionTokensDetails: { reasoningTokens: 7 },
+    })
+  })
+
+  it('keeps spec cachedTokens and leftover extra promptTokensDetails keys', () => {
+    expect(
+      fromSpecTokenUsage(
+        [
+          {
+            inputTokens: 1,
+            outputTokens: 2,
+            totalTokens: 3,
+            cachedInputTokens: 10,
+          },
+        ],
+        { promptTokensDetails: { audioTokens: 5 } },
+      ),
+    ).toEqual({
+      promptTokens: 1,
+      completionTokens: 2,
+      totalTokens: 3,
+      promptTokensDetails: { cachedTokens: 10, audioTokens: 5 },
+    })
+  })
+
+  it('keeps spec reasoningTokens and leftover extra completionTokensDetails keys', () => {
+    expect(
+      fromSpecTokenUsage(
+        [
+          {
+            inputTokens: 1,
+            outputTokens: 2,
+            totalTokens: 3,
+            reasoningTokens: 7,
+          },
+        ],
+        { completionTokensDetails: { audioTokens: 4 } },
+      ),
+    ).toEqual({
+      promptTokens: 1,
+      completionTokens: 2,
+      totalTokens: 3,
+      completionTokensDetails: { reasoningTokens: 7, audioTokens: 4 },
+    })
+  })
+
+  it('returns undefined for missing or empty spec usage without leftover', () => {
+    expect(fromSpecTokenUsage(undefined)).toBeUndefined()
+    expect(fromSpecTokenUsage([])).toBeUndefined()
+  })
+})
+
+describe('rebuildTokenUsage', () => {
+  it('returns TokenUsage objects unchanged', () => {
+    expect(
+      rebuildTokenUsage({
+        promptTokens: 1,
+        completionTokens: 2,
+        totalTokens: 3,
+        cost: 0.01,
+      }),
+    ).toEqual({
+      promptTokens: 1,
+      completionTokens: 2,
+      totalTokens: 3,
+      cost: 0.01,
+    })
+  })
+
+  it('rebuilds spec usage[] plus leftover into TokenUsage', () => {
+    expect(
+      rebuildTokenUsage([{ inputTokens: 4, outputTokens: 5, totalTokens: 9 }], {
+        cost: 0.02,
+      }),
+    ).toEqual({
+      promptTokens: 4,
+      completionTokens: 5,
+      totalTokens: 9,
+      cost: 0.02,
+    })
+  })
+})

package/src/utilities/ag-ui-usage.ts

@@ -0,0 +1,148 @@
+import type { TokenUsage } from '../types'
+
+/** AG-UI spec `usage[]` item (provider/model labels + token counts only). */
+export interface SpecTokenUsage {
+  provider?: string
+  model?: string
+  inputTokens?: number
+  outputTokens?: number
+  totalTokens?: number
+  reasoningTokens?: number
+  cachedInputTokens?: number
+}
+
+export interface ToSpecTokenUsageOptions {
+  provider?: string
+  model?: string
+}
+
+/** TokenUsage fields that have no AG-UI `usage[]` equivalent. */
+export type TokenUsageLeftover = Omit<
+  TokenUsage,
+  'promptTokens' | 'completionTokens' | 'totalTokens'
+>
+
+function definedDetails<T extends object>(value: T): T | undefined {
+  return Object.keys(value).length > 0 ? value : undefined
+}
+
+function withoutKey<T extends object, K extends keyof T>(
+  value: T,
+  key: K,
+): Omit<T, K> {
+  const next = { ...value }
+  delete next[key]
+  return next
+}
+
+export function isTanstackUsage(usage: unknown): usage is TokenUsage {
+  return (
+    typeof usage === 'object' &&
+    usage != null &&
+    !Array.isArray(usage) &&
+    'promptTokens' in usage
+  )
+}
+
+export function toSpecTokenUsage(
+  usage: TokenUsage,
+  options?: ToSpecTokenUsageOptions,
+): { usage: Array<SpecTokenUsage>; leftover?: TokenUsageLeftover } {
+  const {
+    promptTokens,
+    completionTokens,
+    totalTokens,
+    promptTokensDetails,
+    completionTokensDetails,
+    ...rest
+  } = usage
+
+  const spec: SpecTokenUsage = {
+    ...(options?.provider !== undefined ? { provider: options.provider } : {}),
+    ...(options?.model !== undefined ? { model: options.model } : {}),
+    inputTokens: promptTokens,
+    outputTokens: completionTokens,
+    totalTokens,
+  }
+  const cachedInputTokens = promptTokensDetails?.cachedTokens
+  if (cachedInputTokens !== undefined) {
+    spec.cachedInputTokens = cachedInputTokens
+  }
+  const reasoningTokens = completionTokensDetails?.reasoningTokens
+  if (reasoningTokens !== undefined) {
+    spec.reasoningTokens = reasoningTokens
+  }
+
+  const leftoverPrompt = promptTokensDetails
+    ? definedDetails(withoutKey(promptTokensDetails, 'cachedTokens'))
+    : undefined
+  const leftoverCompletion = completionTokensDetails
+    ? definedDetails(withoutKey(completionTokensDetails, 'reasoningTokens'))
+    : undefined
+
+  return {
+    usage: [spec],
+    leftover: definedDetails({
+      ...rest,
+      ...(leftoverPrompt !== undefined
+        ? { promptTokensDetails: leftoverPrompt }
+        : {}),
+      ...(leftoverCompletion !== undefined
+        ? { completionTokensDetails: leftoverCompletion }
+        : {}),
+    }),
+  }
+}
+
+export function rebuildTokenUsage(
+  usage: unknown,
+  leftover?: TokenUsageLeftover,
+): TokenUsage | undefined {
+  if (isTanstackUsage(usage)) {
+    return usage
+  }
+  if (Array.isArray(usage)) {
+    return fromSpecTokenUsage(usage, leftover)
+  }
+  return fromSpecTokenUsage(undefined, leftover)
+}
+
+export function fromSpecTokenUsage(
+  usage: ReadonlyArray<SpecTokenUsage> | undefined,
+  leftover?: TokenUsageLeftover,
+): TokenUsage | undefined {
+  const spec = usage?.[0]
+  if (spec == null && leftover == null) {
+    return undefined
+  }
+
+  const {
+    promptTokensDetails: leftoverPromptDetails,
+    completionTokensDetails: leftoverCompletionDetails,
+    ...leftoverRest
+  } = leftover ?? {}
+
+  const promptTokensDetails = definedDetails({
+    ...(spec?.cachedInputTokens !== undefined
+      ? { cachedTokens: spec.cachedInputTokens }
+      : {}),
+    ...leftoverPromptDetails,
+  })
+  const completionTokensDetails = definedDetails({
+    ...(spec?.reasoningTokens !== undefined
+      ? { reasoningTokens: spec.reasoningTokens }
+      : {}),
+    ...leftoverCompletionDetails,
+  })
+
+  return {
+    promptTokens: spec?.inputTokens ?? 0,
+    completionTokens: spec?.outputTokens ?? 0,
+    totalTokens: spec?.totalTokens ?? 0,
+    ...leftoverRest,
+    ...(promptTokensDetails !== undefined ? { promptTokensDetails } : {}),
+    ...(completionTokensDetails !== undefined
+      ? { completionTokensDetails }
+      : {}),
+  }
+}

package/src/utilities/ag-ui-wire.ts

@@ -1,4 +1,13 @@
-import type { ContentPart, MessagePart, UIMessage } from '../types'
+import type {
+  ContentPart,
+  MessagePart,
+  StructuredOutputPart,
+  TanStackMessageMetadata,
+  UIMessage,
+  UIResourcePart,
+} from '../types'
+import type { MetadataRecord } from './merge-metadata'
+import { tanstackMetadata } from './merge-metadata'
 
 type AGUITextInputContent = { type: 'text'; text: string }
 type AGUIInputContent =
@@ -9,6 +18,7 @@ type AGUIToolCallMirror = {
   id: string
   type: 'function'
   function: { name: string; arguments: string }
+  encryptedValue?: string
 }
 
 type AGUIToolMessage = {
@@ -23,11 +33,18 @@ type AGUIReasoningMessage = {
   role: 'reasoning'
   id: string
   content: string
+  encryptedValue?: string
+  metadata?: MetadataRecord
 }
 
-type WireAnchorMessage = UIMessage & {
+/** Spec AG-UI message. No `parts`, no `createdAt` Date. */
+type WireAnchorMessage = {
+  id: string
+  role: UIMessage['role']
+  name?: string
   content?: string | Array<AGUIInputContent>
   toolCalls?: Array<AGUIToolCallMirror>
+  metadata?: MetadataRecord
 }
 
 export type WireMessage =
@@ -37,12 +54,10 @@ export type WireMessage =
 
 /**
  * Serialize TanStack `UIMessage`s into the AG-UI `RunAgentInput.messages`
- * wire shape. Each anchor (system/user/assistant) carries the canonical
- * `parts` array verbatim plus AG-UI mirror fields (`content`, `toolCalls`)
- * so AG-UI Zod parsing succeeds. Tool results and thinking parts on
- * assistant messages are additionally emitted as fan-out
- * `{role:'tool',...}` and `{role:'reasoning',...}` entries for strict
- * AG-UI server consumers.
+ * wire shape. Anchors are spec-only (`id`, `role`, `name`, `content`,
+ * `toolCalls`, `metadata`). Tool results and thinking parts on assistant
+ * messages are additionally emitted as fan-out `{role:'tool',...}` and
+ * `{role:'reasoning',...}` entries for strict AG-UI server consumers.
  */
 export function uiMessagesToWire(
   messages: Array<UIMessage>,
@@ -50,52 +65,70 @@ export function uiMessagesToWire(
   const wire: Array<WireMessage> = []
 
   for (const msg of messages) {
-    // Defensive: if parts is missing (ModelMessage-shaped input), pass through as-is.
-    // UIMessage always has parts; ModelMessage uses content directly.
+    // Defensive: ModelMessage-shaped input has no `parts`; fall back to `content`.
     const parts: ReadonlyArray<MessagePart> =
       // eslint-disable-next-line @typescript-eslint/no-unnecessary-type-assertion -- runtime input may be ModelMessage-shaped (no `parts`); cast forces the optional-chain fallback below to remain in scope
       (msg.parts as ReadonlyArray<MessagePart> | undefined) ?? []
 
     if (msg.role === 'system') {
-      wire.push({
-        ...msg,
-        content:
-          parts.length > 0
-            ? collectText(parts)
-            : ((msg as { content?: string }).content ?? ''),
-      })
+      wire.push(
+        toAnchor(
+          msg,
+          {
+            content:
+              parts.length > 0
+                ? collectText(parts)
+                : ((msg as { content?: string }).content ?? ''),
+          },
+          parts,
+        ),
+      )
       continue
     }
 
     if (msg.role === 'user') {
-      wire.push({
-        ...msg,
-        content:
-          parts.length > 0
-            ? collectUserContent(parts)
-            : ((msg as { content?: string }).content ?? ''),
-      })
+      wire.push(
+        toAnchor(
+          msg,
+          {
+            content:
+              parts.length > 0
+                ? collectUserContent(parts)
+                : ((msg as { content?: string }).content ?? ''),
+          },
+          parts,
+        ),
+      )
       continue
     }
 
     // assistant: emit reasoning fan-outs first, then anchor, then tool fan-outs
     for (const part of parts) {
       if (part.type === 'thinking') {
-        wire.push({
+        const reasoning: AGUIReasoningMessage = {
           role: 'reasoning',
           id: deriveReasoningId(msg.id, part),
           content: part.content,
-        })
+        }
+        if (part.signature) {
+          reasoning.encryptedValue = part.signature
+        }
+        wire.push(reasoning)
       }
     }
 
     const text = collectText(parts)
     const toolCalls = collectToolCalls(parts)
-    wire.push({
-      ...msg,
-      ...(text !== '' && { content: text }),
-      ...(toolCalls && { toolCalls }),
-    })
+    wire.push(
+      toAnchor(
+        msg,
+        {
+          ...(text !== '' && { content: text }),
+          ...(toolCalls && { toolCalls }),
+        },
+        parts,
+      ),
+    )
 
     for (const part of parts) {
       if (part.type === 'tool-result') {
@@ -116,6 +149,76 @@ export function uiMessagesToWire(
   return wire
 }
 
+function toAnchor(
+  msg: UIMessage,
+  extras: {
+    content?: string | Array<AGUIInputContent>
+    toolCalls?: Array<AGUIToolCallMirror>
+  },
+  parts: ReadonlyArray<MessagePart>,
+): WireAnchorMessage {
+  const metadata = messageMetadata(msg, parts)
+  const name = (msg as { name?: string }).name
+  return {
+    id: msg.id,
+    role: msg.role,
+    ...(name !== undefined && { name }),
+    ...extras,
+    ...(metadata !== undefined && { metadata }),
+  }
+}
+
+function messageMetadata(
+  msg: UIMessage,
+  parts: ReadonlyArray<MessagePart>,
+): MetadataRecord | undefined {
+  const base: MetadataRecord = { ...(msg.metadata ?? {}) }
+  const tanstack: MetadataRecord = { ...(tanstackMetadata(msg) ?? {}) }
+  if (msg.createdAt) tanstack.createdAt = msg.createdAt.toISOString()
+
+  const leftover = unfinishedStructuredOutput(parts)
+  if (leftover) tanstack.structuredOutput = leftover
+
+  const toolCallMetadata: Record<string, unknown> = {}
+  for (const part of parts) {
+    if (part.type === 'tool-call' && part.metadata !== undefined) {
+      toolCallMetadata[part.id] = part.metadata
+    }
+  }
+  if (Object.keys(toolCallMetadata).length > 0) {
+    tanstack.toolCallMetadata = toolCallMetadata
+  }
+
+  const uiResources = parts.filter(
+    (p): p is UIResourcePart => p.type === 'ui-resource',
+  )
+  if (uiResources.length > 0) tanstack.uiResources = uiResources
+
+  if (Object.keys(tanstack).length > 0) base.tanstack = tanstack
+  return Object.keys(base).length > 0 ? base : undefined
+}
+
+function unfinishedStructuredOutput(
+  parts: ReadonlyArray<MessagePart>,
+): TanStackMessageMetadata['structuredOutput'] | undefined {
+  for (const p of parts) {
+    if (p.type === 'structured-output' && p.status !== 'complete') {
+      return structuredOutputLeftover(p)
+    }
+  }
+  return undefined
+}
+
+function structuredOutputLeftover(
+  part: StructuredOutputPart,
+): NonNullable<TanStackMessageMetadata['structuredOutput']> {
+  return {
+    status: part.status,
+    raw: part.raw,
+    ...(part.errorMessage !== undefined && { errorMessage: part.errorMessage }),
+  }
+}
+
 function collectText(parts: ReadonlyArray<MessagePart>): string {
   // The streamed JSON of a completed structured-output part is the source of
   // truth for multi-turn coherence — emitting it back as assistant content
@@ -171,16 +274,31 @@ function collectUserContent(
   return out
 }
 
+function thoughtSignatureFromMetadata(metadata: unknown): string | undefined {
+  if (
+    metadata == null ||
+    typeof metadata !== 'object' ||
+    Array.isArray(metadata)
+  ) {
+    return undefined
+  }
+  if (!('thoughtSignature' in metadata)) return undefined
+  const value = metadata.thoughtSignature
+  return typeof value === 'string' && value !== '' ? value : undefined
+}
+
 function collectToolCalls(
   parts: ReadonlyArray<MessagePart>,
 ): Array<AGUIToolCallMirror> | undefined {
   const calls: Array<AGUIToolCallMirror> = []
   for (const p of parts) {
     if (p.type === 'tool-call') {
+      const encryptedValue = thoughtSignatureFromMetadata(p.metadata)
       calls.push({
         id: p.id,
         type: 'function',
         function: { name: p.name, arguments: p.arguments },
+        ...(encryptedValue !== undefined ? { encryptedValue } : {}),
       })
     }
   }