@assistant-ui/react-ai-sdk 1.3.31 → 1.3.32

package/src/generativeTools.ts

@@ -1,7 +1,21 @@
 import { jsonSchema, type ToolSet } from "ai";
-import { toJSONSchema, type ToolJSONSchema } from "assistant-stream";
-import type { Toolkit, ToolkitDeclaration } from "@assistant-ui/core/react";
-import { defaultToModelOutput, frontendTools } from "./frontendTools";
+import type { MCPClient, MCPClientConfig } from "@ai-sdk/mcp";
+import { createMCPClient } from "@ai-sdk/mcp";
+import { Experimental_StdioMCPTransport } from "@ai-sdk/mcp/mcp-stdio";
+import {
+  toJSONSchema,
+  type Tool,
+  type McpServerConfig,
+  type ToolJSONSchema,
+  type ToolModelOutputFunction,
+} from "assistant-stream";
+import type { Toolkit, ToolkitDefinition } from "@assistant-ui/core/react";
+import { frontendTools } from "./frontendTools";
+import { toAISDKContent, toAISDKDefaultOutput } from "./toolOutputConversion";
+import {
+  unwrapModelContentEnvelope,
+  type ModelContentEnvelope,
+} from "./modelContentEnvelope";
 
 const EMPTY_SCHEMA = { type: "object" as const, properties: {} };
 
