@assistant-ui/react-mcp 0.1.16 → 0.1.18

package/src/auth/createOAuthProvider.ts

@@ -7,6 +7,11 @@ import type {
 } from "@modelcontextprotocol/client";
 import type { MCPStorage } from "../resources/storage/types";
 import type { MCPAuthConfig } from "../mcp-scope";
+import type { MCPPersistedAuthState } from "./types";
+import {
+  isAuthStateForServerUrl,
+  normalizeMcpServerUrl,
+} from "../utils/serverUrl";
 
 const STATE_PREFIX = "aui-mcp:";
 
@@ -49,6 +54,7 @@ export function decodeServerIdFromState(state: string): string | null {
 
 export type CreateOAuthProviderOptions = {
   serverId: string;
+  serverUrl: string;
   /** Must be `auth.type === "oauth"`. */
   config: Extract<MCPAuthConfig, { type: "oauth" }>;
   storage: MCPStorage;
@@ -58,16 +64,64 @@ export type CreateOAuthProviderOptions = {
 };
 
 type OAuthProviderCache = {
+  token?: string | undefined;
   tokens?: OAuthTokens | undefined;
+  tokensClientId?: string | undefined;
   clientInformation?: OAuthClientInformationFull | undefined;
+  clientInformationSource?: MCPPersistedAuthState["clientInformationSource"];
   codeVerifier?: string | undefined;
   state?: string | undefined;
   discoveryState?: OAuthDiscoveryState | undefined;
 };
 
-type OAuthProviderPersistence = {
+type OAuthConfig = Extract<MCPAuthConfig, { type: "oauth" }>;
+
+type OAuthCredentialState = {
+  tokens?: OAuthTokens | undefined;
+  tokensClientId?: string | undefined;
+  clientInformation?: OAuthClientInformationFull | undefined;
+  clientInformationSource?: "registered" | undefined;
+};
+
+const registeredClientId = (
+  state: OAuthCredentialState | null | undefined,
+): string | undefined =>
+  state?.clientInformationSource === "registered"
+    ? state.clientInformation?.client_id
+    : undefined;
+
+export const hasUsableOAuthTokens = (
+  state: OAuthCredentialState | null | undefined,
+  config: OAuthConfig,
+): boolean => {
+  const clientId = config.clientId ?? registeredClientId(state);
+  return (
+    clientId !== undefined &&
+    state?.tokens !== undefined &&
+    state.tokensClientId === clientId
+  );
+};
+
+const hasUsableRegisteredClientInformation = (
+  state: OAuthCredentialState | null | undefined,
+  config: OAuthConfig,
+): boolean => {
+  const clientId = registeredClientId(state);
+  return (
+    clientId !== undefined &&
+    (config.clientId === undefined || config.clientId === clientId)
+  );
+};
+
+type OAuthProviderEndpointCache = {
+  serverUrl: string;
   cached: OAuthProviderCache | null;
   cachePromise: Promise<OAuthProviderCache> | null;
+  invalidated: boolean;
+};
+
+type OAuthProviderPersistence = {
+  endpoint: OAuthProviderEndpointCache | null;
   queue: Promise<void>;
   invalidated: boolean;
 };
@@ -111,7 +165,11 @@ const persistenceByIdentity = new WeakMap<
 const getPersistence = (
   storage: MCPStorage,
   serverId: string,
-): OAuthProviderPersistence => {
+  serverUrl: string,
+): {
+  persistence: OAuthProviderPersistence;
+  endpoint: OAuthProviderEndpointCache;
+} => {
   const identity = getStorageIdentity(storage);
   let byServerId = persistenceByIdentity.get(identity);
   if (!byServerId) {
@@ -122,14 +180,25 @@ const getPersistence = (
   let persistence = byServerId.get(serverId);
   if (!persistence) {
     persistence = {
-      cached: null,
-      cachePromise: null,
+      endpoint: null,
       queue: Promise.resolve(),
       invalidated: false,
     };
     byServerId.set(serverId, persistence);
   }
-  return persistence;
+
+  let endpoint = persistence.endpoint;
+  if (endpoint?.serverUrl !== serverUrl) {
+    if (endpoint) endpoint.invalidated = true;
+    endpoint = {
+      serverUrl,
+      cached: null,
+      cachePromise: null,
+      invalidated: false,
+    };
+    persistence.endpoint = endpoint;
+  }
+  return { persistence, endpoint };
 };
 
 /**
@@ -155,7 +224,10 @@ export const clearOAuthProviderAuthState = async (
   byServerId.delete(serverId);
   if (byServerId.size === 0) persistenceByIdentity.delete(identity);
 
-  await Promise.allSettled([persistence.cachePromise, persistence.queue]);
+  if (persistence.endpoint) persistence.endpoint.invalidated = true;
+  const cachePromise = persistence.endpoint?.cachePromise;
+  if (cachePromise) await Promise.allSettled([cachePromise]);
+  await persistence.queue;
   await storage.clearAuthState(serverId);
 };
 
@@ -167,15 +239,28 @@ export const clearOAuthProviderAuthState = async (
 export function createOAuthProvider(
   opts: CreateOAuthProviderOptions,
 ): OAuthClientProvider {
-  const { serverId, config, storage, redirectUri, onAuthorizationUrl } = opts;
-  const persistence = getPersistence(storage, serverId);
+  const {
+    serverId,
+    serverUrl,
+    config,
+    storage,
+    redirectUri,
+    onAuthorizationUrl,
+  } = opts;
+  const normalizedServerUrl = normalizeMcpServerUrl(serverUrl);
+  const { persistence, endpoint } = getPersistence(
+    storage,
+    serverId,
+    normalizedServerUrl,
+  );
   let pendingState: string | undefined;
 
-  // The cache is shared with every other provider for this (storage, serverId),
-  // so a statically configured client stays a read-time overlay owned by this
-  // provider. Writing it into the cache would leak this provider's registration
-  // to a replacement built for a different, or absent, clientId.
-  const staticClientInformation = (():
+  // The cache is shared with every other provider for this storage, server id,
+  // and server URL, so a statically configured client stays a read-time overlay
+  // owned by this provider. Writing it into the cache would leak this provider's
+  // registration to a replacement built for a different, or absent, clientId.
+  // The SDK's write-backs, its issuer stamp included, replace the overlay.
+  const configuredClientInformation = ():
     | OAuthClientInformationFull
     | undefined => {
     if (!config.clientId) return undefined;
@@ -185,50 +270,92 @@ export function createOAuthProvider(
     };
     if (config.clientSecret) ci.client_secret = config.clientSecret;
     return ci;
-  })();
+  };
+  let clientInformationOverlay = configuredClientInformation();
+
+  const activeClientId = (cache: OAuthProviderCache): string | undefined =>
+    clientInformationOverlay?.client_id ?? registeredClientId(cache);
 
   const loadCache = (): Promise<OAuthProviderCache> => {
-    if (persistence.cached) return Promise.resolve(persistence.cached);
-    if (persistence.cachePromise) return persistence.cachePromise;
-
-    persistence.cachePromise = storage.loadAuthState(serverId).then(
-      (persisted) => {
-        const initial: OAuthProviderCache = {};
-        if (persisted?.tokens) initial.tokens = persisted.tokens;
-        if (persisted?.clientInformation)
-          initial.clientInformation = persisted.clientInformation;
-        if (persisted?.codeVerifier)
-          initial.codeVerifier = persisted.codeVerifier;
-        if (persisted?.state) initial.state = persisted.state;
-        if (persisted?.discoveryState)
-          initial.discoveryState = persisted.discoveryState;
-        persistence.cached = initial;
-        return initial;
-      },
-      (error) => {
-        persistence.cachePromise = null;
-        throw error;
-      },
-    );
-    return persistence.cachePromise;
+    if (endpoint.invalidated) return Promise.resolve({});
+    if (endpoint.cached) return Promise.resolve(endpoint.cached);
+    if (endpoint.cachePromise) return endpoint.cachePromise;
+
+    endpoint.cachePromise = persistence.queue
+      .then(() => storage.loadAuthState(serverId))
+      .then(
+        async (persisted) => {
+          const initial: OAuthProviderCache = {};
+          let needsMigration = false;
+          if (endpoint.invalidated) return initial;
+          if (
+            persisted &&
+            isAuthStateForServerUrl(persisted, normalizedServerUrl)
+          ) {
+            if (
+              hasUsableRegisteredClientInformation(persisted, config) &&
+              persisted.clientInformation
+            ) {
+              initial.clientInformation = persisted.clientInformation;
+              initial.clientInformationSource = "registered";
+            } else if (
+              persisted?.clientInformation ||
+              persisted?.clientInformationSource !== undefined
+            ) {
+              needsMigration = true;
+            }
+            if (hasUsableOAuthTokens(persisted, config)) {
+              initial.tokens = persisted.tokens;
+              initial.tokensClientId = persisted.tokensClientId;
+            } else if (
+              persisted?.tokens ||
+              persisted?.tokensClientId !== undefined
+            ) {
+              needsMigration = true;
+            }
+            if (persisted?.token) initial.token = persisted.token;
+            if (persisted?.codeVerifier)
+              initial.codeVerifier = persisted.codeVerifier;
+            if (persisted?.state) initial.state = persisted.state;
+            if (persisted?.discoveryState)
+              initial.discoveryState = persisted.discoveryState;
+          }
+          endpoint.cached = initial;
+          if (needsMigration) await persist().catch(() => {});
+          return initial;
+        },
+        (error) => {
+          endpoint.cachePromise = null;
+          throw error;
+        },
+      );
+    return endpoint.cachePromise;
   };
 
-  const persist = () => {
+  function persist() {
     const task = persistence.queue.then(async () => {
-      if (persistence.invalidated) return;
-      const c = persistence.cached;
+      if (persistence.invalidated || endpoint.invalidated) return;
+      const c = endpoint.cached;
       if (!c) return;
       const next: Parameters<typeof storage.saveAuthState>[1] = {};
-      if (c.tokens) next.tokens = c.tokens;
-      if (c.clientInformation) next.clientInformation = c.clientInformation;
+      if (hasUsableOAuthTokens(c, config) && c.tokens && c.tokensClientId) {
+        next.tokens = c.tokens;
+        next.tokensClientId = c.tokensClientId;
+      }
+      if (c.clientInformation && c.clientInformationSource === "registered") {
+        next.clientInformation = c.clientInformation;
+        next.clientInformationSource = "registered";
+      }
+      if (c.token) next.token = c.token;
       if (c.codeVerifier) next.codeVerifier = c.codeVerifier;
       if (c.state) next.state = c.state;
       if (c.discoveryState) next.discoveryState = c.discoveryState;
+      next.serverUrl = normalizedServerUrl;
       await storage.saveAuthState(serverId, next);
     });
     persistence.queue = task.catch(() => {});
     return task;
-  };
+  }
 
   const clientMetadata: OAuthClientMetadata = {
     client_name: "assistant-ui",
@@ -258,20 +385,37 @@ export function createOAuthProvider(
     },
     async clientInformation() {
       const c = await loadCache();
-      return staticClientInformation ?? c.clientInformation;
+      if (clientInformationOverlay) return clientInformationOverlay;
+      if (c.clientInformationSource !== "registered") return undefined;
+      return c.clientInformation;
     },
     async saveClientInformation(info) {
+      if (clientInformationOverlay) {
+        clientInformationOverlay = info as OAuthClientInformationFull;
+        return;
+      }
       const c = await loadCache();
       c.clientInformation = info as OAuthClientInformationFull;
+      c.clientInformationSource = "registered";
+      if (c.tokensClientId !== c.clientInformation.client_id) {
+        delete c.tokens;
+        delete c.tokensClientId;
+      }
       await persist();
     },
     async tokens() {
       const c = await loadCache();
+      const clientId = activeClientId(c);
+      if (clientId === undefined || c.tokensClientId !== clientId)
+        return undefined;
       return c.tokens;
     },
     async saveTokens(tokens) {
       const c = await loadCache();
       c.tokens = tokens;
+      const clientId = activeClientId(c);
+      if (clientId) c.tokensClientId = clientId;
+      else delete c.tokensClientId;
       delete c.state;
       await persist();
     },
@@ -305,8 +449,19 @@ export function createOAuthProvider(
     },
     async invalidateCredentials(scope) {
       const c = await loadCache();
-      if (scope === "all" || scope === "tokens") delete c.tokens;
-      if (scope === "all" || scope === "client") delete c.clientInformation;
+      if (scope === "all" || scope === "tokens") {
+        delete c.tokens;
+        delete c.tokensClientId;
+      }
+      if (scope === "all" || scope === "client") {
+        delete c.clientInformation;
+        delete c.clientInformationSource;
+        if (!config.clientId) {
+          delete c.tokens;
+          delete c.tokensClientId;
+        }
+        clientInformationOverlay = configuredClientInformation();
+      }
       if (scope === "all" || scope === "verifier") {
         delete c.codeVerifier;
         delete c.state;

package/src/auth/types.ts

@@ -5,11 +5,15 @@ import type {
 } from "@modelcontextprotocol/client";
 
 export type MCPPersistedAuthState = {
+  /** MCP server URL this authentication state belongs to. Required with credentials. */
+  serverUrl?: string;
   tokens?: OAuthTokens;
+  tokensClientId?: string;
   clientInformation?: OAuthClientInformationFull;
+  clientInformationSource?: "registered";
   codeVerifier?: string;
   state?: string;
   discoveryState?: OAuthDiscoveryState;
-  /** Bearer token (entered at add-form time). */
+  /** Host-persisted bearer token. Must be paired with serverUrl. */
   token?: string;
 };

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

@@ -0,0 +1,261 @@
+// @vitest-environment jsdom
+
+import {
+  cleanup,
+  fireEvent,
+  render,
+  screen,
+  waitFor,
+  within,
+} from "@testing-library/react";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { createRef } from "react";
+
+const mocks = vi.hoisted(() => ({ addCustomServer: vi.fn() }));
+vi.mock("@assistant-ui/store", async (importOriginal) => ({
+  ...(await importOriginal()),
+  useAui: () => ({ mcp: { addCustomServer: mocks.addCustomServer } }),
+}));
+
+import { McpAddFormPrimitiveRoot as Root } from "./McpAddFormRoot";
+import { McpAddFormPrimitiveNameField as NameField } from "./McpAddFormNameField";
+import { McpAddFormPrimitiveUrlField as UrlField } from "./McpAddFormUrlField";
+import { McpAddFormPrimitiveAuthSelect as AuthSelect } from "./McpAddFormAuthSelect";
+import { McpAddFormPrimitiveAuthFields as AuthFields } from "./McpAddFormAuthFields";
+import { McpAddFormPrimitiveError as ErrorMessage } from "./McpAddFormError";
+
+const Form = ({ name = "Add server" }: { name?: string }) => (
+  <Root aria-label={name}>
+    <label>
+      Name
+      <NameField />
+    </label>
+    <label>
+      URL
+      <UrlField />
+    </label>
+    <label>
+      Auth
+      <AuthSelect />
+    </label>
+    <AuthFields />
+    <ErrorMessage />
+  </Root>
+);
+
+beforeEach(() => {
+  mocks.addCustomServer.mockReset().mockResolvedValue("server");
+});
+afterEach(cleanup);
+
+describe("MCP add form accessibility", () => {
+  it("does not add a wrapper around an asChild error without a custom ID", () => {
+    render(
+      <Root aria-label="Add server">
+        <NameField aria-label="Name" />
+        <ErrorMessage asChild>
+          <div>Enter the server name.</div>
+        </ErrorMessage>
+      </Root>,
+    );
+    fireEvent.submit(screen.getByRole("form"));
+    const error = screen.getByRole("alert");
+    expect(error.parentElement).toBe(screen.getByRole("form"));
+    expect(error.id).not.toBe("");
+    expect(
+      screen
+        .getByRole("textbox", { name: "Name" })
+        .getAttribute("aria-describedby"),
+    ).toBe(error.id);
+  });
+
+  it.each(["element", "slot", "child"] as const)(
+    "keeps a custom error ID on the %s associated with fields",
+    (target) => {
+      const ref = createRef<HTMLDivElement>();
+      render(
+        <Root aria-label="Add server">
+          <NameField aria-label="Name" />
+          <ErrorMessage
+            {...(target !== "child" && { id: "custom-error" })}
+            asChild={target !== "element"}
+            ref={ref}
+          >
+            {target !== "element" ? (
+              <div {...(target === "child" && { id: "custom-error" })}>
+                Enter the server name.
+              </div>
+            ) : (
+              "Enter the server name."
+            )}
+          </ErrorMessage>
+        </Root>,
+      );
+      fireEvent.submit(screen.getByRole("form"));
+      const error = screen.getByRole("alert");
+      expect(error.id).toBe("custom-error");
+      expect(ref.current).toBe(error);
+      const descriptionId = screen
+        .getByRole("textbox", { name: "Name" })
+        .getAttribute("aria-describedby")!;
+      expect(document.getElementById(descriptionId)?.textContent).toBe(
+        "Enter the server name.",
+      );
+      expect(document.querySelectorAll(`[id="${descriptionId}"]`)).toHaveLength(
+        1,
+      );
+    },
+  );
+
+  it("keeps enclosing labels as the accessible names", () => {
+    render(
+      <Root>
+        <label>
+          Server name
+          <NameField />
+        </label>
+        <label>
+          Server address
+          <UrlField />
+        </label>
+        <label>
+          Authentication method
+          <AuthSelect />
+        </label>
+      </Root>,
+    );
+    expect(screen.getByRole("textbox", { name: "Server name" })).toBeTruthy();
+    expect(
+      screen.getByRole("textbox", { name: "Server address" }),
+    ).toBeTruthy();
+    expect(
+      screen.getByRole("combobox", { name: "Authentication method" }),
+    ).toBeTruthy();
+  });
+  it("preserves field labels and only marks the field with an error", () => {
+    render(<Form />);
+    const name = screen.getByRole("textbox", { name: "Name" });
+    const url = screen.getByRole("textbox", { name: "URL" });
+    const auth = screen.getByRole("combobox", { name: "Auth" });
+    expect(screen.getByRole("textbox", { name: "OAuth scopes" })).toBeTruthy();
+
+    fireEvent.submit(screen.getByRole("form"));
+    const error = screen.getByRole("alert");
+    expect(error.textContent).toBe("Name is required");
+    expect(name.getAttribute("aria-invalid")).toBe("true");
+    expect(name.getAttribute("aria-describedby")).toBe(error.id);
+    expect(url.hasAttribute("aria-invalid")).toBe(false);
+
+    fireEvent.change(name, { target: { value: "Docs" } });
+    expect(screen.queryByRole("alert")).toBeNull();
+    expect(name.hasAttribute("aria-invalid")).toBe(false);
+    fireEvent.submit(screen.getByRole("form"));
+    expect(url.getAttribute("aria-describedby")).toBe(
+      screen.getByRole("alert").id,
+    );
+    fireEvent.change(url, { target: { value: "https://example.com/mcp" } });
+    fireEvent.change(auth, { target: { value: "bearer" } });
+    fireEvent.submit(screen.getByRole("form"));
+    const token = screen.getByLabelText("Bearer token");
+    expect(token.getAttribute("aria-invalid")).toBe("true");
+    expect(token.getAttribute("aria-describedby")).toBe(
+      screen.getByRole("alert").id,
+    );
+    fireEvent.change(auth, { target: { value: "none" } });
+    expect(screen.queryByRole("alert")).toBeNull();
+    expect(screen.queryByLabelText("Bearer token")).toBeNull();
+  });
+
+  it("announces server errors without marking unrelated fields invalid", async () => {
+    mocks.addCustomServer.mockRejectedValue(new Error("Connection failed"));
+    render(<Form />);
+    fireEvent.change(screen.getByRole("textbox", { name: "Name" }), {
+      target: { value: "Docs" },
+    });
+    fireEvent.change(screen.getByRole("textbox", { name: "URL" }), {
+      target: { value: "https://example.com/mcp" },
+    });
+    fireEvent.submit(screen.getByRole("form"));
+    await waitFor(() =>
+      expect(screen.getByRole("alert").textContent).toBe("Connection failed"),
+    );
+    expect(document.querySelector('[aria-invalid="true"]')).toBeNull();
+  });
+
+  it("uses distinct IDs for simultaneously mounted forms", () => {
+    render(
+      <>
+        <Form name="First" />
+        <Form name="Second" />
+      </>,
+    );
+    const first = screen.getByRole("form", { name: "First" });
+    const second = screen.getByRole("form", { name: "Second" });
+    fireEvent.submit(first);
+    fireEvent.submit(second);
+    const firstError = within(first).getByRole("alert");
+    const secondError = within(second).getByRole("alert");
+    expect(firstError.id).not.toBe(secondError.id);
+    expect(
+      within(first)
+        .getByRole("textbox", { name: "Name" })
+        .getAttribute("aria-describedby"),
+    ).toBe(firstError.id);
+    expect(
+      within(second)
+        .getByRole("textbox", { name: "Name" })
+        .getAttribute("aria-describedby"),
+    ).toBe(secondError.id);
+  });
+
+  it("preserves caller IDs and ARIA overrides while associating errors", () => {
+    render(
+      <Root aria-label="Add server">
+        <NameField
+          id="server-name"
+          aria-label="Server name"
+          aria-describedby="hint"
+          aria-invalid={false}
+        />
+        <p id="hint">Use a recognizable name.</p>
+        <ErrorMessage />
+      </Root>,
+    );
+    const field = screen.getByRole("textbox", { name: "Server name" });
+    expect(field.id).toBe("server-name");
+    expect(field.getAttribute("aria-describedby")).toBe("hint");
+    fireEvent.submit(screen.getByRole("form"));
+    expect(field.getAttribute("aria-invalid")).toBe("false");
+    expect(field.getAttribute("aria-describedby")).toBe(
+      `hint ${screen.getByRole("alert").id}`,
+    );
+    fireEvent.change(field, { target: { value: "Docs" } });
+    expect(screen.queryByRole("alert")).toBeNull();
+    expect(field.getAttribute("aria-describedby")).toBe("hint");
+  });
+
+  it("keeps URL help text associated when validation errors appear and clear", () => {
+    render(
+      <Root aria-label="Add server">
+        <NameField aria-label="Name" />
+        <UrlField aria-label="URL" aria-describedby="url-hint url-format" />
+        <p id="url-hint">Use the server endpoint.</p>
+        <p id="url-format">An HTTP or HTTPS URL is required.</p>
+        <ErrorMessage />
+      </Root>,
+    );
+    fireEvent.change(screen.getByRole("textbox", { name: "Name" }), {
+      target: { value: "Docs" },
+    });
+    const field = screen.getByRole("textbox", { name: "URL" });
+    fireEvent.submit(screen.getByRole("form"));
+    expect(field.getAttribute("aria-invalid")).toBe("true");
+    expect(field.getAttribute("aria-describedby")).toBe(
+      `url-hint url-format ${screen.getByRole("alert").id}`,
+    );
+    fireEvent.change(field, { target: { value: "https://example.com/mcp" } });
+    expect(screen.queryByRole("alert")).toBeNull();
+    expect(field.hasAttribute("aria-invalid")).toBe(false);
+    expect(field.getAttribute("aria-describedby")).toBe("url-hint url-format");
+  });
+});