@tanstack/ai 0.27.0 → 0.28.0

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/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 ? 'input-complete' : '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/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/extend-adapter.ts

@@ -144,6 +144,14 @@ type ExtractCustomModelNames<TDefs extends ReadonlyArray<ExtendedModelDef>> =
 // Factory Type Inference
 // ===========================
 
+/**
+ * The widest factory shape `extendAdapter` accepts: any function taking a
+ * model as its first parameter. Parameters are contravariant, so `never`
+ * params and an `unknown` return accept every factory without resorting
+ * to `any`.
+ */
+type AnyAdapterFactory = (model: never, ...args: Array<never>) => unknown
+
 /**
  * Infer the model parameter type from an adapter factory function.
  * For generic functions like `<T extends Union>(model: T)`, this gets `T` which
@@ -151,32 +159,44 @@ type ExtractCustomModelNames<TDefs extends ReadonlyArray<ExtendedModelDef>> =
  */
 type InferFactoryModels<TFactory> = TFactory extends (
   model: infer TModel,
-  ...args: Array<any>
-) => any
+  ...args: Array<never>
+) => unknown
   ? TModel extends string
     ? TModel
     : string
   : string
 
-/**
- * Infer the config parameter type from an adapter factory function.
- */
-type InferConfig<TFactory> = TFactory extends (
-  model: any,
-  config?: infer TConfig,
-) => any
-  ? TConfig
-  : undefined
-
 /**
  * Infer the adapter return type from a factory function.
  */
 type InferAdapterReturn<TFactory> = TFactory extends (
-  ...args: Array<any>
+  ...args: Array<never>
 ) => infer TReturn
   ? TReturn
   : never
 
+/**
+ * Extracts all parameter types after the model parameter from a factory,
+ * preserving labels and optionality (e.g. `[apiKey: string, config?: C]`).
+ * Note: overloaded factories resolve against their last overload (a
+ * `Parameters` limitation).
+ */
+type InferRestArgs<TFactory extends AnyAdapterFactory> =
+  Parameters<TFactory> extends [unknown?, ...infer TRest] ? TRest : []
+
+/**
+ * The factory signature produced by `extendAdapter`: accepts both original
+ * and custom model names while preserving all remaining parameters and the
+ * return type of the original factory.
+ */
+type ExtendedFactory<
+  TFactory extends AnyAdapterFactory,
+  TDefs extends ReadonlyArray<ExtendedModelDef>,
+> = (
+  model: InferFactoryModels<TFactory> | ExtractCustomModelNames<TDefs>,
+  ...args: InferRestArgs<TFactory>
+) => InferAdapterReturn<TFactory>
+
 // ===========================
 // extendAdapter Function
 // ===========================
@@ -225,19 +245,17 @@ type InferAdapterReturn<TFactory> = TFactory extends (
  * ```
  */
 export function extendAdapter<
-  TFactory extends (...args: Array<any>) => any,
+  TFactory extends AnyAdapterFactory,
   const TDefs extends ReadonlyArray<ExtendedModelDef>,
->(
-  factory: TFactory,
-  _customModels: TDefs,
-): (
-  model: InferFactoryModels<TFactory> | ExtractCustomModelNames<TDefs>,
-  ...args: InferConfig<TFactory> extends undefined
-    ? []
-    : [config?: InferConfig<TFactory>]
-) => InferAdapterReturn<TFactory> {
+>(factory: TFactory, _customModels: TDefs): ExtendedFactory<TFactory, TDefs>
+// The implementation signature stays at the honest `AnyAdapterFactory` width;
+// the overload above performs the deliberate model-union widening.
+export function extendAdapter(
+  factory: AnyAdapterFactory,
+  _customModels: ReadonlyArray<ExtendedModelDef>,
+): AnyAdapterFactory {
   // At runtime, we simply pass through to the original factory.
   // The _customModels parameter is only used for type inference.
   // No runtime validation - users are trusted to pass valid model names.
-  return factory as any
+  return factory
 }

package/src/index.ts

@@ -52,6 +52,16 @@ export {
   type InferToolOutput,
 } from './activities/chat/tools/tool-definition'
 
+// MCP chat option types
+export type {
+  MCPToolSource,
+  ChatMCPOptions,
+  MCPConnectionPolicy,
+} from './activities/chat/mcp/types'
+
+// MCP error classes (value exports — usable with instanceof)
+export { MCPDuplicateToolNameError } from './activities/chat/mcp/manager'
+
 // Schema conversion (Standard JSON Schema compliant)
 export {
   convertSchemaToJsonSchema,

package/src/types.ts

@@ -490,6 +490,12 @@ export type ToolExecutionContext<TContext = unknown> =
   RuntimeContextField<TContext> & {
     /** The ID of the tool call being executed */
     toolCallId?: string
+    /**
+     * Abort signal for the current chat run. Aborts when the run's
+     * `abortController` fires (or middleware aborts). Long-running tools —
+     * e.g. MCP `callTool` — should forward this to cancel in-flight work.
+     */
+    abortSignal?: AbortSignal
     /**
      * Emit a custom event during tool execution.
      * Events are streamed to the client in real-time as AG-UI CUSTOM events.