@assistant-ui/react-mcp 0.1.15 → 0.1.17

package/src/resources/McpServerResource.ts

@@ -1,6 +1,7 @@
 import { useState, useRef, useEffect, useMemo, useEffectEvent } from "react";
 import { resource, useResource, withKey } from "@assistant-ui/tap";
 import type { ClientOutput } from "@assistant-ui/store";
+import { shallowEqual } from "@assistant-ui/store/internal";
 import {
   Client,
   StreamableHTTPClientTransport,
@@ -10,9 +11,17 @@ import {
   type ElicitResult,
   type StreamableHTTPClientTransportOptions,
 } from "@modelcontextprotocol/client";
-import { createOAuthProvider } from "../auth/createOAuthProvider";
+import {
+  clearOAuthProviderAuthState,
+  createOAuthProvider,
+  hasUsableOAuthTokens,
+} from "../auth/createOAuthProvider";
 import { buildHeaders } from "../auth/buildHeaders";
 import { assertValidServerId } from "../utils/serverId";
+import {
+  hasPersistedCredentials,
+  isAuthStateForServerUrl,
+} from "../utils/serverUrl";
 import { validateElicitationContent } from "./validateElicitationContent";
 import type { MCPStorage } from "./storage/types";
 import type {
@@ -46,13 +55,13 @@ type McpServerResourceInstanceProps = McpServerResourceProps & {
   transportCloseQueueRef: { current: Promise<void> };
 };
 
-const getConnectionDependencies = (
+export const getConnectionDependencies = (
   props: McpServerResourceProps,
 ): readonly unknown[] => {
   const auth = props.auth;
   const authDependencies =
     auth.type === "bearer"
-      ? [auth.type, auth.token]
+      ? [auth.type, auth.token, props.storage.scopeId]
       : auth.type === "oauth"
         ? [
             auth.type,
@@ -63,6 +72,7 @@ const getConnectionDependencies = (
             auth.registrationEndpoint,
             auth.clientId,
             auth.clientSecret,
+            props.storage.scopeId,
           ]
         : [auth.type];
 
@@ -76,13 +86,6 @@ const getConnectionDependencies = (
   ];
 };
 
-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"> => {
@@ -102,6 +105,11 @@ const useMcpServerResourceInstance = (
     null,
   );
   const connectionGenerationRef = useRef(0);
+  const pendingAuthValidationRef = useRef<{
+    count: number;
+    promise: Promise<void>;
+    resolve: () => void;
+  } | null>(null);
   const elicitationResolversRef = useRef(
     new Map<
       string,
@@ -233,11 +241,23 @@ const useMcpServerResourceInstance = (
     },
   );
 
+  const unboundAuthMessage = () =>
+    `MCP server "${props.id}" has saved authentication for a different URL. Authenticate again to connect to ${props.url}.`;
+
+  const loadAuthState = useEffectEvent(async () => {
+    const state = await props.storage.loadAuthState(props.id);
+    if (isAuthStateForServerUrl(state, props.url)) {
+      return { state, unbound: false };
+    }
+    return { state: null, unbound: hasPersistedCredentials(state) };
+  });
+
   const buildTransport = useEffectEvent(
     async (): Promise<StreamableHTTPClientTransport> => {
       if (props.auth.type === "oauth") {
         const authProvider = createOAuthProvider({
           serverId: props.id,
+          serverUrl: props.url,
           config: props.auth,
           storage: props.storage,
           redirectUri: props.redirectUri,
@@ -248,8 +268,9 @@ const useMcpServerResourceInstance = (
         });
       }
       if (props.auth.type === "bearer") {
-        const persisted = await props.storage.loadAuthState(props.id);
-        const headers = buildHeaders(props.auth, persisted);
+        const { state, unbound } = await loadAuthState();
+        const headers = buildHeaders(props.auth, state);
+        if (!headers && unbound) throw new Error(unboundAuthMessage());
         const transportOpts: StreamableHTTPClientTransportOptions = {};
         if (headers) transportOpts.requestInit = { headers };
         return new StreamableHTTPClientTransport(
@@ -464,6 +485,46 @@ const useMcpServerResourceInstance = (
   });
 
   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 { state: persisted, unbound } = await loadAuthState();
+      if (!isCurrentConnection(validationGeneration)) {
+        throw createInterruptedAuthError();
+      }
+      if (unbound) throw new Error(unboundAuthMessage());
+      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();
@@ -472,9 +533,6 @@ const useMcpServerResourceInstance = (
     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();
@@ -486,7 +544,7 @@ const useMcpServerResourceInstance = (
       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);
@@ -518,9 +576,9 @@ const useMcpServerResourceInstance = (
         return;
       }
       const generation = connectionGenerationRef.current;
-      let persisted: Awaited<ReturnType<MCPStorage["loadAuthState"]>>;
+      let loaded: Awaited<ReturnType<typeof loadAuthState>>;
       try {
-        persisted = await props.storage.loadAuthState(props.id);
+        loaded = await loadAuthState();
       } catch (error) {
         if (signal.cancelled || !isCurrentConnection(generation)) return;
         const message = error instanceof Error ? error.message : String(error);
@@ -531,11 +589,21 @@ const useMcpServerResourceInstance = (
         return;
       }
       if (signal.cancelled || !isCurrentConnection(generation)) return;
+      if (loaded.unbound) {
+        setLastError({ message: unboundAuthMessage() });
+        return;
+      }
+      const persisted = loaded.state;
       if (props.auth.type === "oauth") {
-        if (!persisted?.tokens) return;
+        if (!hasUsableOAuthTokens(persisted, props.auth)) return;
       } else if (!persisted?.token) {
         return;
       }
+      const pendingAuthValidation = pendingAuthValidationRef.current;
+      if (pendingAuthValidation) {
+        await pendingAuthValidation.promise;
+        if (signal.cancelled || !isCurrentConnection(generation)) return;
+      }
       void doConnect();
     },
   );
@@ -549,6 +617,9 @@ const useMcpServerResourceInstance = (
     pendingDisposalRef.current = pendingDisposal;
     mountedRef.current = true;
     const signal = { cancelled: false };
+    // Auto-connect opens a transport, so it belongs to the same effect as the
+    // disposal that closes it.
+    // eslint-disable-next-line react-hooks/set-state-in-effect
     void tryAutoConnect(signal);
     return () => {
       mountedRef.current = false;
@@ -597,7 +668,7 @@ const useMcpServerResourceInstance = (
     remove: async () => {
       await doDisconnect();
       try {
-        await props.storage.clearAuthState(props.id);
+        await clearOAuthProviderAuthState(props.storage, props.id);
         await props.onRemove();
       } catch (err) {
         setLastError({
@@ -694,7 +765,7 @@ export const McpServerResource = resource(function useMcpServerResource(
   const dependencies = getConnectionDependencies(props);
   const [connection, setConnection] = useState({ dependencies, generation: 0 });
   let currentConnection = connection;
-  if (!areConnectionDependenciesEqual(connection.dependencies, dependencies)) {
+  if (!shallowEqual(connection.dependencies, dependencies)) {
     currentConnection = {
       dependencies,
       generation: connection.generation + 1,
@@ -702,8 +773,9 @@ export const McpServerResource = resource(function useMcpServerResource(
     setConnection(currentConnection);
   }
 
-  // Storage resources do not expose a stable scope identity and may return a
-  // fresh client on ordinary renders, so storage changes cannot key remounts.
+  // 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,

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,13 +127,36 @@ 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",
     });
   });
 
+  it("keeps valid server URL bindings", () => {
+    expect(
+      normalizePersistedAuthState({
+        serverUrl: "http://mcp.example.com/docs",
+        token: "bearer-token",
+      }),
+    ).toEqual({
+      serverUrl: "http://mcp.example.com/docs",
+      token: "bearer-token",
+    });
+  });
+
+  it("rejects auth state with an unsafe server URL binding", () => {
+    expect(
+      normalizePersistedAuthState({
+        serverUrl: "javascript:alert(1)",
+        token: "bearer-token",
+      }),
+    ).toBeNull();
+  });
+
   it("keeps valid OAuth tokens and client information", () => {
     const tokens = {
       access_token: "access-token",
@@ -149,11 +174,15 @@ describe("normalizePersistedAuthState", () => {
     expect(
       normalizePersistedAuthState({
         tokens,
+        tokensClientId: "client-id",
         clientInformation,
+        clientInformationSource: "registered",
       }),
     ).toEqual({
       tokens,
+      tokensClientId: "client-id",
       clientInformation,
+      clientInformationSource: "registered",
     });
   });
 
@@ -373,6 +402,7 @@ describe("McpLocalStorage auth state", () => {
     const createProvider = () =>
       createOAuthProvider({
         serverId: "docs",
+        serverUrl: "https://mcp.example.com/mcp",
         config: { type: "oauth", clientId: "client-id" },
         storage: loadStorage(storage),
         redirectUri: "http://localhost/callback",
@@ -394,6 +424,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 +449,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,10 +1,12 @@
 import { resource } from "@assistant-ui/tap";
+import { useMemo } from "react";
 import {
   OAuthMetadataSchema,
   OAuthClientInformationFullSchema,
   OAuthProtectedResourceMetadataSchema,
   OAuthTokensSchema,
 } from "@modelcontextprotocol/core";
+import { normalizeMcpServerUrl } from "../../utils/serverUrl";
 import type { MCPAuthConfig, MCPCustomServerRecord } from "../../mcp-scope";
 import type { MCPPersistedAuthState } from "../../auth/types";
 import { assertValidServerId } from "../../utils/serverId";
@@ -15,6 +17,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 {
@@ -152,6 +161,16 @@ const isSecureNetworkUrl = (value: unknown): value is string => {
   }
 };
 
+const isMcpServerUrl = (value: unknown): value is string => {
+  if (typeof value !== "string") return false;
+  try {
+    const url = new URL(value);
+    return url.protocol === "https:" || url.protocol === "http:";
+  } catch {
+    return false;
+  }
+};
+
 const normalizeDiscoveryState = (
   value: unknown,
 ): MCPPersistedAuthState["discoveryState"] | undefined => {
@@ -187,18 +206,29 @@ export const normalizePersistedAuthState = (
   value: unknown,
 ): MCPPersistedAuthState | null => {
   if (!isRecord(value)) return null;
+  if ("serverUrl" in value && !isMcpServerUrl(value.serverUrl)) return null;
 
   const state: MCPPersistedAuthState = {};
+  if (isMcpServerUrl(value.serverUrl)) {
+    state.serverUrl = normalizeMcpServerUrl(value.serverUrl);
+  }
   if (isNonEmptyString(value.token)) state.token = value.token;
+  if (isNonEmptyString(value.tokensClientId)) {
+    state.tokensClientId = value.tokensClientId;
+  }
   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;
 
   const clientInformation = normalizeClientInformation(value.clientInformation);
   if (clientInformation) state.clientInformation = clientInformation;
+  if (value.clientInformationSource === "registered") {
+    state.clientInformationSource = value.clientInformationSource;
+  }
 
   const discoveryState = normalizeDiscoveryState(value.discoveryState);
   if (discoveryState) state.discoveryState = discoveryState;
@@ -208,54 +238,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);
+  });
+});