@assistant-ui/react-mcp 0.0.0

package/src/primitives/server/McpServerDisconnectButton.tsx

@@ -0,0 +1,38 @@
+import {
+  type ComponentPropsWithoutRef,
+  type ComponentRef,
+  forwardRef,
+} from "react";
+import { Primitive } from "@radix-ui/react-primitive";
+import { useAui, useAuiState } from "@assistant-ui/store";
+
+export namespace McpServerPrimitiveDisconnectButton {
+  export type Element = ComponentRef<typeof Primitive.button>;
+  export type Props = ComponentPropsWithoutRef<typeof Primitive.button>;
+}
+
+const DISCONNECTABLE = new Set(["connected", "connecting", "authPending"]);
+
+export const McpServerPrimitiveDisconnectButton = forwardRef<
+  McpServerPrimitiveDisconnectButton.Element,
+  McpServerPrimitiveDisconnectButton.Props
+>((props, ref) => {
+  const state = useAuiState((s) => s.mcpServer.connectionState);
+  const aui = useAui();
+  if (!DISCONNECTABLE.has(state)) return null;
+  return (
+    <Primitive.button
+      {...props}
+      type="button"
+      ref={ref}
+      onClick={(e) => {
+        props.onClick?.(e);
+        if (e.defaultPrevented) return;
+        void aui.mcpServer().disconnect();
+      }}
+    />
+  );
+});
+
+McpServerPrimitiveDisconnectButton.displayName =
+  "McpServerPrimitive.DisconnectButton";

package/src/primitives/server/McpServerError.tsx

@@ -0,0 +1,27 @@
+import {
+  type ComponentPropsWithoutRef,
+  type ComponentRef,
+  forwardRef,
+} from "react";
+import { Primitive } from "@radix-ui/react-primitive";
+import { useAuiState } from "@assistant-ui/store";
+
+export namespace McpServerPrimitiveError {
+  export type Element = ComponentRef<typeof Primitive.div>;
+  export type Props = ComponentPropsWithoutRef<typeof Primitive.div>;
+}
+
+export const McpServerPrimitiveError = forwardRef<
+  McpServerPrimitiveError.Element,
+  McpServerPrimitiveError.Props
+>((props, ref) => {
+  const message = useAuiState((s) => s.mcpServer.lastError?.message ?? null);
+  if (message === null) return null;
+  return (
+    <Primitive.div {...props} ref={ref}>
+      {props.children ?? message}
+    </Primitive.div>
+  );
+});
+
+McpServerPrimitiveError.displayName = "McpServerPrimitive.Error";

package/src/primitives/server/McpServerIcon.tsx

@@ -0,0 +1,35 @@
+import {
+  type ComponentPropsWithoutRef,
+  type ComponentRef,
+  forwardRef,
+} from "react";
+import { Primitive } from "@radix-ui/react-primitive";
+import { useAuiState } from "@assistant-ui/store";
+
+export namespace McpServerPrimitiveIcon {
+  export type Element = ComponentRef<typeof Primitive.img>;
+  export type Props = Omit<
+    ComponentPropsWithoutRef<typeof Primitive.img>,
+    "src" | "alt"
+  > & {
+    /** Override the icon URL. Defaults to the server's icon prop. */
+    src?: string;
+    alt?: string;
+  };
+}
+
+export const McpServerPrimitiveIcon = forwardRef<
+  McpServerPrimitiveIcon.Element,
+  McpServerPrimitiveIcon.Props
+>(({ src, alt, ...props }, ref) => {
+  const iconUrl = useAuiState((s) => s.mcpServer.icon ?? null);
+  const name = useAuiState((s) => s.mcpServer.name);
+  const effective = src ?? iconUrl;
+  if (!effective) return null;
+  return (
+    // biome-ignore lint/performance/noImgElement: unstyled primitive — consumers can swap via asChild
+    <Primitive.img {...props} ref={ref} src={effective} alt={alt ?? name} />
+  );
+});
+
+McpServerPrimitiveIcon.displayName = "McpServerPrimitive.Icon";

package/src/primitives/server/McpServerName.tsx

