@assistant-ui/react-mcp 0.1.15 → 0.1.16

package/src/resources/McpServerResource.test.ts

@@ -4,6 +4,7 @@ import { useEffect, useState } from "react";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import type { MCPAuthConfig } from "../mcp-scope";
 import type { MCPStorage } from "./storage/types";
+import type { McpServerResourceProps } from "./McpServerResource";
 
 const mocks = vi.hoisted(() => {
   const clients: any[] = [];
@@ -70,7 +71,8 @@ vi.mock("@modelcontextprotocol/client", async (importOriginal) => ({
   StreamableHTTPClientTransport: mocks.StreamableHTTPClientTransport,
 }));
 
-const { McpServerResource } = await import("./McpServerResource");
+const { McpServerResource, getConnectionDependencies } =
+  await import("./McpServerResource");
 
 const never = <T>() => new Promise<T>(() => {});
 
@@ -559,14 +561,107 @@ describe("McpServerResource connection lifecycle", () => {
 describe("McpServerResource completeAuth", () => {
   beforeEach(resetMocks);
 
+  it("lets callback validation win over mount-time auto-connect", async () => {
+    const pendingLoads: Array<
+      (value: { state?: string; tokens?: { access_token: string } }) => void
+    > = [];
+    const storage = createStorage();
+    vi.mocked(storage.loadAuthState).mockImplementation(
+      () =>
+        new Promise((resolve) => {
+          pendingLoads.push(resolve);
+        }),
+    );
+    const root = mount({
+      auth: { type: "oauth" },
+      storage,
+      autoConnect: true,
+    });
+
+    try {
+      await waitFor(() => pendingLoads.length > 1);
+      const callbackLoadIndex = pendingLoads.length;
+      const completeAuth = root
+        .getValue()
+        .completeAuth("https://example.com/callback?code=abc&state=expected");
+      await waitFor(() => pendingLoads.length > callbackLoadIndex);
+
+      for (const resolve of pendingLoads.slice(0, callbackLoadIndex)) {
+        resolve({ tokens: { access_token: "persisted" } });
+      }
+      await flushMacrotask();
+
+      expect(mocks.StreamableHTTPClientTransport).not.toHaveBeenCalled();
+
+      pendingLoads[callbackLoadIndex]!({ state: "expected" });
+      await expect(completeAuth).resolves.toBeUndefined();
+      await flushMacrotask();
+
+      expect(mocks.transports).toHaveLength(1);
+      expect(mocks.transports[0].finishAuth).toHaveBeenCalledTimes(1);
+      expect(root.getValue().getState().connectionState).toBe("connected");
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("resumes auto-connect when callback validation fails", async () => {
+    const pendingLoads: Array<
+      (value: { state?: string; tokens?: { access_token: string } }) => void
+    > = [];
+    const storage = createStorage();
+    vi.mocked(storage.loadAuthState).mockImplementation(
+      () =>
+        new Promise((resolve) => {
+          pendingLoads.push(resolve);
+        }),
+    );
+    const root = mount({
+      auth: { type: "oauth" },
+      storage,
+      autoConnect: true,
+    });
+
+    try {
+      await waitFor(() => pendingLoads.length > 1);
+      const callbackLoadIndex = pendingLoads.length;
+      const completeAuth = root
+        .getValue()
+        .completeAuth("https://example.com/callback?code=abc&state=expected");
+      await waitFor(() => pendingLoads.length > callbackLoadIndex);
+
+      for (const resolve of pendingLoads.slice(0, callbackLoadIndex)) {
+        resolve({ tokens: { access_token: "persisted" } });
+      }
+      await flushMacrotask();
+
+      expect(mocks.StreamableHTTPClientTransport).not.toHaveBeenCalled();
+
+      pendingLoads[callbackLoadIndex]!({ state: "different" });
+      await expect(completeAuth).rejects.toThrow(
+        "OAuth state does not match the authorization request",
+      );
+      await waitFor(() => mocks.transports.length === 1);
+      await waitForResourceUpdate(
+        () => root.getValue().getState().connectionState === "connected",
+      );
+    } finally {
+      root.unmount();
+    }
+  });
+
   it("completes auth across the StrictMode effect replay", async () => {
+    const storage = createStorage();
+    vi.mocked(storage.loadAuthState).mockResolvedValue({
+      state: "aui-mcp:ZG9jcw.nonce",
+    });
     let completeAuth: Promise<void> | undefined;
     let started = false;
-    const root = mount({ auth: { type: "oauth" } }, (server) => {
+    const root = mount({ auth: { type: "oauth" }, storage }, (server) => {
       if (started) return;
       started = true;
       completeAuth = server.completeAuth(
-        "https://example.com/callback?code=abc",
+        "https://example.com/callback?code=abc&state=aui-mcp%3AZG9jcw.nonce&iss=https%3A%2F%2Fauth.example.com",
       );
     });
 
@@ -575,26 +670,149 @@ describe("McpServerResource completeAuth", () => {
       await flushMacrotask();
 
       expect(mocks.transports[0].finishAuth).toHaveBeenCalledTimes(1);
+      const params = mocks.transports[0].finishAuth.mock.calls[0][0];
+      expect(params).toBeInstanceOf(URLSearchParams);
+      expect(params.get("code")).toBe("abc");
+      expect(params.get("iss")).toBe("https://auth.example.com");
       expect(root.getValue().getState().connectionState).toBe("connected");
     } finally {
       root.unmount();
     }
   });
 
-  it("rejects when the callback URL has no authorization code", async () => {
-    const root = mount();
+  it("rejects before transport setup without a usable authorization code", async () => {
+    const storage = createStorage();
+    vi.mocked(storage.loadAuthState).mockResolvedValue({ state: "abc" });
+    const root = mount({ auth: { type: "oauth" }, storage });
 
     try {
       await expect(
-        root.getValue().completeAuth("https://example.com/callback?state=abc"),
+        root
+          .getValue()
+          .completeAuth("https://example.com/callback?state=abc&code="),
       ).rejects.toThrow("missing authorization code in callback URL");
       await flushMacrotask();
 
+      expect(root.getValue().getState()).toMatchObject({
+        connectionState: "disconnected",
+        lastError: null,
+      });
+      expect(mocks.transports).toHaveLength(0);
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("forwards OAuth error callbacks to the transport", async () => {
+    const storage = createStorage();
+    vi.mocked(storage.loadAuthState).mockResolvedValue({ state: "expected" });
+    mocks.finishAuthResults.push(() =>
+      Promise.reject(new Error("access_denied: Denied")),
+    );
+    const root = mount({ auth: { type: "oauth" }, storage });
+
+    try {
+      await expect(
+        root
+          .getValue()
+          .completeAuth(
+            "https://example.com/callback?error=access_denied&error_description=Denied&state=expected&iss=https%3A%2F%2Fauth.example.com",
+          ),
+      ).rejects.toThrow("access_denied: Denied");
+      await flushMacrotask();
+
+      const params = mocks.transports[0].finishAuth.mock.calls[0][0];
+      expect(params).toBeInstanceOf(URLSearchParams);
+      expect(params.get("error")).toBe("access_denied");
+      expect(params.get("error_description")).toBe("Denied");
+      expect(params.get("iss")).toBe("https://auth.example.com");
       expect(root.getValue().getState()).toMatchObject({
         connectionState: "error",
-        lastError: {
-          message: "missing authorization code in callback URL",
-        },
+        lastError: { message: "access_denied: Denied" },
+      });
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("rejects callbacks without a pending authorization request", async () => {
+    const storage = createStorage();
+    vi.mocked(storage.loadAuthState).mockResolvedValue(null);
+    const root = mount({ auth: { type: "oauth" }, storage });
+
+    try {
+      await expect(
+        root
+          .getValue()
+          .completeAuth("https://example.com/callback?code=abc&state=expected"),
+      ).rejects.toThrow(
+        "no pending OAuth authorization request for this server",
+      );
+
+      expect(mocks.transports).toHaveLength(0);
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("rejects callbacks whose state does not match the authorization request", async () => {
+    const storage = createStorage();
+    vi.mocked(storage.loadAuthState).mockResolvedValue({
+      state: "aui-mcp:ZG9jcw.expected",
+    });
+    const root = mount({ auth: { type: "oauth" }, storage });
+
+    try {
+      await root.getValue().connect();
+      const transport = mocks.transports[0];
+
+      await expect(
+        root
+          .getValue()
+          .completeAuth(
+            "https://example.com/callback?code=abc&state=aui-mcp%3AZG9jcw.forged",
+          ),
+      ).rejects.toThrow("OAuth state does not match the authorization request");
+      await flushMacrotask();
+
+      expect(mocks.transports).toHaveLength(1);
+      expect(transport.close).not.toHaveBeenCalled();
+      expect(root.getValue().getState()).toMatchObject({
+        connectionState: "connected",
+        lastError: null,
+      });
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("does not resume authorization after disconnecting during validation", async () => {
+    let resolveAuthState!: (state: { state: string }) => void;
+    const storage = createStorage();
+    vi.mocked(storage.loadAuthState).mockImplementation(
+      () =>
+        new Promise((resolve) => {
+          resolveAuthState = resolve;
+        }),
+    );
+    const root = mount({ auth: { type: "oauth" }, storage });
+
+    try {
+      const completeAuth = root
+        .getValue()
+        .completeAuth("https://example.com/callback?code=abc&state=expected");
+      await waitFor(() => resolveAuthState !== undefined);
+
+      await root.getValue().disconnect();
+      resolveAuthState({ state: "expected" });
+
+      await expect(completeAuth).rejects.toThrow(
+        'MCP server "docs" authorization was interrupted before completion.',
+      );
+      expect(mocks.transports).toHaveLength(0);
+      expect(root.getValue().getState()).toMatchObject({
+        connectionState: "disconnected",
+        lastError: null,
       });
     } finally {
       root.unmount();
@@ -605,11 +823,15 @@ describe("McpServerResource completeAuth", () => {
     mocks.finishAuthResults.push(() =>
       Promise.reject(new Error("invalid_grant")),
     );
-    const root = mount({ auth: { type: "oauth" } });
+    const storage = createStorage();
+    vi.mocked(storage.loadAuthState).mockResolvedValue({ state: "expected" });
+    const root = mount({ auth: { type: "oauth" }, storage });
 
     try {
       await expect(
-        root.getValue().completeAuth("https://example.com/callback?code=abc"),
+        root
+          .getValue()
+          .completeAuth("https://example.com/callback?code=abc&state=expected"),
       ).rejects.toThrow("invalid_grant");
       await flushMacrotask();
 
@@ -619,7 +841,9 @@ describe("McpServerResource completeAuth", () => {
           message: "invalid_grant",
         },
       });
-      expect(mocks.transports[0].finishAuth).toHaveBeenCalledWith("abc");
+      expect(mocks.transports[0].finishAuth).toHaveBeenCalledWith(
+        expect.any(URLSearchParams),
+      );
       expect(mocks.transports[0].close).toHaveBeenCalledTimes(1);
     } finally {
       root.unmount();
@@ -634,13 +858,15 @@ describe("McpServerResource completeAuth", () => {
           resolveFinishAuth = resolve;
         }),
     );
-    const root = mount({ auth: { type: "oauth" } });
+    const storage = createStorage();
+    vi.mocked(storage.loadAuthState).mockResolvedValue({ state: "expected" });
+    const root = mount({ auth: { type: "oauth" }, storage });
     let didUnmount = false;
 
     try {
       const completeAuth = root
         .getValue()
-        .completeAuth("https://example.com/callback?code=abc");
+        .completeAuth("https://example.com/callback?code=abc&state=expected");
       await waitFor(
         () => mocks.transports[0]?.finishAuth.mock.calls.length === 1,
       );
@@ -665,13 +891,15 @@ describe("McpServerResource completeAuth", () => {
           rejectFinishAuth = reject;
         }),
     );
-    const root = mount({ auth: { type: "oauth" } });
+    const storage = createStorage();
+    vi.mocked(storage.loadAuthState).mockResolvedValue({ state: "expected" });
+    const root = mount({ auth: { type: "oauth" }, storage });
     let didUnmount = false;
 
     try {
       const completeAuth = root
         .getValue()
-        .completeAuth("https://example.com/callback?code=abc");
+        .completeAuth("https://example.com/callback?code=abc&state=expected");
       await waitFor(
         () => mocks.transports[0]?.finishAuth.mock.calls.length === 1,
       );
@@ -1481,3 +1709,136 @@ describe("McpServerResource resource methods", () => {
     }
   });
 });
+
+describe("getConnectionDependencies storage scope", () => {
+  const propsWith = (storage: MCPStorage): McpServerResourceProps => ({
+    id: "docs",
+    kind: "connector",
+    name: "Docs",
+    url: "https://example.com/mcp",
+    auth: { type: "oauth" },
+    storage,
+    redirectUri: "https://example.com/callback",
+    autoConnect: false,
+    onRemove: async () => {},
+  });
+
+  it("keys the connection on a declared storage scopeId", () => {
+    const a = { ...createStorage(), scopeId: "local-storage:a" };
+    const b = { ...createStorage(), scopeId: "local-storage:b" };
+
+    expect(getConnectionDependencies(propsWith(a))).not.toEqual(
+      getConnectionDependencies(propsWith(b)),
+    );
+  });
+
+  it("treats storages sharing a scopeId as the same connection target", () => {
+    const a = { ...createStorage(), scopeId: "local-storage:same" };
+    const b = { ...createStorage(), scopeId: "local-storage:same" };
+
+    expect(getConnectionDependencies(propsWith(a))).toEqual(
+      getConnectionDependencies(propsWith(b)),
+    );
+  });
+
+  it("does not key the connection on storage identity when no scopeId is declared", () => {
+    expect(getConnectionDependencies(propsWith(createStorage()))).toEqual(
+      getConnectionDependencies(propsWith(createStorage())),
+    );
+  });
+
+  it("ignores the storage scope for none-auth servers", () => {
+    const a = { ...createStorage(), scopeId: "local-storage:a" };
+    const b = { ...createStorage(), scopeId: "local-storage:b" };
+    const noneProps = (storage: MCPStorage): McpServerResourceProps => ({
+      ...propsWith(storage),
+      auth: { type: "none" },
+    });
+
+    expect(getConnectionDependencies(noneProps(a))).toEqual(
+      getConnectionDependencies(noneProps(b)),
+    );
+  });
+});
+
+describe("McpServerResource oauth storage swap", () => {
+  beforeEach(resetMocks);
+
+  it("reconnects onto the replacement storage when the scope changes", async () => {
+    const persisted = {
+      tokens: { access_token: "tok", token_type: "bearer" },
+    };
+    const storageA = {
+      ...createStorage(),
+      scopeId: "scope:a",
+      loadAuthState: vi.fn(async () => persisted),
+    };
+    const storageB = {
+      ...createStorage(),
+      scopeId: "scope:b",
+      loadAuthState: vi.fn(async () => persisted),
+    };
+    let setStorage!: (s: MCPStorage) => void;
+
+    const Host = resource(function useHost() {
+      const [storage, set] = useState<MCPStorage>(storageA);
+      setStorage = set;
+      return useResource(
+        McpServerResource({
+          id: "docs",
+          kind: "connector",
+          name: "Docs",
+          url: "https://example.com/mcp",
+          auth: { type: "oauth" },
+          storage,
+          redirectUri: "https://example.com/callback",
+          autoConnect: true,
+          connectionTimeout: 10_000,
+          onRemove: vi.fn(async () => {}),
+        }),
+      );
+    });
+
+    const root = createTapRoot(function SwapRoot() {
+      return useResource(Host());
+    });
+
+    try {
+      await waitFor(() => mocks.transports.length === 1);
+      const authProviderA =
+        mocks.StreamableHTTPClientTransport.mock.calls[0]?.[1]?.authProvider;
+      await authProviderA.tokens();
+      expect(storageA.loadAuthState).toHaveBeenCalledWith("docs");
+      expect(storageB.loadAuthState).not.toHaveBeenCalled();
+
+      setStorage(storageB);
+      await waitForResourceUpdate(() => mocks.transports.length === 2);
+      await waitForResourceUpdate(
+        () => vi.mocked(mocks.transports[0]!.close).mock.calls.length > 0,
+      );
+
+      const authProviderB =
+        mocks.StreamableHTTPClientTransport.mock.calls[1]?.[1]?.authProvider;
+      const callsBefore = vi.mocked(storageA.loadAuthState).mock.calls.length;
+      await authProviderB.tokens();
+      expect(storageB.loadAuthState).toHaveBeenCalledWith("docs");
+      expect(vi.mocked(storageA.loadAuthState).mock.calls.length).toBe(
+        callsBefore,
+      );
+
+      await authProviderB.saveTokens({
+        access_token: "fresh",
+        token_type: "bearer",
+      });
+      expect(storageB.saveAuthState).toHaveBeenCalledWith(
+        "docs",
+        expect.objectContaining({
+          tokens: expect.objectContaining({ access_token: "fresh" }),
+        }),
+      );
+      expect(storageA.saveAuthState).not.toHaveBeenCalled();
+    } finally {
+      root.unmount();
+    }
+  });
+});

package/src/resources/McpServerResource.ts

@@ -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";
@@ -46,13 +49,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 +66,7 @@ const getConnectionDependencies = (
             auth.registrationEndpoint,
             auth.clientId,
             auth.clientSecret,
+            props.storage.scopeId,
           ]
         : [auth.type];
 
@@ -102,6 +106,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,
@@ -464,6 +473,45 @@ 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 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();
@@ -472,9 +520,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 +531,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);
@@ -536,6 +581,11 @@ const useMcpServerResourceInstance = (
       } else if (!persisted?.token) {
         return;
       }
+      const pendingAuthValidation = pendingAuthValidationRef.current;
+      if (pendingAuthValidation) {
+        await pendingAuthValidation.promise;
+        if (signal.cancelled || !isCurrentConnection(generation)) return;
+      }
       void doConnect();
     },
   );
@@ -597,7 +647,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({
@@ -702,8 +752,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,