@assistant-ui/react-mcp 0.1.17 → 0.1.19

package/src/resources/McpManagerResource.ts

@@ -1,4 +1,11 @@
-import { useState, useEffect, useMemo, useEffectEvent, useRef } from "react";
+import {
+  useState,
+  useEffect,
+  useMemo,
+  useEffectEvent,
+  useRef,
+  useCallback,
+} from "react";
 import { useResource, resource, withKey } from "@assistant-ui/tap";
 import {
   useClientLookup,
@@ -12,6 +19,7 @@ import { createMcpId } from "../utils/createMcpId";
 import { clearOAuthProviderAuthState } from "../auth/createOAuthProvider";
 import type { Tool } from "assistant-stream";
 import { McpServerResource } from "./McpServerResource";
+import { withMcpServerRemovalFence } from "./McpServerRemovalFence";
 import { McpLocalStorage } from "./storage/McpLocalStorage";
 import type { MCPStorage, MCPStorageElement } from "./storage/types";
 import { assertUniqueServerIds } from "../utils/serverId";
@@ -52,6 +60,12 @@ const reportCustomStorageFailure = (
   );
 };
 
+const reportBlockedCustomServerPersistence = () => {
+  console.error(
+    "[assistant-ui/react-mcp] custom server changes remain in memory because loading the persisted list failed; remount the manager to retry",
+  );
+};
+
 const persistCustomServers = async (
   storage: MCPStorage,
   records: MCPCustomServerRecord[],
@@ -63,59 +77,139 @@ const persistCustomServers = async (
   }
 };
 
