@assistant-ui/react-mcp 0.1.14 → 0.1.16

package/src/resources/McpServerResource.test.ts

@@ -1,9 +1,10 @@
-import { createTapRoot, useResource } from "@assistant-ui/tap";
+import { createTapRoot, resource, useResource } from "@assistant-ui/tap";
 import type { ClientOutput } from "@assistant-ui/store";
-import { useEffect } from "react";
+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>(() => {});
 
@@ -444,6 +446,47 @@ describe("McpServerResource connectionTimeout", () => {
 describe("McpServerResource connection lifecycle", () => {
   beforeEach(resetMocks);
 
+  it("replaces direct resource connections when the server id changes", async () => {
+    const storage = createStorage();
+    let updateId = (_id: string) => {};
+    const DynamicServer = resource(function useDynamicServer() {
+      const [id, setId] = useState("docs");
+      updateId = setId;
+      return useResource(
+        McpServerResource({
+          id,
+          kind: "connector",
+          name: "Docs",
+          url: "https://example.com/mcp",
+          auth: { type: "none" },
+          storage,
+          redirectUri: "https://example.com/callback",
+          autoConnect: true,
+          onRemove: vi.fn(async () => {}),
+        }),
+      );
+    });
+    const root = createTapRoot(function Root() {
+      return useResource(DynamicServer());
+    });
+
+    try {
+      await waitForResourceUpdate(() => mocks.transports.length === 1);
+      const firstTransport = mocks.transports[0];
+
+      updateId("internal-docs");
+
+      await waitForResourceUpdate(
+        () =>
+          firstTransport.close.mock.calls.length === 1 &&
+          mocks.transports.length === 2,
+      );
+      expect(root.getValue().getState().id).toBe("internal-docs");
+    } finally {
+      root.unmount();
+    }
+  });
+
   it("closes a pending connection when the resource unmounts", async () => {
     let resolveConnect!: () => void;
     mocks.connectResults.push(
@@ -518,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",
       );
     });
 
@@ -534,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();
@@ -564,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();
 
@@ -578,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();
@@ -593,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,
       );
@@ -624,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,
       );
@@ -1440,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();
+    }
+  });
+});