@tanstack/ai 0.27.0 → 0.29.0

package/src/activities/chat/index.ts

@@ -25,6 +25,7 @@ import {
 import { maxIterations as maxIterationsStrategy } from './agent-loop-strategies'
 import { convertMessagesToModelMessages, generateMessageId } from './messages'
 import { MiddlewareRunner } from './middleware/compose'
+import { MCPManager } from './mcp/manager'
 import type {
   ApprovalRequest,
   ClientToolRequest,
@@ -69,6 +70,7 @@ import type {
   MergeContext,
   UnionToIntersection,
 } from './runtime-context-types'
+import type { ChatMCPOptions } from './mcp/types'
 
 // ===========================
 // Activity Kind
@@ -208,6 +210,12 @@ export interface TextActivityOptions<
         | ProviderTool<string, TAdapter['~types']['toolCapabilities'][number]>
       >
     | undefined
+  /**
+   * Hand MCP clients/pools to chat(): their tools are discovered at run start
+   * and merged into the run; `connection` controls whether chat() closes them
+   * when the run ends. See docs/tools/mcp.md "Managing MCP clients with chat()".
+   */
+  mcp?: ChatMCPOptions
   /** Additional metadata to attach to the request. */
   metadata?: TextOptions['metadata']
   /** Model-specific provider options (type comes from adapter) */
@@ -432,6 +440,28 @@ interface TextEngineConfig<
 type ToolPhaseResult = 'continue' | 'stop' | 'wait'
 type CyclePhase = 'processText' | 'executeToolCalls'
 
+/**
+ * Combine two optional AbortSignals into one that aborts when either does.
+ * Returns the other signal directly when one is absent or already aborted.
+ * (Manual implementation — `AbortSignal.any` requires Node >= 20.3.)
+ */
+function combineAbortSignals(
+  a: AbortSignal | undefined,
+  b: AbortSignal | undefined,
+): AbortSignal | undefined {
+  if (!a) return b
+  if (!b) return a
+  if (a.aborted) return a
+  if (b.aborted) return b
+  const controller = new AbortController()
+  const onAbort = (source: AbortSignal) => () => {
+    controller.abort(source.reason)
+  }
+  a.addEventListener('abort', onAbort(a), { once: true })
+  b.addEventListener('abort', onAbort(b), { once: true })
+  return controller.signal
+}
+
 class TextEngine<
   TAdapter extends AnyTextAdapter,
   TContext = unknown,
@@ -486,6 +516,9 @@ class TextEngine<
   private readonly deferredPromises: Array<Promise<unknown>> = []
   private abortReason?: string
   private readonly middlewareAbortController?: AbortController
+  // Combines the caller's signal with middleware abort() so running tools
+  // observe both cancellation sources via ctx.abortSignal.
+  private readonly toolAbortSignal?: AbortSignal
   private terminalHookCalled = false
 
   private readonly logger: InternalLogger
@@ -583,6 +616,10 @@ class TextEngine<
     ]
     this.middlewareRunner = new MiddlewareRunner(allMiddleware, logger)
     this.middlewareAbortController = new AbortController()
+    this.toolAbortSignal = combineAbortSignals(
+      this.effectiveSignal,
+      this.middlewareAbortController.signal,
+    )
     this.middlewareCtx = {
       requestId: this.requestId,
       streamId: this.streamId,
@@ -1225,6 +1262,7 @@ class TextEngine<
         },
       },
       this.middlewareCtx.context,
+      this.toolAbortSignal,
     )
 
     // Consume the async generator, yielding custom events and collecting the return value
@@ -1386,6 +1424,7 @@ class TextEngine<
         },
       },
       this.middlewareCtx.context,
+      this.toolAbortSignal,
     )
 
     // Consume the async generator, yielding custom events and collecting the return value
@@ -1639,8 +1678,9 @@ class TextEngine<
       const wireContent =
         typeof content === 'string' ? content : JSON.stringify(content)
 
-      // Emit TOOL_CALL_START + TOOL_CALL_ARGS before TOOL_CALL_END so that
-      // the client can reconstruct the full tool call during continuations.
+      // argsMap is set only on continuation re-executions, where the adapter
+      // never streamed these calls. Otherwise it already emitted END, so a
+      // second one here would be an orphan that fails verifyEvents (#519).
       if (argsMap) {
         chunks.push({
           type: 'TOOL_CALL_START',
@@ -1660,18 +1700,18 @@ class TextEngine<
           delta: args,
           args,
         } as StreamChunk)
-      }
 