-const useMcpManagerResource = (
-  props: McpManagerResourceProps,
-): ClientOutput<"mcp"> => {
-  const connectors = props.connectors ?? NO_CONNECTORS;
-  const autoConnect = props.autoConnect ?? true;
-  const redirectUri = props.oauthRedirectUri ?? defaultRedirectUri();
-  const connectionTimeout = props.connectionTimeout;
+type CustomServerPersistenceQueues = Map<string, Promise<void>>;
 
-  const storageElement = props.storage ?? McpLocalStorage();
-  const storage = useResource(storageElement);
+const enqueueCustomServerTask = (
+  persistenceQueues: CustomServerPersistenceQueues,
+  scopeKey: string,
+  task: () => Promise<void>,
+) => {
+  const previous = persistenceQueues.get(scopeKey);
+  const next = (previous ?? Promise.resolve()).then(task);
+  persistenceQueues.set(scopeKey, next);
+  void next.then(() => {
+    if (persistenceQueues.get(scopeKey) === next) {
+      persistenceQueues.delete(scopeKey);
+    }
+  });
+};
+
+const enqueueCustomServerPersistence = (
+  persistenceQueues: CustomServerPersistenceQueues,
+  scopeKey: string,
+  storage: MCPStorage,
+  records: MCPCustomServerRecord[],
+) =>
+  enqueueCustomServerTask(persistenceQueues, scopeKey, () =>
+    persistCustomServers(storage, records),
+  );
+
+const holdCustomServerPersistence = (
+  persistenceQueues: CustomServerPersistenceQueues,
+  scopeKey: string,
+) => {
+  let release = () => {};
+  const gate = new Promise<void>((resolve) => {
+    release = resolve;
+  });
+  enqueueCustomServerTask(persistenceQueues, scopeKey, () => gate);
+  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;
+  persistenceQueues: CustomServerPersistenceQueues;
+};
 
+const useMcpCustomServersResource = ({
+  storage,
+  scopeKey,
+  persistenceQueues,
+}: McpCustomServersResourceProps) => {
   const [customServers, setCustomServers] = useState<MCPCustomServerRecord[]>(
     [],
   );
   const [isHydrated, setIsHydrated] = useState(false);
 
-  const hydratedRef = useRef(false);
-  const storageRef = useRef(storage);
-  const persistenceQueueRef = useRef(Promise.resolve());
-
-  useEffect(() => {
-    storageRef.current = storage;
-  }, [storage]);
+  const customServersRef = useRef<MCPCustomServerRecord[]>([]);
+  const hydrationStateRef = useRef<"pending" | "succeeded" | "failed">(
+    "pending",
+  );
+  const hasPendingMutationRef = useRef(false);
+  const [removedBeforeHydration] = useState(() => new Set<string>());
+  const reportedBlockedPersistenceRef = useRef(false);
 
   const hydrate = useEffectEvent(async (signal: { cancelled: boolean }) => {
-    const markHydrated = () => {
-      if (!signal.cancelled) {
-        hydratedRef.current = true;
-        setIsHydrated(true);
-      }
-    };
+    // A revisited scope must not read behind writes still queued against it.
+    while (true) {
+      const pendingPersistence = persistenceQueues.get(scopeKey);
+      if (!pendingPersistence) break;
+      await pendingPersistence;
+      if (signal.cancelled) return;
+      if (persistenceQueues.get(scopeKey) === pendingPersistence) break;
+    }
 
     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);
+        hydrationStateRef.current = "failed";
+        if (hasPendingMutationRef.current) {
+          reportBlockedCustomServerPersistence();
+          reportedBlockedPersistenceRef.current = true;
+        }
+        setIsHydrated(true);
       }
-      markHydrated();
       return;
     }
-
-    if (signal.cancelled) 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.
-    setCustomServers((prev) => {
-      if (prev.length === 0) return records;
-      const persistedIds = new Set(records.map((r) => r.id));
-      return [...records, ...prev.filter((r) => !persistedIds.has(r.id))];
-    });
-    markHydrated();
+    const hadPendingMutation = hasPendingMutationRef.current;
+    const hydratedRecords = records.filter(
+      (record) => !removedBeforeHydration.has(record.id),
+    );
+    const mergedRecords = (() => {
+      const prev = customServersRef.current;
+      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";
+    hasPendingMutationRef.current = false;
+    if (hadPendingMutation) {
+      enqueueCustomServerPersistence(
+        persistenceQueues,
+        scopeKey,
+        storage,
+        mergedRecords,
+      );
+    }
+    if (signal.cancelled) return;
+    setCustomServers(mergedRecords);
+    setIsHydrated(true);
   });
 
   useEffect(() => {
@@ -129,13 +223,82 @@ const useMcpManagerResource = (
     };
   }, []);
 
-  useEffect(() => {
-    if (!hydratedRef.current) return;
-    const targetStorage = storageRef.current;
-    persistenceQueueRef.current = persistenceQueueRef.current.then(() =>
-      persistCustomServers(targetStorage, customServers),
+  const updateCustomServers = useCallback(
+    (
+      updater: (records: MCPCustomServerRecord[]) => MCPCustomServerRecord[],
+    ) => {
+      const next = updater(customServersRef.current);
+      customServersRef.current = next;
+      setCustomServers(next);
+
+      if (hydrationStateRef.current === "succeeded") {
+        enqueueCustomServerPersistence(
+          persistenceQueues,
+          scopeKey,
+          storage,
+          next,
+        );
+        return;
+      }
+
+      hasPendingMutationRef.current = true;
+      if (
+        hydrationStateRef.current === "failed" &&
+        !reportedBlockedPersistenceRef.current
+      ) {
+        reportBlockedCustomServerPersistence();
+        reportedBlockedPersistenceRef.current = true;
+      }
+    },
+    [persistenceQueues, scopeKey, storage],
+  );
+
+  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);
+
+const useMcpManagerResource = (
+  props: McpManagerResourceProps,
+): ClientOutput<"mcp"> => {
+  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);
+  const [persistenceQueues] = useState<CustomServerPersistenceQueues>(
+    () => new Map(),
+  );
+  const storageScopeKey =
+    storage.scopeId === undefined ? "unscoped" : `scoped:${storage.scopeId}`;
+  const { customServers, isHydrated, updateCustomServers, removeCustomServer } =
+    useResource(
+      withKey(
+        storageScopeKey,
+        McpCustomServersResource({
+          storage,
+          scopeKey: storageScopeKey,
+          persistenceQueues,
+        }),
+      ),
     );
-  }, [customServers]);
 
   const serverElements = useMemo(() => {
     assertUniqueServerIds([
@@ -170,24 +333,30 @@ const useMcpManagerResource = (
     const customElements = customServers.map((s) =>
       withKey(
         s.id,
-        McpServerResource({
-          id: s.id,
-          kind: "custom",
-          name: s.name,
-          url: s.url,
-          auth: s.auth,
-          storage,
-          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));
-          },
-        }),
+        McpServerResource(
+          withMcpServerRemovalFence(
+            {
+              id: s.id,
+              kind: "custom",
+              name: s.name,
+              url: s.url,
+              auth: s.auth,
+              storage,
+              redirectUri,
+              autoConnect,
+              connectionTimeout: s.connectionTimeout ?? connectionTimeout,
+              ...(s.cache !== undefined ? { cache: s.cache } : {}),
+              ...(s.elicitation !== undefined
+                ? { elicitation: s.elicitation }
+                : {}),
+              onRemove: async () => {
+                removeCustomServer(s.id);
+              },
+            },
+            () =>
+              holdCustomServerPersistence(persistenceQueues, storageScopeKey),
+          ),
+        ),
       ),
     );
     return [...connectorElements, ...customElements];
@@ -198,6 +367,9 @@ const useMcpManagerResource = (
     redirectUri,
     autoConnect,
     connectionTimeout,
+    removeCustomServer,
+    persistenceQueues,
+    storageScopeKey,
   ]);
 
   const lookup = useClientLookup(serverElements);
@@ -289,7 +461,7 @@ const useMcpManagerResource = (
         ...(elicitation !== undefined ? { elicitation } : {}),
         createdAt: Date.now(),
       };
-      setCustomServers((prev) => [...prev, record]);
+      updateCustomServers((prev) => [...prev, record]);
       return record.id;
     },
     removeServer: async (id) => {
@@ -307,8 +479,18 @@ const useMcpManagerResource = (
       try {
         await lookup.get({ key: id }).remove();
       } catch {
-        await clearOAuthProviderAuthState(storage, id);
-        setCustomServers((prev) => prev.filter((s) => s.id !== id));
+        const releasePersistence = holdCustomServerPersistence(
+          persistenceQueues,
+          storageScopeKey,
+        );
+        try {
+          await clearOAuthProviderAuthState(storage, id);
+          removeCustomServer(id);
+        } catch (error) {
+          releasePersistence();
+          throw error;
+        }
+        releasePersistence();
       }
     },
   };

package/src/resources/McpServerRemovalFence.ts

@@ -0,0 +1,21 @@
+import type { McpServerResourceProps } from "./McpServerResource";
+
+const beginRemovalFenceSymbol = Symbol();
+
+type McpServerResourceInternalProps = McpServerResourceProps & {
+  [beginRemovalFenceSymbol]?: () => () => void;
+};
+
+export const withMcpServerRemovalFence = (
+  props: McpServerResourceProps,
+  beginRemovalFence: () => () => void,
+): McpServerResourceProps =>
+  ({
+    ...props,
+    [beginRemovalFenceSymbol]: beginRemovalFence,
+  }) as McpServerResourceProps;
+
+export const beginMcpServerRemovalFence = (
+  props: McpServerResourceProps,
+): (() => void) | undefined =>
+  (props as McpServerResourceInternalProps)[beginRemovalFenceSymbol]?.();

package/src/resources/McpServerResource.test.ts

@@ -113,8 +113,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 +156,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,
@@ -292,7 +291,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 +339,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({
@@ -1474,7 +1477,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

@@ -34,6 +34,7 @@ import type {
   MCPToolInfo,
 } from "../mcp-scope";
 import { createMcpId } from "../utils/createMcpId";
+import { beginMcpServerRemovalFence } from "./McpServerRemovalFence";
 
 export type McpServerResourceProps = {
   id: string;
@@ -379,7 +380,7 @@ const useMcpServerResourceInstance = (
               };
               elicitationResolversRef.current.set(id, {
                 resolve,
-                signal: context.signal,
+                signal: context.mcpReq.signal,
                 onAbort,
                 requestedSchema,
               });
@@ -394,10 +395,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,
                 });
               }
@@ -666,16 +667,19 @@ const useMcpServerResourceInstance = (
     connect: doConnect,
     disconnect: doDisconnect,
     remove: async () => {
-      await doDisconnect();
+      const releaseRemovalFence = beginMcpServerRemovalFence(props);
       try {
+        await doDisconnect();
         await clearOAuthProviderAuthState(props.storage, props.id);
         await props.onRemove();
       } catch (err) {
+        releaseRemovalFence?.();
         setLastError({
           message: err instanceof Error ? err.message : String(err),
         });
         throw err;
       }
+      releaseRemovalFence?.();
     },
     callTool: async (name, args) => {
       const client = clientRef.current;
@@ -708,7 +712,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" ||
@@ -748,10 +752,11 @@ const useMcpServerResourceInstance = (
           content: response.content as ElicitResult["content"],
         };
         resolvePendingElicitation(id, result);
-        return;
+        return undefined;
       }
 
       resolvePendingElicitation(id, { action: response.action });
+      return undefined;
     },
   };
 };

