@assistant-ui/react-mcp 0.1.18 → 0.1.20

package/src/resources/McpManagerResource.test.ts

@@ -92,6 +92,51 @@ describe("McpManagerResource server ids", () => {
     expect(() => assertUniqueServerIds(["docs", "linear"])).not.toThrow();
   });
 
+  it("keeps the first persisted custom server when ids are duplicated", async () => {
+    const docsServer: MCPCustomServerRecord = {
+      id: "docs",
+      name: "Docs",
+      url: "https://example.com/docs/mcp",
+      auth: { type: "none" },
+      createdAt: 1,
+    };
+    const saveCustomServers = vi.fn(async () => {});
+    const consoleError = vi
+      .spyOn(console, "error")
+      .mockImplementation(() => {});
+    const root = mount(
+      [],
+      McpCustomStorage({
+        loadCustomServers: vi.fn(async () => [
+          docsServer,
+          { ...docsServer, name: "Duplicate Docs", createdAt: 2 },
+        ]),
+        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),
+      );
+      expect(root.getValue().getState().customServers).toHaveLength(1);
+      expect(root.getValue().getState().customServers[0]).toMatchObject({
+        id: "docs",
+        name: "Docs",
+      });
+      expect(saveCustomServers).not.toHaveBeenCalled();
+      expect(consoleError).toHaveBeenCalledWith(
+        '[assistant-ui/react-mcp] ignored duplicate custom server id "docs" loaded from storage',
+      );
+    } finally {
+      consoleError.mockRestore();
+      root.unmount();
+    }
+  });
+
   it("passes connector cache configuration to its client", async () => {
     mocks.Client.mockClear();
     const root = mount([
@@ -374,6 +419,46 @@ describe("McpManagerResource storage failures", () => {
 });
 
 describe("McpManagerResource storage ordering", () => {
+  it("preserves a removal made before custom server hydration finishes", async () => {
+    const docsServer: MCPCustomServerRecord = {
+      id: "docs",
+      name: "Docs",
+      url: "https://example.com/docs/mcp",
+      auth: { type: "none" },
+      createdAt: 1,
+    };
+    let resolveLoad!: (records: MCPCustomServerRecord[]) => void;
+    const load = new Promise<MCPCustomServerRecord[]>((resolve) => {
+      resolveLoad = resolve;
+    });
+    const saveCustomServers = vi.fn(async () => {});
+    const root = mount(
+      [],
+      McpCustomStorage({
+        loadCustomServers: vi.fn(() => load),
+        saveCustomServers,
+        loadAuthState: vi.fn(async () => null),
+        saveAuthState: vi.fn(async () => {}),
+        clearAuthState: vi.fn(async () => {}),
+      }),
+    );
+
+    try {
+      await root.getValue().removeServer("docs");
+      resolveLoad([docsServer]);
+
+      await vi.waitFor(() =>
+        expect(root.getValue().getState().isHydrated).toBe(true),
+      );
+      expect(root.getValue().getState().customServers).toHaveLength(0);
+      await vi.waitFor(() =>
+        expect(saveCustomServers).toHaveBeenCalledWith([]),
+      );
+    } finally {
+      root.unmount();
+    }
+  });
+
   it("persists custom server updates in invocation order", async () => {
     let resolveFirstSave: (() => void) | undefined;
     const firstSave = new Promise<void>((resolve) => {

package/src/resources/McpManagerResource.ts

@@ -116,6 +116,20 @@ const holdCustomServerPersistence = (
   return release;
 };
 
+const deduplicateCustomServers = (records: MCPCustomServerRecord[]) => {
+  const seen = new Set<string>();
+  return records.filter((record) => {
+    if (seen.has(record.id)) {
+      console.error(
+        `[assistant-ui/react-mcp] ignored duplicate custom server id "${record.id}" loaded from storage`,
+      );
+      return false;
+    }
+    seen.add(record.id);
+    return true;
+  });
+};
+
 type McpCustomServersResourceProps = {
   storage: MCPStorage;
   scopeKey: string;
@@ -137,6 +151,7 @@ const useMcpCustomServersResource = ({
     "pending",
   );
   const hasPendingMutationRef = useRef(false);
+  const [removedBeforeHydration] = useState(() => new Set<string>());
   const reportedBlockedPersistenceRef = useRef(false);
 
   const hydrate = useEffectEvent(async (signal: { cancelled: boolean }) => {
@@ -151,7 +166,8 @@ const useMcpCustomServersResource = ({
 
     let records: Awaited<ReturnType<typeof storage.loadCustomServers>>;
     try {
-      records = await storage.loadCustomServers();
+      const loadedRecords = await storage.loadCustomServers();
+      records = deduplicateCustomServers(loadedRecords);
     } catch (error) {
       if (!signal.cancelled) {
         reportCustomStorageFailure("load", error);
@@ -164,16 +180,21 @@ const useMcpCustomServersResource = ({
       }
       return;
     }
-
     // Merge rather than replace so any addCustomServer calls that
     // happened before hydration resolved aren't silently overwritten.
     // Persisted order wins; pre-hydration locals append.
     const hadPendingMutation = hasPendingMutationRef.current;
+    const hydratedRecords = records.filter(
+      (record) => !removedBeforeHydration.has(record.id),
+    );
     const mergedRecords = (() => {
       const prev = customServersRef.current;
-      if (prev.length === 0) return records;
-      const persistedIds = new Set(records.map((r) => r.id));
-      return [...records, ...prev.filter((r) => !persistedIds.has(r.id))];
+      if (prev.length === 0) return hydratedRecords;
+      const persistedIds = new Set(hydratedRecords.map((r) => r.id));
+      return [
+        ...hydratedRecords,
+        ...prev.filter((r) => !persistedIds.has(r.id)),
+      ];
     })();
     customServersRef.current = mergedRecords;
     hydrationStateRef.current = "succeeded";
@@ -232,7 +253,22 @@ const useMcpCustomServersResource = ({
     [persistenceQueues, scopeKey, storage],
   );
 
-  return { customServers, isHydrated, updateCustomServers };
+  const removeCustomServer = useCallback(
+    (id: string) => {
+      if (hydrationStateRef.current === "pending") {
+        removedBeforeHydration.add(id);
+      }
+      updateCustomServers((prev) => prev.filter((record) => record.id !== id));
+    },
+    [removedBeforeHydration, updateCustomServers],
+  );
+
+  return {
+    customServers,
+    isHydrated,
+    updateCustomServers,
+    removeCustomServer,
+  };
 };
 
 const McpCustomServersResource = resource(useMcpCustomServersResource);
@@ -252,16 +288,17 @@ const useMcpManagerResource = (
   );
   const storageScopeKey =
     storage.scopeId === undefined ? "unscoped" : `scoped:${storage.scopeId}`;
-  const { customServers, isHydrated, updateCustomServers } = useResource(
-    withKey(
-      storageScopeKey,
-      McpCustomServersResource({
-        storage,
-        scopeKey: storageScopeKey,
-        persistenceQueues,
-      }),
-    ),
-  );
+  const { customServers, isHydrated, updateCustomServers, removeCustomServer } =
+    useResource(
+      withKey(
+        storageScopeKey,
+        McpCustomServersResource({
+          storage,
+          scopeKey: storageScopeKey,
+          persistenceQueues,
+        }),
+      ),
+    );
 
   const serverElements = useMemo(() => {
     assertUniqueServerIds([
@@ -313,9 +350,7 @@ const useMcpManagerResource = (
                 ? { elicitation: s.elicitation }
                 : {}),
               onRemove: async () => {
-                updateCustomServers((prev) =>
-                  prev.filter((record) => record.id !== s.id),
-                );
+                removeCustomServer(s.id);
               },
             },
             () =>
@@ -332,7 +367,7 @@ const useMcpManagerResource = (
     redirectUri,
     autoConnect,
     connectionTimeout,
-    updateCustomServers,
+    removeCustomServer,
     persistenceQueues,
     storageScopeKey,
   ]);
@@ -450,9 +485,7 @@ const useMcpManagerResource = (
         );
         try {
           await clearOAuthProviderAuthState(storage, id);
-          updateCustomServers((prev) =>
-            prev.filter((record) => record.id !== id),
-          );
+          removeCustomServer(id);
         } catch (error) {
           releasePersistence();
           throw error;

package/src/resources/McpServerResource.test.ts

@@ -1,5 +1,6 @@
 import { createTapRoot, resource, useResource } from "@assistant-ui/tap";
 import type { ClientOutput } from "@assistant-ui/store";
+import { UnauthorizedError } from "@modelcontextprotocol/client";
 import { useEffect, useState } from "react";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import type { MCPAuthConfig } from "../mcp-scope";
@@ -113,8 +114,8 @@ const requestElicitation = (
   client: any,
   message: string,
   requestedSchema: unknown,
-  context: { signal: AbortSignal } = {
-    signal: new AbortController().signal,
+  context: { mcpReq: { signal: AbortSignal } } = {
+    mcpReq: { signal: new AbortController().signal },
   },
 ) => {
   const handler = client.requestHandlers.get("elicitation/create");
@@ -156,7 +157,6 @@ const mount = (
     cache?: { readonly defaultTtlMs?: number } | undefined;
     elicitation?: boolean | undefined;
     kind?: "connector" | "custom" | undefined;
-    storage?: MCPStorage | undefined;
     onRemove?: (() => Promise<void>) | undefined;
   },
   onMount?: (server: ClientOutput<"mcpServer">) => void,
@@ -193,6 +193,15 @@ const mount = (
 const unboundAuthMessage =
   'MCP server "docs" has saved authentication for a different URL. Authenticate again to connect to https://example.com/mcp.';
 
+const getOAuthProvider = (index: number) => {
+  const provider =
+    mocks.StreamableHTTPClientTransport.mock.calls[index]?.[1]?.authProvider;
+  if (!provider) throw new Error("Expected OAuth provider");
+  return provider as {
+    redirectToAuthorization: (url: URL) => Promise<void>;
+  };
+};
+
 describe("McpServerResource automatic authentication", () => {
   beforeEach(resetMocks);
 
@@ -292,7 +301,9 @@ describe("McpServerResource automatic authentication", () => {
     });
 
     try {
-      await waitFor(() => storage.loadAuthState.mock.calls.length > 0);
+      await waitFor(
+        () => vi.mocked(storage.loadAuthState).mock.calls.length > 0,
+      );
       await flushMacrotask();
 
       expect(root.getValue().getState()).toMatchObject({
@@ -338,7 +349,9 @@ describe("McpServerResource automatic authentication", () => {
     });
 
     try {
-      await waitFor(() => storage.loadAuthState.mock.calls.length > 0);
+      await waitFor(
+        () => vi.mocked(storage.loadAuthState).mock.calls.length > 0,
+      );
       await flushMacrotask();
 
       expect(root.getValue().getState()).toMatchObject({
@@ -629,6 +642,70 @@ describe("McpServerResource connectionTimeout", () => {
 describe("McpServerResource connection lifecycle", () => {
   beforeEach(resetMocks);
 
+  it("publishes authorization URLs from the current connection", async () => {
+    const root = mount({ auth: { type: "oauth" } });
+
+    try {
+      await root.getValue().connect();
+      await getOAuthProvider(0).redirectToAuthorization(
+        new URL("https://auth.example.com/current"),
+      );
+      await waitForResourceUpdate(
+        () => root.getValue().getState().authorizationUrl !== null,
+      );
+
+      expect(root.getValue().getState().authorizationUrl).toBe(
+        "https://auth.example.com/current",
+      );
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("ignores authorization URLs after disconnect", async () => {
+    const root = mount({ auth: { type: "oauth" } });
+
+    try {
+      await root.getValue().connect();
+      const provider = getOAuthProvider(0);
+      await root.getValue().disconnect();
+
+      await provider.redirectToAuthorization(
+        new URL("https://auth.example.com/stale"),
+      );
+      await flushMacrotask();
+
+      expect(root.getValue().getState()).toMatchObject({
+        connectionState: "disconnected",
+        authorizationUrl: null,
+      });
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("ignores authorization URLs from a superseded connection", async () => {
+    const root = mount({ auth: { type: "oauth" } });
+
+    try {
+      await root.getValue().connect();
+      const staleProvider = getOAuthProvider(0);
+      await root.getValue().connect();
+
+      await staleProvider.redirectToAuthorization(
+        new URL("https://auth.example.com/stale"),
+      );
+      await flushMacrotask();
+
+      expect(root.getValue().getState()).toMatchObject({
+        connectionState: "connected",
+        authorizationUrl: null,
+      });
+    } finally {
+      root.unmount();
+    }
+  });
+
   it("replaces direct resource connections when the server id changes", async () => {
     const storage = createStorage();
     let updateId = (_id: string) => {};
@@ -744,6 +821,44 @@ describe("McpServerResource connection lifecycle", () => {
 describe("McpServerResource completeAuth", () => {
   beforeEach(resetMocks);
 
+  it("transfers authorization URL ownership when reusing the auth transport", async () => {
+    const storage = createStorage();
+    vi.mocked(storage.loadAuthState).mockResolvedValue({
+      serverUrl: "https://example.com/mcp",
+      state: "expected",
+    });
+    mocks.connectResults.push(() =>
+      Promise.reject(new UnauthorizedError("authorization required")),
+    );
+    const root = mount({ auth: { type: "oauth" }, storage });
+
+    try {
+      await root.getValue().connect();
+      await waitForResourceUpdate(
+        () => root.getValue().getState().connectionState === "authRequired",
+      );
+      const provider = getOAuthProvider(0);
+
+      await root
+        .getValue()
+        .completeAuth("https://example.com/callback?code=abc&state=expected");
+      await provider.redirectToAuthorization(
+        new URL("https://auth.example.com/reauthorize"),
+      );
+      await waitForResourceUpdate(
+        () => root.getValue().getState().authorizationUrl !== null,
+      );
+
+      expect(root.getValue().getState()).toMatchObject({
+        connectionState: "connected",
+        authorizationUrl: "https://auth.example.com/reauthorize",
+      });
+      expect(mocks.transports).toHaveLength(1);
+    } finally {
+      root.unmount();
+    }
+  });
+
   it("lets callback validation win over mount-time auto-connect", async () => {
     const pendingLoads: Array<(value: MCPPersistedAuthState | null) => void> =
       [];
@@ -1474,7 +1589,7 @@ describe("McpServerResource elicitation", () => {
           type: "object",
           properties: {},
         },
-        { signal: controller.signal },
+        { mcpReq: { signal: controller.signal } },
       );
       await waitForResourceUpdate(
         () => root.getValue().getState().pendingElicitations.length === 1,

package/src/resources/McpServerResource.ts

@@ -105,6 +105,9 @@ const useMcpServerResourceInstance = (
   const pendingTransportRef = useRef<StreamableHTTPClientTransport | null>(
     null,
   );
+  const transportGenerationRef = useRef(
+    new WeakMap<StreamableHTTPClientTransport, { current: number }>(),
+  );
   const connectionGenerationRef = useRef(0);
   const pendingAuthValidationRef = useRef<{
     count: number;
@@ -254,19 +257,29 @@ const useMcpServerResourceInstance = (
   });
 
   const buildTransport = useEffectEvent(
-    async (): Promise<StreamableHTTPClientTransport> => {
+    async (generation: number): Promise<StreamableHTTPClientTransport> => {
       if (props.auth.type === "oauth") {
+        const generationOwner = { current: generation };
         const authProvider = createOAuthProvider({
           serverId: props.id,
           serverUrl: props.url,
           config: props.auth,
           storage: props.storage,
           redirectUri: props.redirectUri,
-          onAuthorizationUrl: (url) => setAuthorizationUrl(url.toString()),
-        });
-        return new StreamableHTTPClientTransport(new URL(props.url), {
-          authProvider,
+          onAuthorizationUrl: (url) => {
+            if (isCurrentConnection(generationOwner.current)) {
+              setAuthorizationUrl(url.toString());
+            }
+          },
         });
+        const transport = new StreamableHTTPClientTransport(
+          new URL(props.url),
+          {
+            authProvider,
+          },
+        );
+        transportGenerationRef.current.set(transport, generationOwner);
+        return transport;
       }
       if (props.auth.type === "bearer") {
         const { state, unbound } = await loadAuthState();
@@ -380,7 +393,7 @@ const useMcpServerResourceInstance = (
               };
               elicitationResolversRef.current.set(id, {
                 resolve,
-                signal: context.signal,
+                signal: context.mcpReq.signal,
                 onAbort,
                 requestedSchema,
               });
@@ -395,10 +408,10 @@ const useMcpServerResourceInstance = (
             ]);
             const entry = elicitationResolversRef.current.get(id);
             if (entry) {
-              if (context.signal.aborted) {
+              if (context.mcpReq.signal.aborted) {
                 entry.onAbort();
               } else {
-                context.signal.addEventListener("abort", entry.onAbort, {
+                context.mcpReq.signal.addEventListener("abort", entry.onAbort, {
                   once: true,
                 });
               }
@@ -443,7 +456,7 @@ const useMcpServerResourceInstance = (
     setTools([]);
     let transport: StreamableHTTPClientTransport | null = null;
     try {
-      transport = await buildTransport();
+      transport = await buildTransport(generation);
       if (!isCurrentConnection(generation)) {
         await closeQueuedTransports([transport]);
         return;
@@ -536,12 +549,14 @@ const useMcpServerResourceInstance = (
     try {
       let transport = transportRef.current;
       if (!transport) {
-        transport = await buildTransport();
+        transport = await buildTransport(generation);
         if (!isCurrentConnection(generation)) {
           await closeQueuedTransports([transport]);
           throw createInterruptedAuthError();
         }
       }
+      const generationOwner = transportGenerationRef.current.get(transport);
+      if (generationOwner) generationOwner.current = generation;
       transportRef.current = null;
       clientRef.current = null;
       pendingTransportRef.current = transport;
@@ -712,7 +727,7 @@ const useMcpServerResourceInstance = (
     ): readonly { property: string; message: string }[] | undefined => {
       if (response.action === "accept") {
         const entry = elicitationResolversRef.current.get(id);
-        if (!entry) return;
+        if (!entry) return undefined;
 
         if (
           typeof response.content !== "object" ||
@@ -752,10 +767,11 @@ const useMcpServerResourceInstance = (
           content: response.content as ElicitResult["content"],
         };
         resolvePendingElicitation(id, result);
-        return;
+        return undefined;
       }
 
       resolvePendingElicitation(id, { action: response.action });
+      return undefined;
     },
   };
 };