-      chunks.push({
-        type: 'TOOL_CALL_END',
-        timestamp: Date.now(),
-        model: finishEvent.model,
-        toolCallId: result.toolCallId,
-        toolCallName: result.toolName,
-        toolName: result.toolName,
-        result: wireContent,
-        ...(result.state !== undefined && { state: result.state }),
-      } as StreamChunk)
+        chunks.push({
+          type: 'TOOL_CALL_END',
+          timestamp: Date.now(),
+          model: finishEvent.model,
+          toolCallId: result.toolCallId,
+          toolCallName: result.toolName,
+          toolName: result.toolName,
+          result: wireContent,
+          ...(result.state !== undefined && { state: result.state }),
+        } as StreamChunk)
+      }
 
       // AG-UI spec TOOL_CALL_RESULT event (content is string-only per spec)
       chunks.push({
@@ -2568,10 +2608,16 @@ export function chat<
 async function* runStreamingText<TContext = unknown>(
   options: TextActivityOptions<AnyTextAdapter, undefined, true, TContext>,
 ): AsyncIterable<StreamChunk> {
-  const { adapter, middleware, context, debug, ...textOptions } = options
+  const { adapter, middleware, context, debug, mcp, ...textOptions } = options
   const model = adapter.model
   const logger = resolveDebugOption(debug)
 
+  const mcpManager = MCPManager.from(mcp)
+  const mcpTools = await mcpManager.discover()
+  if (mcpTools.length > 0) {
+    textOptions.tools = [...(textOptions.tools ?? []), ...mcpTools]
+  }
+
   const engine = new TextEngine(
     {
       adapter,
@@ -2586,8 +2632,12 @@ async function* runStreamingText<TContext = unknown>(
     logger,
   )
 
-  for await (const chunk of engine.run()) {
-    yield chunk
+  try {
+    for await (const chunk of engine.run()) {
+      yield chunk
+    }
+  } finally {
+    await mcpManager.dispose()
   }
 }
 
@@ -2624,8 +2674,15 @@ async function runAgenticStructuredOutput<
 >(
   options: TextActivityOptions<AnyTextAdapter, TSchema, boolean, TContext>,
 ): Promise<InferSchemaType<TSchema>> {
-  const { adapter, outputSchema, middleware, context, debug, ...textOptions } =
-    options
+  const {
+    adapter,
+    outputSchema,
+    middleware,
+    context,
+    debug,
+    mcp,
+    ...textOptions
+  } = options
   const model = adapter.model
   const logger = resolveDebugOption(debug)
 
@@ -2659,6 +2716,12 @@ async function runAgenticStructuredOutput<
   const nativeCombined =
     adapter.supportsCombinedToolsAndSchema?.(options.modelOptions) === true
 
+  const mcpManager = MCPManager.from(mcp)
+  const mcpTools = await mcpManager.discover()
+  if (mcpTools.length > 0) {
+    textOptions.tools = [...(textOptions.tools ?? []), ...mcpTools]
+  }
+
   const engine = new TextEngine(
     {
       adapter,
@@ -2679,9 +2742,13 @@ async function runAgenticStructuredOutput<
     logger,
   )
 
-  // Consume the stream — chunks pipe through middleware but are not yielded externally
-  for await (const _chunk of engine.run()) {
-    // intentionally empty
+  try {
+    // Consume the stream — chunks pipe through middleware but are not yielded externally
+    for await (const _chunk of engine.run()) {
+      // intentionally empty
+    }
+  } finally {
+    await mcpManager.dispose()
   }
 
   const finalizationError = engine.getFinalizationError()
@@ -2912,8 +2979,15 @@ async function* runStreamingStructuredOutputImpl<
   options: TextActivityOptions<AnyTextAdapter, TSchema, true, TContext>,
   jsonSchema: NonNullable<ReturnType<typeof convertSchemaToJsonSchema>>,
 ): StructuredOutputStreamInternal<InferSchemaType<TSchema>> {
-  const { adapter, outputSchema, middleware, context, debug, ...textOptions } =
-    options
+  const {
+    adapter,
+    outputSchema,
+    middleware,
+    context,
+    debug,
+    mcp,
+    ...textOptions
+  } = options
   const model = adapter.model
   const logger = resolveDebugOption(debug)
 
@@ -2927,6 +3001,12 @@ async function* runStreamingStructuredOutputImpl<
   const nativeCombined =
     adapter.supportsCombinedToolsAndSchema?.(options.modelOptions) === true
 
+  const mcpManager = MCPManager.from(mcp)
+  const mcpTools = await mcpManager.discover()
+  if (mcpTools.length > 0) {
+    textOptions.tools = [...(textOptions.tools ?? []), ...mcpTools]
+  }
+
   // Inputs may be UIMessages (from useChat) or ModelMessages (from server-side
   // callers). TextEngine handles the conversion uniformly.
   const engine = new TextEngine(
@@ -2948,8 +3028,12 @@ async function* runStreamingStructuredOutputImpl<
     logger,
   )
 
-  for await (const chunk of engine.run()) {
-    yield chunk
+  try {
+    for await (const chunk of engine.run()) {
+      yield chunk
+    }
+  } finally {
+    await mcpManager.dispose()
   }
 
   // Schema validation for the streaming variant remains the consumer's

package/src/activities/chat/mcp/manager.ts

@@ -0,0 +1,85 @@
+import type { ServerTool } from '../tools/tool-definition'
+import type { ChatMCPOptions, MCPToolSource } from './types'
+
+export class MCPDuplicateToolNameError extends Error {
+  constructor(public readonly toolName: string) {
+    super(
+      `Duplicate MCP tool name "${toolName}" in chat({ mcp.clients }). ` +
+        `Set a unique \`prefix\` on one of the MCP clients (or use a pool, ` +
+        `which auto-prefixes) to disambiguate.`,
+    )
+    this.name = 'MCPDuplicateToolNameError'
+  }
+}
+
+/**
+ * Encapsulates MCP tool discovery + connection lifecycle for chat().
+ * Built from chat()'s `mcp` option; runners only call `discover()` then
+ * `dispose()`. A manager built from `undefined` is an inert no-op
+ * (`discover()` → `[]`, `dispose()` → no-op), so runners need no branching.
+ */
+export class MCPManager {
+  static from(options: ChatMCPOptions | undefined): MCPManager {
+    return new MCPManager(options)
+  }
+
+  readonly #sources: ReadonlyArray<MCPToolSource>
+  readonly #shouldClose: boolean
+  readonly #lazyTools: boolean
+  readonly #onDiscoveryError?: (
+    error: unknown,
+    source: MCPToolSource,
+  ) => void | Promise<void>
+
+  private constructor(options: ChatMCPOptions | undefined) {
+    this.#sources = options?.clients ?? []
+    // default 'close'; only 'keep-alive' disables closing
+    this.#shouldClose = options ? options.connection !== 'keep-alive' : false
+    this.#lazyTools = options?.lazyTools ?? false
+    this.#onDiscoveryError = options?.onDiscoveryError
+  }
+
+  /**
+   * Discover + merge tools from all sources. Throws on a fatal discovery error
+   * (no `onDiscoveryError`, or it re-threw) or a duplicate tool name; in that
+   * case it first closes any connected sources when the policy is 'close'.
+   */
+  async discover(): Promise<Array<ServerTool>> {
+    if (this.#sources.length === 0) return []
+    try {
+      const settled = await Promise.allSettled(
+        this.#sources.map((s) => s.tools({ lazy: this.#lazyTools })),
+      )
+      const tools: Array<ServerTool> = []
+      const zipped = this.#sources.map(
+        (source, i) => [source, settled[i]] as const,
+      )
+      for (const [source, result] of zipped) {
+        if (result === undefined) continue
+        if (result.status === 'fulfilled') {
+          tools.push(...result.value)
+        } else if (this.#onDiscoveryError) {
+          // throw/reject inside handler ⇒ propagate (fail-fast); return ⇒ skip
+          await this.#onDiscoveryError(result.reason, source)
+        } else {
+          throw result.reason
+        }
+      }
+      const seen = new Set<string>()
+      for (const t of tools) {
+        if (seen.has(t.name)) throw new MCPDuplicateToolNameError(t.name)
+        seen.add(t.name)
+      }
+      return tools
+    } catch (err) {
+      await this.dispose() // cleanup-on-failure (no-op if keep-alive)
+      throw err
+    }
+  }
+
+  /** Close sources iff policy is 'close'. Idempotent; never throws. */
+  async dispose(): Promise<void> {
+    if (!this.#shouldClose || this.#sources.length === 0) return
+    await Promise.allSettled(this.#sources.map((s) => s.close()))
+  }
+}

package/src/activities/chat/mcp/types.ts

@@ -0,0 +1,66 @@
+import type { ServerTool } from '../tools/tool-definition'
+
+/**
+ * Minimal structural shape that `chat({ mcp })` needs from an MCP client.
+ *
+ * `@tanstack/ai-mcp`'s `MCPClient` and `MCPClients` satisfy this interface by
+ * shape — the core `@tanstack/ai` package does NOT import `@tanstack/ai-mcp`
+ * (ai-mcp depends on ai, not the reverse).
+ */
+export interface MCPToolSource {
+  // Keep the options shape in sync with ai-mcp's `ToolsOptions` — extra
+  // optional fields added there still match structurally, but chat() only
+  // forwards what is declared here.
+  tools: (options?: { lazy?: boolean }) => Promise<Array<ServerTool>>
+  close: () => Promise<void>
+}
+
+/**
+ * Controls what happens to MCP connections when the chat run ends.
+ *
+ * - `'close'` (default) — `chat()` closes each connection when the run ends
+ *   (after the agent loop completes and the stream is drained), so tools can
+ *   still execute throughout the run.
+ * - `'keep-alive'` — `chat()` never closes the connections; the caller owns
+ *   their lifecycle (e.g. keep them warm across requests).
+ */
+export type MCPConnectionPolicy = 'close' | 'keep-alive'
+
+/**
+ * Options controlling MCP tool discovery and lifecycle for a `chat()` call.
+ */
+export interface ChatMCPOptions {
+  /**
+   * The MCP clients or client pools to discover tools from and manage.
+   */
+  clients: Array<MCPToolSource>
+
+  /**
+   * Connection lifecycle policy applied to all clients when the run ends.
+   *
+   * Defaults to `'close'`.
+   */
+  connection?: MCPConnectionPolicy
+
+  /**
+   * When `true`, tool schemas are fetched lazily (forwarded to
+   * `tools({ lazy: true })`).
+   *
+   * Defaults to `false`.
+   */
+  lazyTools?: boolean
+
+  /**
+   * Called when tool discovery fails for a single source.
+   *
+   * - Throw (or re-throw) from this handler to fail the entire chat call fast.
+   * - Return normally to skip that source and continue with remaining clients.
+   * - Omit this handler entirely to rethrow the error (fail-fast by default).
+   *
+   * Async handlers are awaited, so a rejected promise also fails fast.
+   */
+  onDiscoveryError?: (
+    error: unknown,
+    source: MCPToolSource,
+  ) => void | Promise<void>
+}

package/src/activities/chat/messages.ts

@@ -211,6 +211,7 @@ function isToolCallIncluded(part: ToolCallPart): boolean {
     part.state === 'input-complete' ||
     part.state === 'complete' ||
     part.state === 'approval-responded' ||
+    part.state === 'error' ||
     part.output !== undefined
   )
 }

package/src/activities/chat/stream/message-updaters.ts

@@ -161,10 +161,14 @@ export function updateToolCallApproval(
     )
 
     if (toolCallPart) {
-      toolCallPart.state = 'approval-requested'
-      toolCallPart.approval = {
-        id: approvalId,
-        needsApproval: true,
+      const index = parts.indexOf(toolCallPart)
+      parts[index] = {
+        ...toolCallPart,
+        state: 'approval-requested',
+        approval: {
+          id: approvalId,
+          needsApproval: true,
+        },
       }
     }
 
@@ -192,7 +196,8 @@ export function updateToolCallState(
     )
 
     if (toolCallPart) {
-      toolCallPart.state = state
+      const index = parts.indexOf(toolCallPart)
+      parts[index] = { ...toolCallPart, state }
     }
 
     return { ...msg, parts }
@@ -217,8 +222,12 @@ export function updateToolCallWithOutput(
     )
 
     if (toolCallPart) {
-      toolCallPart.output = errorText ? { error: errorText } : output
-      toolCallPart.state = state ?? (errorText ? 'input-complete' : 'complete')
+      const index = parts.indexOf(toolCallPart)
+      parts[index] = {
+        ...toolCallPart,
+        output: errorText ? { error: errorText } : output,
+        state: state ?? (errorText ? 'error' : 'complete'),
+      }
     }
 
     return { ...msg, parts }
@@ -242,8 +251,12 @@ export function updateToolCallApprovalResponse(
     )
 
     if (toolCallPart && toolCallPart.approval) {
-      toolCallPart.approval.approved = approved
-      toolCallPart.state = 'approval-responded'
+      const index = parts.indexOf(toolCallPart)
+      parts[index] = {
+        ...toolCallPart,
+        approval: { ...toolCallPart.approval, approved },
+        state: 'approval-responded',
+      }
     }
 
     return { ...msg, parts }

package/src/activities/chat/stream/processor.ts

@@ -317,7 +317,7 @@ export class StreamProcessor {
       this.messages,
       toolCallId,
       output,
-      error ? 'input-complete' : undefined,
+      error ? 'error' : undefined,
       error,
     )
 
@@ -1184,7 +1184,7 @@ export class StreamProcessor {
         this.messages,
         chunk.toolCallId,
         output,
-        chunk.state === 'output-error' ? 'input-complete' : undefined,
+        chunk.state === 'output-error' ? 'error' : undefined,
       )
 
       // Step 2: Create/update the tool-result part (for LLM conversation history)
@@ -1240,7 +1240,7 @@ export class StreamProcessor {
       this.messages,
       chunk.toolCallId,
       output,
-      chunk.state === 'output-error' ? 'input-complete' : undefined,
+      chunk.state === 'output-error' ? 'error' : undefined,
     )
 
     // Step 2: Create/update the tool-result part
@@ -1690,11 +1690,22 @@ export class StreamProcessor {
     _index: number,
     toolCall: InternalToolCallState,
   ): void {
+    // Finalize the internal bookkeeping: the call's input arguments ARE
+    // complete regardless of whether execution later failed, so the call still
+    // counts as a completed tool call in getCompletedToolCalls()/getState().
     toolCall.state = 'input-complete'
 
     // Try final parse
     toolCall.parsedArguments = this.jsonParser.parse(toolCall.arguments)
 
+    // Don't downgrade the rendered part of a call that already reached the
+    // terminal 'error' state (e.g. an output-error TOOL_CALL_RESULT arrived
+    // without a preceding TOOL_CALL_END). The RUN_FINISHED / finalizeStream
+    // safety net must not clobber a failed call back to 'input-complete'.
+    if (this.isToolCallPartErrored(toolCall.id)) {
+      return
+    }
+
     // Update UIMessage
     this.messages = updateToolCallPart(this.messages, messageId, {
       id: toolCall.id,
@@ -1714,6 +1725,22 @@ export class StreamProcessor {
     )
   }
 
+  /**
+   * Whether the rendered tool-call part for the given id has reached the
+   * terminal 'error' state. Used to prevent the completion safety net from
+   * downgrading a failed call back to 'input-complete'.
+   */
+  private isToolCallPartErrored(toolCallId: string): boolean {
+    return this.messages.some((msg) =>
+      msg.parts.some(
+        (part) =>
+          part.type === 'tool-call' &&
+          part.id === toolCallId &&
+          part.state === 'error',
+      ),
+    )
+  }
+
   /**
    * Emit pending text update for a specific message.
    *

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

@@ -599,6 +599,7 @@ export async function* executeToolCalls<TContext = unknown>(
   ) => CustomEvent,
   middlewareHooks?: ToolExecutionMiddlewareHooks,
   userContext?: TContext,
+  abortSignal?: AbortSignal,
 ): AsyncGenerator<CustomEvent, ExecuteToolCallsResult, void> {
   const results: Array<ToolResult> = []
   const needsApproval: Array<ApprovalRequest> = []
@@ -679,6 +680,7 @@ export async function* executeToolCalls<TContext = unknown>(
     const context = {
       toolCallId: toolCall.id,
       context: userContext,
+      abortSignal,
       emitCustomEvent: (eventName: string, value: Record<string, any>) => {
         if (createCustomEventChunk) {
           pendingEvents.push(

package/src/activities/generateVideo/index.ts

@@ -15,6 +15,7 @@ import type { DebugOption } from '../../logger/types'
 import type { VideoAdapter } from './adapter'
 import type {
   StreamChunk,
+  TokenUsage,
   VideoJobResult,
   VideoStatusResult,
   VideoUrlResult,
@@ -380,6 +381,7 @@ async function* runStreamingVideoGeneration<
             status: 'completed',
             url: urlResult.url,
             expiresAt: urlResult.expiresAt,
+            ...(urlResult.usage ? { usage: urlResult.usage } : {}),
           },
           timestamp: Date.now(),
         } as StreamChunk
@@ -454,6 +456,7 @@ export async function getVideoJobStatus<
   progress?: number
   url?: string
   error?: string
+  usage?: TokenUsage
 }> {
   const { adapter, jobId } = options
   const requestId = createId('video-status')
@@ -487,10 +490,19 @@ export async function getVideoJobStatus<
         duration: Date.now() - startTime,
         timestamp: Date.now(),
       })
+      if (urlResult.usage) {
+        aiEventClient.emit('video:usage', {
+          requestId,
+          model: adapter.model,
+          usage: urlResult.usage,
+          timestamp: Date.now(),
+        })
+      }
       return {
         status: statusResult.status,
         progress: statusResult.progress,
         url: urlResult.url,
+        ...(urlResult.usage ? { usage: urlResult.usage } : {}),
       }
     } catch (error) {
       const errorMessage =