@assistant-ui/react-mcp 0.1.2 → 0.1.4

package/src/resources/McpManagerResource.ts

@@ -10,7 +10,7 @@ import { ModelContext } from "@assistant-ui/core/store";
 import type { Tool } from "assistant-stream";
 import { McpServerResource } from "./McpServerResource";
 import { McpLocalStorage } from "./storage/McpLocalStorage";
-import type { MCPStorageElement } from "./storage/types";
+import type { MCPStorage, MCPStorageElement } from "./storage/types";
 import { assertUniqueServerIds } from "../utils/serverId";
 import type {
   MCPAuthConfig,
@@ -49,6 +49,17 @@ const reportCustomStorageFailure = (
   );
 };
 
+const persistCustomServers = async (
+  storage: MCPStorage,
+  records: MCPCustomServerRecord[],
+) => {
+  try {
+    await storage.saveCustomServers(records);
+  } catch (error) {
+    reportCustomStorageFailure("save", error);
+  }
+};
+
 const useMcpManagerResource = (
   props: McpManagerResourceProps,
 ): ClientOutput<"mcp"> => {
@@ -66,6 +77,12 @@ const useMcpManagerResource = (
   const [isHydrated, setIsHydrated] = useState(false);
 
   const hydratedRef = useRef(false);
+  const storageRef = useRef(storage);
+  const persistenceQueueRef = useRef(Promise.resolve());
+
+  useEffect(() => {
+    storageRef.current = storage;
+  }, [storage]);
 
   const hydrate = useEffectEvent(async (signal: { cancelled: boolean }) => {
     const markHydrated = () => {
@@ -106,19 +123,12 @@ const useMcpManagerResource = (
     };
   }, []);
 
-  const persistCustomServers = useEffectEvent(
-    async (records: MCPCustomServerRecord[]) => {
-      if (!hydratedRef.current) return;
-      try {
-        await storage.saveCustomServers(records);
-      } catch (error) {
-        reportCustomStorageFailure("save", error);
-      }
-    },
-  );
-
   useEffect(() => {
-    void persistCustomServers(customServers);
+    if (!hydratedRef.current) return;
+    const targetStorage = storageRef.current;
+    persistenceQueueRef.current = persistenceQueueRef.current.then(() =>
+      persistCustomServers(targetStorage, customServers),
+    );
   }, [customServers]);
 
   const serverElements = useMemo(() => {

package/src/resources/McpServerResource.test.ts

@@ -150,6 +150,9 @@ const mount = (
     connectionTimeout?: number | undefined;
     cache?: { readonly defaultTtlMs?: number } | undefined;
     elicitation?: boolean | undefined;
+    kind?: "connector" | "custom" | undefined;
+    storage?: MCPStorage | undefined;
+    onRemove?: (() => Promise<void>) | undefined;
   },
   onMount?: (server: ClientOutput<"mcpServer">) => void,
 ) => {
@@ -160,11 +163,11 @@ const mount = (
     const server = useResource(
       McpServerResource({
         id: "docs",
-        kind: "connector",
+        kind: props?.kind ?? "connector",
         name: "Docs",
         url: "https://example.com/mcp",
         auth: props?.auth ?? { type: "none" },
-        storage: createStorage(),
+        storage: props?.storage ?? createStorage(),
         redirectUri: "https://example.com/callback",
         autoConnect: false,
         connectionTimeout,
@@ -172,7 +175,7 @@ const mount = (
         ...(props?.elicitation !== undefined
           ? { elicitation: props.elicitation }
           : {}),
-        onRemove: vi.fn(async () => {}),
+        onRemove: props?.onRemove ?? vi.fn(async () => {}),
       }),
     );
     useEffect(() => {
@@ -1200,6 +1203,37 @@ describe("McpServerResource tools listChanged", () => {
   });
 });
 
+describe("McpServerResource removal", () => {
+  beforeEach(() => {
+    resetMocks();
+  });
+
+  it("exposes auth storage failures and preserves the server", async () => {
+    const error = new Error("storage unavailable");
+    const storage = createStorage();
+    vi.mocked(storage.clearAuthState).mockRejectedValueOnce(error);
+    const onRemove = vi.fn(async () => {});
+    const root = mount({ kind: "custom", storage, onRemove });
+
+    try {
+      await expect(root.getValue().remove()).rejects.toBe(error);
+      await waitForResourceUpdate(
+        () =>
+          root.getValue().getState().lastError?.message ===
+          "storage unavailable",
+      );
+
+      expect(root.getValue().getState()).toMatchObject({
+        connectionState: "disconnected",
+        lastError: { message: "storage unavailable" },
+      });
+      expect(onRemove).not.toHaveBeenCalled();
+    } finally {
+      root.unmount();
+    }
+  });
+});
+
 describe("McpServerResource resource methods", () => {
   beforeEach(() => {
     mocks.clients.length = 0;

package/src/resources/McpServerResource.ts

@@ -539,8 +539,15 @@ const useMcpServerResource = (
     disconnect: doDisconnect,
     remove: async () => {
       await doDisconnect();
-      await props.storage.clearAuthState(props.id);
-      await props.onRemove();
+      try {
+        await props.storage.clearAuthState(props.id);
+        await props.onRemove();
+      } catch (err) {
+        setLastError({
+          message: err instanceof Error ? err.message : String(err),
+        });
+        throw err;
+      }
     },
     callTool: async (name, args) => {
       const client = clientRef.current;