@assistant-ui/react-mcp 0.1.14 → 0.1.16

package/src/resources/McpServerResource.ts

@@ -1,5 +1,5 @@
 import { useState, useRef, useEffect, useMemo, useEffectEvent } from "react";
-import { resource } from "@assistant-ui/tap";
+import { resource, useResource, withKey } from "@assistant-ui/tap";
 import type { ClientOutput } from "@assistant-ui/store";
 import {
   Client,
@@ -10,7 +10,10 @@ import {
   type ElicitResult,
   type StreamableHTTPClientTransportOptions,
 } from "@modelcontextprotocol/client";
-import { createOAuthProvider } from "../auth/createOAuthProvider";
+import {
+  clearOAuthProviderAuthState,
+  createOAuthProvider,
+} from "../auth/createOAuthProvider";
 import { buildHeaders } from "../auth/buildHeaders";
 import { assertValidServerId } from "../utils/serverId";
 import { validateElicitationContent } from "./validateElicitationContent";
@@ -24,6 +27,7 @@ import type {
   MCPServerState,
   MCPToolInfo,
 } from "../mcp-scope";
+import { createMcpId } from "../utils/createMcpId";
 
 export type McpServerResourceProps = {
   id: string;
@@ -41,8 +45,50 @@ export type McpServerResourceProps = {
   onRemove: () => Promise<void>;
 };
 
-const useMcpServerResource = (
+type McpServerResourceInstanceProps = McpServerResourceProps & {
+  transportCloseQueueRef: { current: Promise<void> };
+};
+
+export const getConnectionDependencies = (
   props: McpServerResourceProps,
+): readonly unknown[] => {
+  const auth = props.auth;
+  const authDependencies =
+    auth.type === "bearer"
+      ? [auth.type, auth.token, props.storage.scopeId]
+      : auth.type === "oauth"
+        ? [
+            auth.type,
+            auth.scopes?.length,
+            ...(auth.scopes ?? []),
+            auth.authorizationEndpoint,
+            auth.tokenEndpoint,
+            auth.registrationEndpoint,
+            auth.clientId,
+            auth.clientSecret,
+            props.storage.scopeId,
+          ]
+        : [auth.type];
+
+  return [
+    props.id,
+    props.url,
+    ...authDependencies,
+    props.redirectUri,
+    props.cache?.defaultTtlMs,
+    props.elicitation !== false,
+  ];
+};
+
+const areConnectionDependenciesEqual = (
+  left: readonly unknown[],
+  right: readonly unknown[],
+) =>
+  left.length === right.length &&
+  left.every((value, index) => Object.is(value, right[index]));
+
+const useMcpServerResourceInstance = (
+  props: McpServerResourceInstanceProps,
 ): ClientOutput<"mcpServer"> => {
   assertValidServerId(props.id);
   const [connectionState, setConnectionState] =
@@ -59,8 +105,12 @@ const useMcpServerResource = (
   const pendingTransportRef = useRef<StreamableHTTPClientTransport | null>(
     null,
   );
-  const transportCloseQueueRef = useRef(Promise.resolve());
   const connectionGenerationRef = useRef(0);
+  const pendingAuthValidationRef = useRef<{
+    count: number;
+    promise: Promise<void>;
+    resolve: () => void;
+  } | null>(null);
   const elicitationResolversRef = useRef(
     new Map<
       string,
@@ -88,10 +138,10 @@ const useMcpServerResource = (
   const closeQueuedTransports = (
     transports: StreamableHTTPClientTransport[],
   ): Promise<void> => {
-    const task = transportCloseQueueRef.current.then(async () => {
+    const task = props.transportCloseQueueRef.current.then(async () => {
       await Promise.all(transports.map(closeTransportSafely));
     });
-    transportCloseQueueRef.current = task;
+    props.transportCloseQueueRef.current = task;
     return task;
   };
 
@@ -132,7 +182,7 @@ const useMcpServerResource = (
     }
   };
 
-  const closeTransports = async () => {
+  const detachTransports = () => {
     cancelPendingElicitations();
     const pendingTransport = pendingTransportRef.current;
     const activeTransport = transportRef.current;
@@ -140,13 +190,18 @@ const useMcpServerResource = (
     transportRef.current = null;
     clientRef.current = null;
 
-    const transports = new Set(
-      [pendingTransport, activeTransport].filter(
-        (transport): transport is StreamableHTTPClientTransport =>
-          transport !== null,
+    return [
+      ...new Set(
+        [pendingTransport, activeTransport].filter(
+          (transport): transport is StreamableHTTPClientTransport =>
+            transport !== null,
+        ),
       ),
-    );
-    await closeQueuedTransports([...transports]);
+    ];
+  };
+
+  const closeTransports = async () => {
+    await closeQueuedTransports(detachTransports());
   };
 
   const isCurrentConnection = (generation: number) =>
@@ -305,10 +360,7 @@ const useMcpServerResource = (
             }
             const { message, requestedSchema } = request.params;
 
-            const id =
-              typeof crypto !== "undefined" && "randomUUID" in crypto
-                ? crypto.randomUUID()
-                : `mcp-${Date.now()}-${Math.random().toString(36).slice(2)}`;
+            const id = createMcpId();
             const promise = new Promise<ElicitResult>((resolve) => {
               const onAbort = () => {
                 resolvePendingElicitation(id, { action: "cancel" });
@@ -421,6 +473,45 @@ const useMcpServerResource = (
   });
 
   const doCompleteAuth = useEffectEvent(async (callbackUrl: string) => {
+    const validationGeneration = connectionGenerationRef.current;
+    const url = new URL(callbackUrl);
+    const state = url.searchParams.get("state");
+    if (!state) throw new Error('missing "state" parameter');
+    let pendingAuthValidation = pendingAuthValidationRef.current;
+    if (!pendingAuthValidation) {
+      let resolve!: () => void;
+      const promise = new Promise<void>((resolvePromise) => {
+        resolve = resolvePromise;
+      });
+      pendingAuthValidation = { count: 0, promise, resolve };
+      pendingAuthValidationRef.current = pendingAuthValidation;
+    }
+    pendingAuthValidation.count += 1;
+    try {
+      const persisted = await props.storage.loadAuthState(props.id);
+      if (!isCurrentConnection(validationGeneration)) {
+        throw createInterruptedAuthError();
+      }
+      if (!persisted?.state) {
+        throw new Error(
+          "no pending OAuth authorization request for this server",
+        );
+      }
+      if (persisted.state !== state) {
+        throw new Error("OAuth state does not match the authorization request");
+      }
+      if (!url.searchParams.get("code") && !url.searchParams.get("error")) {
+        throw new Error("missing authorization code in callback URL");
+      }
+    } finally {
+      pendingAuthValidation.count -= 1;
+      if (pendingAuthValidation.count === 0) {
+        pendingAuthValidationRef.current = null;
+        pendingAuthValidation.resolve();
+      }
+    }
+
+    // Claim the generation before a waiting auto-connect can resume.
     const generation = ++connectionGenerationRef.current;
     cancelPendingElicitations();
     await closePendingTransport();
@@ -429,9 +520,6 @@ const useMcpServerResource = (
     setConnectionState("authPending");
     setLastError(null);
     try {
-      const url = new URL(callbackUrl);
-      const code = url.searchParams.get("code");
-      if (!code) throw new Error("missing authorization code in callback URL");
       let transport = transportRef.current;
       if (!transport) {
         transport = await buildTransport();
@@ -443,7 +531,7 @@ const useMcpServerResource = (
       transportRef.current = null;
       clientRef.current = null;
       pendingTransportRef.current = transport;
-      await transport.finishAuth(code);
+      await transport.finishAuth(url.searchParams);
       if (!isCurrentConnection(generation)) throw createInterruptedAuthError();
       setAuthorizationUrl(null);
       const connected = await finalizeConnect(transport, generation);
@@ -493,12 +581,18 @@ const useMcpServerResource = (
       } else if (!persisted?.token) {
         return;
       }
+      const pendingAuthValidation = pendingAuthValidationRef.current;
+      if (pendingAuthValidation) {
+        await pendingAuthValidation.promise;
+        if (signal.cancelled || !isCurrentConnection(generation)) return;
+      }
       void doConnect();
     },
   );
 
   // Auto-connect on mount when usable auth exists.
   useEffect(() => {
+    const transportCloseQueueRef = props.transportCloseQueueRef;
     const previousDisposal = pendingDisposalRef.current;
     if (previousDisposal) previousDisposal.cancelled = true;
     const pendingDisposal = { cancelled: false };
@@ -509,14 +603,15 @@ const useMcpServerResource = (
     return () => {
       mountedRef.current = false;
       signal.cancelled = true;
-      // Defer disposal so StrictMode can replay setup before closing the transport.
-      queueMicrotask(() => {
+      const task = transportCloseQueueRef.current.then(async () => {
+        await new Promise<void>((resolve) => queueMicrotask(resolve));
         if (pendingDisposal.cancelled) return;
         connectionGenerationRef.current += 1;
-        void closeTransports();
+        await Promise.all(detachTransports().map(closeTransportSafely));
       });
+      transportCloseQueueRef.current = task;
     };
-  }, []);
+  }, [props.transportCloseQueueRef]);
 
   const state = useMemo<MCPServerState>(
     () => ({
@@ -552,7 +647,7 @@ const useMcpServerResource = (
     remove: async () => {
       await doDisconnect();
       try {
-        await props.storage.clearAuthState(props.id);
+        await clearOAuthProviderAuthState(props.storage, props.id);
         await props.onRemove();
       } catch (err) {
         setLastError({
@@ -640,4 +735,30 @@ const useMcpServerResource = (
   };
 };
 
-export const McpServerResource = resource(useMcpServerResource);
+const McpServerResourceInstance = resource(useMcpServerResourceInstance);
+
+export const McpServerResource = resource(function useMcpServerResource(
+  props: McpServerResourceProps,
+): ClientOutput<"mcpServer"> {
+  const transportCloseQueueRef = useRef(Promise.resolve());
+  const dependencies = getConnectionDependencies(props);
+  const [connection, setConnection] = useState({ dependencies, generation: 0 });
+  let currentConnection = connection;
+  if (!areConnectionDependenciesEqual(connection.dependencies, dependencies)) {
+    currentConnection = {
+      dependencies,
+      generation: connection.generation + 1,
+    };
+    setConnection(currentConnection);
+  }
+
+  // Storage keys remounts through its optional scopeId rather than object
+  // identity, because a defaulted storage element is rebuilt on ordinary
+  // renders.
+  return useResource(
+    withKey(
+      currentConnection.generation,
+      McpServerResourceInstance({ ...props, transportCloseQueueRef }),
+    ),
+  );
+});

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

@@ -1,8 +1,10 @@
-import { createTapRoot, useResource } from "@assistant-ui/tap";
+import { createTapRoot, resource, useResource } from "@assistant-ui/tap";
+import { useState } from "react";
 import { auth, type FetchLike } from "@modelcontextprotocol/client";
 import { describe, expect, it } from "vitest";
 import { createOAuthProvider } from "../../auth/createOAuthProvider";
 
+import type { MCPStorage } from "./types";
 import {
   McpLocalStorage,
   normalizeCustomServerRecords,
@@ -125,10 +127,12 @@ describe("normalizePersistedAuthState", () => {
       normalizePersistedAuthState({
         token: "bearer-token",
         codeVerifier: "pkce-verifier",
+        state: "aui-mcp:ZG9jcw.nonce",
       }),
     ).toEqual({
       token: "bearer-token",
       codeVerifier: "pkce-verifier",
+      state: "aui-mcp:ZG9jcw.nonce",
     });
   });
 
@@ -394,6 +398,7 @@ describe("McpLocalStorage auth state", () => {
       JSON.parse(storage.getItem("test-mcp:auth:docs") ?? "null"),
     ).toMatchObject({
       codeVerifier: expect.any(String),
+      state: authorizationUrls[0]!.searchParams.get("state"),
       discoveryState: {
         authorizationServerUrl: "https://auth.example.com",
       },
@@ -418,3 +423,68 @@ describe("McpLocalStorage auth state", () => {
     });
   });
 });
+
+describe("McpLocalStorage instance identity", () => {
+  it("derives a scope from the prefix for the shared default backing", () => {
+    let storage!: MCPStorage;
+
+    createTapRoot(function McpStorageScopeRoot() {
+      storage = useResource(McpLocalStorage({ keyPrefix: "test-mcp" }));
+      return storage;
+    });
+
+    expect(storage.scopeId).toBe("local-storage:test-mcp");
+  });
+
+  it("declares no scope for a custom backing store unless one is named", () => {
+    const backing = createStorage();
+    let unnamed!: MCPStorage;
+    let named!: MCPStorage;
+
+    createTapRoot(function McpStorageCustomScopeRoot() {
+      unnamed = useResource(
+        McpLocalStorage({ keyPrefix: "test-mcp", storage: backing }),
+      );
+      return unnamed;
+    });
+    createTapRoot(function McpStorageNamedScopeRoot() {
+      named = useResource(
+        McpLocalStorage({
+          keyPrefix: "test-mcp",
+          storage: backing,
+          scopeId: "session:alpha",
+        }),
+      );
+      return named;
+    });
+
+    expect(unnamed.scopeId).toBeUndefined();
+    expect(named.scopeId).toBe("session:alpha");
+  });
+
+  it("returns the same instance across re-renders", () => {
+    const backing = createStorage();
+    const seen: MCPStorage[] = [];
+    let rerender!: () => void;
+
+    const useHost = () => {
+      const [, setTick] = useState(0);
+      rerender = () => setTick((n) => n + 1);
+      const storage = useResource(
+        McpLocalStorage({ keyPrefix: "test-mcp", storage: backing }),
+      );
+      seen.push(storage);
+      return storage;
+    };
+    const Host = resource(useHost);
+
+    createTapRoot(function McpStorageIdentityRoot() {
+      return useResource(Host());
+    });
+    rerender();
+    rerender();
+
+    expect(seen.length).toBeGreaterThan(1);
+    expect(new Set(seen).size).toBe(1);
+  });
+});

package/src/resources/storage/McpLocalStorage.ts

@@ -1,4 +1,5 @@
 import { resource } from "@assistant-ui/tap";
+import { useMemo } from "react";
 import {
   OAuthMetadataSchema,
   OAuthClientInformationFullSchema,
@@ -15,6 +16,13 @@ export type McpLocalStorageOptions = {
   keyPrefix?: string;
   /** Override the underlying Storage. Defaults to globalThis.localStorage. */
   storage?: Storage;
+  /**
+   * Stable identity for the backing data, used to key server reconnects.
+   * Required to get reconnect-on-swap behavior when `storage` is overridden;
+   * without it a custom backing store declares no scope, since a prefix
+   * alone cannot distinguish two different stores.
+   */
+  scopeId?: string;
 };
 
 function resolveStorage(opts: McpLocalStorageOptions): Storage | null {
@@ -193,6 +201,7 @@ export const normalizePersistedAuthState = (
   if (isNonEmptyString(value.codeVerifier)) {
     state.codeVerifier = value.codeVerifier;
   }
+  if (isNonEmptyString(value.state)) state.state = value.state;
 
   const tokens = normalizeOAuthTokens(value.tokens);
   if (tokens) state.tokens = tokens;
@@ -208,54 +217,67 @@ export const normalizePersistedAuthState = (
 
 const useMcpLocalStorage = (opts: McpLocalStorageOptions = {}): MCPStorage => {
   const prefix = opts.keyPrefix ?? "aui-mcp";
-  const customServersKey = `${prefix}:custom-servers`;
-  const authKey = (id: string) => `${prefix}:auth:${id}`;
   const storage = resolveStorage(opts);
-
-  const read = <T>(key: string, fallback: T): T => {
-    if (!storage) return fallback;
-    try {
-      const raw = storage.getItem(key);
-      if (raw == null) return fallback;
-      return JSON.parse(raw) as T;
-    } catch {
-      return fallback;
-    }
-  };
-
-  const write = (key: string, value: unknown): void => {
-    if (!storage) return;
-    try {
-      storage.setItem(key, JSON.stringify(value));
-    } catch {
-      // quota or serialization failure — silently drop
-    }
-  };
-
-  const remove = (key: string): void => {
-    if (!storage) return;
-    try {
-      storage.removeItem(key);
-    } catch {
-      // ignore
-    }
-  };
-
-  return {
-    loadCustomServers: async () =>
-      normalizeCustomServerRecords(read<unknown>(customServersKey, [])),
-    saveCustomServers: async (records) => {
-      write(customServersKey, records);
-    },
-    loadAuthState: async (id) =>
-      normalizePersistedAuthState(read<unknown>(authKey(id), null)),
-    saveAuthState: async (id, state) => {
-      write(authKey(id), state);
-    },
-    clearAuthState: async (id) => {
-      remove(authKey(id));
-    },
-  };
+  // Deriving a scope from the prefix is only honest for the shared
+  // globalThis.localStorage; two custom backing stores under one prefix hold
+  // different data, so an overridden backing declares no scope unless the
+  // caller names one.
+  const scopeId =
+    opts.scopeId ??
+    (opts.storage === undefined ? `local-storage:${prefix}` : undefined);
+
+  // Callers key per-server coordination state on this instance, so it has to
+  // stay referentially stable for as long as the underlying store does.
+  return useMemo(() => {
+    const customServersKey = `${prefix}:custom-servers`;
+    const authKey = (id: string) => `${prefix}:auth:${id}`;
+
+    const read = <T>(key: string, fallback: T): T => {
+      if (!storage) return fallback;
+      try {
+        const raw = storage.getItem(key);
+        if (raw == null) return fallback;
+        return JSON.parse(raw) as T;
+      } catch {
+        return fallback;
+      }
+    };
+
+    const write = (key: string, value: unknown): void => {
+      if (!storage) return;
+      try {
+        storage.setItem(key, JSON.stringify(value));
+      } catch {
+        // quota or serialization failure — silently drop
+      }
+    };
+
+    const remove = (key: string): void => {
+      if (!storage) return;
+      try {
+        storage.removeItem(key);
+      } catch {
+        // ignore
+      }
+    };
+
+    return {
+      ...(scopeId !== undefined ? { scopeId } : {}),
+      loadCustomServers: async () =>
+        normalizeCustomServerRecords(read<unknown>(customServersKey, [])),
+      saveCustomServers: async (records) => {
+        write(customServersKey, records);
+      },
+      loadAuthState: async (id) =>
+        normalizePersistedAuthState(read<unknown>(authKey(id), null)),
+      saveAuthState: async (id, state) => {
+        write(authKey(id), state);
+      },
+      clearAuthState: async (id) => {
+        remove(authKey(id));
+      },
+    };
+  }, [prefix, storage, scopeId]);
 };
 
 export const McpLocalStorage = resource(useMcpLocalStorage);

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

@@ -0,0 +1,99 @@
+import { createTapRoot, resource, useResource } from "@assistant-ui/tap";
+import { describe, expect, it } from "vitest";
+import { useState } from "react";
+import type { MCPStorage } from "./types";
+import { McpMemoryStorage } from "./McpMemoryStorage";
+
+const mountWithRerender = () => {
+  const seen: MCPStorage[] = [];
+  let setTick!: (update: (value: number) => number) => void;
+
+  const useHost = () => {
+    const [, setValue] = useState(0);
+    setTick = setValue;
+    const storage = useResource(McpMemoryStorage());
+    seen.push(storage);
+    return storage;
+  };
+  const Host = resource(useHost);
+
+  createTapRoot(function MemoryStorageRoot() {
+    return useResource(Host());
+  });
+
+  return {
+    seen,
+    latest: () => seen[seen.length - 1]!,
+    rerender: () => setTick((value) => value + 1),
+  };
+};
+
+describe("McpMemoryStorage", () => {
+  it("returns the same instance across re-renders", () => {
+    const { seen, rerender } = mountWithRerender();
+
+    rerender();
+    rerender();
+
+    expect(seen.length).toBeGreaterThan(1);
+    expect(new Set(seen).size).toBe(1);
+  });
+
+  it("keeps persisted auth state across re-renders", async () => {
+    const { seen, latest, rerender } = mountWithRerender();
+    await seen[0]!.saveAuthState("docs", { codeVerifier: "pkce-verifier" });
+
+    rerender();
+
+    await expect(latest().loadAuthState("docs")).resolves.toEqual({
+      codeVerifier: "pkce-verifier",
+    });
+  });
+
+  it("keeps custom servers across re-renders", async () => {
+    const { seen, latest, rerender } = mountWithRerender();
+    await seen[0]!.saveCustomServers([
+      {
+        id: "docs",
+        name: "Docs",
+        url: "https://docs.example.com/mcp",
+        auth: { type: "none" },
+        createdAt: 1,
+      },
+    ]);
+
+    rerender();
+
+    await expect(latest().loadCustomServers()).resolves.toHaveLength(1);
+  });
+});
+
+describe("McpMemoryStorage scope identity", () => {
+  it("scopes each instance uniquely and keeps it stable per instance", () => {
+    let first!: MCPStorage;
+    let second!: MCPStorage;
+    const seen: (string | undefined)[] = [];
+    let rerender!: () => void;
+
+    const Host = resource(function useHost() {
+      const [, setTick] = useState(0);
+      rerender = () => setTick((n) => n + 1);
+      first = useResource(McpMemoryStorage());
+      seen.push(first.scopeId);
+      return first;
+    });
+    createTapRoot(function MemoryScopeRootA() {
+      return useResource(Host());
+    });
+    createTapRoot(function MemoryScopeRootB() {
+      second = useResource(McpMemoryStorage());
+      return second;
+    });
+    rerender();
+
+    expect(first.scopeId).toMatch(/^memory:/);
+    expect(second.scopeId).toMatch(/^memory:/);
+    expect(first.scopeId).not.toBe(second.scopeId);
+    expect(new Set(seen).size).toBe(1);
+  });
+});

package/src/resources/storage/McpMemoryStorage.ts

@@ -1,24 +1,30 @@
 import { resource } from "@assistant-ui/tap";
+import { useMemo } from "react";
+import { generateId } from "@assistant-ui/core";
 import type { MCPCustomServerRecord } from "../../mcp-scope";
 import type { MCPPersistedAuthState } from "../../auth/types";
 import type { MCPStorage } from "./types";
 
-const useMcpMemoryStorage = (): MCPStorage => {
-  let servers: MCPCustomServerRecord[] = [];
-  const auth = new Map<string, MCPPersistedAuthState>();
-  return {
-    loadCustomServers: async () => [...servers],
-    saveCustomServers: async (records) => {
-      servers = [...records];
-    },
-    loadAuthState: async (id) => auth.get(id) ?? null,
-    saveAuthState: async (id, state) => {
-      auth.set(id, state);
-    },
-    clearAuthState: async (id) => {
-      auth.delete(id);
-    },
-  };
-};
+const useMcpMemoryStorage = (): MCPStorage =>
+  useMemo(() => {
+    let servers: MCPCustomServerRecord[] = [];
+    const auth = new Map<string, MCPPersistedAuthState>();
+    return {
+      // Each memory store is its own private data: a distinct instance is a
+      // distinct scope, so replacing one keys a reconnect.
+      scopeId: `memory:${generateId()}`,
+      loadCustomServers: async () => [...servers],
+      saveCustomServers: async (records) => {
+        servers = [...records];
+      },
+      loadAuthState: async (id) => auth.get(id) ?? null,
+      saveAuthState: async (id, state) => {
+        auth.set(id, state);
+      },
+      clearAuthState: async (id) => {
+        auth.delete(id);
+      },
+    };
+  }, []);
 
 export const McpMemoryStorage = resource(useMcpMemoryStorage);