@assistant-ui/react-mcp 0.0.22 → 0.1.1

package/src/resources/McpManagerResource.test.ts

@@ -1,10 +1,39 @@
 import { createTapRoot, useResource } from "@assistant-ui/tap";
-import { describe, expect, it } from "vitest";
+import { describe, expect, it, vi } from "vitest";
 import { defineConnector } from "../connector";
 import type { MCPConnector } from "../mcp-scope";
 import { assertUniqueServerIds } from "../utils/serverId";
 import { McpManagerResource } from "./McpManagerResource";
+import { McpCustomStorage } from "./storage/McpCustomStorage";
 import { McpMemoryStorage } from "./storage/McpMemoryStorage";
+import type { MCPStorageElement } from "./storage/types";
+
+const mocks = vi.hoisted(() => {
+  const Client = vi.fn().mockImplementation(function Client(this: any) {
+    this.connect = vi.fn(async () => {});
+    this.listTools = vi.fn(async () => ({ tools: [] }));
+    this.setRequestHandler = vi.fn();
+    this.setNotificationHandler = vi.fn();
+  });
+  const StreamableHTTPClientTransport = vi
+    .fn()
+    .mockImplementation(function StreamableHTTPClientTransport(this: any) {
+      this.close = vi.fn(async () => {});
+    });
+
+  return { Client, StreamableHTTPClientTransport };
+});
+
+vi.mock("@modelcontextprotocol/client", async (importOriginal) => ({
+  ...(await importOriginal()),
+  Client: mocks.Client,
+  StreamableHTTPClientTransport: mocks.StreamableHTTPClientTransport,
+}));
+
+vi.mock("@assistant-ui/store", async (importOriginal) => ({
+  ...(await importOriginal()),
+  useAssistantClientRef: () => ({ current: null }),
+}));
 
 const connector = (id: string, name = id): MCPConnector =>
   defineConnector({
@@ -14,12 +43,15 @@ const connector = (id: string, name = id): MCPConnector =>
     auth: { type: "none" },
   });
 
