@assistant-ui/react-mcp 0.1.11 → 0.1.13

package/src/primitives/addForm/McpAddFormRoot.test.tsx

@@ -0,0 +1,80 @@
+// @vitest-environment jsdom
+
+import {
+  cleanup,
+  fireEvent,
+  render,
+  screen,
+  waitFor,
+} from "@testing-library/react";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+const mocks = vi.hoisted(() => ({
+  addCustomServer: vi.fn(),
+}));
+
+vi.mock("@assistant-ui/store", async (importOriginal) => ({
+  ...(await importOriginal()),
+  useAui: () => ({
+    mcp: { addCustomServer: mocks.addCustomServer },
+  }),
+}));
+
+import { McpAddFormPrimitiveError } from "./McpAddFormError";
+import { McpAddFormPrimitiveNameField } from "./McpAddFormNameField";
+import { McpAddFormPrimitiveRoot } from "./McpAddFormRoot";
+import { McpAddFormPrimitiveSubmit } from "./McpAddFormSubmit";
+import { McpAddFormPrimitiveUrlField } from "./McpAddFormUrlField";
+
+describe("McpAddFormPrimitiveRoot", () => {
+  beforeEach(() => {
+    mocks.addCustomServer.mockReset();
+    mocks.addCustomServer.mockResolvedValue("server-1");
+  });
+
+  afterEach(() => {
+    cleanup();
+    vi.restoreAllMocks();
+  });
+
+  it.each(["throws", "rejects"] as const)(
+    "does not turn a successful add into an error when onSubmitted %s",
+    async (mode) => {
+      const callbackError = new Error("navigation failed");
+      const consoleError = vi
+        .spyOn(console, "error")
+        .mockImplementation(() => undefined);
+
+      render(
+        <McpAddFormPrimitiveRoot
+          onSubmitted={() => {
+            if (mode === "throws") throw callbackError;
+            return Promise.reject(callbackError);
+          }}
+        >
+          <McpAddFormPrimitiveNameField aria-label="Name" />
+          <McpAddFormPrimitiveUrlField aria-label="URL" />
+          <McpAddFormPrimitiveError />
+          <McpAddFormPrimitiveSubmit>Submit</McpAddFormPrimitiveSubmit>
+        </McpAddFormPrimitiveRoot>,
+      );
+
+      fireEvent.change(screen.getByLabelText("Name"), {
+        target: { value: "Docs" },
+      });
+      fireEvent.change(screen.getByLabelText("URL"), {
+        target: { value: "https://example.com/mcp" },
+      });
+      fireEvent.click(screen.getByRole("button", { name: "Submit" }));
+
+      await waitFor(() => expect(mocks.addCustomServer).toHaveBeenCalledOnce());
+      expect(screen.queryByText(callbackError.message)).toBeNull();
+      await waitFor(() => {
+        expect(consoleError).toHaveBeenCalledWith(
+          "[react-mcp] onSubmitted callback threw an error",
+          callbackError,
+        );
+      });
+    },
+  );
+});

package/src/primitives/addForm/McpAddFormRoot.tsx

@@ -11,6 +11,7 @@ import { Primitive } from "@radix-ui/react-primitive";
 import { useAui } from "@assistant-ui/store";
 import { AddFormContext, type AddFormState } from "./context";
 import type { MCPAuthConfig } from "../../mcp-scope";
+import { invokeMcpCallback } from "../../utils/invokeMcpCallback";
 
 const INITIAL: AddFormState = {
   name: "",
@@ -106,7 +107,7 @@ export const McpAddFormPrimitiveRoot = forwardRef<
         auth: buildAuth(),
       });
       setState(INITIAL);