@@ -0,0 +1,26 @@
+import {
+  type ComponentPropsWithoutRef,
+  type ComponentRef,
+  forwardRef,
+} from "react";
+import { Primitive } from "@radix-ui/react-primitive";
+import { useAuiState } from "@assistant-ui/store";
+
+export namespace McpServerPrimitiveName {
+  export type Element = ComponentRef<typeof Primitive.span>;
+  export type Props = ComponentPropsWithoutRef<typeof Primitive.span>;
+}
+
+export const McpServerPrimitiveName = forwardRef<
+  McpServerPrimitiveName.Element,
+  McpServerPrimitiveName.Props
+>((props, ref) => {
+  const name = useAuiState((s) => s.mcpServer.name);
+  return (
+    <Primitive.span {...props} ref={ref}>
+      {props.children ?? name}
+    </Primitive.span>
+  );
+});
+
+McpServerPrimitiveName.displayName = "McpServerPrimitive.Name";

package/src/primitives/server/McpServerOAuthLink.tsx

@@ -0,0 +1,38 @@
+import {
+  type ComponentPropsWithoutRef,
+  type ComponentRef,
+  forwardRef,
+} from "react";
+import { Primitive } from "@radix-ui/react-primitive";
+import { useAuiState } from "@assistant-ui/store";
+
+export namespace McpServerPrimitiveOAuthLink {
+  export type Element = ComponentRef<typeof Primitive.a>;
+  export type Props = Omit<
+    ComponentPropsWithoutRef<typeof Primitive.a>,
+    "href"
+  > & {
+    /** Optional href override; defaults to the server's authorizationUrl. */
+    href?: string;
+  };
+}
+
+export const McpServerPrimitiveOAuthLink = forwardRef<
+  McpServerPrimitiveOAuthLink.Element,
+  McpServerPrimitiveOAuthLink.Props
+>(({ href, ...props }, ref) => {
+  const url = useAuiState((s) => s.mcpServer.authorizationUrl);
+  const effective = href ?? url;
+  if (!effective) return null;
+  return (
+    <Primitive.a
+      target="_blank"
+      {...props}
+      ref={ref}
+      href={effective}
+      rel={props.rel ?? "noopener noreferrer"}
+    />
+  );
+});
+
+McpServerPrimitiveOAuthLink.displayName = "McpServerPrimitive.OAuthLink";

package/src/primitives/server/McpServerRemoveButton.tsx

@@ -0,0 +1,35 @@
+import {
+  type ComponentPropsWithoutRef,
+  type ComponentRef,
+  forwardRef,
+} from "react";
+import { Primitive } from "@radix-ui/react-primitive";
+import { useAui, useAuiState } from "@assistant-ui/store";
+
+export namespace McpServerPrimitiveRemoveButton {
+  export type Element = ComponentRef<typeof Primitive.button>;
+  export type Props = ComponentPropsWithoutRef<typeof Primitive.button>;
+}
+
+export const McpServerPrimitiveRemoveButton = forwardRef<
+  McpServerPrimitiveRemoveButton.Element,
+  McpServerPrimitiveRemoveButton.Props
+>((props, ref) => {
+  const kind = useAuiState((s) => s.mcpServer.kind);
+  const aui = useAui();
+  if (kind !== "custom") return null;
+  return (
+    <Primitive.button
+      {...props}
+      type="button"
+      ref={ref}
+      onClick={(e) => {
+        props.onClick?.(e);
+        if (e.defaultPrevented) return;
+        void aui.mcpServer().remove();
+      }}
+    />
+  );
+});
+
+McpServerPrimitiveRemoveButton.displayName = "McpServerPrimitive.RemoveButton";

package/src/primitives/server/McpServerRoot.tsx

@@ -0,0 +1,35 @@
+import {
+  type ComponentPropsWithoutRef,
+  type ComponentRef,
+  forwardRef,
+} from "react";
+import { Primitive } from "@radix-ui/react-primitive";
+import { useAuiState } from "@assistant-ui/store";
+
+export namespace McpServerPrimitiveRoot {
+  export type Element = ComponentRef<typeof Primitive.div>;
+  export type Props = ComponentPropsWithoutRef<typeof Primitive.div>;
+}
+
+export const McpServerPrimitiveRoot = forwardRef<
+  McpServerPrimitiveRoot.Element,
+  McpServerPrimitiveRoot.Props
+>((props, ref) => {
+  const id = useAuiState((s) => s.mcpServer.id);
+  const kind = useAuiState((s) => s.mcpServer.kind);
+  const connectionState = useAuiState((s) => s.mcpServer.connectionState);
+  const hasError = useAuiState((s) => s.mcpServer.lastError !== null);
+
+  return (
+    <Primitive.div
+      {...props}
+      ref={ref}
+      data-server-id={id}
+      data-kind={kind}
+      data-connection-state={connectionState}
+      data-has-error={hasError || undefined}
+    />
+  );
+});
+
+McpServerPrimitiveRoot.displayName = "McpServerPrimitive.Root";

