@assistant-ui/react-mcp 0.0.16 → 0.0.18

package/src/hooks/useMcpOAuthCallback.tsx

@@ -2,6 +2,20 @@ import { type FC, type ReactNode, useEffect, useRef, useState } from "react";
 import { useAui } from "@assistant-ui/store";
 import { decodeServerIdFromState } from "../auth/createOAuthProvider";
 
+export const createMcpOAuthCallbackError = (
+  err: unknown,
+  serverId: string | null,
+): Error => {
+  const message = err instanceof Error ? err.message : String(err);
+  if (serverId) {
+    return new Error(
+      `MCP OAuth callback for server "${serverId}" failed: ${message}`,
+      { cause: err },
+    );
+  }
+  return new Error(`MCP OAuth callback failed: ${message}`, { cause: err });
+};
+
 export type UseMcpOAuthCallbackOptions = {
   /** Defaults to `window.location.href`. */
   url?: string;
@@ -39,27 +53,28 @@ export function useMcpOAuthCallback(
     startedRef.current = url;
 
     (async () => {
+      let serverId: string | null = null;
       try {
         const parsed = new URL(url);
         const state = parsed.searchParams.get("state");
+        if (state) serverId = decodeServerIdFromState(state);
         const error = parsed.searchParams.get("error");
         if (error) {
           throw new Error(
             parsed.searchParams.get("error_description") ?? error,
           );
         }
-        if (!state) throw new Error("missing state parameter in callback URL");
-        const serverId = decodeServerIdFromState(state);
+        if (!state) throw new Error('missing "state" parameter');
         if (!serverId) {
-          throw new Error("callback state does not match an MCP server");
+          throw new Error("state was not created by assistant-ui MCP");
         }
         setResult({ status: "running", serverId, error: null });
         await aui.mcp().server({ id: serverId }).completeAuth(url);
         setResult({ status: "done", serverId, error: null });
         optsRef.current.onComplete?.(serverId);
       } catch (err) {
-        const e = err instanceof Error ? err : new Error(String(err));
-        setResult((prev) => ({ ...prev, status: "error", error: e }));
+        const e = createMcpOAuthCallbackError(err, serverId);
+        setResult({ status: "error", serverId, error: e });
         optsRef.current.onError?.(e);
       }
     })();

package/src/mcp-scope.ts

@@ -19,6 +19,7 @@ export type MCPConnector = {
   url: string;
   icon?: string | undefined;
   auth: MCPAuthConfig;
+  connectionTimeout?: number | undefined;
 };
 
 export type MCPCustomServerRecord = {
@@ -26,6 +27,7 @@ export type MCPCustomServerRecord = {
   name: string;
   url: string;
   auth: MCPAuthConfig;
+  connectionTimeout?: number | undefined;
   createdAt: number;
 };
 
@@ -70,6 +72,8 @@ export type MCPServerMethods = {
   disconnect: () => Promise<void>;
   remove: () => Promise<void>;
   callTool: (name: string, args: unknown) => Promise<unknown>;
+  /** List resources exposed by the server. Returns the raw MCP `ListResourcesResult`. */
+  listResources: (params?: { cursor?: string | undefined }) => Promise<unknown>;
   /** Read a resource by URI. Returns the raw MCP `ReadResourceResult`. */
   readResource: (uri: string) => Promise<unknown>;
   /** OAuth only: pass full callback URL (e.g. window.location.href) */
@@ -93,6 +97,7 @@ export type MCPManagerMethods = {
     name: string;
     url: string;
     auth: MCPAuthConfig;
+    connectionTimeout?: number | undefined;
   }) => Promise<string>;
   removeServer: (id: string) => Promise<void>;
 };

package/src/resources/McpManagerResource.test.ts

@@ -0,0 +1,40 @@
+import { createTapRoot, useResource } from "@assistant-ui/tap";
+import { describe, expect, it } from "vitest";
+import { defineConnector } from "../connector";
+import type { MCPConnector } from "../mcp-scope";
+import { assertUniqueServerIds } from "../utils/serverId";
+import { McpManagerResource } from "./McpManagerResource";
+import { McpMemoryStorage } from "./storage/McpMemoryStorage";
+
+const connector = (id: string, name = id): MCPConnector =>
+  defineConnector({
+    id,
+    name,
+    url: `https://example.com/${id}/mcp`,
+    auth: { type: "none" },
+  });
+
+const mount = (connectors: MCPConnector[]) =>
+  createTapRoot(function Root() {
+    return useResource(
+      McpManagerResource({
+        connectors,
+        storage: McpMemoryStorage(),
+        autoConnect: false,
+      }),
+    );
+  });
+
+describe("McpManagerResource server ids", () => {
+  it("throws when connectors reuse an id", () => {
+    expect(() =>
+      mount([connector("docs", "Docs"), connector("docs", "Internal Docs")]),
+    ).toThrow(
+      'McpManagerResource received duplicate MCP server id "docs". Server ids must be unique because they are used for lookups, OAuth routing, and tool name prefixes.',
+    );
+  });
+
+  it("allows distinct ids", () => {
+    expect(() => assertUniqueServerIds(["docs", "linear"])).not.toThrow();
+  });
+});

package/src/resources/McpManagerResource.ts

@@ -11,6 +11,7 @@ import type { Tool } from "assistant-stream";
 import { McpServerResource } from "./McpServerResource";
 import { McpLocalStorage } from "./storage/McpLocalStorage";
 import type { MCPStorageElement } from "./storage/types";
+import { assertUniqueServerIds } from "../utils/serverId";
 import type {
   MCPAuthConfig,
   MCPConnector,
@@ -25,6 +26,8 @@ export type McpManagerResourceProps = {
   oauthRedirectUri?: string | undefined;
   /** Connect on mount when usable auth exists. Default true. */
   autoConnect?: boolean | undefined;
+  /** Optional timeout in milliseconds for connect/listTools calls. Disabled by default. */
+  connectionTimeout?: number | undefined;
 };
 
 function defaultRedirectUri(): string {
@@ -42,6 +45,7 @@ const useMcpManagerResource = (
   const connectors = props.connectors ?? NO_CONNECTORS;
   const autoConnect = props.autoConnect ?? true;
   const redirectUri = props.oauthRedirectUri ?? defaultRedirectUri();
+  const connectionTimeout = props.connectionTimeout;
 
   const storageElement = props.storage ?? McpLocalStorage();
   const storage = useResource(storageElement);
@@ -101,6 +105,11 @@ const useMcpManagerResource = (
   }, [customServers]);
 
   const serverElements = useMemo(() => {
+    assertUniqueServerIds([
+      ...connectors.map((c) => c.id),
+      ...customServers.map((s) => s.id),
+    ]);
+
     const connectorElements = connectors.map((c) =>
       withKey(
         c.id,
@@ -114,6 +123,7 @@ const useMcpManagerResource = (
           storage,
           redirectUri,
           autoConnect,
+          connectionTimeout: c.connectionTimeout ?? connectionTimeout,
           onRemove: async () => {
             // connectors cannot be removed
           },
@@ -132,6 +142,7 @@ const useMcpManagerResource = (
           storage,
           redirectUri,
           autoConnect,
+          connectionTimeout: s.connectionTimeout ?? connectionTimeout,
           onRemove: async () => {
             setCustomServers((prev) => prev.filter((x) => x.id !== s.id));
           },
@@ -139,7 +150,14 @@ const useMcpManagerResource = (
       ),
     );
     return [...connectorElements, ...customElements];
-  }, [connectors, customServers, storage, redirectUri, autoConnect]);
+  }, [
+    connectors,
+    customServers,
+    storage,
+    redirectUri,
+    autoConnect,
+    connectionTimeout,
+  ]);
 
   const lookup = useClientLookup(serverElements);
 
@@ -208,7 +226,7 @@ const useMcpManagerResource = (
     },
     connector: ({ index }) => serverByKind("connector", index),
     customServer: ({ index }) => serverByKind("custom", index),
-    addCustomServer: async ({ name, url, auth }) => {
+    addCustomServer: async ({ name, url, auth, connectionTimeout }) => {
       const record: MCPCustomServerRecord = {
         id:
           typeof crypto !== "undefined" && "randomUUID" in crypto
@@ -217,6 +235,7 @@ const useMcpManagerResource = (
         name,
         url,
         auth: auth as MCPAuthConfig,
+        connectionTimeout,
         createdAt: Date.now(),
       };
       setCustomServers((prev) => [...prev, record]);

package/src/resources/McpServerResource.test.ts

@@ -0,0 +1,364 @@
+import { createTapRoot, useResource } from "@assistant-ui/tap";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import type { MCPAuthConfig } from "../mcp-scope";
+import type { MCPStorage } from "./storage/types";
+
+const mocks = vi.hoisted(() => {
+  const clients: any[] = [];
+  const transports: any[] = [];
+  const connectResults: Array<() => Promise<void>> = [];
+  const listToolsResults: Array<
+    () => Promise<{
+      tools: Array<{
+        name: string;
+        description?: string;
+        inputSchema: unknown;
+      }>;
+    }>
+  > = [];
+  const finishAuthResults: Array<() => Promise<void>> = [];
+
+  const Client = vi.fn().mockImplementation(function Client(this: any) {
+    const index = clients.length;
+    this.connect = vi.fn(() => connectResults[index]?.() ?? Promise.resolve());
+    this.listTools = vi.fn(
+      () => listToolsResults[index]?.() ?? Promise.resolve({ tools: [] }),
+    );
+    this.callTool = vi.fn();
+    this.listResources = vi.fn(() => Promise.resolve({ resources: [] }));
+    this.readResource = vi.fn();
+    clients.push(this);
+  });
+
+  const StreamableHTTPClientTransport = vi
+    .fn()
+    .mockImplementation(function StreamableHTTPClientTransport(this: any) {
+      const index = transports.length;
+      this.close = vi.fn(() => Promise.resolve());
+      this.finishAuth = vi.fn(
+        () => finishAuthResults[index]?.() ?? Promise.resolve(),
+      );
+      transports.push(this);
+    });
+
+  return {
+    Client,
+    StreamableHTTPClientTransport,
+    clients,
+    transports,
+    connectResults,
+    listToolsResults,
+    finishAuthResults,
+  };
+});
+
+vi.mock("@modelcontextprotocol/sdk/client/index.js", () => ({
+  Client: mocks.Client,
+}));
+
+vi.mock("@modelcontextprotocol/sdk/client/streamableHttp.js", () => ({
+  StreamableHTTPClientTransport: mocks.StreamableHTTPClientTransport,
+}));
+
+const { McpServerResource } = await import("./McpServerResource");
+
+const never = <T>() => new Promise<T>(() => {});
+
+const tick = async () => {
+  await Promise.resolve();
+};
+
+const flushMacrotask = async () => {
+  await new Promise<void>((resolve) => {
+    const channel = new MessageChannel();
+    channel.port1.onmessage = () => {
+      channel.port1.close();
+      channel.port2.close();
+      resolve();
+    };
+    channel.port2.postMessage(null);
+  });
+};
+
+const waitFor = async (predicate: () => boolean) => {
+  for (let i = 0; i < 20; i++) {
+    if (predicate()) return;
+    await tick();
+  }
+  expect(predicate()).toBe(true);
+};
+
+const createStorage = (): MCPStorage => ({
+  loadCustomServers: vi.fn(async () => []),
+  saveCustomServers: vi.fn(async () => {}),
+  loadAuthState: vi.fn(async () => null),
+  saveAuthState: vi.fn(async () => {}),
+  clearAuthState: vi.fn(async () => {}),
+});
+
+const resetMocks = () => {
+  mocks.clients.length = 0;
+  mocks.transports.length = 0;
+  mocks.connectResults.length = 0;
+  mocks.listToolsResults.length = 0;
+  mocks.finishAuthResults.length = 0;
+  mocks.Client.mockClear();
+  mocks.StreamableHTTPClientTransport.mockClear();
+};
+
+const mount = (props?: {
+  auth?: MCPAuthConfig | undefined;
+  connectionTimeout?: number | undefined;
+}) => {
+  const connectionTimeout =
+    props && "connectionTimeout" in props ? props.connectionTimeout : 10_000;
+
+  return createTapRoot(function Root() {
+    return useResource(
+      McpServerResource({
+        id: "docs",
+        kind: "connector",
+        name: "Docs",
+        url: "https://example.com/mcp",
+        auth: props?.auth ?? { type: "none" },
+        storage: createStorage(),
+        redirectUri: "https://example.com/callback",
+        autoConnect: false,
+        connectionTimeout,
+        onRemove: vi.fn(async () => {}),
+      }),
+    );
+  });
+};
+
+describe("McpServerResource connectionTimeout", () => {
+  beforeEach(() => {
+    vi.useFakeTimers();
+    resetMocks();
+  });
+
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
+  it("fails the connection when client.connect hangs", async () => {
+    mocks.connectResults.push(() => never());
+    const root = mount();
+
+    try {
+      const connectPromise = root.getValue().connect();
+      await waitFor(() => mocks.clients[0]?.connect.mock.calls.length === 1);
+
+      await vi.advanceTimersByTimeAsync(10_000);
+      await connectPromise;
+
+      expect(root.getValue().getState()).toMatchObject({
+        connectionState: "error",
+        tools: [],
+        lastError: {
+          message:
+            'MCP server "docs" timed out while connecting after 10000ms.',
+        },
+      });
+      expect(mocks.transports[0].close).toHaveBeenCalledTimes(1);
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("fails the connection when client.listTools hangs", async () => {
+    mocks.listToolsResults.push(() => never());
+    const root = mount();
+
+    try {
+      const connectPromise = root.getValue().connect();
+      await waitFor(() => mocks.clients[0]?.listTools.mock.calls.length === 1);
+
+      await vi.advanceTimersByTimeAsync(10_000);
+      await connectPromise;
+
+      expect(root.getValue().getState()).toMatchObject({
+        connectionState: "error",
+        tools: [],
+        lastError: {
+          message:
+            'MCP server "docs" timed out while listing tools after 10000ms.',
+        },
+      });
+      expect(mocks.transports[0].close).toHaveBeenCalledTimes(1);
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("uses one timeout budget for connect and listTools", async () => {
+    mocks.connectResults.push(
+      () => new Promise<void>((resolve) => setTimeout(resolve, 9_000)),
+    );
+    mocks.listToolsResults.push(() => never());
+    const root = mount();
+
+    try {
+      const connectPromise = root.getValue().connect();
+      await waitFor(() => mocks.clients[0]?.connect.mock.calls.length === 1);
+
+      await vi.advanceTimersByTimeAsync(9_000);
+      await waitFor(() => mocks.clients[0]?.listTools.mock.calls.length === 1);
+      await vi.advanceTimersByTimeAsync(999);
+      await tick();
+      expect(root.getValue().getState().connectionState).toBe("connecting");
+
+      await vi.advanceTimersByTimeAsync(1);
+      await connectPromise;
+
+      expect(root.getValue().getState()).toMatchObject({
+        connectionState: "error",
+        tools: [],
+        lastError: {
+          message:
+            'MCP server "docs" timed out while listing tools after 10000ms.',
+        },
+      });
+      expect(mocks.transports[0].close).toHaveBeenCalledTimes(1);
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("keeps waiting when connectionTimeout is undefined", async () => {
+    mocks.connectResults.push(() => never());
+    const root = mount({ connectionTimeout: undefined });
+
+    try {
+      void root.getValue().connect();
+      await waitFor(() => mocks.clients[0]?.connect.mock.calls.length === 1);
+      await vi.advanceTimersByTimeAsync(10_000);
+      await tick();
+
+      expect(root.getValue().getState()).toMatchObject({
+        connectionState: "connecting",
+        tools: [],
+        lastError: null,
+      });
+      expect(mocks.transports[0].close).not.toHaveBeenCalled();
+    } finally {
+      root.unmount();
+    }
+  });
+});
+
+describe("McpServerResource completeAuth", () => {
+  beforeEach(resetMocks);
+
+  it("rejects when the callback URL has no authorization code", async () => {
+    const root = mount();
+
+    try {
+      await expect(
+        root.getValue().completeAuth("https://example.com/callback?state=abc"),
+      ).rejects.toThrow("missing authorization code in callback URL");
+      await flushMacrotask();
+
+      expect(root.getValue().getState()).toMatchObject({
+        connectionState: "error",
+        lastError: {
+          message: "missing authorization code in callback URL",
+        },
+      });
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("rejects after storing finishAuth failures on the server state", async () => {
+    mocks.finishAuthResults.push(() =>
+      Promise.reject(new Error("invalid_grant")),
+    );
+    const root = mount({ auth: { type: "oauth" } });
+
+    try {
+      await expect(
+        root.getValue().completeAuth("https://example.com/callback?code=abc"),
+      ).rejects.toThrow("invalid_grant");
+      await flushMacrotask();
+
+      expect(root.getValue().getState()).toMatchObject({
+        connectionState: "error",
+        lastError: {
+          message: "invalid_grant",
+        },
+      });
+      expect(mocks.transports[0].finishAuth).toHaveBeenCalledWith("abc");
+      expect(mocks.transports[0].close).toHaveBeenCalledTimes(1);
+    } finally {
+      root.unmount();
+    }
+  });
+});
+
+describe("McpServerResource resource methods", () => {
+  beforeEach(() => {
+    mocks.clients.length = 0;
+    mocks.transports.length = 0;
+    mocks.connectResults.length = 0;
+    mocks.listToolsResults.length = 0;
+    mocks.Client.mockClear();
+    mocks.StreamableHTTPClientTransport.mockClear();
+  });
+
+  it("lists resources from a connected server", async () => {
+    const result = {
+      resources: [
+        {
+          uri: "docs://intro",
+          name: "Intro",
+          mimeType: "text/markdown",
+        },
+      ],
+    };
+    const root = mount();
+
+    try {
+      await root.getValue().connect();
+      mocks.clients[0].listResources.mockResolvedValueOnce(result);
+
+      await expect(root.getValue().listResources()).resolves.toBe(result);
+      expect(mocks.clients[0].listResources).toHaveBeenCalledTimes(1);
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("forwards the resource pagination cursor", async () => {
+    const result = {
+      resources: [{ uri: "docs://page-two", name: "Page two" }],
+    };
+    const root = mount();
+
+    try {
+      await root.getValue().connect();
+      mocks.clients[0].listResources.mockResolvedValueOnce(result);
+
+      await expect(
+        root.getValue().listResources({ cursor: "next-page" }),
+      ).resolves.toBe(result);
+      expect(mocks.clients[0].listResources).toHaveBeenCalledWith({
+        cursor: "next-page",
+      });
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("rejects listResources when the server is disconnected", async () => {
+    const root = mount();
+
+    try {
+      await expect(root.getValue().listResources()).rejects.toThrow(
+        'MCP server "docs" is not connected',
+      );
+    } finally {
+      root.unmount();
+    }
+  });
+});

package/src/resources/McpServerResource.ts

@@ -30,6 +30,7 @@ export type McpServerResourceProps = {
   storage: MCPStorage;
   redirectUri: string;
   autoConnect: boolean;
+  connectionTimeout?: number | undefined;
   onRemove: () => Promise<void>;
 };
 
@@ -46,6 +47,35 @@ const useMcpServerResource = (
   const clientRef = useRef<Client | null>(null);
   const transportRef = useRef<StreamableHTTPClientTransport | null>(null);
 
+  const withConnectionTimeout = useEffectEvent(
+    async <T>(
+      promise: Promise<T>,
+      phase: "connecting" | "listing tools",
+      startedAt: number,
+    ): Promise<T> => {
+      const timeoutMs = props.connectionTimeout;
+      if (timeoutMs === undefined) return await promise;
+      const remainingMs = timeoutMs - (Date.now() - startedAt);
+      const timeoutError = () =>
+        new Error(
+          `MCP server "${props.id}" timed out while ${phase} after ${timeoutMs}ms.`,
+        );
+      if (remainingMs <= 0) throw timeoutError();
+
+      let timeout: ReturnType<typeof setTimeout> | undefined;
+      try {
+        return await Promise.race([
+          promise,
+          new Promise<never>((_, reject) => {
+            timeout = setTimeout(() => reject(timeoutError()), remainingMs);
+          }),
+        ]);
+      } finally {
+        if (timeout !== undefined) clearTimeout(timeout);
+      }
+    },
+  );
+
   const buildTransport = useEffectEvent(
     async (): Promise<StreamableHTTPClientTransport> => {
       if (props.auth.type === "oauth") {
@@ -80,16 +110,25 @@ const useMcpServerResource = (
         name: "assistant-ui-mcp",
         version: "0.0.0",
       });
+      const startedAt = Date.now();
       // SDK's StreamableHTTPClientTransport.sessionId is `string | undefined`
       // but Transport.sessionId is declared `string?` — under
       // exactOptionalPropertyTypes the SDK's own classes don't satisfy its
       // Transport interface. Cast to bridge the gap.
-      await client.connect(transport as unknown as Transport);
+      await withConnectionTimeout(
+        client.connect(transport as unknown as Transport),
+        "connecting",
+        startedAt,
+      );
       // Defer ref assignment until listTools() also succeeds — otherwise a
       // post-connect failure leaves stale refs that `callTool()` would
       // happily walk into, producing confusing SDK errors instead of
       // "not connected".
-      const list = await client.listTools();
+      const list = await withConnectionTimeout(
+        client.listTools(),
+        "listing tools",
+        startedAt,
+      );
       clientRef.current = client;
       transportRef.current = transport;
       setTools(
@@ -182,10 +221,12 @@ const useMcpServerResource = (
       await finalizeConnect(transport);
     } catch (err) {
       await closeTransport();
+      const error = err instanceof Error ? err : new Error(String(err));
       setLastError({
-        message: err instanceof Error ? err.message : String(err),
+        message: error.message,
       });
       setConnectionState("error");
+      throw error;
     }
   });
 
@@ -269,6 +310,13 @@ const useMcpServerResource = (
         arguments: args as Record<string, unknown> | undefined,
       });
     },
+    listResources: async (params) => {
+      const client = clientRef.current;
+      if (!client) {
+        throw new Error(`MCP server "${props.id}" is not connected`);
+      }
+      return await client.listResources(params);
+    },
     readResource: async (uri) => {
       const client = clientRef.current;
       if (!client) {