-const mount = (connectors: MCPConnector[]) =>
+const mount = (
+  connectors: MCPConnector[],
+  storage: MCPStorageElement = McpMemoryStorage(),
+) =>
   createTapRoot(function Root() {
     return useResource(
       McpManagerResource({
         connectors,
-        storage: McpMemoryStorage(),
+        storage,
         autoConnect: false,
       }),
     );
@@ -37,4 +69,142 @@ describe("McpManagerResource server ids", () => {
   it("allows distinct ids", () => {
     expect(() => assertUniqueServerIds(["docs", "linear"])).not.toThrow();
   });
+
+  it("passes connector cache configuration to its client", async () => {
+    mocks.Client.mockClear();
+    const root = mount([
+      defineConnector({
+        id: "docs",
+        name: "Docs",
+        url: "https://example.com/docs/mcp",
+        auth: { type: "none" },
+        cache: { defaultTtlMs: 5_000 },
+      }),
+    ]);
+
+    try {
+      await root.getValue().connector({ index: 0 }).connect();
+
+      expect(mocks.Client).toHaveBeenCalledWith(
+        {
+          name: "assistant-ui-mcp",
+          version: "0.0.0",
+        },
+        expect.objectContaining({ defaultCacheTtlMs: 5_000 }),
+      );
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("passes custom server cache configuration to its client", async () => {
+    mocks.Client.mockClear();
+    const root = mount([]);
+
+    try {
+      const id = await root.getValue().addCustomServer({
+        name: "Docs",
+        url: "https://example.com/docs/mcp",
+        auth: { type: "none" },
+        cache: { defaultTtlMs: 5_000 },
+      });
+
+      await vi.waitFor(() =>
+        expect(root.getValue().getState().customServers).toHaveLength(1),
+      );
+      await root.getValue().server({ id }).connect();
+
+      expect(mocks.Client).toHaveBeenCalledWith(
+        {
+          name: "assistant-ui-mcp",
+          version: "0.0.0",
+        },
+        expect.objectContaining({ defaultCacheTtlMs: 5_000 }),
+      );
+    } finally {
+      root.unmount();
+    }
+  });
+});
+
+describe("McpManagerResource storage failures", () => {
+  it("handles custom server load failures", async () => {
+    const error = new Error("load failed");
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => {});
+    const root = mount(
+      [],
+      McpCustomStorage({
+        loadCustomServers: vi.fn(async () => {
+          throw error;
+        }),
+        saveCustomServers: vi.fn(async () => {}),
+        loadAuthState: vi.fn(async () => null),
+        saveAuthState: vi.fn(async () => {}),
+        clearAuthState: vi.fn(async () => {}),
+      }),
+    );
+
+    try {
+      await vi.waitFor(() =>
+        expect(root.getValue().getState().isHydrated).toBe(true),
+      );
+      expect(root.getValue().getState().customServers).toHaveLength(0);
+      expect(consoleError).toHaveBeenCalledWith(
+        "[assistant-ui/react-mcp] failed to load custom servers:",
+        error,
+      );
+    } finally {
+      root.unmount();
+      consoleError.mockRestore();
+    }
+  });
+
+  it("handles custom server save failures", async () => {
+    const error = new Error("save failed");
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => {});
+    const saveCustomServers = vi.fn(async () => {});
+    const root = mount(
+      [],
+      McpCustomStorage({
+        loadCustomServers: vi.fn(async () => []),
+        saveCustomServers,
+        loadAuthState: vi.fn(async () => null),
+        saveAuthState: vi.fn(async () => {}),
+        clearAuthState: vi.fn(async () => {}),
+      }),
+    );
+
+    try {
+      await vi.waitFor(() =>
+        expect(root.getValue().getState().isHydrated).toBe(true),
+      );
+      await vi.waitFor(() => expect(saveCustomServers).toHaveBeenCalled());
+      saveCustomServers.mockClear();
+      saveCustomServers.mockRejectedValue(error);
+
+      await root.getValue().addCustomServer({
+        name: "Docs",
+        url: "https://example.com/docs/mcp",
+        auth: { type: "none" },
+      });
+
+      await vi.waitFor(() => {
+        expect(saveCustomServers).toHaveBeenCalledWith([
+          expect.objectContaining({ name: "Docs" }),
+        ]);
+        expect(root.getValue().getState().customServers).toHaveLength(1);
+      });
+      expect(consoleError).toHaveBeenCalledWith(
+        "[assistant-ui/react-mcp] failed to save custom servers:",
+        error,
+      );
+    } finally {
+      root.unmount();
+      consoleError.mockRestore();
+    }
+  });
 });

package/src/resources/McpManagerResource.ts

@@ -39,6 +39,16 @@ function defaultRedirectUri(): string {
 // array each render (which would invalidate the serverElements memo below).
 const NO_CONNECTORS: MCPConnector[] = [];
 
+const reportCustomStorageFailure = (
+  operation: "load" | "save",
+  error: unknown,
+) => {
+  console.error(
+    `[assistant-ui/react-mcp] failed to ${operation} custom servers:`,
+    error,
+  );
+};
+
 const useMcpManagerResource = (
   props: McpManagerResourceProps,
 ): ClientOutput<"mcp"> => {
@@ -69,8 +79,11 @@ const useMcpManagerResource = (
     try {
       records = await storage.loadCustomServers();
     } catch (error) {
+      if (!signal.cancelled) {
+        reportCustomStorageFailure("load", error);
+      }
       markHydrated();
-      throw error;
+      return;
     }
 
     if (signal.cancelled) return;
@@ -96,7 +109,11 @@ const useMcpManagerResource = (
   const persistCustomServers = useEffectEvent(
     async (records: MCPCustomServerRecord[]) => {
       if (!hydratedRef.current) return;
-      await storage.saveCustomServers(records);
+      try {
+        await storage.saveCustomServers(records);
+      } catch (error) {
+        reportCustomStorageFailure("save", error);
+      }
     },
   );
 
@@ -124,6 +141,10 @@ const useMcpManagerResource = (
           redirectUri,
           autoConnect,
           connectionTimeout: c.connectionTimeout ?? connectionTimeout,
+          ...(c.cache !== undefined ? { cache: c.cache } : {}),
+          ...(c.elicitation !== undefined
+            ? { elicitation: c.elicitation }
+            : {}),
           onRemove: async () => {
             // connectors cannot be removed
           },
@@ -143,6 +164,10 @@ const useMcpManagerResource = (
           redirectUri,
           autoConnect,
           connectionTimeout: s.connectionTimeout ?? connectionTimeout,
+          ...(s.cache !== undefined ? { cache: s.cache } : {}),
+          ...(s.elicitation !== undefined
+            ? { elicitation: s.elicitation }
+            : {}),
           onRemove: async () => {
             setCustomServers((prev) => prev.filter((x) => x.id !== s.id));
           },
@@ -226,7 +251,14 @@ const useMcpManagerResource = (
     },
     connector: ({ index }) => serverByKind("connector", index),
     customServer: ({ index }) => serverByKind("custom", index),
-    addCustomServer: async ({ name, url, auth, connectionTimeout }) => {
+    addCustomServer: async ({
+      name,
+      url,
+      auth,
+      connectionTimeout,
+      cache,
+      elicitation,
+    }) => {
       const record: MCPCustomServerRecord = {
         id:
           typeof crypto !== "undefined" && "randomUUID" in crypto
@@ -236,6 +268,8 @@ const useMcpManagerResource = (
         url,
         auth: auth as MCPAuthConfig,
         connectionTimeout,
+        ...(cache !== undefined ? { cache } : {}),
+        ...(elicitation !== undefined ? { elicitation } : {}),
         createdAt: Date.now(),
       };
       setCustomServers((prev) => [...prev, record]);