@tanstack/ai 0.58.0 → 0.61.0

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

@@ -1,15 +1,8 @@
+import type { TokenUsage as SpecTokenUsage } from '@ag-ui/core'
 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 type { TokenUsage as SpecTokenUsage } from '@ag-ui/core'
 
 export interface ToSpecTokenUsageOptions {
   provider?: string
@@ -68,11 +61,16 @@ export function toSpecTokenUsage(
   if (cachedInputTokens !== undefined) {
     spec.cachedInputTokens = cachedInputTokens
   }
+  if (promptTokensDetails?.cacheWriteTokens !== undefined) {
+    spec.cacheWriteInputTokens = promptTokensDetails.cacheWriteTokens
+  }
   const reasoningTokens = completionTokensDetails?.reasoningTokens
   if (reasoningTokens !== undefined) {
     spec.reasoningTokens = reasoningTokens
   }
 
+  // cacheWriteTokens stays in the leftover too, so clients built before
+  // cacheWriteInputTokens existed still rebuild it.
   const leftoverPrompt = promptTokensDetails
     ? definedDetails(withoutKey(promptTokensDetails, 'cachedTokens'))
     : undefined
@@ -94,6 +92,80 @@ export function toSpecTokenUsage(
   }
 }
 
+function sumNumbers<T extends object>(
+  current: T | undefined,
+  next: T | undefined,
+): T | undefined {
+  if (!current) return next
+  if (!next) return current
+  const result = { ...current }
+  for (const key of Object.keys(next) as Array<keyof T>) {
+    const value = next[key]
+    if (typeof value !== 'number') continue
+    const previous = current[key]
+    result[key] = ((typeof previous === 'number' ? previous : 0) +
+      value) as T[keyof T]
+  }
+  return result
+}
+
+function sumOptional(current?: number, next?: number): number | undefined {
+  if (current === undefined) return next
+  if (next === undefined) return current
+  return current + next
+}
+
+/**
+ * Add two usage totals, as for a parent run and its children. Numbers add
+ * up. `billed` adds up only in the same unit. `providerUsageDetails` is
+ * opaque, so the latest one stays. `@tanstack/ai-persistence` sums per-run
+ * usage with the same rules.
+ */
+export function addTokenUsage(
+  current: TokenUsage,
+  next: TokenUsage,
+): TokenUsage {
+  const promptTokensDetails = sumNumbers(
+    current.promptTokensDetails,
+    next.promptTokensDetails,
+  )
+  const completionTokensDetails = sumNumbers(
+    current.completionTokensDetails,
+    next.completionTokensDetails,
+  )
+  const costDetails = sumNumbers(current.costDetails, next.costDetails)
+  const cost = sumOptional(current.cost, next.cost)
+  const durationSeconds = sumOptional(
+    current.durationSeconds,
+    next.durationSeconds,
+  )
+  const unitsBilled = sumOptional(current.unitsBilled, next.unitsBilled)
+  const billed =
+    current.billed && next.billed && current.billed.unit === next.billed.unit
+      ? {
+          quantity: current.billed.quantity + next.billed.quantity,
+          unit: current.billed.unit,
+        }
+      : (next.billed ?? current.billed)
+  const providerUsageDetails =
+    next.providerUsageDetails ?? current.providerUsageDetails
+  return {
+    ...current,
+    ...next,
+    promptTokens: current.promptTokens + next.promptTokens,
+    completionTokens: current.completionTokens + next.completionTokens,
+    totalTokens: current.totalTokens + next.totalTokens,
+    ...(promptTokensDetails && { promptTokensDetails }),
+    ...(completionTokensDetails && { completionTokensDetails }),
+    ...(cost !== undefined && { cost }),
+    ...(costDetails && { costDetails }),
+    ...(durationSeconds !== undefined && { durationSeconds }),
+    ...(unitsBilled !== undefined && { unitsBilled }),
+    ...(billed && { billed }),
+    ...(providerUsageDetails && { providerUsageDetails }),
+  }
+}
+
 export function rebuildTokenUsage(
   usage: unknown,
   leftover?: TokenUsageLeftover,
@@ -111,8 +183,20 @@ export function fromSpecTokenUsage(
   usage: ReadonlyArray<SpecTokenUsage> | undefined,
   leftover?: TokenUsageLeftover,
 ): TokenUsage | undefined {
-  const spec = usage?.[0]
-  if (spec == null && leftover == null) {
+  const spec = usage?.reduce<SpecTokenUsage>((total, entry) => {
+    for (const key of [
+      'inputTokens',
+      'outputTokens',
+      'totalTokens',
+      'cachedInputTokens',
+      'cacheWriteInputTokens',
+      'reasoningTokens',
+    ] as const) {
+      if (entry[key] !== undefined) total[key] = (total[key] ?? 0) + entry[key]
+    }
+    return total
+  }, {})
+  if ((usage === undefined || usage.length === 0) && leftover == null) {
     return undefined
   }
 
@@ -126,6 +210,9 @@ export function fromSpecTokenUsage(
     ...(spec?.cachedInputTokens !== undefined
       ? { cachedTokens: spec.cachedInputTokens }
       : {}),
+    ...(spec?.cacheWriteInputTokens !== undefined
+      ? { cacheWriteTokens: spec.cacheWriteInputTokens }
+      : {}),
     ...leftoverPromptDetails,
   })
   const completionTokensDetails = definedDetails({

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

@@ -12,13 +12,17 @@ import type {
   MessagePart,
   ModelMessage,
   StructuredOutputPart,
+  SubagentPart,
   TanStackMessageMetadata,
   UIMessage,
   UIResourcePart,
 } from '../types'
 import type { MetadataRecord } from './merge-metadata'
 import { tanstackMetadata } from './merge-metadata'
+import { isProviderExecutedToolCall } from './provider-executed'
 import { normalizeToolResult } from './tool-result'
+import { wireSubagentInfo, wireSubagentRunId } from './subagent-wire'
+import type { SubagentWireInfo } from './subagent-wire'
 import {
   coerceCreatedAt,
   modelMessageToUIMessage,
@@ -168,9 +172,45 @@ export function uiMessagesToWire(
       continue
     }
 
-    // assistant: emit reasoning fan-outs first, then anchor, then tool fan-outs
+    // assistant: reasoning fan-outs, then anchor, then tool fan-outs.
+    //
+    // Provider-executed tools (Anthropic web_search / web_fetch) run inside one
+    // provider response, and the provider signs every thinking block against
+    // the blocks before it. Emitting all reasoning first and a single anchor
+    // with the joined text and every tool call turns
+    // "thinking, tool, thinking, text, tool" into
+    // "thinking, thinking, text, tool, tool", and the provider rejects the next
+    // turn ("thinking blocks in the latest assistant message cannot be
+    // modified"). Split the message into ordered segments at each thinking
+    // part that follows a provider-executed tool call, the same rule
+    // buildAssistantMessages applies, so the wire keeps the signed order.
+    // Segments after the first get derived anchor ids; strict AG-UI consumers
+    // still see plain anchors.
+    type Segment = {
+      thinking: Array<Extract<MessagePart, { type: 'thinking' }>>
+      parts: Array<MessagePart>
+    }
+    let current: Segment = { thinking: [], parts: [] }
+    const segments: Array<Segment> = [current]
     for (const part of parts) {
       if (part.type === 'thinking') {
+        if (
+          current.parts.some(
+            (p) => p.type === 'tool-call' && isProviderExecutedToolCall(p),
+          )
+        ) {
+          current = { thinking: [part], parts: [] }
+          segments.push(current)
+        } else {
+          current.thinking.push(part)
+        }
+      } else {
+        current.parts.push(part)
+      }
+    }
+
+    segments.forEach((segment, index) => {
+      for (const part of segment.thinking) {
         const reasoning: WireReasoningMessage = {
           role: 'reasoning',
           id: uniqueWireId(deriveReasoningId(uiMessage.id, part), usedWireIds),
@@ -181,22 +221,29 @@ export function uiMessagesToWire(
         }
         wire.push(reasoning)
       }
-    }
 
-    const text = collectText(parts)
-    const toolCalls = collectToolCalls(parts)
-    wire.push(
-      toAnchor(
-        uiMessage,
-        'assistant',
-        {
-          ...(text !== '' && { content: text }),
-          ...(toolCalls && { toolCalls }),
-        },
-        parts,
-        includeSnapshotStructuredOutput,
-      ),
-    )
+      const text = collectText(segment.parts)
+      const toolCalls = collectToolCalls(segment.parts)
+      const anchorMessage: UIMessage =
+        index === 0
+          ? uiMessage
+          : {
+              ...uiMessage,
+              id: uniqueWireId(`${uiMessage.id}-segment-${index}`, usedWireIds),
+            }
+      wire.push(
+        toAnchor(
+          anchorMessage,
+          'assistant',
+          {
+            ...(text !== '' && { content: text }),
+            ...(toolCalls && { toolCalls }),
+          },
+          segment.parts,
+          includeSnapshotStructuredOutput,
+        ),
+      )
+    })
 
     const explicitToolResults = new Set(
       parts.flatMap((part) =>
@@ -263,6 +310,10 @@ export function uiMessagesToWire(
         })
       }
     }
+
+    for (const part of parts) {
+      if (part.type === 'subagent') wire.push(...subagentToWire(part, options))
+    }
   }
 
   return wire
@@ -332,6 +383,8 @@ function messageMetadata(
   const tanstack: TanStackMessageMetadata = {}
   if (previousTanstack?.model !== undefined)
     tanstack.model = previousTanstack.model
+  if (previousTanstack?.runId !== undefined)
+    tanstack.runId = previousTanstack.runId
   if (previousTanstack?.signature !== undefined)
     tanstack.signature = previousTanstack.signature
   const createdAt = coerceCreatedAt(msg.createdAt)
@@ -424,6 +477,71 @@ function collectText(parts: ReadonlyArray<MessagePart>): string {
   return out.join('')
 }
 
+/**
+ * A child's messages, tagged with its AG-UI `subagentRunId`. The card data
+ * rides in `metadata.tanstack.subagent`, so the other side can rebuild the
+ * card and continue a suspended child.
+ */
+function subagentToWire(
+  part: SubagentPart,
+  options?: { includeSnapshotStructuredOutput: boolean },
+): Array<WireMessage> {
+  const { subagent } = part
+  const info: SubagentWireInfo = {
+    name: subagent.name,
+    status: subagent.status,
+    ...(subagent.description !== undefined && {
+      description: subagent.description,
+    }),
+    ...(subagent.error !== undefined && { error: subagent.error }),
+    ...(subagent.interruptIds !== undefined && {
+      interruptIds: subagent.interruptIds,
+    }),
+    ...(subagent.parentSubagentRunId !== undefined && {
+      parentSubagentRunId: subagent.parentSubagentRunId,
+    }),
+    ...(subagent.parentToolCallId !== undefined && {
+      parentToolCallId: subagent.parentToolCallId,
+    }),
+    ...(subagent.metadata !== undefined && { metadata: subagent.metadata }),
+  }
+  const child = uiMessagesToWire(subagent.messages, options)
+  const own = child.some((message) => wireSubagentRunId(message) === undefined)
+  const messages: Array<WireMessage> = own
+    ? child
+    : [{ id: `subagent:${subagent.id}`, role: 'assistant' }, ...child]
+  return messages.map((message, index) => {
+    const nestedId = wireSubagentRunId(message)
+    if (nestedId !== undefined) {
+      const nested = wireSubagentInfo(message)
+      if (!nested || nested.parentSubagentRunId !== undefined) return message
+      return withSubagentInfo(message, nestedId, {
+        ...nested,
+        parentSubagentRunId: subagent.id,
+      })
+    }
+    return withSubagentInfo(
+      message,
+      subagent.id,
+      !own && index === 0 ? { ...info, placeholder: true } : info,
+    )
+  })
+}
+
+function withSubagentInfo(
+  message: WireMessage,
+  subagentRunId: string,
+  info: SubagentWireInfo,
+): WireMessage {
+  const metadata = isRecord(message.metadata) ? message.metadata : {}
+  const tanstack = isRecord(metadata.tanstack) ? metadata.tanstack : {}
+  return {
+    ...message,
+    subagentRunId,
+    metadata: { ...metadata, tanstack: { ...tanstack, subagent: info } },
+  }
+}
+
 function collectUserContent(
   parts: ReadonlyArray<MessagePart>,
 ): string | Array<InputContent> {

package/src/utilities/content-source.ts

@@ -0,0 +1,138 @@
+import type { ContentPartFileSource, ContentPartSource } from '../types'
+
+/**
+ * Narrow a {@link ContentPartSource} to the provider-file-reference arm.
+ *
+ * Issuer adapters use this to route a file source to their native wire field;
+ * everyone else is protected by the core preflight (see
+ * {@link assertMessagesFileSourceSupport}) plus a defensive throw at their own
+ * mapping site.
+ */
+export function isFileSource(
+  source: ContentPartSource,
+): source is ContentPartFileSource {
+  return source.type === 'file'
+}
+
+/**
+ * Resolve the handle `providerName` should send for a file source.
+ *
+ * A file source carries one opaque handle (`value`) and, optionally, the
+ * provider that issued it. An adapter always knows which provider it talks
+ * to, so a source that names no provider is taken as-is.
+ *
+ * @throws when the source names a different issuing provider. A handle only
+ * resolves at the provider that minted it.
+ */
+export function fileReferenceFor(
+  source: ContentPartFileSource,
+  providerName: string,
+): string {
+  if (source.provider !== undefined && source.provider !== providerName) {
+    throw new Error(
+      `${providerName}: file source was issued by ${source.provider}. ` +
+        `A provider file handle only works with the provider that issued ` +
+        `it. Upload the file with ${providerName}Files(), or pass a data or ` +
+        `url source instead.`,
+    )
+  }
+  return source.value
+}
+
+/**
+ * Build the standard error a non-issuer adapter throws when it encounters a
+ * `{ type: 'file' }` source it can't consume — either because the provider has
+ * no file-reference input surface, or because the endpoint requires raw bytes
+ * (image edits, Veo) rather than a reference.
+ *
+ * @param detail Optional context appended to the message (e.g. a modality or
+ * endpoint name, or a pointer to the adapter that does support references).
+ * When provided it replaces the generic remediation tail, so a site-specific
+ * hint ("pass inline bytes") is never contradicted by generic advice.
+ */
+export function unsupportedFileSourceError(
+  providerName: string,
+  detail?: string,
+): Error {
+  return new Error(
+    `${providerName} does not support provider file-handle sources ` +
+      `({ type: 'file' })` +
+      (detail
+        ? ` ${detail}.`
+        : `. Pass a data or url source, or upload via the provider's files ` +
+          `adapter where supported.`),
+  )
+}
+
+/**
+ * The slice of an adapter the file-source preflight reads. Adapters that can
+ * consume `{ type: 'file' }` sources declare `supportsFileSources: true`;
+ * everything else — including adapters written before this arm existed —
+ * fails closed at the activity layer instead of falling through to a
+ * URL/data branch and silently mis-mapping the reference.
+ */
+export interface FileSourceCapable {
+  name: string
+  supportsFileSources?: boolean
+}
+
+/** True when a content-part-like value carries a `{ type: 'file' }` source. */
+function partHasFileSource(part: unknown): boolean {
+  if (typeof part !== 'object' || part === null) return false
+  const source = (part as { source?: unknown }).source
+  return (
+    typeof source === 'object' &&
+    source !== null &&
+    (source as { type?: unknown }).type === 'file'
+  )
+}
+
+/**
+ * True when `value` is a content part with a file source, or an array
+ * (possibly nested — fused embedding items) that contains one.
+ */
+function inputHasFileSource(value: unknown): boolean {
+  if (Array.isArray(value)) return value.some(inputHasFileSource)
+  return partHasFileSource(value)
+}
+
+/**
+ * Fail-closed preflight for media prompts and embedding inputs
+ * (`generateImage` / `generateVideo` / `embed`): throws when the input
+ * carries a `{ type: 'file' }` source and the adapter hasn't declared
+ * `supportsFileSources`. Runs in the activity dispatcher — the same layer
+ * that validates modality — so an adapter that predates the file arm can
+ * never receive one. Walks a single part, an array of parts, and nested
+ * arrays (fused embedding items).
+ */
+export function assertPromptFileSourceSupport(
+  adapter: FileSourceCapable,
+  prompt: unknown,
+): void {
+  if (adapter.supportsFileSources === true) return
+  if (inputHasFileSource(prompt)) {
+    throw unsupportedFileSourceError(adapter.name)
+  }
+}
+
+/**
+ * Fail-closed preflight for chat messages: throws when any message content
+ * part carries a `{ type: 'file' }` source and the adapter hasn't declared
+ * `supportsFileSources`. See {@link assertPromptFileSourceSupport}.
+ */
+export function assertMessagesFileSourceSupport(
+  adapter: FileSourceCapable,
+  messages: ReadonlyArray<unknown>,
+): void {
+  if (adapter.supportsFileSources === true) return
+  for (const message of messages) {
+    if (typeof message !== 'object' || message === null) continue
+    const content = (message as { content?: unknown }).content
+    if (!Array.isArray(content)) continue
+    for (const part of content) {
+      if (partHasFileSource(part)) {
+        throw unsupportedFileSourceError(adapter.name)
+      }
+    }
+  }
+}

package/src/utilities/normalize-stream-chunk.ts

@@ -11,6 +11,14 @@ function stringField(value: unknown): string | undefined {
   return typeof value === 'string' && value !== '' ? value : undefined
 }
 
+function chunkFields(chunk: object) {
+  const fields: Record<string, unknown> = {}
+  for (const key of Object.keys(chunk)) {
+    fields[key] = Reflect.get(chunk, key)
+  }
+  return fields
+}
+
 function encryptedValueExtras(chunk: AdapterYieldChunk): Array<StreamChunk> {
   const extras: Array<StreamChunk> = []
   const timestamp =
@@ -19,7 +27,7 @@ function encryptedValueExtras(chunk: AdapterYieldChunk): Array<StreamChunk> {
       : undefined
 
   if (typeof chunk.signature === 'string' && chunk.signature !== '') {
-    const source = chunk as Record<string, unknown>
+    const source = chunkFields(chunk)
     const toolCallId = stringField(source.toolCallId)
     const entityId =
       stringField(chunk.stepId) ??
@@ -62,7 +70,7 @@ export function normalizeStreamChunk(
   chunk: AdapterYieldChunk,
 ): Array<StreamChunk> {
   const specKeys = specKeysFor(chunk.type)
-  const source = chunk as Record<string, unknown>
+  const source = chunkFields(chunk)
   const specChunk: Record<string, unknown> & {
     metadata?: MetadataRecord | null
   } = {}

package/src/utilities/provider-executed.ts

@@ -30,3 +30,16 @@ export function isProviderExecutedToolCall(
 ): boolean {
   return getProviderExecutedMetadata(toolCall) !== null
 }
+
+/**
+ * True when `id` is a `${parentId}-segment-${n}` message. The wire and the
+ * run loop split one provider turn into segments at thinking that follows a
+ * provider-executed tool call, to keep signed thinking order. Readers fold a
+ * segment back into its parent so the UI still shows one message.
+ */
+export function isAssistantSegmentOf(
+  id: string | undefined,
+  parentId: string,
+): boolean {
+  return id?.startsWith(`${parentId}-segment-`) === true
+}

package/src/utilities/spec-event-keys.ts

@@ -30,11 +30,6 @@ const SPEC_KEYS: ReadonlyMap<string, ReadonlySet<string>> = new Map([
     EventType.TOOL_CALL_RESULT,
     keys('messageId', 'toolCallId', 'content', 'role'),
   ],
-  [EventType.THINKING_START, keys('title')],
-  [EventType.THINKING_END, SHARED],
-  [EventType.THINKING_TEXT_MESSAGE_START, SHARED],
-  [EventType.THINKING_TEXT_MESSAGE_CONTENT, keys('delta')],
-  [EventType.THINKING_TEXT_MESSAGE_END, SHARED],
   [EventType.STATE_SNAPSHOT, keys('snapshot')],
   [EventType.STATE_DELTA, keys('delta')],
   [EventType.MESSAGES_SNAPSHOT, keys('messages')],
@@ -45,7 +40,10 @@ const SPEC_KEYS: ReadonlyMap<string, ReadonlySet<string>> = new Map([
   [EventType.ACTIVITY_DELTA, keys('messageId', 'activityType', 'patch')],
   [EventType.RAW, keys('event', 'source')],
   [EventType.CUSTOM, keys('name', 'value')],
-  [EventType.RUN_STARTED, keys('threadId', 'runId', 'parentRunId', 'input')],
+  [
+    EventType.RUN_STARTED,
+    keys('threadId', 'runId', 'protocolVersion', 'parentRunId', 'input'),
+  ],
   [
     EventType.RUN_FINISHED,
     keys('threadId', 'runId', 'result', 'outcome', 'usage'),
@@ -63,10 +61,39 @@ const SPEC_KEYS: ReadonlyMap<string, ReadonlySet<string>> = new Map([
     EventType.REASONING_ENCRYPTED_VALUE,
     keys('subtype', 'entityId', 'encryptedValue'),
   ],
+  [
+    EventType.SUBAGENT_STARTED,
+    keys(
+      'subagentRunId',
+      'name',
+      'description',
+      'parentSubagentRunId',
+      'parentToolCallId',
+      'parentMessageId',
+    ),
+  ],
+  [EventType.SUBAGENT_FINISHED, keys('subagentRunId', 'result', 'outcome')],
+  [EventType.SUBAGENT_ERROR, keys('subagentRunId', 'message', 'code')],
 ])
 
+const RUN_SCOPED = new Set<string>([
+  EventType.RUN_STARTED,
+  EventType.RUN_FINISHED,
+  EventType.RUN_ERROR,
+  EventType.MESSAGES_SNAPSHOT,
+])
+
+const ATTRIBUTABLE_KEYS = new Map(
+  [...SPEC_KEYS]
+    .filter(([type]) => !RUN_SCOPED.has(type))
+    .map(
+      ([type, fields]) =>
+        [type, new Set([...fields, 'subagentRunId'])] as const,
+    ),
+)
+
 export function specKeysFor(type: string): ReadonlySet<string> {
-  return SPEC_KEYS.get(type) ?? SHARED
+  return ATTRIBUTABLE_KEYS.get(type) ?? SPEC_KEYS.get(type) ?? SHARED
 }
 
 export function isSpecTopLevelKey(type: string, key: string): boolean {