@assistant-ui/react-mcp 0.0.0

package/src/resources/McpServerResource.ts

@@ -0,0 +1,292 @@
+import {
+  resource,
+  tapState,
+  tapRef,
+  tapEffect,
+  tapMemo,
+  tapEffectEvent,
+} from "@assistant-ui/tap";
+import type { ClientOutput } from "@assistant-ui/store";
+import { Client } from "@modelcontextprotocol/sdk/client/index.js";
+import {
+  StreamableHTTPClientTransport,
+  type StreamableHTTPClientTransportOptions,
+} from "@modelcontextprotocol/sdk/client/streamableHttp.js";
+import { UnauthorizedError } from "@modelcontextprotocol/sdk/client/auth.js";
+import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
+import { createOAuthProvider } from "../auth/createOAuthProvider";
+import { buildHeaders } from "../auth/buildHeaders";
+import { assertValidServerId } from "../utils/serverId";
+import type { MCPStorage } from "./storage/types";
+import type {
+  MCPAuthConfig,
+  MCPConnectionState,
+  MCPServerKind,
+  MCPServerState,
+  MCPToolInfo,
+} from "../mcp-scope";
+
+export type McpServerResourceProps = {
+  id: string;
+  kind: MCPServerKind;
+  name: string;
+  url: string;
+  icon?: string | undefined;
+  auth: MCPAuthConfig;
+  storage: MCPStorage;
+  redirectUri: string;
+  autoConnect: boolean;
+  onRemove: () => Promise<void>;
+};
+
+export const McpServerResource = resource(
+  (props: McpServerResourceProps): ClientOutput<"mcpServer"> => {
+    assertValidServerId(props.id);
+    const [connectionState, setConnectionState] =
+      tapState<MCPConnectionState>("disconnected");
+    const [tools, setTools] = tapState<MCPToolInfo[]>([]);
+    const [lastError, setLastError] = tapState<{ message: string } | null>(
+      null,
+    );
+    const [authorizationUrl, setAuthorizationUrl] = tapState<string | null>(
+      null,
+    );
+
+    const clientRef = tapRef<Client | null>(null);
+    const transportRef = tapRef<StreamableHTTPClientTransport | null>(null);
+
+    const buildTransport = tapEffectEvent(
+      async (): Promise<StreamableHTTPClientTransport> => {
+        if (props.auth.type === "oauth") {
+          const authProvider = createOAuthProvider({
+            serverId: props.id,
+            config: props.auth,
+            storage: props.storage,
+            redirectUri: props.redirectUri,
+            onAuthorizationUrl: (url) => setAuthorizationUrl(url.toString()),
+          });
+          return new StreamableHTTPClientTransport(new URL(props.url), {
+            authProvider,
+          });
+        }
+        if (props.auth.type === "bearer") {
+          const persisted = await props.storage.loadAuthState(props.id);
+          const headers = buildHeaders(props.auth, persisted);
+          const transportOpts: StreamableHTTPClientTransportOptions = {};
+          if (headers) transportOpts.requestInit = { headers };
+          return new StreamableHTTPClientTransport(
+            new URL(props.url),
+            transportOpts,
+          );
+        }
+        return new StreamableHTTPClientTransport(new URL(props.url));
+      },
+    );
+
+    const finalizeConnect = tapEffectEvent(
+      async (transport: StreamableHTTPClientTransport) => {
+        const client = new Client({
+          name: "assistant-ui-mcp",
+          version: "0.0.0",
+        });
+        // SDK's StreamableHTTPClientTransport.sessionId is `string | undefined`
+        // but Transport.sessionId is declared `string?` — under
+        // exactOptionalPropertyTypes the SDK's own classes don't satisfy its
+        // Transport interface. Cast to bridge the gap.
+        await client.connect(transport as unknown as Transport);
+        // Defer ref assignment until listTools() also succeeds — otherwise a
+        // post-connect failure leaves stale refs that `callTool()` would
+        // happily walk into, producing confusing SDK errors instead of
+        // "not connected".
+        const list = await client.listTools();
+        clientRef.current = client;
+        transportRef.current = transport;
+        setTools(
+          list.tools.map((t) => {
+            const info: MCPToolInfo = {
+              name: t.name,
+              inputSchema: t.inputSchema,
+            };
+            if (t.description !== undefined) info.description = t.description;
+            return info;
+          }),
+        );
+        setConnectionState("connected");
+      },
+    );
+
+    const closeTransport = async () => {
+      const t = transportRef.current;
+      transportRef.current = null;
+      clientRef.current = null;
+      if (t) {
+        try {
+          await t.close();
+        } catch {
+          // ignore close errors
+        }
+      }
+    };
+
+    const doConnect = tapEffectEvent(async () => {
+      // Close any prior transport/client so a re-connect doesn't leak.
+      await closeTransport();
+      setConnectionState("connecting");
+      setLastError(null);
+      setAuthorizationUrl(null);
+      // Clear tools so a reconnect (error → connect, or authRequired
+      // → connect) doesn't expose the previous attempt's tool list.
+      setTools([]);
+      let transport: StreamableHTTPClientTransport | null = null;
+      try {
+        transport = await buildTransport();
+        // Don't assign to transportRef until connect succeeds — otherwise a
+        // failed `listTools()` leaves an orphaned transport that future
+        // doConnect / doDisconnect calls treat as live.
+        await finalizeConnect(transport);
+      } catch (err) {
+        if (err instanceof UnauthorizedError) {
+          // OAuth: keep the transport alive so completeAuth can call
+          // finishAuth on it. Closing it before storing would leave a
+          // closed transport on transportRef.
+          transportRef.current = transport;
+          setConnectionState("authRequired");
+        } else {
+          if (transport) {
+            try {
+              await transport.close();
+            } catch {
+              // ignore close errors
+            }
+          }
+          setLastError({
+            message: err instanceof Error ? err.message : String(err),
+          });
+          setConnectionState("error");
+        }
+      }
+    });
+
+    const doDisconnect = tapEffectEvent(async () => {
+      setTools([]);
+      setAuthorizationUrl(null);
+      setConnectionState("disconnected");
+      await closeTransport();
+    });
+
+    const doCompleteAuth = tapEffectEvent(async (callbackUrl: string) => {
+      setConnectionState("authPending");
+      setLastError(null);
+      try {
+        const url = new URL(callbackUrl);
+        const code = url.searchParams.get("code");
+        if (!code)
+          throw new Error("missing authorization code in callback URL");
+        let transport = transportRef.current;
+        if (!transport) {
+          transport = await buildTransport();
+          transportRef.current = transport;
+        }
+        await transport.finishAuth(code);
+        setAuthorizationUrl(null);
+        await finalizeConnect(transport);
+      } catch (err) {
+        await closeTransport();
+        setLastError({
+          message: err instanceof Error ? err.message : String(err),
+        });
+        setConnectionState("error");
+      }
+    });
+
+    const tryAutoConnect = tapEffectEvent(
+      async (signal: { cancelled: boolean }) => {
+        if (!props.autoConnect) return;
+        if (props.auth.type === "none") {
+          void doConnect();
+          return;
+        }
+        // Static connector token short-circuits the storage read.
+        if (props.auth.type === "bearer" && props.auth.token) {
+          void doConnect();
+          return;
+        }
+        const persisted = await props.storage.loadAuthState(props.id);
+        if (signal.cancelled) return;
+        if (props.auth.type === "oauth") {
+          if (!persisted?.tokens) return;
+        } else if (!persisted?.token) {
+          return;
+        }
+        void doConnect();
+      },
+    );
+
+    // Auto-connect on mount when usable auth exists.
+    tapEffect(() => {
+      const signal = { cancelled: false };
+      void tryAutoConnect(signal);
+      return () => {
+        signal.cancelled = true;
+        const t = transportRef.current;
+        transportRef.current = null;
+        clientRef.current = null;
+        if (t) t.close().catch(() => {});
+      };
+    }, []);
+
+    const state = tapMemo<MCPServerState>(
+      () => ({
+        id: props.id,
+        kind: props.kind,
+        name: props.name,
+        url: props.url,
+        icon: props.icon,
+        connectionState,
+        lastError,
+        tools,
+        authorizationUrl,
+      }),
+      [
+        props.id,
+        props.kind,
+        props.name,
+        props.url,
+        props.icon,
+        connectionState,
+        lastError,
+        tools,
+        authorizationUrl,
+      ],
+    );
+
+    return {
+      getState: () => state,
+      connect: doConnect,
+      disconnect: doDisconnect,
+      remove: async () => {
+        await doDisconnect();
+        await props.storage.clearAuthState(props.id);
+        await props.onRemove();
+      },
+      callTool: async (name, args) => {
+        const client = clientRef.current;
+        if (!client) {
+          throw new Error(`MCP server "${props.id}" is not connected`);
+        }
+        return await client.callTool({
+          name,
+          arguments: args as Record<string, unknown> | undefined,
+        });
+      },
+      readResource: async (uri) => {
+        const client = clientRef.current;
+        if (!client) {
+          throw new Error(`MCP server "${props.id}" is not connected`);
+        }
+        return await client.readResource({ uri });
+      },
+      completeAuth: doCompleteAuth,
+    };
+  },
+);

