@assistant-ui/react-mcp 0.0.17 → 0.0.19

package/src/hooks/useMcpOAuthCallback.tsx

@@ -2,6 +2,20 @@ import { type FC, type ReactNode, useEffect, useRef, useState } from "react";
 import { useAui } from "@assistant-ui/store";
 import { decodeServerIdFromState } from "../auth/createOAuthProvider";
 
+export const createMcpOAuthCallbackError = (
+  err: unknown,
+  serverId: string | null,
+): Error => {
+  const message = err instanceof Error ? err.message : String(err);
+  if (serverId) {
+    return new Error(
+      `MCP OAuth callback for server "${serverId}" failed: ${message}`,
+      { cause: err },
+    );
+  }
+  return new Error(`MCP OAuth callback failed: ${message}`, { cause: err });
+};
+
 export type UseMcpOAuthCallbackOptions = {
   /** Defaults to `window.location.href`. */
   url?: string;
@@ -39,27 +53,28 @@ export function useMcpOAuthCallback(
     startedRef.current = url;
 
     (async () => {
+      let serverId: string | null = null;
       try {
         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 in callback URL");
-        const serverId = decodeServerIdFromState(state);
+        if (!state) throw new Error('missing "state" parameter');
         if (!serverId) {
-          throw new Error("callback state does not match an MCP server");
+          throw new Error("state was not created by assistant-ui MCP");
         }
         setResult({ status: "running", serverId, error: null });
         await aui.mcp().server({ id: serverId }).completeAuth(url);
         setResult({ status: "done", serverId, error: null });
         optsRef.current.onComplete?.(serverId);
       } catch (err) {
-        const e = err instanceof Error ? err : new Error(String(err));
-        setResult((prev) => ({ ...prev, status: "error", error: e }));
+        const e = createMcpOAuthCallbackError(err, serverId);
+        setResult({ status: "error", serverId, error: e });
         optsRef.current.onError?.(e);
       }
     })();

package/src/mcp-scope.ts

@@ -72,6 +72,8 @@ export type MCPServerMethods = {
   disconnect: () => Promise<void>;
   remove: () => Promise<void>;
   callTool: (name: string, args: unknown) => Promise<unknown>;
+  /** List resources exposed by the server. Returns the raw MCP `ListResourcesResult`. */
+  listResources: (params?: { cursor?: string | undefined }) => Promise<unknown>;
   /** Read a resource by URI. Returns the raw MCP `ReadResourceResult`. */
   readResource: (uri: string) => Promise<unknown>;
   /** OAuth only: pass full callback URL (e.g. window.location.href) */

package/src/resources/McpServerResource.test.ts

@@ -1,5 +1,8 @@
 import { createTapRoot, useResource } from "@assistant-ui/tap";
+import type { ClientOutput } from "@assistant-ui/store";
+import { useEffect } from "react";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import type { MCPAuthConfig } from "../mcp-scope";
 import type { MCPStorage } from "./storage/types";
 
 const mocks = vi.hoisted(() => {
@@ -15,6 +18,8 @@ const mocks = vi.hoisted(() => {
       }>;
     }>
   > = [];
+  const finishAuthResults: Array<() => Promise<void>> = [];
+  const closeResults: Array<() => Promise<void>> = [];
 
   const Client = vi.fn().mockImplementation(function Client(this: any) {
     const index = clients.length;
@@ -23,6 +28,7 @@ const mocks = vi.hoisted(() => {
       () => listToolsResults[index]?.() ?? Promise.resolve({ tools: [] }),
     );
     this.callTool = vi.fn();
+    this.listResources = vi.fn(() => Promise.resolve({ resources: [] }));
     this.readResource = vi.fn();
     clients.push(this);
   });
@@ -30,8 +36,11 @@ const mocks = vi.hoisted(() => {
   const StreamableHTTPClientTransport = vi
     .fn()
     .mockImplementation(function StreamableHTTPClientTransport(this: any) {
-      this.close = vi.fn(() => Promise.resolve());
-      this.finishAuth = vi.fn(() => Promise.resolve());
+      const index = transports.length;
+      this.close = vi.fn(() => closeResults[index]?.() ?? Promise.resolve());
+      this.finishAuth = vi.fn(
+        () => finishAuthResults[index]?.() ?? Promise.resolve(),
+      );
       transports.push(this);
     });
 
@@ -42,6 +51,8 @@ const mocks = vi.hoisted(() => {
     transports,
     connectResults,
     listToolsResults,
+    finishAuthResults,
+    closeResults,
   };
 });
 
@@ -61,6 +72,18 @@ const tick = async () => {
   await Promise.resolve();
 };
 
+const flushMacrotask = async () => {
+  await new Promise<void>((resolve) => {
+    const channel = new MessageChannel();
+    channel.port1.onmessage = () => {
+      channel.port1.close();
+      channel.port2.close();
+      resolve();
+    };
+    channel.port2.postMessage(null);
+  });
+};
+
 const waitFor = async (predicate: () => boolean) => {
   for (let i = 0; i < 20; i++) {
     if (predicate()) return;
@@ -77,18 +100,35 @@ const createStorage = (): MCPStorage => ({
   clearAuthState: vi.fn(async () => {}),
 });
 
-const mount = (props?: { connectionTimeout?: number | undefined }) => {
+const resetMocks = () => {
+  mocks.clients.length = 0;
+  mocks.transports.length = 0;
+  mocks.connectResults.length = 0;
+  mocks.listToolsResults.length = 0;
+  mocks.finishAuthResults.length = 0;
+  mocks.closeResults.length = 0;
+  mocks.Client.mockClear();
+  mocks.StreamableHTTPClientTransport.mockClear();
+};
+
+const mount = (
+  props?: {
+    auth?: MCPAuthConfig | undefined;
+    connectionTimeout?: number | undefined;
+  },
+  onMount?: (server: ClientOutput<"mcpServer">) => void,
+) => {
   const connectionTimeout =
     props && "connectionTimeout" in props ? props.connectionTimeout : 10_000;
 
   return createTapRoot(function Root() {
-    return useResource(
+    const server = useResource(
       McpServerResource({
         id: "docs",
         kind: "connector",
         name: "Docs",
         url: "https://example.com/mcp",
-        auth: { type: "none" },
+        auth: props?.auth ?? { type: "none" },
         storage: createStorage(),
         redirectUri: "https://example.com/callback",
         autoConnect: false,
@@ -96,18 +136,17 @@ const mount = (props?: { connectionTimeout?: number | undefined }) => {
         onRemove: vi.fn(async () => {}),
       }),
     );
+    useEffect(() => {
+      onMount?.(server);
+    }, [server]);
+    return server;
   });
 };
 
 describe("McpServerResource connectionTimeout", () => {
   beforeEach(() => {
     vi.useFakeTimers();
-    mocks.clients.length = 0;
-    mocks.transports.length = 0;
-    mocks.connectResults.length = 0;
-    mocks.listToolsResults.length = 0;
-    mocks.Client.mockClear();
-    mocks.StreamableHTTPClientTransport.mockClear();
+    resetMocks();
   });
 
   afterEach(() => {
@@ -219,3 +258,280 @@ describe("McpServerResource connectionTimeout", () => {
     }
   });
 });
+
+describe("McpServerResource connection lifecycle", () => {
+  beforeEach(resetMocks);
+
+  it("closes a pending connection when the resource unmounts", async () => {
+    let resolveConnect!: () => void;
+    mocks.connectResults.push(
+      () =>
+        new Promise<void>((resolve) => {
+          resolveConnect = resolve;
+        }),
+    );
+    const root = mount({ connectionTimeout: undefined });
+    let didUnmount = false;
+    try {
+      const connectPromise = root.getValue().connect();
+      await waitFor(() => mocks.clients[0]?.connect.mock.calls.length === 1);
+
+      root.unmount();
+      didUnmount = true;
+      await flushMacrotask();
+
+      expect(mocks.transports[0].close).toHaveBeenCalledTimes(1);
+
+      resolveConnect();
+      await connectPromise;
+
+      expect(mocks.clients[0].listTools).not.toHaveBeenCalled();
+      expect(mocks.transports[0].close).toHaveBeenCalledTimes(1);
+    } finally {
+      if (!didUnmount) root.unmount();
+    }
+  });
+
+  it("waits for pending transports to close before a newer reconnect", async () => {
+    let resolveFirstConnect!: () => void;
+    let resolveFirstClose!: () => void;
+    mocks.connectResults.push(
+      () =>
+        new Promise<void>((resolve) => {
+          resolveFirstConnect = resolve;
+        }),
+    );
+    mocks.closeResults.push(
+      () =>
+        new Promise<void>((resolve) => {
+          resolveFirstClose = resolve;
+        }),
+    );
+    const root = mount({ connectionTimeout: undefined });
+
+    try {
+      const firstConnect = root.getValue().connect();
+      await waitFor(() => mocks.clients[0]?.connect.mock.calls.length === 1);
+
+      const supersededReconnect = root.getValue().connect();
+      await waitFor(() => mocks.transports[0]?.close.mock.calls.length === 1);
+
+      const latestReconnect = root.getValue().connect();
+      await flushMacrotask();
+      expect(mocks.transports).toHaveLength(1);
+
+      resolveFirstClose();
+      await supersededReconnect;
+      await waitFor(() => mocks.clients[1]?.connect.mock.calls.length === 1);
+      await latestReconnect;
+
+      resolveFirstConnect();
+      await firstConnect;
+    } finally {
+      root.unmount();
+    }
+  });
+});
+
+describe("McpServerResource completeAuth", () => {
+  beforeEach(resetMocks);
+
+  it("completes auth across the StrictMode effect replay", async () => {
+    let completeAuth: Promise<void> | undefined;
+    let started = false;
+    const root = mount({ auth: { type: "oauth" } }, (server) => {
+      if (started) return;
+      started = true;
+      completeAuth = server.completeAuth(
+        "https://example.com/callback?code=abc",
+      );
+    });
+
+    try {
+      await expect(completeAuth).resolves.toBeUndefined();
+      await flushMacrotask();
+
+      expect(mocks.transports[0].finishAuth).toHaveBeenCalledTimes(1);
+      expect(root.getValue().getState().connectionState).toBe("connected");
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("rejects when the callback URL has no authorization code", async () => {
+    const root = mount();
+
+    try {
+      await expect(
+        root.getValue().completeAuth("https://example.com/callback?state=abc"),
+      ).rejects.toThrow("missing authorization code in callback URL");
+      await flushMacrotask();
+
+      expect(root.getValue().getState()).toMatchObject({
+        connectionState: "error",
+        lastError: {
+          message: "missing authorization code in callback URL",
+        },
+      });
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("rejects after storing finishAuth failures on the server state", async () => {
+    mocks.finishAuthResults.push(() =>
+      Promise.reject(new Error("invalid_grant")),
+    );
+    const root = mount({ auth: { type: "oauth" } });
+
+    try {
+      await expect(
+        root.getValue().completeAuth("https://example.com/callback?code=abc"),
+      ).rejects.toThrow("invalid_grant");
+      await flushMacrotask();
+
+      expect(root.getValue().getState()).toMatchObject({
+        connectionState: "error",
+        lastError: {
+          message: "invalid_grant",
+        },
+      });
+      expect(mocks.transports[0].finishAuth).toHaveBeenCalledWith("abc");
+      expect(mocks.transports[0].close).toHaveBeenCalledTimes(1);
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("rejects when the resource unmounts during auth completion", async () => {
+    let resolveFinishAuth!: () => void;
+    mocks.finishAuthResults.push(
+      () =>
+        new Promise<void>((resolve) => {
+          resolveFinishAuth = resolve;
+        }),
+    );
+    const root = mount({ auth: { type: "oauth" } });
+    let didUnmount = false;
+
+    try {
+      const completeAuth = root
+        .getValue()
+        .completeAuth("https://example.com/callback?code=abc");
+      await waitFor(
+        () => mocks.transports[0]?.finishAuth.mock.calls.length === 1,
+      );
+
+      root.unmount();
+      didUnmount = true;
+      resolveFinishAuth();
+
+      await expect(completeAuth).rejects.toThrow(
+        'MCP server "docs" authorization was interrupted before completion.',
+      );
+    } finally {
+      if (!didUnmount) root.unmount();
+    }
+  });
+
+  it("normalizes late auth failures after unmount", async () => {
+    let rejectFinishAuth!: (error: Error) => void;
+    mocks.finishAuthResults.push(
+      () =>
+        new Promise<void>((_, reject) => {
+          rejectFinishAuth = reject;
+        }),
+    );
+    const root = mount({ auth: { type: "oauth" } });
+    let didUnmount = false;
+
+    try {
+      const completeAuth = root
+        .getValue()
+        .completeAuth("https://example.com/callback?code=abc");
+      await waitFor(
+        () => mocks.transports[0]?.finishAuth.mock.calls.length === 1,
+      );
+
+      root.unmount();
+      didUnmount = true;
+      const finishAuthError = new Error("Connection closed");
+      rejectFinishAuth(finishAuthError);
+
+      await expect(completeAuth).rejects.toMatchObject({
+        message:
+          'MCP server "docs" authorization was interrupted before completion.',
+        cause: finishAuthError,
+      });
+    } finally {
+      if (!didUnmount) root.unmount();
+    }
+  });
+});
+
+describe("McpServerResource resource methods", () => {
+  beforeEach(() => {
+    mocks.clients.length = 0;
+    mocks.transports.length = 0;
+    mocks.connectResults.length = 0;
+    mocks.listToolsResults.length = 0;
+    mocks.Client.mockClear();
+    mocks.StreamableHTTPClientTransport.mockClear();
+  });
+
+  it("lists resources from a connected server", async () => {
+    const result = {
+      resources: [
+        {
+          uri: "docs://intro",
+          name: "Intro",
+          mimeType: "text/markdown",
+        },
+      ],
+    };
+    const root = mount();
+
+    try {
+      await root.getValue().connect();
+      mocks.clients[0].listResources.mockResolvedValueOnce(result);
+
+      await expect(root.getValue().listResources()).resolves.toBe(result);
+      expect(mocks.clients[0].listResources).toHaveBeenCalledTimes(1);
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("forwards the resource pagination cursor", async () => {
+    const result = {
+      resources: [{ uri: "docs://page-two", name: "Page two" }],
+    };
+    const root = mount();
+
+    try {
+      await root.getValue().connect();
+      mocks.clients[0].listResources.mockResolvedValueOnce(result);
+
+      await expect(
+        root.getValue().listResources({ cursor: "next-page" }),
+      ).resolves.toBe(result);
+      expect(mocks.clients[0].listResources).toHaveBeenCalledWith({
+        cursor: "next-page",
+      });
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("rejects listResources when the server is disconnected", async () => {
+    const root = mount();
+
+    try {
+      await expect(root.getValue().listResources()).rejects.toThrow(
+        'MCP server "docs" is not connected',
+      );
+    } finally {
+      root.unmount();
+    }
+  });
+});