@@ -14,11 +28,14 @@ const humanNotSupported = (): never => {
 // AI SDK leaves `abortSignal` optional; assistant-ui's execute requires one.
 const neverAbort = new AbortController().signal;
 
+const parametersToInputSchema = (parameters: Tool["parameters"] | undefined) =>
+  jsonSchema(parameters ? toJSONSchema(parameters) : EMPTY_SCHEMA);
+
 export interface GenerativeToolsOptions {
   /**
    * The server build of a generative toolkit (schema + server `execute`). Typed
    * as the canonical {@link Toolkit} so callers don't need to cast; the server
-   * build carries `execute`, recovered internally as {@link ToolkitDeclaration}.
+   * build carries `execute`, recovered internally as {@link ToolkitDefinition}.
    */
   toolkit: Toolkit;
   /**
@@ -29,6 +46,17 @@ export interface GenerativeToolsOptions {
   frontendTools?: Record<string, ToolJSONSchema>;
 }
 
+export type AISDKToolkitOptions = {
+  toolkit: Toolkit;
+};
+
+export type AISDKToolkitToolsOptions = {
+  /**
+   * Tools uploaded by the frontend request body.
+   */
+  frontend?: Record<string, ToolJSONSchema>;
+};
+
 /**
  * Builds an AI SDK `ToolSet` for server-side use with `streamText` /
  * `generateText` from a generative `toolkit` and the frontend-uploaded tools.
@@ -38,7 +66,8 @@ export interface GenerativeToolsOptions {
  * resolves to the server build — schema + `execute`, with `render` stripped) and
  * pass it here. Tools without an `execute` are still exposed to the model but
  * left for the client to fulfill. `frontendTools` lets the client contribute
- * tools that aren't in the static toolkit.
+ * tools that aren't in the static toolkit. Use {@link AISDKToolkit} when the
+ * toolkit contains MCP entries.
  *
  * @example
  * ```ts
@@ -50,28 +79,190 @@ export interface GenerativeToolsOptions {
  * });
  * ```
  */
-export const generativeTools = (options: GenerativeToolsOptions): ToolSet => ({
-  ...(options.frontendTools ? frontendTools(options.frontendTools) : {}),
-  // `toolkit` last so its server-side `execute` wins over an uploaded entry of
-  // the same name. The cast recovers the declaration shape — the server build
-  // carries `execute`, which the canonical `Toolkit` type erases.
-  ...toServerToolSet(options.toolkit as ToolkitDeclaration),
-});
-
-const toServerToolSet = (toolkit: ToolkitDeclaration): ToolSet =>
+export const generativeTools = (options: GenerativeToolsOptions): ToolSet => {
+  assertNoMcpToolkitTools(options.toolkit);
+  return {
+    ...(options.frontendTools ? frontendTools(options.frontendTools) : {}),
+    // `toolkit` last so its server-side `execute` wins over an uploaded entry of
+    // the same name. The cast recovers the declaration shape — the server build
+    // carries `execute`, which the canonical `Toolkit` type erases.
+    ...toProviderToolSet(options.toolkit),
+    ...toServerToolSet(options.toolkit as ToolkitDefinition),
+  };
+};
+
+export class AISDKToolkit {
+  readonly #toolkit: Toolkit;
+  readonly #mcpClients = new Map<string, Promise<MCPClient>>();
+
+  constructor(options: AISDKToolkitOptions) {
+    this.#toolkit = options.toolkit;
+  }
+
+  async tools(options: AISDKToolkitToolsOptions = {}): Promise<ToolSet> {
+    return {
+      ...(options.frontend ? frontendTools(options.frontend) : {}),
+      ...(await this.#mcpTools()),
+      ...toProviderToolSet(this.#toolkit),
+      ...toServerToolSet(this.#toolkit as ToolkitDefinition),
+    };
+  }
+
+  async close(): Promise<void> {
+    const clientPromises = [...this.#mcpClients.values()];
+    this.#mcpClients.clear();
+    const results = await Promise.allSettled(clientPromises);
+    const clients = results.flatMap((result) =>
+      result.status === "fulfilled" ? [result.value] : [],
+    );
+    const closeResults = await Promise.allSettled(
+      clients.map((client) => client.close()),
+    );
+    const errors = [
+      ...results.flatMap((result) =>
+        result.status === "rejected" ? [result.reason] : [],
+      ),
+      ...closeResults.flatMap((result) =>
+        result.status === "rejected" ? [result.reason] : [],
+      ),
+    ];
+    if (errors.length === 1) throw errors[0];
+    if (errors.length > 1) {
+      throw new AggregateError(
+        errors,
+        "Failed to close one or more MCP clients",
+      );
+    }
+  }
+
+  async #mcpTools(): Promise<ToolSet> {
+    const toolSets = await Promise.all(
+      Object.entries(this.#toolkit)
+        .filter((entry): entry is [string, McpToolkitTool] =>
+          isMcpToolkitTool(entry[1]),
+        )
+        .map(async ([name, tool]) => {
+          const client = await this.#mcpClient(name, tool.server);
+          return [name, await client.tools()] as const;
+        }),
+    );
+
+    const tools: ToolSet = {};
+    const toolSources = new Map<string, string>();
+    for (const [serverName, toolSet] of toolSets) {
+      for (const [toolName, tool] of Object.entries(toolSet)) {
+        const existingServerName = toolSources.get(toolName);
+        if (existingServerName) {
+          throw new Error(
+            `MCP tool name collision: "${toolName}" is exposed by both "${existingServerName}" and "${serverName}". Rename one of the toolkit entries or expose distinct MCP tool names.`,
+          );
+        }
+        toolSources.set(toolName, serverName);
+        tools[toolName] = tool;
+      }
+    }
+    return tools;
+  }
+
+  #mcpClient(name: string, config: McpServerConfig): Promise<MCPClient> {
+    const existing = this.#mcpClients.get(name);
+    if (existing) return existing;
+    let next: Promise<MCPClient>;
+    next = createMCPClient(toMCPClientConfig(config)).catch((error) => {
+      if (this.#mcpClients.get(name) === next) {
+        this.#mcpClients.delete(name);
+      }
+      throw error;
+    });
+    this.#mcpClients.set(name, next);
+    return next;
+  }
+}
+
+const toMCPClientConfig = (config: McpServerConfig): MCPClientConfig => {
+  if (config.type === "stdio") {
+    return {
+      transport: new Experimental_StdioMCPTransport({
+        command: config.command,
+        ...(config.args && { args: [...config.args] }),
+        ...(config.env && { env: config.env }),
+        ...(config.cwd && { cwd: config.cwd }),
+      }),
+    };
+  }
+
+  return {
+    transport: {
+      type: config.type,
+      url: config.url,
+      ...(config.headers && { headers: config.headers }),
+      ...(config.redirect && { redirect: config.redirect }),
+    },
+  };
+};
+
+type ToolkitTool = Toolkit[string];
+
+type McpToolkitTool = ToolkitTool & {
+  type: "mcp";
+  server: McpServerConfig;
+};
+
+const isMcpToolkitTool = (tool: ToolkitTool): tool is McpToolkitTool =>
+  tool.type === "mcp" && !tool.disabled;
+
+const assertNoMcpToolkitTools = (toolkit: Toolkit): void => {
+  const mcpToolName = Object.entries(toolkit).find(([, tool]) =>
+    isMcpToolkitTool(tool),
+  )?.[0];
+  if (!mcpToolName) return;
+
+  throw new Error(
+    `MCP toolkit entry "${mcpToolName}" requires AISDKToolkit. Use new AISDKToolkit({ toolkit }).tools(...) instead of generativeTools(...).`,
+  );
+};
+
+type AISDKToModelOutputOptions<TArgs, TResult> = Omit<
+  Parameters<ToolModelOutputFunction<TArgs, TResult>>[0],
+  "output"
+> & {
+  output: TResult | ModelContentEnvelope<TResult>;
+};
+
+const toAISDKToModelOutput =
+  <TArgs, TResult>(toModelOutput?: ToolModelOutputFunction<TArgs, TResult>) =>
+  async (options: AISDKToModelOutputOptions<TArgs, TResult>) => {
+    const { result, modelContent } = unwrapModelContentEnvelope(options.output);
+
+    if (modelContent !== undefined) {
+      return toAISDKContent(modelContent);
+    }
+
+    if (!toModelOutput) {
+      return toAISDKDefaultOutput(result);
+    }
+
+    const parts = await toModelOutput({
+      ...options,
+      output: result,
+    });
+    return toAISDKContent(parts);
+  };
+
+const toServerToolSet = (toolkit: ToolkitDefinition): ToolSet =>
   Object.fromEntries(
     Object.entries(toolkit)
-      .filter(([, t]) => !t.disabled)
+      .filter(
+        ([, t]) => t.type !== "mcp" && t.type !== "provider" && !t.disabled,
+      )
       .map(([name, t]) => {
         const execute = t.execute;
         return [
           name,
           {
             ...(t.description !== undefined && { description: t.description }),
-            inputSchema: jsonSchema(
-              t.parameters ? toJSONSchema(t.parameters) : EMPTY_SCHEMA,
-            ),
-            toModelOutput: t.toModelOutput ?? defaultToModelOutput,
+            inputSchema: parametersToInputSchema(t.parameters),
+            toModelOutput: toAISDKToModelOutput(t.toModelOutput),
             ...(t.providerOptions && { providerOptions: t.providerOptions }),
             ...(execute && {
               execute: (
@@ -88,3 +279,32 @@ const toServerToolSet = (toolkit: ToolkitDeclaration): ToolSet =>
         ];
       }),
   ) as ToolSet;
+
+const toProviderToolSet = (toolkit: Toolkit): ToolSet =>
+  Object.fromEntries(
+    Object.entries(toolkit)
+      .filter((entry): entry is [string, ProviderToolkitTool] =>
+        isProviderToolkitTool(entry[1]),
+      )
+      .map(([name, t]) => [
+        name,
+        {
+          type: "provider",
+          id: t.providerId,
+          args: t.args,
+          ...(t.parameters && {
+            inputSchema: parametersToInputSchema(t.parameters),
+          }),
+          ...(t.providerOptions && { providerOptions: t.providerOptions }),
+          ...(t.supportsDeferredResults !== undefined && {
+            supportsDeferredResults: t.supportsDeferredResults,
+          }),
+        },
+      ]),
+  ) as ToolSet;
+
+type ProviderToolkitTool = Extract<Toolkit[string], { type: "provider" }>;
+
+const isProviderToolkitTool = (
+  tool: Toolkit[string],
+): tool is ProviderToolkitTool => tool.type === "provider" && !tool.disabled;

package/src/index.ts

@@ -14,7 +14,10 @@ export type {
 } from "./ui/resumable";
 export { frontendTools } from "./frontendTools";
 export {
+  AISDKToolkit,
   generativeTools,
+  type AISDKToolkitOptions,
+  type AISDKToolkitToolsOptions,
   type GenerativeToolsOptions,
 } from "./generativeTools";
 export { injectQuoteContext } from "./injectQuoteContext";

package/src/modelContentEnvelope.ts

@@ -7,9 +7,9 @@ export type ModelContentEnvelope<TResult = unknown> = {
   readonly value: TResult;
 };
 
-export function isModelContentEnvelope(
-  value: unknown,
-): value is ModelContentEnvelope {
+export function isModelContentEnvelope<TResult = unknown>(
+  value: TResult | ModelContentEnvelope<TResult>,
+): value is ModelContentEnvelope<TResult> {
   return (
     value != null &&
     typeof value === "object" &&
@@ -25,8 +25,10 @@ export function wrapModelContentEnvelope<TResult>(
   return { [ENVELOPE_KEY]: modelContent, value: result };
 }
 
-export function unwrapModelContentEnvelope(output: unknown): {
-  result: unknown;
+export function unwrapModelContentEnvelope<TResult>(
+  output: TResult | ModelContentEnvelope<TResult>,
+): {
+  result: TResult;
   modelContent?: readonly ToolModelContentPart[];
 } {
   if (isModelContentEnvelope(output)) {

package/src/toolOutputConversion.ts

@@ -0,0 +1,29 @@
+import type { JSONValue } from "ai";
+import type { ToolModelContentPart } from "assistant-stream";
+
+export const toAISDKContent = (parts: readonly ToolModelContentPart[]) => ({
+  type: "content" as const,
+  value: parts.map((part) => {
+    if (part.type === "text") {
+      return { type: "text" as const, text: part.text };
+    }
+    const isImage = part.mediaType.startsWith("image/");
+    return isImage
+      ? {
+          type: "image-data" as const,
+          data: part.data,
+          mediaType: part.mediaType,
+        }
+      : {
+          type: "file-data" as const,
+          data: part.data,
+          mediaType: part.mediaType,
+          ...(part.filename !== undefined && { filename: part.filename }),
+        };
+  }),
+});
+
+export const toAISDKDefaultOutput = (output: unknown) =>
+  typeof output === "string"
+    ? { type: "text" as const, value: output }
+    : { type: "json" as const, value: (output ?? null) as JSONValue };

package/src/ui/use-chat/useAISDKRuntime.test.ts

@@ -2,6 +2,7 @@
 
 import { act, renderHook, waitFor } from "@testing-library/react";
 import { beforeEach, describe, expect, it, vi } from "vitest";
+import { validateUIMessages } from "ai";
 
 // Mock only the sibling module that requires AUI store context (not available
 // in isolation). Every other dependency — useExternalStoreRuntime,
@@ -140,6 +141,54 @@ describe("useAISDKRuntime", () => {
     expect(chat.messages[0].parts[1].state).toBe("output-available");
   });
 
+  it("strips stale approval when cancelling a tool pending approval so history stays valid (#4195)", async () => {
+    const chat = createChatHelpers([
+      {
+        id: "a1",
+        role: "assistant",
+        parts: [
+          {
+            type: "dynamic-tool",
+            toolName: "mcp_search",
+            toolCallId: "tc-1",
+            state: "approval-requested",
+            input: { q: "hi" },
+            approval: { id: "appr-1" },
+          },
+        ],
+      },
+    ]);
+
+    const { result } = renderHook(() => useAISDKRuntime(chat));
+
+    await waitFor(() => {
+      expect(result.current.thread.getState().messages.length).toBeGreaterThan(
+        0,
+      );
+    });
+
+    act(() => {
+      result.current.thread.append({
+        role: "user",
+        content: [{ type: "text", text: "what" }],
+      });
+    });
+
+    await waitFor(() => {
+      expect(chat.sendMessage).toHaveBeenCalledTimes(1);
+    });
+
+    const part = chat.messages[0].parts[0];
+    expect(part.state).toBe("output-error");
+    // The pending-approval object must not survive into the terminal state,
+    // otherwise AI SDK's validateUIMessages rejects the next request.
+    expect(part.approval).toBeUndefined();
+
+    await expect(
+      validateUIMessages({ messages: chat.messages }),
+    ).resolves.toBeDefined();
+  });
+
   it("appends a new user message without sending when startRun is false", async () => {
     const chat = createChatHelpers([
       { id: "u1", role: "user", parts: [{ type: "text", text: "earlier" }] },

package/src/ui/use-chat/useAISDKRuntime.ts

@@ -180,8 +180,9 @@ export const useAISDKRuntime = <UI_MESSAGE extends UIMessage = UIMessage>(
           return part;
 
         hasChanges = true;
+        const { approval: _approval, ...rest } = part;
         return {
-          ...part,
+          ...rest,
           state: "output-error" as const,
           errorText: "User cancelled tool call by sending a new message.",
         };

package/src/ui/utils/toCreateMessage.test.ts

@@ -0,0 +1,54 @@
+import { describe, expect, it } from "vitest";
+import type { AppendMessage } from "@assistant-ui/core";
+import { toCreateMessage } from "./toCreateMessage";
+
+const baseMessage = {
+  role: "user",
+  parentId: null,
+  sourceId: null,
+  runConfig: undefined,
+  metadata: undefined,
+} as const;
+
+describe("toCreateMessage", () => {
+  it("converts a data part in message content into a data-<name> part", () => {
+    const message = {
+      ...baseMessage,
+      content: [{ type: "data", name: "workflow", data: { field: 1 } }],
+    } as unknown as AppendMessage;
+
+    const result = toCreateMessage(message);
+
+    expect(result.parts).toEqual([
+      { type: "data-workflow", data: { field: 1 } },
+    ]);
+  });
+
+  it("converts a data part inside an attachment without throwing", () => {
+    const message = {
+      ...baseMessage,
+      content: [],
+      attachments: [
+        {
+          id: "some-id",
+          type: "document",
+          name: "some-name",
+          status: { type: "complete" },
+          content: [
+            {
+              type: "data",
+              name: "some-content-name",
+              data: { field: 1 },
+            },
+          ],
+        },
+      ],
+    } as unknown as AppendMessage;
+
+    const result = toCreateMessage(message);
+
+    expect(result.parts).toEqual([
+      { type: "data-some-content-name", data: { field: 1 } },
+    ]);
+  });
+});

package/src/ui/utils/toCreateMessage.ts

@@ -41,6 +41,11 @@ export const toCreateMessage = <UI_MESSAGE extends UIMessage = UIMessage>(
           mediaType: part.mimeType,
           ...(part.filename && { filename: part.filename }),
         };
+      case "data":
+        return {
+          type: `data-${part.name}`,
+          data: part.data,
+        };
       default:
         throw new Error(`Unsupported part type: ${part.type}`);
     }