@assistant-ui/react-mcp 0.1.15 → 0.1.17

package/src/auth/createOAuthProvider.ts

@@ -7,6 +7,11 @@ import type {
 } from "@modelcontextprotocol/client";
 import type { MCPStorage } from "../resources/storage/types";
 import type { MCPAuthConfig } from "../mcp-scope";
+import type { MCPPersistedAuthState } from "./types";
+import {
+  isAuthStateForServerUrl,
+  normalizeMcpServerUrl,
+} from "../utils/serverUrl";
 
 const STATE_PREFIX = "aui-mcp:";
 
@@ -49,6 +54,7 @@ export function decodeServerIdFromState(state: string): string | null {
 
 export type CreateOAuthProviderOptions = {
   serverId: string;
+  serverUrl: string;
   /** Must be `auth.type === "oauth"`. */
   config: Extract<MCPAuthConfig, { type: "oauth" }>;
   storage: MCPStorage;
@@ -57,6 +63,174 @@ export type CreateOAuthProviderOptions = {
   onAuthorizationUrl: (url: URL) => void;
 };
 
+type OAuthProviderCache = {
+  token?: string | undefined;
+  tokens?: OAuthTokens | undefined;
+  tokensClientId?: string | undefined;
+  clientInformation?: OAuthClientInformationFull | undefined;
+  clientInformationSource?: MCPPersistedAuthState["clientInformationSource"];
+  codeVerifier?: string | undefined;
+  state?: string | undefined;
+  discoveryState?: OAuthDiscoveryState | undefined;
+};
+
+type OAuthConfig = Extract<MCPAuthConfig, { type: "oauth" }>;
+
+type OAuthCredentialState = {
+  tokens?: OAuthTokens | undefined;
+  tokensClientId?: string | undefined;
+  clientInformation?: OAuthClientInformationFull | undefined;
+  clientInformationSource?: "registered" | undefined;
+};
+
+const registeredClientId = (
+  state: OAuthCredentialState | null | undefined,
+): string | undefined =>
+  state?.clientInformationSource === "registered"
+    ? state.clientInformation?.client_id
+    : undefined;
+
+export const hasUsableOAuthTokens = (
+  state: OAuthCredentialState | null | undefined,
+  config: OAuthConfig,
+): boolean => {
+  const clientId = config.clientId ?? registeredClientId(state);
+  return (
+    clientId !== undefined &&
+    state?.tokens !== undefined &&
+    state.tokensClientId === clientId
+  );
+};
+
+const hasUsableRegisteredClientInformation = (
+  state: OAuthCredentialState | null | undefined,
+  config: OAuthConfig,
+): boolean => {
+  const clientId = registeredClientId(state);
+  return (
+    clientId !== undefined &&
+    (config.clientId === undefined || config.clientId === clientId)
+  );
+};
+
+type OAuthProviderEndpointCache = {
+  serverUrl: string;
+  cached: OAuthProviderCache | null;
+  cachePromise: Promise<OAuthProviderCache> | null;
+  invalidated: boolean;
+};
+
+type OAuthProviderPersistence = {
+  endpoint: OAuthProviderEndpointCache | null;
+  queue: Promise<void>;
+  invalidated: boolean;
+};
+
+// scopeId, not object identity, is what addresses the same persisted data, so
+// storages sharing one share an anchor and an unscoped storage is its own
+// identity. Every storage declaring a scope holds that scope's anchor, so the
+// coordination state below is collected once the last of them is gone.
+const anchorByStorage = new WeakMap<MCPStorage, object>();
+const anchorByScope = new Map<string, WeakRef<object>>();
+const anchorRegistry = new FinalizationRegistry<string>((scopeId) => {
+  if (!anchorByScope.get(scopeId)?.deref()) anchorByScope.delete(scopeId);
+});
+
+const getStorageIdentity = (storage: MCPStorage): object => {
+  const existing = anchorByStorage.get(storage);
+  if (existing) return existing;
+
+  const { scopeId } = storage;
+  if (scopeId === undefined) return storage;
+
+  let anchor = anchorByScope.get(scopeId)?.deref();
+  if (!anchor) {
+    anchor = {};
+    anchorByScope.set(scopeId, new WeakRef(anchor));
+    anchorRegistry.register(anchor, scopeId);
+  }
+  anchorByStorage.set(storage, anchor);
+  return anchor;
+};
+
+// McpServerResource builds a fresh provider for every transport, so the cache,
+// the in-flight load, and the write queue have to outlive any one provider.
+// saveAuthState replaces the whole record, so two providers writing their own
+// snapshots concurrently would drop whichever field the loser had added.
+const persistenceByIdentity = new WeakMap<
+  object,
+  Map<string, OAuthProviderPersistence>
+>();
+
+const getPersistence = (
+  storage: MCPStorage,
+  serverId: string,
+  serverUrl: string,
+): {
+  persistence: OAuthProviderPersistence;
+  endpoint: OAuthProviderEndpointCache;
+} => {
+  const identity = getStorageIdentity(storage);
+  let byServerId = persistenceByIdentity.get(identity);
+  if (!byServerId) {
+    byServerId = new Map();
+    persistenceByIdentity.set(identity, byServerId);
+  }
+
+  let persistence = byServerId.get(serverId);
+  if (!persistence) {
+    persistence = {
+      endpoint: null,
+      queue: Promise.resolve(),
+      invalidated: false,
+    };
+    byServerId.set(serverId, persistence);
+  }
+
+  let endpoint = persistence.endpoint;
+  if (endpoint?.serverUrl !== serverUrl) {
+    if (endpoint) endpoint.invalidated = true;
+    endpoint = {
+      serverUrl,
+      cached: null,
+      cachePromise: null,
+      invalidated: false,
+    };
+    persistence.endpoint = endpoint;
+  }
+  return { persistence, endpoint };
+};
+
+/**
+ * Clears persisted OAuth state after the in-flight load and every queued write
+ * for that server have settled, so a discarded provider cannot recreate the
+ * record it was mid-save on.
+ */
+export const clearOAuthProviderAuthState = async (
+  storage: MCPStorage,
+  serverId: string,
+): Promise<void> => {
+  const identity = getStorageIdentity(storage);
+  const byServerId = persistenceByIdentity.get(identity);
+  const persistence = byServerId?.get(serverId);
+  if (!byServerId || !persistence) {
+    await storage.clearAuthState(serverId);
+    return;
+  }
+
+  // Detaching the entry before awaiting keeps a provider built during the clear
+  // on a fresh generation instead of inheriting the fenced one.
+  persistence.invalidated = true;
+  byServerId.delete(serverId);
+  if (byServerId.size === 0) persistenceByIdentity.delete(identity);
+
+  if (persistence.endpoint) persistence.endpoint.invalidated = true;
+  const cachePromise = persistence.endpoint?.cachePromise;
+  if (cachePromise) await Promise.allSettled([cachePromise]);
+  await persistence.queue;
+  await storage.clearAuthState(serverId);
+};
+
 /**
  * Builds an OAuthClientProvider for the MCP SDK, backed by MCPStorage.
  * Token refresh and DCR are handled by the SDK; this provider only mediates
@@ -65,49 +239,124 @@ export type CreateOAuthProviderOptions = {
 export function createOAuthProvider(
   opts: CreateOAuthProviderOptions,
 ): OAuthClientProvider {
-  const { serverId, config, storage, redirectUri, onAuthorizationUrl } = opts;
+  const {
+    serverId,
+    serverUrl,
+    config,
+    storage,
+    redirectUri,
+    onAuthorizationUrl,
+  } = opts;
+  const normalizedServerUrl = normalizeMcpServerUrl(serverUrl);
+  const { persistence, endpoint } = getPersistence(
+    storage,
+    serverId,
+    normalizedServerUrl,
+  );
+  let pendingState: string | undefined;
 
-  type Cache = {
-    tokens?: OAuthTokens | undefined;
-    clientInformation?: OAuthClientInformationFull | undefined;
-    codeVerifier?: string | undefined;
-    discoveryState?: OAuthDiscoveryState | undefined;
-  };
-  let cached: Cache | null = null;
-
-  const loadCache = async (): Promise<Cache> => {
-    if (cached) return cached;
-    const persisted = await storage.loadAuthState(serverId);
-    const initial: Cache = {};
-    if (persisted?.tokens) initial.tokens = persisted.tokens;
-    if (config.clientId) {
-      const ci: OAuthClientInformationFull = {
-        client_id: config.clientId,
-        redirect_uris: [redirectUri],
-      };
-      if (config.clientSecret) ci.client_secret = config.clientSecret;
-      initial.clientInformation = ci;
-    } else if (persisted?.clientInformation) {
-      initial.clientInformation = persisted.clientInformation;
-    }
-    if (persisted?.codeVerifier) initial.codeVerifier = persisted.codeVerifier;
-    if (persisted?.discoveryState)
-      initial.discoveryState = persisted.discoveryState;
-    cached = initial;
-    return cached;
+  // The cache is shared with every other provider for this storage, server id,
+  // and server URL, so a statically configured client stays a read-time overlay
+  // owned by this provider. Writing it into the cache would leak this provider's
+  // registration to a replacement built for a different, or absent, clientId.
+  // The SDK's write-backs, its issuer stamp included, replace the overlay.
+  const configuredClientInformation = ():
+    | OAuthClientInformationFull
+    | undefined => {
+    if (!config.clientId) return undefined;
+    const ci: OAuthClientInformationFull = {
+      client_id: config.clientId,
+      redirect_uris: [redirectUri],
+    };
+    if (config.clientSecret) ci.client_secret = config.clientSecret;
+    return ci;
   };
+  let clientInformationOverlay = configuredClientInformation();
 
-  const persist = async () => {
-    const c = cached;
-    if (!c) return;
-    const next: Parameters<typeof storage.saveAuthState>[1] = {};
-    if (c.tokens) next.tokens = c.tokens;
-    if (c.clientInformation) next.clientInformation = c.clientInformation;
-    if (c.codeVerifier) next.codeVerifier = c.codeVerifier;
-    if (c.discoveryState) next.discoveryState = c.discoveryState;
-    await storage.saveAuthState(serverId, next);
+  const activeClientId = (cache: OAuthProviderCache): string | undefined =>
+    clientInformationOverlay?.client_id ?? registeredClientId(cache);
+
+  const loadCache = (): Promise<OAuthProviderCache> => {
+    if (endpoint.invalidated) return Promise.resolve({});
+    if (endpoint.cached) return Promise.resolve(endpoint.cached);
+    if (endpoint.cachePromise) return endpoint.cachePromise;
+
+    endpoint.cachePromise = persistence.queue
+      .then(() => storage.loadAuthState(serverId))
+      .then(
+        async (persisted) => {
+          const initial: OAuthProviderCache = {};
+          let needsMigration = false;
+          if (endpoint.invalidated) return initial;
+          if (
+            persisted &&
+            isAuthStateForServerUrl(persisted, normalizedServerUrl)
+          ) {
+            if (
+              hasUsableRegisteredClientInformation(persisted, config) &&
+              persisted.clientInformation
+            ) {
+              initial.clientInformation = persisted.clientInformation;
+              initial.clientInformationSource = "registered";
+            } else if (
+              persisted?.clientInformation ||
+              persisted?.clientInformationSource !== undefined
+            ) {
+              needsMigration = true;
+            }
+            if (hasUsableOAuthTokens(persisted, config)) {
+              initial.tokens = persisted.tokens;
+              initial.tokensClientId = persisted.tokensClientId;
+            } else if (
+              persisted?.tokens ||
+              persisted?.tokensClientId !== undefined
+            ) {
+              needsMigration = true;
+            }
+            if (persisted?.token) initial.token = persisted.token;
+            if (persisted?.codeVerifier)
+              initial.codeVerifier = persisted.codeVerifier;
+            if (persisted?.state) initial.state = persisted.state;
+            if (persisted?.discoveryState)
+              initial.discoveryState = persisted.discoveryState;
+          }
+          endpoint.cached = initial;
+          if (needsMigration) await persist().catch(() => {});
+          return initial;
+        },
+        (error) => {
+          endpoint.cachePromise = null;
+          throw error;
+        },
+      );
+    return endpoint.cachePromise;
   };
 
+  function persist() {
+    const task = persistence.queue.then(async () => {
+      if (persistence.invalidated || endpoint.invalidated) return;
+      const c = endpoint.cached;
+      if (!c) return;
+      const next: Parameters<typeof storage.saveAuthState>[1] = {};
+      if (hasUsableOAuthTokens(c, config) && c.tokens && c.tokensClientId) {
+        next.tokens = c.tokens;
+        next.tokensClientId = c.tokensClientId;
+      }
+      if (c.clientInformation && c.clientInformationSource === "registered") {
+        next.clientInformation = c.clientInformation;
+        next.clientInformationSource = "registered";
+      }
+      if (c.token) next.token = c.token;
+      if (c.codeVerifier) next.codeVerifier = c.codeVerifier;
+      if (c.state) next.state = c.state;
+      if (c.discoveryState) next.discoveryState = c.discoveryState;
+      next.serverUrl = normalizedServerUrl;
+      await storage.saveAuthState(serverId, next);
+    });
+    persistence.queue = task.catch(() => {});
+    return task;
+  }
+
   const clientMetadata: OAuthClientMetadata = {
     client_name: "assistant-ui",
     redirect_uris: [redirectUri],
@@ -131,24 +380,43 @@ export function createOAuthProvider(
         typeof crypto !== "undefined" && "randomUUID" in crypto
           ? crypto.randomUUID()
           : `${Date.now()}.${Math.random()}`;
-      return `${encodeServerIdInState(serverId)}.${nonce}`;
+      pendingState = `${encodeServerIdInState(serverId)}.${nonce}`;
+      return pendingState;
     },
     async clientInformation() {
       const c = await loadCache();
+      if (clientInformationOverlay) return clientInformationOverlay;
+      if (c.clientInformationSource !== "registered") return undefined;
       return c.clientInformation;
     },
     async saveClientInformation(info) {
+      if (clientInformationOverlay) {
+        clientInformationOverlay = info as OAuthClientInformationFull;
+        return;
+      }
       const c = await loadCache();
       c.clientInformation = info as OAuthClientInformationFull;
+      c.clientInformationSource = "registered";
+      if (c.tokensClientId !== c.clientInformation.client_id) {
+        delete c.tokens;
+        delete c.tokensClientId;
+      }
       await persist();
     },
     async tokens() {
       const c = await loadCache();
+      const clientId = activeClientId(c);
+      if (clientId === undefined || c.tokensClientId !== clientId)
+        return undefined;
       return c.tokens;
     },
     async saveTokens(tokens) {
       const c = await loadCache();
       c.tokens = tokens;
+      const clientId = activeClientId(c);
+      if (clientId) c.tokensClientId = clientId;
+      else delete c.tokensClientId;
+      delete c.state;
       await persist();
     },
     async redirectToAuthorization(url) {
@@ -157,6 +425,10 @@ export function createOAuthProvider(
     async saveCodeVerifier(codeVerifier) {
       const c = await loadCache();
       c.codeVerifier = codeVerifier;
+      if (pendingState) {
+        c.state = pendingState;
+        pendingState = undefined;
+      }
       await persist();
     },
     async codeVerifier() {
@@ -177,9 +449,23 @@ export function createOAuthProvider(
     },
     async invalidateCredentials(scope) {
       const c = await loadCache();
-      if (scope === "all" || scope === "tokens") delete c.tokens;
-      if (scope === "all" || scope === "client") delete c.clientInformation;
-      if (scope === "all" || scope === "verifier") delete c.codeVerifier;
+      if (scope === "all" || scope === "tokens") {
+        delete c.tokens;
+        delete c.tokensClientId;
+      }
+      if (scope === "all" || scope === "client") {
+        delete c.clientInformation;
+        delete c.clientInformationSource;
+        if (!config.clientId) {
+          delete c.tokens;
+          delete c.tokensClientId;
+        }
+        clientInformationOverlay = configuredClientInformation();
+      }
+      if (scope === "all" || scope === "verifier") {
+        delete c.codeVerifier;
+        delete c.state;
+      }
       if (scope === "all" || scope === "discovery") delete c.discoveryState;
       await persist();
     },

package/src/auth/types.ts

@@ -5,10 +5,15 @@ import type {
 } from "@modelcontextprotocol/client";
 
 export type MCPPersistedAuthState = {
+  /** MCP server URL this authentication state belongs to. Required with credentials. */
+  serverUrl?: string;
   tokens?: OAuthTokens;
+  tokensClientId?: string;
   clientInformation?: OAuthClientInformationFull;
+  clientInformationSource?: "registered";
   codeVerifier?: string;
+  state?: string;
   discoveryState?: OAuthDiscoveryState;
-  /** Bearer token (entered at add-form time). */
+  /** Host-persisted bearer token. Must be paired with serverUrl. */
   token?: string;
 };

package/src/hooks/useMcpOAuthCallback.test.ts

@@ -1,6 +1,12 @@
 // @vitest-environment jsdom
 
-import { renderHook, waitFor } from "@testing-library/react";
+import {
+  createElement,
+  Suspense,
+  startTransition,
+  type PropsWithChildren,
+} from "react";
+import { act, renderHook, waitFor } from "@testing-library/react";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 
 const mocks = vi.hoisted(() => {
@@ -98,6 +104,73 @@ describe("useMcpOAuthCallback", () => {
       });
     },
   );
+
+  it("lets the server validate error callback parameters", async () => {
+    const authError = new Error("authorization response issuer mismatch");
+    const url = `${callbackUrl}&error=access_denied&error_description=untrusted`;
+    mocks.completeAuth.mockRejectedValueOnce(authError);
+
+    const { result } = renderHook(() => useMcpOAuthCallback({ url }));
+
+    await waitFor(() => expect(result.current.status).toBe("error"));
+    expect(mocks.completeAuth).toHaveBeenCalledWith(url);
+    expect(result.current.error).toMatchObject({
+      message:
+        'MCP OAuth callback for server "docs" failed: authorization response issuer mismatch',
+      cause: authError,
+    });
+  });
+
+  it("keeps callbacks scoped to committed renders", async () => {
+    let resolveAuth!: () => void;
+    mocks.completeAuth.mockReturnValueOnce(
+      new Promise<void>((resolve) => {
+        resolveAuth = resolve;
+      }),
+    );
+    const onCompleteA = vi.fn();
+    const onCompleteB = vi.fn();
+    const interruptedRender = vi.fn();
+    const pending = new Promise<never>(() => {});
+    let blocked = false;
+    const Blocker = () => {
+      if (blocked) {
+        interruptedRender();
+        throw pending;
+      }
+      return null;
+    };
+    const Wrapper = ({ children }: PropsWithChildren) =>
+      createElement(
+        Suspense,
+        { fallback: null },
+        children,
+        createElement(Blocker),
+      );
+
+    const { rerender } = renderHook(
+      ({ onComplete }) => useMcpOAuthCallback({ url: callbackUrl, onComplete }),
+      {
+        initialProps: { onComplete: onCompleteA },
+        wrapper: Wrapper,
+      },
+    );
+    await waitFor(() => expect(mocks.completeAuth).toHaveBeenCalledOnce());
+
+    act(() => {
+      blocked = true;
+      startTransition(() => rerender({ onComplete: onCompleteB }));
+    });
+    expect(interruptedRender).toHaveBeenCalled();
+
+    await act(async () => {
+      resolveAuth();
+      await Promise.resolve();
+    });
+
+    expect(onCompleteB).not.toHaveBeenCalled();
+    expect(onCompleteA).toHaveBeenCalledWith("docs");
+  });
 });
 
 describe("createMcpOAuthCallbackError", () => {

package/src/hooks/useMcpOAuthCallback.tsx

@@ -1,4 +1,11 @@
-import { type FC, type ReactNode, useEffect, useRef, useState } from "react";
+import {
+  type FC,
+  type ReactNode,
+  useEffect,
+  useInsertionEffect,
+  useRef,
+  useState,
+} from "react";
 import { useAui } from "@assistant-ui/store";
 import { decodeServerIdFromState } from "../auth/createOAuthProvider";
 import { invokeMcpCallback } from "../utils/invokeMcpCallback";
@@ -44,7 +51,9 @@ export function useMcpOAuthCallback(
   // single-use OAuth code is double-redeemed and the second attempt 4xxs.
   const startedRef = useRef<string | null>(null);
   const optsRef = useRef(opts);
-  optsRef.current = opts;
+  useInsertionEffect(() => {
+    optsRef.current = opts;
+  });
 
   useEffect(() => {
     const url =
@@ -59,12 +68,6 @@ export function useMcpOAuthCallback(
         const parsed = new URL(url);
         const state = parsed.searchParams.get("state");
         if (state) serverId = decodeServerIdFromState(state);
-        const error = parsed.searchParams.get("error");
-        if (error) {
-          throw new Error(
-            parsed.searchParams.get("error_description") ?? error,
-          );
-        }
         if (!state) throw new Error('missing "state" parameter');
         if (!serverId) {
           throw new Error("state was not created by assistant-ui MCP");

package/src/resources/McpManagerResource.ts

@@ -9,6 +9,7 @@ import {
 import { useAssistantScopeEffect } from "@assistant-ui/store/client";
 import { ModelContext } from "@assistant-ui/core/store";
 import { createMcpId } from "../utils/createMcpId";
+import { clearOAuthProviderAuthState } from "../auth/createOAuthProvider";
 import type { Tool } from "assistant-stream";
 import { McpServerResource } from "./McpServerResource";
 import { McpLocalStorage } from "./storage/McpLocalStorage";
@@ -119,6 +120,9 @@ const useMcpManagerResource = (
 
   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;
@@ -303,7 +307,7 @@ const useMcpManagerResource = (
       try {
         await lookup.get({ key: id }).remove();
       } catch {
-        await storage.clearAuthState(id);
+        await clearOAuthProviderAuthState(storage, id);
         setCustomServers((prev) => prev.filter((s) => s.id !== id));
       }
     },