@assistant-ui/react-mcp 0.1.15 → 0.1.16

package/src/auth/createOAuthProvider.ts

@@ -57,6 +57,108 @@ export type CreateOAuthProviderOptions = {
   onAuthorizationUrl: (url: URL) => void;
 };
 
+type OAuthProviderCache = {
+  tokens?: OAuthTokens | undefined;
+  clientInformation?: OAuthClientInformationFull | undefined;
+  codeVerifier?: string | undefined;
+  state?: string | undefined;
+  discoveryState?: OAuthDiscoveryState | undefined;
+};
+
+type OAuthProviderPersistence = {
+  cached: OAuthProviderCache | null;
+  cachePromise: Promise<OAuthProviderCache> | 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,
+): OAuthProviderPersistence => {
+  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 = {
+      cached: null,
+      cachePromise: null,
+      queue: Promise.resolve(),
+      invalidated: false,
+    };
+    byServerId.set(serverId, persistence);
+  }
+  return persistence;
+};
+
+/**
+ * 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);
+
+  await Promise.allSettled([persistence.cachePromise, 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
@@ -66,46 +168,66 @@ export function createOAuthProvider(
   opts: CreateOAuthProviderOptions,
 ): OAuthClientProvider {
   const { serverId, config, storage, redirectUri, onAuthorizationUrl } = opts;
+  const persistence = getPersistence(storage, serverId);
+  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, serverId),
+  // 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.
+  const staticClientInformation = (():
+    | 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;
+  })();
+
+  const loadCache = (): Promise<OAuthProviderCache> => {
+    if (persistence.cached) return Promise.resolve(persistence.cached);
+    if (persistence.cachePromise) return persistence.cachePromise;
+
+    persistence.cachePromise = storage.loadAuthState(serverId).then(
+      (persisted) => {
+        const initial: OAuthProviderCache = {};
+        if (persisted?.tokens) initial.tokens = persisted.tokens;
+        if (persisted?.clientInformation)
+          initial.clientInformation = persisted.clientInformation;
+        if (persisted?.codeVerifier)
+          initial.codeVerifier = persisted.codeVerifier;
+        if (persisted?.state) initial.state = persisted.state;
+        if (persisted?.discoveryState)
+          initial.discoveryState = persisted.discoveryState;
+        persistence.cached = initial;
+        return initial;
+      },
+      (error) => {
+        persistence.cachePromise = null;
+        throw error;
+      },
+    );
+    return persistence.cachePromise;
   };
 
-  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 persist = () => {
+    const task = persistence.queue.then(async () => {
+      if (persistence.invalidated) return;
+      const c = persistence.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.state) next.state = c.state;
+      if (c.discoveryState) next.discoveryState = c.discoveryState;
+      await storage.saveAuthState(serverId, next);
+    });
+    persistence.queue = task.catch(() => {});
+    return task;
   };
 
   const clientMetadata: OAuthClientMetadata = {
@@ -131,11 +253,12 @@ 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();
-      return c.clientInformation;
+      return staticClientInformation ?? c.clientInformation;
     },
     async saveClientInformation(info) {
       const c = await loadCache();
@@ -149,6 +272,7 @@ export function createOAuthProvider(
     async saveTokens(tokens) {
       const c = await loadCache();
       c.tokens = tokens;
+      delete c.state;
       await persist();
     },
     async redirectToAuthorization(url) {
@@ -157,6 +281,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() {
@@ -179,7 +307,10 @@ export function createOAuthProvider(
       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 === "verifier") {
+        delete c.codeVerifier;
+        delete c.state;
+      }
       if (scope === "all" || scope === "discovery") delete c.discoveryState;
       await persist();
     },

package/src/auth/types.ts

@@ -8,6 +8,7 @@ export type MCPPersistedAuthState = {
   tokens?: OAuthTokens;
   clientInformation?: OAuthClientInformationFull;
   codeVerifier?: string;
+  state?: string;
   discoveryState?: OAuthDiscoveryState;
   /** Bearer token (entered at add-form time). */
   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";
@@ -303,7 +304,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));
       }
     },