package/src/resources/storage/McpCustomStorage.ts

@@ -0,0 +1,6 @@
+import { resource } from "@assistant-ui/tap";
+import type { MCPStorage } from "./types";
+
+export const McpCustomStorage = resource(
+  (impl: MCPStorage): MCPStorage => impl,
+);

package/src/resources/storage/McpLocalStorage.ts

@@ -0,0 +1,77 @@
+import { resource } from "@assistant-ui/tap";
+import type { MCPCustomServerRecord } from "../../mcp-scope";
+import type { MCPPersistedAuthState } from "../../auth/types";
+import type { MCPStorage } from "./types";
+
+export type McpLocalStorageOptions = {
+  /** Namespace prefix for keys. Default "aui-mcp". */
+  keyPrefix?: string;
+  /** Override the underlying Storage. Defaults to globalThis.localStorage. */
+  storage?: Storage;
+};
+
+function resolveStorage(opts: McpLocalStorageOptions): Storage | null {
+  if (opts.storage) return opts.storage;
+  if (typeof globalThis !== "undefined" && "localStorage" in globalThis) {
+    try {
+      return (globalThis as { localStorage: Storage }).localStorage;
+    } catch {
+      return null;
+    }
+  }
+  return null;
+}
+
+export const McpLocalStorage = resource(
+  (opts: McpLocalStorageOptions = {}): MCPStorage => {
+    const prefix = opts.keyPrefix ?? "aui-mcp";
+    const customServersKey = `${prefix}:custom-servers`;
+    const authKey = (id: string) => `${prefix}:auth:${id}`;
+    const storage = resolveStorage(opts);
+
+    const read = <T>(key: string, fallback: T): T => {
+      if (!storage) return fallback;
+      try {
+        const raw = storage.getItem(key);
+        if (raw == null) return fallback;
+        return JSON.parse(raw) as T;
+      } catch {
+        return fallback;
+      }
+    };
+
+    const write = (key: string, value: unknown): void => {
+      if (!storage) return;
+      try {
+        storage.setItem(key, JSON.stringify(value));
+      } catch {
+        // quota or serialization failure — silently drop
+      }
+    };
+
+    const remove = (key: string): void => {
+      if (!storage) return;
+      try {
+        storage.removeItem(key);
+      } catch {
+        // ignore
+      }
+    };
+
+    return {
+      loadCustomServers: async () =>
+        read<MCPCustomServerRecord[]>(customServersKey, []),
+      saveCustomServers: async (records) => {
+        write(customServersKey, records);
+      },
+      loadAuthState: async (id) =>
+        read<MCPPersistedAuthState | null>(authKey(id), null),
+      saveAuthState: async (id, state) => {
+        write(authKey(id), state);
+      },
+      clearAuthState: async (id) => {
+        remove(authKey(id));
+      },
+    };
+  },
+);

