@assistant-ui/react-mcp 0.1.14 → 0.1.16

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.test.ts

@@ -147,6 +147,134 @@ describe("McpManagerResource server ids", () => {
       root.unmount();
     }
   });
+
+  it("replaces a connected transport when connector settings change", async () => {
+    mocks.StreamableHTTPClientTransport.mockClear();
+    let updateConnector = (_connector: MCPConnector) => {};
+    const DynamicManager = resource(function useDynamicManager() {
+      const [currentConnector, setCurrentConnector] = useState(
+        connector("docs"),
+      );
+      updateConnector = setCurrentConnector;
+
+      return useResource(
+        McpManagerResource({
+          connectors: [currentConnector],
+          storage: McpMemoryStorage(),
+        }),
+      );
+    });
+    const root = createTapRoot(function Root() {
+      return useResource(DynamicManager());
+    });
+    let resolveFirstClose = () => {};
+
+    try {
+      await vi.waitFor(() =>
+        expect(mocks.StreamableHTTPClientTransport).toHaveBeenCalledOnce(),
+      );
+      const firstTransport = mocks.StreamableHTTPClientTransport.mock
+        .instances[0] as { close: ReturnType<typeof vi.fn> };
+      firstTransport.close.mockImplementation(
+        () =>
+          new Promise<void>((resolve) => {
+            resolveFirstClose = resolve;
+          }),
+      );
+
+      updateConnector(
+        defineConnector({
+          id: "docs",
+          name: "Docs",
+          url: "https://other.example.com/docs/mcp",
+          auth: { type: "none" },
+        }),
+      );
+
+      await vi.waitFor(() =>
+        expect(firstTransport.close).toHaveBeenCalledOnce(),
+      );
+      expect(mocks.StreamableHTTPClientTransport).toHaveBeenCalledOnce();
+
+      resolveFirstClose();
+      await vi.waitFor(() =>
+        expect(mocks.StreamableHTTPClientTransport).toHaveBeenCalledTimes(2),
+      );
+
+      expect(mocks.StreamableHTTPClientTransport).toHaveBeenLastCalledWith(
+        new URL("https://other.example.com/docs/mcp"),
+      );
+    } finally {
+      resolveFirstClose();
+      root.unmount();
+    }
+  });
+
+  it("keeps a connection across equivalent and cosmetic connector updates", async () => {
+    mocks.StreamableHTTPClientTransport.mockClear();
+    let rerenderEquivalent = () => {};
+    let updatePresentation = () => {};
+    const DynamicManager = resource(function useDynamicManager() {
+      const [, setVersion] = useState(0);
+      const [presentation, setPresentation] = useState({
+        name: "Docs",
+        icon: "docs.svg",
+      });
+      rerenderEquivalent = () => setVersion((version) => version + 1);
+      updatePresentation = () =>
+        setPresentation({ name: "Documentation", icon: "book.svg" });
+
+      return useResource(
+        McpManagerResource({
+          connectors: [
+            defineConnector({
+              id: "docs",
+              name: presentation.name,
+              icon: presentation.icon,
+              url: "https://example.com/docs/mcp",
+              auth: { type: "none" },
+            }),
+          ],
+          storage: McpCustomStorage({
+            loadCustomServers: vi.fn(async () => []),
+            saveCustomServers: vi.fn(async () => {}),
+            loadAuthState: vi.fn(async () => null),
+            saveAuthState: vi.fn(async () => {}),
+            clearAuthState: vi.fn(async () => {}),
+          }),
+          autoConnect: false,
+        }),
+      );
+    });
+    const root = createTapRoot(function Root() {
+      return useResource(DynamicManager());
+    });
+
+    try {
+      await root.getValue().connector({ index: 0 }).connect();
+      const transport = mocks.StreamableHTTPClientTransport.mock
+        .instances[0] as { close: ReturnType<typeof vi.fn> };
+
+      rerenderEquivalent();
+      await new Promise((resolve) => setTimeout(resolve, 0));
+      expect(mocks.StreamableHTTPClientTransport).toHaveBeenCalledOnce();
+      expect(transport.close).not.toHaveBeenCalled();
+
+      updatePresentation();
+      await vi.waitFor(() =>
+        expect(
+          root.getValue().connector({ index: 0 }).getState(),
+        ).toMatchObject({
+          name: "Documentation",
+          icon: "book.svg",
+        }),
+      );
+      expect(mocks.StreamableHTTPClientTransport).toHaveBeenCalledOnce();
+      expect(transport.close).not.toHaveBeenCalled();
+    } finally {
+      root.unmount();
+    }
+  });
 });
 
 describe("McpManagerResource storage failures", () => {

package/src/resources/McpManagerResource.ts

@@ -8,6 +8,8 @@ import {
 } from "@assistant-ui/store";
 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";
@@ -275,10 +277,7 @@ const useMcpManagerResource = (
       elicitation,
     }) => {
       const record: MCPCustomServerRecord = {
-        id:
-          typeof crypto !== "undefined" && "randomUUID" in crypto
-            ? crypto.randomUUID()
-            : `mcp-${Date.now()}-${Math.random().toString(36).slice(2)}`,
+        id: createMcpId(),
         name,
         url,
         auth: auth as MCPAuthConfig,
@@ -305,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));
       }
     },