-      onSubmitted?.(id);
+      invokeMcpCallback("onSubmitted", onSubmitted, id);
     } catch (err) {
       setState((p) => ({
         ...p,

package/src/resources/McpManagerResource.test.ts

@@ -36,6 +36,27 @@ vi.mock("@assistant-ui/store", async (importOriginal) => ({
   useAssistantClientRef: () => ({ current: null }),
 }));
 
+vi.mock("@assistant-ui/store/client", async (importOriginal) => {
+  const actual =
+    await importOriginal<typeof import("@assistant-ui/store/client")>();
+  const { useEffect } = await import("react");
+  const useScopeEffectShim = (
+    _scope: string,
+    effect: () => (() => void) | void,
+    deps: readonly unknown[],
+  ) => {
+    useEffect(() => {
+      const cleanup = effect();
+      return typeof cleanup === "function" ? cleanup : undefined;
+      // oxlint-disable-next-line react-hooks/exhaustive-deps -- caller-provided deps, mirrors the real hook
+    }, deps);
+  };
+  return {
+    ...actual,
+    useAssistantScopeEffect: useScopeEffectShim,
+  };
+});
+
 const connector = (id: string, name = id): MCPConnector =>
   defineConnector({
     id,

package/src/resources/McpManagerResource.ts

@@ -6,6 +6,7 @@ import {
   attachTransformScopes,
   type ClientOutput,
 } from "@assistant-ui/store";
+import { useAssistantScopeEffect } from "@assistant-ui/store/client";
 import { ModelContext } from "@assistant-ui/core/store";
 import type { Tool } from "assistant-stream";
 import { McpServerResource } from "./McpServerResource";
@@ -234,13 +235,17 @@ const useMcpManagerResource = (
 
   const clientRef = useAssistantClientRef();
 
-  useEffect(() => {
-    const client = clientRef.current;
-    if (!client) return;
-    return client.modelContext.register({
-      getModelContext: () => ({ tools: toolkit }),
-    });
-  }, [toolkit, clientRef]);
+  useAssistantScopeEffect(
+    "modelContext",
+    () => {
+      const client = clientRef.current;
+      if (!client) return;
+      return client.modelContext.register({
+        getModelContext: () => ({ tools: toolkit }),
+      });
+    },
+    [toolkit],
+  );
 
   const serverByKind = (kind: "connector" | "custom", index: number) => {
     const list = kind === "connector" ? state.connectors : state.customServers;

package/src/resources/McpServerResource.test.ts

@@ -147,6 +147,8 @@ const resetMocks = () => {
 const mount = (
   props?: {
     auth?: MCPAuthConfig | undefined;
+    storage?: MCPStorage | undefined;
+    autoConnect?: boolean | undefined;
     connectionTimeout?: number | undefined;
     cache?: { readonly defaultTtlMs?: number } | undefined;
     elicitation?: boolean | undefined;
@@ -169,7 +171,7 @@ const mount = (
         auth: props?.auth ?? { type: "none" },
         storage: props?.storage ?? createStorage(),
         redirectUri: "https://example.com/callback",
-        autoConnect: false,
+        autoConnect: props?.autoConnect ?? false,
         connectionTimeout,
         cache: props?.cache,
         ...(props?.elicitation !== undefined
@@ -185,6 +187,144 @@ const mount = (
   });
 };
 
+describe("McpServerResource automatic authentication", () => {
+  beforeEach(resetMocks);
+
+  it("reports auth storage load failures", async () => {
+    const storage = createStorage();
+    vi.mocked(storage.loadAuthState).mockRejectedValue(
+      new Error("auth storage unavailable"),
+    );
+    const root = mount({
+      auth: { type: "oauth" },
+      storage,
+      autoConnect: true,
+    });
+
+    try {
+      await waitForResourceUpdate(
+        () => root.getValue().getState().connectionState === "error",
+      );
+
+      expect(storage.loadAuthState).toHaveBeenCalledWith("docs");
+      expect(root.getValue().getState()).toMatchObject({
+        connectionState: "error",
+        tools: [],
+        lastError: {
+          message:
+            'MCP server "docs" failed to load saved authentication: auth storage unavailable',
+        },
+      });
+      expect(mocks.StreamableHTTPClientTransport).not.toHaveBeenCalled();
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("ignores auth storage failures from the cancelled StrictMode setup", async () => {
+    let rejectCancelledLoad!: (error: Error) => void;
+    const storage = createStorage();
+    vi.mocked(storage.loadAuthState)
+      .mockImplementationOnce(
+        () =>
+          new Promise((_, reject) => {
+            rejectCancelledLoad = reject;
+          }),
+      )
+      .mockResolvedValueOnce(null);
+    const root = mount({
+      auth: { type: "oauth" },
+      storage,
+      autoConnect: true,
+    });
+
+    try {
+      await waitFor(
+        () => vi.mocked(storage.loadAuthState).mock.calls.length > 1,
+      );
+      rejectCancelledLoad(new Error("cancelled auth storage failure"));
+      await flushMacrotask();
+
+      expect(root.getValue().getState()).toMatchObject({
+        connectionState: "disconnected",
+        lastError: null,
+      });
+      expect(mocks.StreamableHTTPClientTransport).not.toHaveBeenCalled();
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("ignores auth storage failures superseded by a manual connection", async () => {
+    let rejectLoad!: (error: Error) => void;
+    const storage = createStorage();
+    vi.mocked(storage.loadAuthState).mockImplementation(
+      () =>
+        new Promise((_, reject) => {
+          rejectLoad = reject;
+        }),
+    );
+    const root = mount({
+      auth: { type: "oauth" },
+      storage,
+      autoConnect: true,
+    });
+
+    try {
+      await waitFor(
+        () => vi.mocked(storage.loadAuthState).mock.calls.length > 0,
+      );
+      await root.getValue().connect();
+      await waitForResourceUpdate(
+        () => root.getValue().getState().connectionState === "connected",
+      );
+
+      rejectLoad(new Error("superseded auth storage failure"));
+      await flushMacrotask();
+
+      expect(root.getValue().getState()).toMatchObject({
+        connectionState: "connected",
+        lastError: null,
+      });
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("ignores successful auth storage loads superseded by disconnect", async () => {
+    let resolveLoad!: (value: { token: string }) => void;
+    const storage = createStorage();
+    vi.mocked(storage.loadAuthState).mockImplementation(
+      () =>
+        new Promise((resolve) => {
+          resolveLoad = resolve;
+        }),
+    );
+    const root = mount({
+      auth: { type: "bearer" },
+      storage,
+      autoConnect: true,
+    });
+
+    try {
+      await waitFor(
+        () => vi.mocked(storage.loadAuthState).mock.calls.length > 0,
+      );
+      await root.getValue().disconnect();
+      resolveLoad({ token: "secret" });
+      await flushMacrotask();
+
+      expect(root.getValue().getState()).toMatchObject({
+        connectionState: "disconnected",
+        lastError: null,
+      });
+      expect(mocks.StreamableHTTPClientTransport).not.toHaveBeenCalled();
+    } finally {
+      root.unmount();
+    }
+  });
+});
+
 describe("McpServerResource connectionTimeout", () => {
   beforeEach(() => {
     vi.useFakeTimers();

package/src/resources/McpServerResource.ts

@@ -474,8 +474,20 @@ const useMcpServerResource = (
         void doConnect();
         return;
       }
-      const persisted = await props.storage.loadAuthState(props.id);
-      if (signal.cancelled) return;
+      const generation = connectionGenerationRef.current;
+      let persisted: Awaited<ReturnType<MCPStorage["loadAuthState"]>>;
+      try {
+        persisted = await props.storage.loadAuthState(props.id);
+      } catch (error) {
+        if (signal.cancelled || !isCurrentConnection(generation)) return;
+        const message = error instanceof Error ? error.message : String(error);
+        setLastError({
+          message: `MCP server "${props.id}" failed to load saved authentication: ${message}`,
+        });
+        setConnectionState("error");
+        return;
+      }
+      if (signal.cancelled || !isCurrentConnection(generation)) return;
       if (props.auth.type === "oauth") {
         if (!persisted?.tokens) return;
       } else if (!persisted?.token) {

package/src/resources/storage/McpLocalStorage.test.ts

@@ -1,5 +1,7 @@
 import { createTapRoot, useResource } from "@assistant-ui/tap";
+import { auth, type FetchLike } from "@modelcontextprotocol/client";
 import { describe, expect, it } from "vitest";
+import { createOAuthProvider } from "../../auth/createOAuthProvider";
 
 import {
   McpLocalStorage,
@@ -155,6 +157,103 @@ describe("normalizePersistedAuthState", () => {
     });
   });
 
+  it("keeps valid OAuth discovery state", () => {
+    const discoveryState = {
+      authorizationServerUrl: "https://auth.example.com",
+      resourceMetadataUrl:
+        "https://mcp.example.com/.well-known/oauth-protected-resource",
+      authorizationServerMetadata: {
+        issuer: "https://auth.example.com",
+        authorization_endpoint: "https://auth.example.com/authorize",
+        token_endpoint: "https://auth.example.com/token",
+        response_types_supported: ["code"],
+      },
+      resourceMetadata: {
+        resource: "https://mcp.example.com",
+        authorization_servers: ["https://auth.example.com"],
+      },
+    };
+
+    expect(normalizePersistedAuthState({ discoveryState })).toEqual({
+      discoveryState,
+    });
+  });
+
+  it("drops malformed OAuth discovery state", () => {
+    expect(
+      normalizePersistedAuthState({
+        token: "bearer-token",
+        discoveryState: {
+          authorizationServerUrl: "not-a-url",
+        },
+      }),
+    ).toEqual({ token: "bearer-token" });
+  });
+
+  it("keeps the URL binding when optional discovery fields are malformed", () => {
+    expect(
+      normalizePersistedAuthState({
+        token: "bearer-token",
+        discoveryState: {
+          authorizationServerUrl: "https://auth.example.com",
+          authorizationServerMetadata: { issuer: 123 },
+          resourceMetadata: { resource: 42 },
+        },
+      }),
+    ).toEqual({
+      token: "bearer-token",
+      discoveryState: {
+        authorizationServerUrl: "https://auth.example.com",
+      },
+    });
+  });
+
+  it.each([
+    "http://auth.example.com",
+    "data:text/plain,auth",
+    "file:///tmp/auth",
+  ])("drops discovery state with an unsafe URL: %s", (url) => {
+    expect(
+      normalizePersistedAuthState({
+        token: "bearer-token",
+        discoveryState: { authorizationServerUrl: url },
+      }),
+    ).toEqual({ token: "bearer-token" });
+  });
+
+  it.each([
+    "http://localhost:3000",
+    "http://foo.localhost:3000",
+    "http://127.0.0.1:3000",
+    "http://127.0.0.2:3000",
+    "http://[::1]:3000",
+  ])("keeps loopback HTTP discovery URLs: %s", (url) => {
+    expect(
+      normalizePersistedAuthState({
+        discoveryState: { authorizationServerUrl: url },
+      }),
+    ).toEqual({
+      discoveryState: { authorizationServerUrl: url },
+    });
+  });
+
+  it("drops an unsafe resource metadata URL but keeps the binding", () => {
+    expect(
+      normalizePersistedAuthState({
+        token: "bearer-token",
+        discoveryState: {
+          authorizationServerUrl: "https://auth.example.com",
+          resourceMetadataUrl: "http://mcp.example.com/oauth-resource",
+        },
+      }),
+    ).toEqual({
+      token: "bearer-token",
+      discoveryState: {
+        authorizationServerUrl: "https://auth.example.com",
+      },
+    });
+  });
+
   it("keeps valid fields when neighboring fields are malformed", () => {
     expect(
       normalizePersistedAuthState({
@@ -237,4 +336,85 @@ describe("McpLocalStorage auth state", () => {
       token: "bearer-token",
     });
   });
+
+  it("completes an SDK callback with discovery state from localStorage", async () => {
+    const storage = createStorage();
+    const requests: Array<{ method: string; url: string }> = [];
+    const fetchFn: FetchLike = async (input, init) => {
+      const url = new URL(
+        input instanceof Request ? input.url : input.toString(),
+      );
+      requests.push({ method: init?.method ?? "GET", url: url.toString() });
+
+      if (url.pathname === "/.well-known/oauth-protected-resource") {
+        return Response.json({
+          resource: "https://mcp.example.com/mcp",
+          authorization_servers: ["https://auth.example.com"],
+        });
+      }
+      if (url.pathname === "/.well-known/oauth-authorization-server") {
+        return Response.json({
+          issuer: "https://auth.example.com",
+          authorization_endpoint: "https://auth.example.com/authorize",
+          token_endpoint: "https://auth.example.com/token",
+          response_types_supported: ["code"],
+          code_challenge_methods_supported: ["S256"],
+        });
+      }
+      if (url.pathname === "/token") {
+        return Response.json({
+          access_token: "access-token",
+          token_type: "Bearer",
+        });
+      }
+      throw new Error(`Unexpected OAuth request: ${url}`);
+    };
+    const authorizationUrls: URL[] = [];
+    const createProvider = () =>
+      createOAuthProvider({
+        serverId: "docs",
+        config: { type: "oauth", clientId: "client-id" },
+        storage: loadStorage(storage),
+        redirectUri: "http://localhost/callback",
+        onAuthorizationUrl: (url) => authorizationUrls.push(url),
+      });
+    const redirectProvider = createProvider();
+
+    await expect(
+      auth(redirectProvider, {
+        serverUrl: "https://mcp.example.com/mcp",
+        resourceMetadataUrl: new URL(
+          "https://mcp.example.com/.well-known/oauth-protected-resource",
+        ),
+        fetchFn,
+      }),
+    ).resolves.toBe("REDIRECT");
+    expect(authorizationUrls).toHaveLength(1);
+    expect(
+      JSON.parse(storage.getItem("test-mcp:auth:docs") ?? "null"),
+    ).toMatchObject({
+      codeVerifier: expect.any(String),
+      discoveryState: {
+        authorizationServerUrl: "https://auth.example.com",
+      },
+    });
+
+    requests.length = 0;
+    const callbackProvider = createProvider();
+    await expect(
+      auth(callbackProvider, {
+        serverUrl: "https://mcp.example.com/mcp",
+        authorizationCode: "authorization-code",
+        fetchFn,
+      }),
+    ).resolves.toBe("AUTHORIZED");
+    expect(requests).toEqual([
+      { method: "POST", url: "https://auth.example.com/token" },
+    ]);
+    await expect(
+      loadStorage(storage).loadAuthState("docs"),
+    ).resolves.toMatchObject({
+      tokens: { access_token: "access-token" },
+    });
+  });
 });

package/src/resources/storage/McpLocalStorage.ts

@@ -1,6 +1,8 @@
 import { resource } from "@assistant-ui/tap";
 import {
+  OAuthMetadataSchema,
   OAuthClientInformationFullSchema,
+  OAuthProtectedResourceMetadataSchema,
   OAuthTokensSchema,
 } from "@modelcontextprotocol/core";
 import type { MCPAuthConfig, MCPCustomServerRecord } from "../../mcp-scope";
@@ -133,6 +135,54 @@ const normalizeClientInformation = (
   return result.success ? result.data : undefined;
 };
 
+const isSecureNetworkUrl = (value: unknown): value is string => {
+  if (!isNonEmptyString(value)) return false;
+  try {
+    const url = new URL(value);
+    return (
+      url.protocol === "https:" ||
+      (url.protocol === "http:" &&
+        (url.hostname === "localhost" ||
+          url.hostname.endsWith(".localhost") ||
+          url.hostname.startsWith("127.") ||
+          url.hostname === "[::1]"))
+    );
+  } catch {
+    return false;
+  }
+};
+
+const normalizeDiscoveryState = (
+  value: unknown,
+): MCPPersistedAuthState["discoveryState"] | undefined => {
+  if (!isRecord(value) || !isSecureNetworkUrl(value.authorizationServerUrl)) {
+    return undefined;
+  }
+
+  // A malformed optional field is dropped alone: keeping the validated
+  // authorization server URL preserves the redirect-time binding, and the SDK
+  // re-discovers whatever metadata is missing.
+  const state: NonNullable<MCPPersistedAuthState["discoveryState"]> = {
+    authorizationServerUrl: value.authorizationServerUrl,
+  };
+
+  if (isSecureNetworkUrl(value.resourceMetadataUrl)) {
+    state.resourceMetadataUrl = value.resourceMetadataUrl;
+  }
+
+  const metadata = OAuthMetadataSchema.safeParse(
+    value.authorizationServerMetadata,
+  );
+  if (metadata.success) state.authorizationServerMetadata = metadata.data;
+
+  const resourceMetadata = OAuthProtectedResourceMetadataSchema.safeParse(
+    value.resourceMetadata,
+  );
+  if (resourceMetadata.success) state.resourceMetadata = resourceMetadata.data;
+
+  return state;
+};
+
 export const normalizePersistedAuthState = (
   value: unknown,
 ): MCPPersistedAuthState | null => {
@@ -150,6 +200,9 @@ export const normalizePersistedAuthState = (
   const clientInformation = normalizeClientInformation(value.clientInformation);
   if (clientInformation) state.clientInformation = clientInformation;
 
+  const discoveryState = normalizeDiscoveryState(value.discoveryState);
+  if (discoveryState) state.discoveryState = discoveryState;
+
   return Object.keys(state).length > 0 ? state : null;
 };
 

package/src/utils/invokeMcpCallback.ts

@@ -0,0 +1,27 @@
+const reportCallbackError = (name: string, error: unknown) => {
+  console.error(`[react-mcp] ${name} callback threw an error`, error);
+};
+
+export const invokeMcpCallback = <TArgs extends unknown[]>(
+  name: string,
+  callback: ((...args: TArgs) => void) | undefined,
+  ...args: TArgs
+) => {
+  if (!callback) return;
+
+  try {
+    const result = callback(...args) as unknown;
+    if (
+      result !== null &&
+      (typeof result === "object" || typeof result === "function") &&
+      "then" in result &&
+      typeof result.then === "function"
+    ) {
+      void Promise.resolve(result).catch((error) => {
+        reportCallbackError(name, error);
+      });
+    }
+  } catch (error) {
+    reportCallbackError(name, error);
+  }
+};