package/src/resources/storage/McpMemoryStorage.ts

@@ -0,0 +1,22 @@
+import { resource } from "@assistant-ui/tap";
+import type { MCPCustomServerRecord } from "../../mcp-scope";
+import type { MCPPersistedAuthState } from "../../auth/types";
+import type { MCPStorage } from "./types";
+
+export const McpMemoryStorage = resource((): MCPStorage => {
+  let servers: MCPCustomServerRecord[] = [];
+  const auth = new Map<string, MCPPersistedAuthState>();
+  return {
+    loadCustomServers: async () => [...servers],
+    saveCustomServers: async (records) => {
+      servers = [...records];
+    },
+    loadAuthState: async (id) => auth.get(id) ?? null,
+    saveAuthState: async (id, state) => {
+      auth.set(id, state);
+    },
+    clearAuthState: async (id) => {
+      auth.delete(id);
+    },
+  };
+});

package/src/resources/storage/types.ts

@@ -0,0 +1,16 @@
+import type { ResourceElement } from "@assistant-ui/tap";
+import type { MCPCustomServerRecord } from "../../mcp-scope";
+import type { MCPPersistedAuthState } from "../../auth/types";
+
+export type MCPStorage = {
+  loadCustomServers: () => Promise<MCPCustomServerRecord[]>;
+  saveCustomServers: (records: MCPCustomServerRecord[]) => Promise<void>;
+  loadAuthState: (serverId: string) => Promise<MCPPersistedAuthState | null>;
+  saveAuthState: (
+    serverId: string,
+    state: MCPPersistedAuthState,
+  ) => Promise<void>;
+  clearAuthState: (serverId: string) => Promise<void>;
+};
+
+export type MCPStorageElement = ResourceElement<MCPStorage>;

package/src/utils/serverId.ts

@@ -0,0 +1,19 @@
+// Server ids are joined with tool names via `__` to form unique tool
+// names for the AI SDK. The separator must not appear inside ids,
+// otherwise two distinct `(serverId, toolName)` pairs could produce
+// the same composite name.
+export function assertValidServerId(id: string): void {
+  if (id.length === 0) {
+    throw new Error("MCP server id must not be empty.");
+  }
+  if (id.includes("__")) {
+    throw new Error(
+      `MCP server id "${id}" must not contain "__" — that sequence is reserved for separating server id from tool name.`,
+    );
+  }
+  if (id.includes("\x1f")) {
+    throw new Error(
+      `MCP server id "${id}" must not contain the ASCII Unit Separator (\\x1f) — it is reserved as the list delimiter in primitive id-list selectors.`,
+    );
+  }
+}