package/src/primitives/server/McpServerStatus.tsx

@@ -0,0 +1,26 @@
+import {
+  type ComponentPropsWithoutRef,
+  type ComponentRef,
+  forwardRef,
+} from "react";
+import { Primitive } from "@radix-ui/react-primitive";
+import { useAuiState } from "@assistant-ui/store";
+
+export namespace McpServerPrimitiveStatus {
+  export type Element = ComponentRef<typeof Primitive.span>;
+  export type Props = ComponentPropsWithoutRef<typeof Primitive.span>;
+}
+
+export const McpServerPrimitiveStatus = forwardRef<
+  McpServerPrimitiveStatus.Element,
+  McpServerPrimitiveStatus.Props
+>((props, ref) => {
+  const state = useAuiState((s) => s.mcpServer.connectionState);
+  return (
+    <Primitive.span {...props} ref={ref} data-state={state}>
+      {props.children ?? state}
+    </Primitive.span>
+  );
+});
+
+McpServerPrimitiveStatus.displayName = "McpServerPrimitive.Status";

package/src/primitives/server/McpServerToolName.tsx

@@ -0,0 +1,26 @@
+import {
+  type ComponentPropsWithoutRef,
+  type ComponentRef,
+  forwardRef,
+} from "react";
+import { Primitive } from "@radix-ui/react-primitive";
+import { useMcpServerTool } from "./McpServerTools";
+
+export namespace McpServerPrimitiveToolName {
+  export type Element = ComponentRef<typeof Primitive.span>;
+  export type Props = ComponentPropsWithoutRef<typeof Primitive.span>;
+}
+
+export const McpServerPrimitiveToolName = forwardRef<
+  McpServerPrimitiveToolName.Element,
+  McpServerPrimitiveToolName.Props
+>((props, ref) => {
+  const tool = useMcpServerTool();
+  return (
+    <Primitive.span {...props} ref={ref} data-tool-name={tool.name}>
+      {props.children ?? tool.name}
+    </Primitive.span>
+  );
+});
+
+McpServerPrimitiveToolName.displayName = "McpServerPrimitive.ToolName";

package/src/primitives/server/McpServerTools.tsx

@@ -0,0 +1,39 @@
+import { type FC, type ReactNode, createContext, useContext } from "react";
+import { useAuiState } from "@assistant-ui/store";
+import type { MCPToolInfo } from "../../mcp-scope";
+
+const ToolContext = createContext<MCPToolInfo | null>(null);
+
+export const useMcpServerTool = (): MCPToolInfo => {
+  const tool = useContext(ToolContext);
+  if (!tool) {
+    throw new Error(
+      "useMcpServerTool must be used inside <McpServerPrimitive.Tools>",
+    );
+  }
+  return tool;
+};
+
+export namespace McpServerPrimitiveTools {
+  export type Props = {
+    children: ReactNode | ((tool: MCPToolInfo) => ReactNode);
+  };
+}
+
+export const McpServerPrimitiveTools: FC<McpServerPrimitiveTools.Props> = ({
+  children,
+}) => {
+  const tools = useAuiState((s) => s.mcpServer.tools);
+  if (tools.length === 0) return null;
+  return (
+    <>
+      {tools.map((tool) => (
+        <ToolContext.Provider key={tool.name} value={tool}>
+          {typeof children === "function" ? children(tool) : children}
+        </ToolContext.Provider>
+      ))}
+    </>
+  );
+};
+
+McpServerPrimitiveTools.displayName = "McpServerPrimitive.Tools";

package/src/primitives/server.ts

@@ -0,0 +1,14 @@
+export { McpServerPrimitiveRoot as Root } from "./server/McpServerRoot";
+export { McpServerPrimitiveName as Name } from "./server/McpServerName";
+export { McpServerPrimitiveStatus as Status } from "./server/McpServerStatus";
+export { McpServerPrimitiveError as Error } from "./server/McpServerError";
+export { McpServerPrimitiveIcon as Icon } from "./server/McpServerIcon";
+export { McpServerPrimitiveConnectButton as ConnectButton } from "./server/McpServerConnectButton";
+export { McpServerPrimitiveDisconnectButton as DisconnectButton } from "./server/McpServerDisconnectButton";
+export { McpServerPrimitiveRemoveButton as RemoveButton } from "./server/McpServerRemoveButton";
+export { McpServerPrimitiveOAuthLink as OAuthLink } from "./server/McpServerOAuthLink";
+export {
+  McpServerPrimitiveTools as Tools,
+  useMcpServerTool,
+} from "./server/McpServerTools";
+export { McpServerPrimitiveToolName as ToolName } from "./server/McpServerToolName";

