@assistant-ui/react-mcp 0.1.20 → 0.1.22

package/src/resources/McpManagerResource.ts

@@ -1,11 +1,4 @@
-import {
-  useState,
-  useEffect,
-  useMemo,
-  useEffectEvent,
-  useRef,
-  useCallback,
-} from "react";
+import { useState, useMemo } from "react";
 import { useResource, resource, withKey } from "@assistant-ui/tap";
 import {
   useClientLookup,
@@ -19,9 +12,14 @@ import { createMcpId } from "../utils/createMcpId";
 import { clearOAuthProviderAuthState } from "../auth/createOAuthProvider";
 import type { Tool } from "assistant-stream";
 import { McpServerResource } from "./McpServerResource";
+import {
+  McpCustomServersResource,
+  holdCustomServerPersistence,
+  type CustomServerPersistenceQueues,
+} from "./McpCustomServerPersistence";
 import { withMcpServerRemovalFence } from "./McpServerRemovalFence";
 import { McpLocalStorage } from "./storage/McpLocalStorage";
-import type { MCPStorage, MCPStorageElement } from "./storage/types";
+import type { MCPStorageElement } from "./storage/types";
 import { assertUniqueServerIds } from "../utils/serverId";
 import type {
   MCPAuthConfig,
@@ -50,229 +48,6 @@ 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 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[],
-) => {
-  try {
-    await storage.saveCustomServers(records);
-  } catch (error) {
-    reportCustomStorageFailure("save", error);
-  }
-};
-
-type CustomServerPersistenceQueues = Map<string, Promise<void>>;
-
-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 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 }) => {
-    // 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 {
-      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);
-      }
-      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 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(() => {
-    const signal = { cancelled: false };
-    // Hydration reads persisted records asynchronously; there is no earlier
-    // point than mount at which to start it.
-    // eslint-disable-next-line react-hooks/set-state-in-effect
-    void hydrate(signal);
-    return () => {
-      signal.cancelled = true;
-    };
-  }, []);
-
-  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"> => {
@@ -472,26 +247,28 @@ const useMcpManagerResource = (
           `Cannot remove connector "${id}" — connectors are app-defined and not removable. Use a custom server id instead.`,
         );
       }
-      // Delegate to McpServerResource.remove() which disconnects,
-      // clears auth state, and unregisters from customServers in one
-      // place. Fallback to manual cleanup if the lookup is empty
-      // (server already gone).
+      let server;
       try {
-        await lookup.get({ key: id }).remove();
+        server = lookup.get({ key: id });
       } catch {
-        const releasePersistence = holdCustomServerPersistence(
-          persistenceQueues,
-          storageScopeKey,
-        );
-        try {
-          await clearOAuthProviderAuthState(storage, id);
-          removeCustomServer(id);
-        } catch (error) {
-          releasePersistence();
-          throw error;
-        }
+        server = undefined;
+      }
+      if (server) {
+        await server.remove();
+        return;
+      }
+      const releasePersistence = holdCustomServerPersistence(
+        persistenceQueues,
+        storageScopeKey,
+      );
+      try {
+        await clearOAuthProviderAuthState(storage, id);
+        removeCustomServer(id);
+      } catch (error) {
         releasePersistence();
+        throw error;
       }
+      releasePersistence();
     },
   };
 };

package/src/resources/McpServerResource.test.ts

@@ -157,6 +157,7 @@ const mount = (
     cache?: { readonly defaultTtlMs?: number } | undefined;
     elicitation?: boolean | undefined;
     kind?: "connector" | "custom" | undefined;
+    url?: string | undefined;
     onRemove?: (() => Promise<void>) | undefined;
   },
   onMount?: (server: ClientOutput<"mcpServer">) => void,