package/src/resources/storage/McpLocalStorage.test.ts

@@ -239,6 +239,8 @@ describe("normalizePersistedAuthState", () => {
 
   it.each([
     "http://auth.example.com",
+    "http://127.example.com",
+    "http://127.0.0.1.example.com",
     "data:text/plain,auth",
     "file:///tmp/auth",
   ])("drops discovery state with an unsafe URL: %s", (url) => {

package/src/resources/storage/McpLocalStorage.ts

@@ -148,12 +148,13 @@ const isSecureNetworkUrl = (value: unknown): value is string => {
   if (!isNonEmptyString(value)) return false;
   try {
     const url = new URL(value);
+    const isIpv4Loopback = /^127(?:\.\d{1,3}){3}$/.test(url.hostname);
     return (
       url.protocol === "https:" ||
       (url.protocol === "http:" &&
         (url.hostname === "localhost" ||
           url.hostname.endsWith(".localhost") ||
-          url.hostname.startsWith("127.") ||
+          isIpv4Loopback ||
           url.hostname === "[::1]"))
     );
   } catch {

package/src/resources/storage/types.ts

@@ -5,14 +5,16 @@ import type { MCPPersistedAuthState } from "../../auth/types";
 export type MCPStorage = {
   /**
    * Stable identity of the backing store. Two storages with the same scopeId
-   * must read and write the same persisted data. When present, server
-   * connections and the OAuth write fence key on it, so swapping to a
-   * differently-scoped storage reconnects instead of leaving a live OAuth flow
-   * on the replaced store, and clearing through a same-scoped replacement still
-   * waits for writes queued against the storage it replaced. When absent, the
-   * fence falls back to object identity while connections never re-key at all,
-   * so a storage rebuilt on every render has to declare a scopeId; without one
-   * a clear runs unfenced against the writes queued by the object it replaced.
+   * must read and write the same persisted data. When present, custom servers,
+   * server connections, and the OAuth write fence key on it. Swapping to a
+   * differently-scoped storage rehydrates its custom servers and reconnects
+   * authenticated servers instead of retaining data or a live OAuth flow from
+   * the replaced store. Clearing through a same-scoped replacement still waits
+   * for writes queued against the storage it replaced. When absent, the fence
+   * falls back to object identity while custom servers and connections never
+   * re-key. Replacing an unscoped storage can therefore persist the previous
+   * store's custom servers into the replacement, and a clear runs unfenced
+   * against writes queued by the object it replaced.
    */
   scopeId?: string;
   loadCustomServers: () => Promise<MCPCustomServerRecord[]>;