package/src/resources/McpManagerResource.ts

@@ -0,0 +1,253 @@
+import {
+  resource,
+  tapState,
+  tapResource,
+  tapEffect,
+  tapMemo,
+  tapEffectEvent,
+  tapRef,
+  withKey,
+} from "@assistant-ui/tap";
+import {
+  tapClientLookup,
+  tapAssistantClientRef,
+  attachTransformScopes,
+  type ClientOutput,
+} from "@assistant-ui/store";
+import { ModelContext } from "@assistant-ui/core/store";
+import type { Tool } from "assistant-stream";
+import { McpServerResource } from "./McpServerResource";
+import { McpLocalStorage } from "./storage/McpLocalStorage";
+import type { MCPStorageElement } from "./storage/types";
+import type {
+  MCPAuthConfig,
+  MCPConnector,
+  MCPCustomServerRecord,
+  MCPManagerState,
+} from "../mcp-scope";
+
+export type McpManagerResourceProps = {
+  connectors?: MCPConnector[] | undefined;
+  storage?: MCPStorageElement | undefined;
+  /** OAuth redirect target. Defaults to `${origin}/mcp/callback`. */
+  oauthRedirectUri?: string | undefined;
+  /** Connect on mount when usable auth exists. Default true. */
+  autoConnect?: boolean | undefined;
+};
+
+function defaultRedirectUri(): string {
+  if (typeof window === "undefined") return "";
+  return `${window.location.origin}/mcp/callback`;
+}
+
+export const McpManagerResource = resource(
+  (props: McpManagerResourceProps): ClientOutput<"mcp"> => {
+    const connectors = props.connectors ?? [];
+    const autoConnect = props.autoConnect ?? true;
+    const redirectUri = props.oauthRedirectUri ?? defaultRedirectUri();
+
+    const storageElement = props.storage ?? McpLocalStorage();
+    const storage = tapResource(storageElement);
+
+    const [customServers, setCustomServers] = tapState<MCPCustomServerRecord[]>(
+      [],
+    );
+    const [isHydrated, setIsHydrated] = tapState(false);
+
+    const hydratedRef = tapRef(false);
+
+    const hydrate = tapEffectEvent(async (signal: { cancelled: boolean }) => {
+      try {
+        const records = await storage.loadCustomServers();
+        if (signal.cancelled) return;
+        // Merge rather than replace so any addCustomServer calls that
+        // happened before hydration resolved aren't silently overwritten.
+        // Persisted order wins; pre-hydration locals append.
+        setCustomServers((prev) => {
+          if (prev.length === 0) return records;
+          const persistedIds = new Set(records.map((r) => r.id));
+          return [...records, ...prev.filter((r) => !persistedIds.has(r.id))];
+        });
+      } finally {
+        if (!signal.cancelled) {
+          hydratedRef.current = true;
+          setIsHydrated(true);
+        }
+      }
+    });
+
+    tapEffect(() => {
+      const signal = { cancelled: false };
+      void hydrate(signal);
+      return () => {
+        signal.cancelled = true;
+      };
+    }, []);
+
+    const persistCustomServers = tapEffectEvent(
+      async (records: MCPCustomServerRecord[]) => {
+        if (!hydratedRef.current) return;
+        await storage.saveCustomServers(records);
+      },
+    );
+
+    tapEffect(() => {
+      void persistCustomServers(customServers);
+    }, [customServers]);
+
+    const serverElements = tapMemo(() => {
+      const connectorElements = connectors.map((c) =>
+        withKey(
+          c.id,
+          McpServerResource({
+            id: c.id,
+            kind: "connector",
+            name: c.name,
+            url: c.url,
+            icon: c.icon,
+            auth: c.auth,
+            storage,
+            redirectUri,
+            autoConnect,
+            onRemove: async () => {
+              // connectors cannot be removed
+            },
+          }),
+        ),
+      );
+      const customElements = customServers.map((s) =>
+        withKey(
+          s.id,
+          McpServerResource({
+            id: s.id,
+            kind: "custom",
+            name: s.name,
+            url: s.url,
+            auth: s.auth,
+            storage,
+            redirectUri,
+            autoConnect,
+            onRemove: async () => {
+              setCustomServers((prev) => prev.filter((x) => x.id !== s.id));
+            },
+          }),
+        ),
+      );
+      return [...connectorElements, ...customElements];
+    }, [connectors, customServers, storage, redirectUri, autoConnect]);
+
+    const lookup = tapClientLookup(() => serverElements, [serverElements]);
+
+    const state = tapMemo<MCPManagerState>(() => {
+      const all = lookup.state;
+      return {
+        servers: all,
+        connectors: all.filter((s) => s.kind === "connector"),
+        customServers: all.filter((s) => s.kind === "custom"),
+        isHydrated,
+      };
+    }, [lookup.state, isHydrated]);
+
+    // ─── Auto-register MCP tools as frontend tools in modelContext ─────
+    // Build the toolkit from connected servers; re-register when the visible
+    // tool surface changes. Tool names are prefixed with the server id to
+    // avoid collisions across connected servers — connector ids must not
+    // contain `__` (enforced by `defineConnector`); `addCustomServer`
+    // generates UUIDs that satisfy the constraint by construction.
+    const toolkit = tapMemo<Record<string, Tool<any, any>>>(() => {
+      const out: Record<string, Tool<any, any>> = {};
+      for (const server of state.servers) {
+        if (server.connectionState !== "connected") continue;
+        for (const tool of server.tools) {
+          const fullName = `${server.id}__${tool.name}`;
+          out[fullName] = {
+            type: "frontend",
+            ...(tool.description !== undefined
+              ? { description: tool.description }
+              : {}),
+            parameters: tool.inputSchema as never,
+            execute: (args) =>
+              lookup
+                .get({ key: server.id })
+                .callTool(tool.name, args as unknown),
+          };
+        }
+      }
+      return out;
+    }, [state, lookup]);
+
+    const clientRef = tapAssistantClientRef();
+
+    tapEffect(() => {
+      const client = clientRef.current;
+      if (!client) return;
+      return client.modelContext().register({
+        getModelContext: () => ({ tools: toolkit }),
+      });
+    }, [toolkit, clientRef]);
+
+    const serverByKind = (kind: "connector" | "custom", index: number) => {
+      const list =
+        kind === "connector" ? state.connectors : state.customServers;
+      const entry = list[index];
+      if (!entry) {
+        throw new Error(
+          `McpManagerResource: no ${kind} at index ${index} (length ${list.length})`,
+        );
+      }
+      return lookup.get({ key: entry.id });
+    };
+
+    return {
+      getState: () => state,
+      server: (query) => {
+        if ("id" in query) return lookup.get({ key: query.id });
+        return serverByKind(query.kind, query.index);
+      },
+      connector: ({ index }) => serverByKind("connector", index),
+      customServer: ({ index }) => serverByKind("custom", index),
+      addCustomServer: async ({ name, url, auth }) => {
+        const record: MCPCustomServerRecord = {
+          id:
+            typeof crypto !== "undefined" && "randomUUID" in crypto
+              ? crypto.randomUUID()
+              : `mcp-${Date.now()}-${Math.random().toString(36).slice(2)}`,
+          name,
+          url,
+          auth: auth as MCPAuthConfig,
+          createdAt: Date.now(),
+        };
+        setCustomServers((prev) => [...prev, record]);
+        return record.id;
+      },
+      removeServer: async (id) => {
+        // removeServer is custom-server only — connectors are app-defined
+        // and not user-removable. Refuse rather than silently no-op.
+        if (state.connectors.some((c) => c.id === id)) {
+          throw new Error(
+            `Cannot remove connector "${id}" — connectors are app-defined and not removable. Use a custom server id instead.`,
+          );
+        }
+        // Delegate to McpServerResource.remove() which disconnects,
+        // clears auth state, and unregisters from customServers in one
+        // place. Fallback to manual cleanup if the lookup is empty
+        // (server already gone).
+        try {
+          await lookup.get({ key: id }).remove();
+        } catch {
+          await storage.clearAuthState(id);
+          setCustomServers((prev) => prev.filter((s) => s.id !== id));
+        }
+      },
+    };
+  },
+);
+
+// Ensure modelContext exists as a sibling when the manager mounts. If an
+// ancestor (e.g. a chat runtime) already provides modelContext, this is a
+// no-op; otherwise it's auto-mounted alongside `mcp`.
+attachTransformScopes(McpManagerResource, (scopes, parent) => {
+  if (!scopes.modelContext && parent.modelContext.source === null) {
+    scopes.modelContext = ModelContext();
+  }
+});