@@ -170,7 +171,7 @@ const mount = (
         id: "docs",
         kind: props?.kind ?? "connector",
         name: "Docs",
-        url: "https://example.com/mcp",
+        url: props?.url ?? "https://example.com/mcp",
         auth: props?.auth ?? { type: "none" },
         storage: props?.storage ?? createStorage(),
         redirectUri: "https://example.com/callback",
@@ -205,6 +206,75 @@ const getOAuthProvider = (index: number) => {
 describe("McpServerResource automatic authentication", () => {
   beforeEach(resetMocks);
 
+  it.each(["bearer", "oauth"] as const)(
+    "rejects %s authentication over remote HTTP",
+    async (type) => {
+      const root = mount({
+        auth: type === "bearer" ? { type, token: "secret" } : { type },
+        url: "http://mcp.example.com/mcp",
+      });
+
+      try {
+        await root.getValue().connect();
+        await waitForResourceUpdate(
+          () => root.getValue().getState().connectionState === "error",
+        );
+
+        expect(root.getValue().getState()).toMatchObject({
+          connectionState: "error",
+          lastError: {
+            message:
+              'Authenticated MCP server "docs" must use HTTPS or loopback HTTP.',
+          },
+        });
+        expect(mocks.StreamableHTTPClientTransport).not.toHaveBeenCalled();
+      } finally {
+        root.unmount();
+      }
+    },
+  );
+
+  it.each([
+    "http://localhost:3000/mcp",
+    "http://dev.localhost:3000/mcp",
+    "http://127.0.0.1:3000/mcp",
+    "http://127.0.0.2:3000/mcp",
+    "http://[::1]:3000/mcp",
+    "http://localhost.:3000/mcp",
+    "http://dev.localhost.:3000/mcp",
+    "http://[::ffff:127.0.0.1]:3000/mcp",
+  ])("allows bearer authentication over loopback HTTP at %s", async (url) => {
+    const root = mount({
+      auth: { type: "bearer", token: "secret" },
+      url,
+      autoConnect: true,
+    });
+
+    try {
+      await waitFor(() => mocks.transports.length > 0);
+      expect(mocks.StreamableHTTPClientTransport).toHaveBeenCalledWith(
+        new URL(url),
+        { requestInit: { headers: { Authorization: "Bearer secret" } } },
+      );
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("allows unauthenticated MCP servers over remote HTTP", async () => {
+    const url = "http://mcp.example.com/mcp";
+    const root = mount({ auth: { type: "none" }, url, autoConnect: true });
+
+    try {
+      await waitFor(() => mocks.transports.length > 0);
+      expect(mocks.StreamableHTTPClientTransport).toHaveBeenCalledWith(
+        new URL(url),
+      );
+    } finally {
+      root.unmount();
+    }
+  });
+
   it("does not auto-connect with authentication from another server URL", async () => {
     const storage = createStorage();
     vi.mocked(storage.loadAuthState).mockResolvedValue({
@@ -642,26 +712,69 @@ describe("McpServerResource connectionTimeout", () => {
 describe("McpServerResource connection lifecycle", () => {
   beforeEach(resetMocks);
 
-  it("publishes authorization URLs from the current connection", async () => {
+  it.each([
+    "https://auth.example.com/current",
+    "http://localhost:3000/current",
+  ])(
+    "publishes authorization URLs from the current connection at %s",
+    async (url) => {
+      const root = mount({ auth: { type: "oauth" } });
+
+      try {
+        await root.getValue().connect();
+        await getOAuthProvider(0).redirectToAuthorization(new URL(url));
+        await waitForResourceUpdate(
+          () => root.getValue().getState().authorizationUrl !== null,
+        );
+
+        expect(root.getValue().getState().authorizationUrl).toBe(url);
+      } finally {
+        root.unmount();
+      }
+    },
+  );
+
+  it("rejects remote HTTP OAuth 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",
+      await expect(
+        getOAuthProvider(0).redirectToAuthorization(
+          new URL("http://auth.example.com/authorize"),
+        ),
+      ).rejects.toThrow(
+        "MCP OAuth authorization URL must use HTTPS or loopback HTTP.",
       );
+      expect(root.getValue().getState().authorizationUrl).toBeNull();
     } finally {
       root.unmount();
     }
   });
 
+  it.each(["javascript:", "data:", "file:"])(
+    "rejects %s OAuth authorization URLs from the current connection",
+    async (protocol) => {
+      const root = mount({ auth: { type: "oauth" } });
+
+      try {
+        await root.getValue().connect();
+
+        await expect(
+          getOAuthProvider(0).redirectToAuthorization(
+            new URL(`${protocol}unsafe`),
+          ),
+        ).rejects.toThrow(
+          `Unsupported MCP OAuth authorization URL protocol: ${protocol}`,
+        );
+        expect(root.getValue().getState().authorizationUrl).toBeNull();
+      } finally {
+        root.unmount();
+      }
+    },
+  );
+
   it("ignores authorization URLs after disconnect", async () => {
     const root = mount({ auth: { type: "oauth" } });