@assistant-ui/ai-sdk 0.0.1

package/src/tools/generativeTools.test.ts

@@ -0,0 +1,785 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { defineMcpToolkit } from "@assistant-ui/core/react";
+import { AISDKToolkit } from "./generativeTools";
+import { wrapModelContentEnvelope } from "../converters/modelContentEnvelope";
+
+const mocks = vi.hoisted(() => ({
+  close: vi.fn(),
+  tools: vi.fn(),
+  createMCPClient: vi.fn(),
+}));
+
+vi.mock("@ai-sdk/mcp", () => ({
+  createMCPClient: mocks.createMCPClient,
+}));
+
+vi.mock("@ai-sdk/mcp/mcp-stdio", () => ({
+  Experimental_StdioMCPTransport: vi.fn((config) => ({
+    type: "stdio",
+    config,
+  })),
+}));
+
+const never = <T>() => new Promise<T>(() => {});
+
+describe("AISDKToolkit.tools()", () => {
+  it("merges frontend tools with toolkit tools", async () => {
+    const toolSet = await new AISDKToolkit({
+      toolkit: {
+        serverTool: {
+          type: "backend",
+          description: "Server tool",
+          parameters: { type: "object", properties: {} },
+          execute: async () => "ok",
+        } as never,
+      },
+    }).tools({
+      frontend: {
+        clientTool: {
+          parameters: { type: "object", properties: {} },
+        },
+      },
+    });
+
+    expect(toolSet.clientTool).toBeDefined();
+    expect(toolSet.serverTool?.description).toBe("Server tool");
+    expect(toolSet.serverTool?.execute).toBeTypeOf("function");
+  });
+
+  it("keeps a flat toolkit tool named tools", async () => {
+    const toolSet = await new AISDKToolkit({
+      toolkit: {
+        tools: {
+          type: "backend",
+          description: "Actually a tool, not config",
+          parameters: { type: "object", properties: {} },
+          execute: async () => "ok",
+        } as never,
+      },
+    }).tools();
+
+    expect(toolSet.tools?.description).toBe("Actually a tool, not config");
+    expect(toolSet.tools?.execute).toBeTypeOf("function");
+  });
+
+  it("converts provider tools without an execute function", async () => {
+    const toolSet = await new AISDKToolkit({
+      toolkit: {
+        web_search: {
+          type: "provider",
+          providerId: "openai.web_search_preview",
+          args: { searchContextSize: "low" },
+        },
+      },
+    }).tools();
+
+    expect(toolSet.web_search).toMatchObject({
+      type: "provider",
+      id: "openai.web_search_preview",
+      args: { searchContextSize: "low" },
+    });
+    expect(toolSet.web_search).not.toHaveProperty("inputSchema");
+    expect(toolSet.web_search).not.toHaveProperty("execute");
+  });
+
+  it("forwards provider tool parameters and providerOptions when present", async () => {
+    const toolSet = await new AISDKToolkit({
+      toolkit: {
+        web_search: {
+          type: "provider",
+          providerId: "openai.web_search_preview",
+          args: { searchContextSize: "low" },
+          parameters: {
+            type: "object",
+            properties: {
+              query: { type: "string" },
+            },
+            required: ["query"],
+          },
+          providerOptions: {
+            openai: { rankingOptions: { scoreThreshold: 0.5 } },
+          },
+        },
+      },
+    }).tools();
+
+    expect(toolSet.web_search).toMatchObject({
+      type: "provider",
+      id: "openai.web_search_preview",
+      args: { searchContextSize: "low" },
+      providerOptions: {
+        openai: { rankingOptions: { scoreThreshold: 0.5 } },
+      },
+    });
+    expect(toolSet.web_search).toHaveProperty("inputSchema");
+  });
+
+  it("forwards explicit false supportsDeferredResults", async () => {
+    const toolSet = await new AISDKToolkit({
+      toolkit: {
+        web_search: {
+          type: "provider",
+          providerId: "openai.web_search_preview",
+          args: { searchContextSize: "low" },
+          supportsDeferredResults: false,
+        },
+      },
+    }).tools();
+
+    expect(toolSet.web_search).toMatchObject({
+      supportsDeferredResults: false,
+    });
+  });
+});
+
+describe("AISDKToolkit", () => {
+  beforeEach(() => {
+    mocks.close.mockReset();
+    mocks.tools.mockReset();
+    mocks.createMCPClient.mockReset();
+  });
+
+  it("loads MCP tools through pooled clients", async () => {
+    mocks.tools.mockResolvedValue({ echo: { inputSchema: {} } });
+    mocks.createMCPClient.mockResolvedValue({
+      tools: mocks.tools,
+      close: mocks.close,
+    });
+
+    const toolkit = new AISDKToolkit({
+      toolkit: {
+        local: {
+          type: "mcp",
+          server: { type: "http", url: "http://localhost:3001/mcp" },
+        },
+      },
+    });
+
+    await expect(toolkit.tools()).resolves.toHaveProperty("echo");
+    await toolkit.tools();
+
+    expect(mocks.createMCPClient).toHaveBeenCalledTimes(1);
+    expect(mocks.createMCPClient).toHaveBeenCalledWith({
+      transport: {
+        type: "http",
+        url: "http://localhost:3001/mcp",
+      },
+    });
+    expect(mocks.tools).toHaveBeenCalledTimes(2);
+  });
+
+  it("does not connect disabled MCP toolkit entries", async () => {
+    const toolkit = new AISDKToolkit({
+      toolkit: {
+        local: {
+          type: "mcp",
+          server: { type: "http", url: "http://localhost:3001/mcp" },
+          disabled: true,
+        },
+      },
+    });
+
+    await expect(toolkit.tools()).resolves.toEqual({});
+
+    expect(mocks.createMCPClient).not.toHaveBeenCalled();
+    expect(mocks.tools).not.toHaveBeenCalled();
+  });
+
+  it("filters disabled MCP tools from enabled toolkit entries", async () => {
+    mocks.tools.mockResolvedValue({
+      publicSearch: { inputSchema: {} },
+      privateSearch: { inputSchema: {} },
+    });
+    mocks.createMCPClient.mockResolvedValue({
+      tools: mocks.tools,
+      close: mocks.close,
+    });
+
+    const toolkit = new AISDKToolkit({
+      toolkit: {
+        local: {
+          type: "mcp",
+          server: { type: "http", url: "http://localhost:3001/mcp" },
+          tools: {
+            privateSearch: {
+              disabled: true,
+            },
+          },
+        } as never,
+      },
+    });
+
+    await expect(toolkit.tools()).resolves.toEqual({
+      publicSearch: { inputSchema: {} },
+    });
+
+    expect(mocks.createMCPClient).toHaveBeenCalledTimes(1);
+    expect(mocks.tools).toHaveBeenCalledTimes(1);
+  });
+
+  it("closes pooled MCP clients", async () => {
+    mocks.tools.mockResolvedValue({});
+    mocks.createMCPClient.mockResolvedValue({
+      tools: mocks.tools,
+      close: mocks.close,
+    });
+
+    const toolkit = new AISDKToolkit({
+      toolkit: {
+        local: {
+          type: "mcp",
+          server: { type: "sse", url: "http://localhost:3001/sse" },
+        },
+      },
+    });
+
+    await toolkit.tools();
+    await toolkit.close();
+
+    expect(mocks.close).toHaveBeenCalledTimes(1);
+  });
+
+  it("clears pooled MCP clients even when initialization fails", async () => {
+    const error = new Error("connect failed");
+    const closeError = new Error("close failed");
+    const close = vi.fn().mockRejectedValue(closeError);
+    mocks.tools.mockResolvedValue({});
+    mocks.createMCPClient
+      .mockResolvedValueOnce({
+        tools: mocks.tools,
+        close,
+      })
+      .mockRejectedValueOnce(error);
+
+    const toolkit = new AISDKToolkit({
+      toolkit: {
+        first: {
+          type: "mcp",
+          server: { type: "http", url: "http://localhost:3001/mcp" },
+        },
+        second: {
+          type: "mcp",
+          server: { type: "http", url: "http://localhost:3002/mcp" },
+        },
+      },
+    });
+
+    const toolsPromise = toolkit.tools();
+    await expect(toolkit.close()).rejects.toMatchObject({
+      errors: [
+        {
+          message:
+            'MCP toolkit entry "second" failed to connect: connect failed',
+          cause: error,
+        },
+        {
+          message: 'MCP toolkit entry "first" failed to close: close failed',
+          cause: closeError,
+        },
+      ],
+    });
+    await expect(toolsPromise).rejects.toMatchObject({
+      message: 'MCP toolkit entry "second" failed to connect: connect failed',
+      cause: error,
+    });
+    expect(close).toHaveBeenCalledTimes(1);
+
+    await expect(toolkit.close()).resolves.toBeUndefined();
+  });
+
+  it("includes the MCP toolkit entry name when closing a client fails", async () => {
+    const closeError = new Error("close failed");
+    mocks.tools.mockResolvedValue({});
+    mocks.createMCPClient.mockResolvedValue({
+      tools: mocks.tools,
+      close: vi.fn().mockRejectedValue(closeError),
+    });
+
+    const toolkit = new AISDKToolkit({
+      toolkit: {
+        github: {
+          type: "mcp",
+          server: { type: "http", url: "http://localhost:3001/mcp" },
+        },
+      },
+    });
+
+    await toolkit.tools();
+
+    await expect(toolkit.close()).rejects.toMatchObject({
+      message: 'MCP toolkit entry "github" failed to close: close failed',
+      cause: closeError,
+    });
+  });
+
+  it("evicts failed MCP client initialization so later calls can retry", async () => {
+    const error = new Error("connect failed");
+    mocks.createMCPClient.mockRejectedValueOnce(error).mockResolvedValueOnce({
+      tools: vi.fn().mockResolvedValue({ echo: { inputSchema: {} } }),
+      close: mocks.close,
+    });
+
+    const toolkit = new AISDKToolkit({
+      toolkit: {
+        local: {
+          type: "mcp",
+          server: { type: "http", url: "http://localhost:3001/mcp" },
+        },
+      },
+    });
+
+    await expect(toolkit.tools()).rejects.toMatchObject({
+      message: 'MCP toolkit entry "local" failed to connect: connect failed',
+      cause: error,
+    });
+    await expect(toolkit.tools()).resolves.toHaveProperty("echo");
+    expect(mocks.createMCPClient).toHaveBeenCalledTimes(2);
+  });
+
+  it("times out MCP client creation", async () => {
+    vi.useFakeTimers();
+    mocks.createMCPClient.mockReturnValue(never());
+
+    const toolkit = new AISDKToolkit({
+      toolkit: {
+        docs: {
+          type: "mcp",
+          server: {
+            type: "http",
+            url: "http://localhost:3001/mcp",
+            connectionTimeout: 10_000,
+          },
+        },
+      },
+    });
+
+    try {
+      const toolsPromise = toolkit.tools();
+      const expectedRejection = expect(toolsPromise).rejects.toThrow(
+        'MCP toolkit entry "docs" timed out while connecting after 10000ms.',
+      );
+      await vi.advanceTimersByTimeAsync(10_000);
+      await expectedRejection;
+
+      mocks.createMCPClient.mockResolvedValueOnce({
+        tools: vi.fn().mockResolvedValue({ echo: { inputSchema: {} } }),
+        close: mocks.close,
+      });
+      await expect(toolkit.tools()).resolves.toHaveProperty("echo");
+      expect(mocks.createMCPClient).toHaveBeenCalledTimes(2);
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+
+  it("closes an MCP client that resolves after its connection times out", async () => {
+    vi.useFakeTimers();
+    const close = vi.fn().mockResolvedValue(undefined);
+    let resolveClient!: (client: {
+      tools: typeof mocks.tools;
+      close: typeof close;
+    }) => void;
+    mocks.createMCPClient.mockReturnValue(
+      new Promise((resolve) => {
+        resolveClient = resolve;
+      }),
+    );
+
+    const toolkit = new AISDKToolkit({
+      toolkit: {
+        docs: {
+          type: "mcp",
+          server: {
+            type: "http",
+            url: "http://localhost:3001/mcp",
+            connectionTimeout: 10_000,
+          },
+        },
+      },
+    });
+
+    try {
+      const toolsPromise = toolkit.tools();
+      const expectedRejection = expect(toolsPromise).rejects.toThrow(
+        'MCP toolkit entry "docs" timed out while connecting after 10000ms.',
+      );
+      await vi.advanceTimersByTimeAsync(10_000);
+      await expectedRejection;
+
+      resolveClient({ tools: mocks.tools, close });
+      await vi.waitFor(() => expect(close).toHaveBeenCalledTimes(1));
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+
+  it("times out MCP tool listing", async () => {
+    vi.useFakeTimers();
+    const close = vi.fn().mockResolvedValue(undefined);
+    mocks.createMCPClient.mockResolvedValue({
+      tools: vi.fn(() => never()),
+      close,
+    });
+
+    const toolkit = new AISDKToolkit({
+      toolkit: {
+        docs: {
+          type: "mcp",
+          server: {
+            type: "http",
+            url: "http://localhost:3001/mcp",
+            connectionTimeout: 10_000,
+          },
+        },
+      },
+    });
+
+    try {
+      const toolsPromise = toolkit.tools();
+      const expectedRejection = expect(toolsPromise).rejects.toThrow(
+        'MCP toolkit entry "docs" timed out while listing tools after 10000ms.',
+      );
+      await vi.advanceTimersByTimeAsync(10_000);
+      await expectedRejection;
+      expect(close).toHaveBeenCalledTimes(1);
+
+      mocks.createMCPClient.mockResolvedValueOnce({
+        tools: vi.fn().mockResolvedValue({ echo: { inputSchema: {} } }),
+        close: mocks.close,
+      });
+      await expect(toolkit.tools()).resolves.toHaveProperty("echo");
+      expect(mocks.createMCPClient).toHaveBeenCalledTimes(2);
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+
+  it("includes the MCP toolkit entry name when listing tools fails", async () => {
+    const error = new Error("list failed");
+    mocks.tools.mockRejectedValue(error);
+    mocks.createMCPClient.mockResolvedValue({
+      tools: mocks.tools,
+      close: mocks.close,
+    });
+
+    const toolkit = new AISDKToolkit({
+      toolkit: {
+        docs: {
+          type: "mcp",
+          server: { type: "http", url: "http://localhost:3001/mcp" },
+        },
+      },
+    });
+
+    await expect(toolkit.tools()).rejects.toMatchObject({
+      message: 'MCP toolkit entry "docs" failed to list tools: list failed',
+      cause: error,
+    });
+  });
+
+  it("rejects duplicate MCP tool names", async () => {
+    mocks.createMCPClient
+      .mockResolvedValueOnce({
+        tools: vi.fn().mockResolvedValue({ echo: { inputSchema: {} } }),
+        close: mocks.close,
+      })
+      .mockResolvedValueOnce({
+        tools: vi.fn().mockResolvedValue({ echo: { inputSchema: {} } }),
+        close: mocks.close,
+      });
+
+    const toolkit = new AISDKToolkit({
+      toolkit: {
+        first: {
+          type: "mcp",
+          server: { type: "http", url: "http://localhost:3001/mcp" },
+        },
+        second: {
+          type: "mcp",
+          server: { type: "http", url: "http://localhost:3002/mcp" },
+        },
+      },
+    });
+
+    await expect(toolkit.tools()).rejects.toThrow(
+      /MCP tool name collision: "echo"/,
+    );
+  });
+
+  it("prefixes MCP tool names when configured", async () => {
+    const docsExecute = vi.fn().mockResolvedValue("docs result");
+    mocks.createMCPClient
+      .mockResolvedValueOnce({
+        tools: vi.fn().mockResolvedValue({
+          search: { inputSchema: {}, execute: docsExecute },
+        }),
+        close: mocks.close,
+      })
+      .mockResolvedValueOnce({
+        tools: vi.fn().mockResolvedValue({ search: { inputSchema: {} } }),
+        close: mocks.close,
+      });
+
+    const toolkit = new AISDKToolkit({
+      toolkit: defineMcpToolkit({
+        docs: {
+          server: { type: "http", url: "http://localhost:3001/mcp" },
+          prefix: "docs_",
+        },
+        github: {
+          server: { type: "http", url: "http://localhost:3002/mcp" },
+          prefix: "github_",
+        },
+      }),
+    });
+
+    const toolSet = await toolkit.tools();
+
+    expect(toolSet).toHaveProperty("docs_search");
+    expect(toolSet).toHaveProperty("github_search");
+    expect(toolSet).not.toHaveProperty("search");
+
+    expect(toolSet.docs_search?.execute).toBeTypeOf("function");
+    const executeOptions = {
+      toolCallId: "call-docs-search",
+      messages: [],
+    };
+    await expect(
+      toolSet.docs_search?.execute?.({ query: "assistant-ui" }, executeOptions),
+    ).resolves.toBe("docs result");
+    expect(docsExecute).toHaveBeenCalledWith(
+      { query: "assistant-ui" },
+      executeOptions,
+    );
+  });
+
+  it("rejects MCP tool names that collide with toolkit tools", async () => {
+    mocks.tools.mockResolvedValue({ search: { inputSchema: {} } });
+    mocks.createMCPClient.mockResolvedValue({
+      tools: mocks.tools,
+      close: mocks.close,
+    });
+
+    const toolkit = new AISDKToolkit({
+      toolkit: {
+        docs: {
+          type: "mcp",
+          server: { type: "http", url: "http://localhost:3001/mcp" },
+        },
+        search: {
+          type: "backend",
+          parameters: { type: "object", properties: {} },
+          execute: async () => "local search",
+        } as never,
+      },
+    });
+
+    await expect(toolkit.tools()).rejects.toThrow(
+      'MCP tool "search" from "docs" conflicts with toolkit tool "search"',
+    );
+  });
+
+  it("rejects MCP tool names that collide with provider tools", async () => {
+    mocks.tools.mockResolvedValue({ web_search: { inputSchema: {} } });
+    mocks.createMCPClient.mockResolvedValue({
+      tools: mocks.tools,
+      close: mocks.close,
+    });
+
+    const toolkit = new AISDKToolkit({
+      toolkit: {
+        docs: {
+          type: "mcp",
+          server: { type: "http", url: "http://localhost:3001/mcp" },
+        },
+        web_search: {
+          type: "provider",
+          providerId: "openai.web_search_preview",
+          args: { searchContextSize: "low" },
+        },
+      },
+    });
+
+    await expect(toolkit.tools()).rejects.toThrow(
+      'MCP tool "web_search" from "docs" conflicts with provider tool "web_search"',
+    );
+  });
+
+  it("rejects MCP tool names that collide with uploaded frontend tools", async () => {
+    mocks.tools.mockResolvedValue({ clientTool: { inputSchema: {} } });
+    mocks.createMCPClient.mockResolvedValue({
+      tools: mocks.tools,
+      close: mocks.close,
+    });
+
+    const toolkit = new AISDKToolkit({
+      toolkit: {
+        docs: {
+          type: "mcp",
+          server: { type: "http", url: "http://localhost:3001/mcp" },
+        },
+      },
+    });
+
+    await expect(
+      toolkit.tools({
+        frontend: {
+          clientTool: {
+            parameters: { type: "object", properties: {} },
+          },
+        },
+      }),
+    ).rejects.toThrow(
+      'MCP tool "clientTool" from "docs" conflicts with frontend tool "clientTool"',
+    );
+  });
+
+  it("ignores disabled MCP tools during name collision checks", async () => {
+    mocks.createMCPClient
+      .mockResolvedValueOnce({
+        tools: vi.fn().mockResolvedValue({ echo: { inputSchema: {} } }),
+        close: mocks.close,
+      })
+      .mockResolvedValueOnce({
+        tools: vi.fn().mockResolvedValue({ echo: { inputSchema: {} } }),
+        close: mocks.close,
+      });
+
+    const toolkit = new AISDKToolkit({
+      toolkit: {
+        first: {
+          type: "mcp",
+          server: { type: "http", url: "http://localhost:3001/mcp" },
+          tools: {
+            echo: {
+              disabled: true,
+            },
+          },
+        } as never,
+        second: {
+          type: "mcp",
+          server: { type: "http", url: "http://localhost:3002/mcp" },
+        },
+      },
+    });
+
+    await expect(toolkit.tools()).resolves.toEqual({
+      echo: { inputSchema: {} },
+    });
+  });
+
+  it("includes provider tools alongside MCP tools", async () => {
+    mocks.tools.mockResolvedValue({ echo: { inputSchema: {} } });
+    mocks.createMCPClient.mockResolvedValue({
+      tools: mocks.tools,
+      close: mocks.close,
+    });
+
+    const toolkit = new AISDKToolkit({
+      toolkit: {
+        local: {
+          type: "mcp",
+          server: { type: "http", url: "http://localhost:3001/mcp" },
+        },
+        web_search: {
+          type: "provider",
+          providerId: "openai.web_search_preview",
+          args: { searchContextSize: "low" },
+          supportsDeferredResults: false,
+        },
+      },
+    });
+
+    await expect(toolkit.tools()).resolves.toMatchObject({
+      echo: { inputSchema: {} },
+      web_search: {
+        type: "provider",
+        id: "openai.web_search_preview",
+        args: { searchContextSize: "low" },
+        supportsDeferredResults: false,
+      },
+    });
+  });
+});
+
+describe("AISDKToolkit toModelOutput", () => {
+  const createWeatherTools = (toModelOutput?: any) =>
+    new AISDKToolkit({
+      toolkit: {
+        get_weather: {
+          ...(toModelOutput && { toModelOutput }),
+        },
+      } as any,
+    }).tools();
+
+  it("adapts assistant-ui model content parts to the AI SDK tool output shape", async () => {
+    const tools = await createWeatherTools(({ output }: any) => [
+      { type: "text", text: `Weather card displayed: ${output.location}` },
+    ]);
+
+    const output = await tools.get_weather!.toModelOutput!({
+      toolCallId: "tc-weather",
+      input: {},
+      output: { location: "San Francisco" },
+    });
+
+    expect(output).toEqual({
+      type: "content",
+      value: [{ type: "text", text: "Weather card displayed: San Francisco" }],
+    });
+  });
+
+  it("uses stored model content envelopes without re-running the custom projector", async () => {
+    let called = false;
+    const tools = await createWeatherTools(() => {
+      called = true;
+      return [{ type: "text", text: "recomputed" }];
+    });
+
+    const output = await tools.get_weather!.toModelOutput!({
+      toolCallId: "tc-weather",
+      input: {},
+      output: wrapModelContentEnvelope({ location: "San Francisco" }, [
+        { type: "text", text: "cached weather receipt" },
+      ]),
+    });
+
+    expect(called).toBe(false);
+    expect(output).toEqual({
+      type: "content",
+      value: [{ type: "text", text: "cached weather receipt" }],
+    });
+  });
+
+  it("falls back to default model output when no custom projector is defined", async () => {
+    const tools = await createWeatherTools();
+
+    const output = await tools.get_weather!.toModelOutput!({
+      toolCallId: "tc-weather",
+      input: {},
+      output: { location: "San Francisco" },
+    });
+
+    expect(output).toEqual({
+      type: "json",
+      value: { location: "San Francisco" },
+    });
+  });
+
+  it("uses stored model content envelopes when no custom projector is defined", async () => {
+    const tools = await createWeatherTools();
+
+    const output = await tools.get_weather!.toModelOutput!({
+      toolCallId: "tc-weather",
+      input: {},
+      output: wrapModelContentEnvelope({ location: "San Francisco" }, [
+        { type: "text", text: "cached weather receipt" },
+      ]),
+    });
+
+    expect(output).toEqual({
+      type: "content",
+      value: [{ type: "text", text: "cached weather receipt" }],
+    });
+  });
+});