@assistant-ui/react-mcp 0.0.6 → 0.0.7

package/src/resources/McpServerResource.ts

@@ -1,11 +1,5 @@
-import {
-  resource,
-  tapState,
-  tapRef,
-  tapEffect,
-  tapMemo,
-  tapEffectEvent,
-} from "@assistant-ui/tap";
+import { useState, useRef, useEffect, useMemo, useEffectEvent } from "react";
+import { resource } from "@assistant-ui/tap";
 import type { ClientOutput } from "@assistant-ui/store";
 import { Client } from "@modelcontextprotocol/sdk/client/index.js";
 import {
@@ -39,254 +33,249 @@ export type McpServerResourceProps = {
   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);
+export const McpServerResource = resource(function McpServerResource(
+  props: McpServerResourceProps,
+): ClientOutput<"mcpServer"> {
+  assertValidServerId(props.id);
+  const [connectionState, setConnectionState] =
+    useState<MCPConnectionState>("disconnected");
+  const [tools, setTools] = useState<MCPToolInfo[]>([]);
+  const [lastError, setLastError] = useState<{ message: string } | null>(null);
+  const [authorizationUrl, setAuthorizationUrl] = useState<string | 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 clientRef = useRef<Client | null>(null);
+  const transportRef = useRef<StreamableHTTPClientTransport | null>(null);
 
-    const finalizeConnect = tapEffectEvent(
-      async (transport: StreamableHTTPClientTransport) => {
-        const client = new Client({
-          name: "assistant-ui-mcp",
-          version: "0.0.0",
+  const buildTransport = useEffectEvent(
+    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()),
         });
-        // 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;
-          }),
+        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,
         );
-        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
-        }
       }
-    };
+      return new StreamableHTTPClientTransport(new URL(props.url));
+    },
+  );
 
-    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;
+  const finalizeConnect = useEffectEvent(
+    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 {
-        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");
-        }
+        await t.close();
+      } catch {
+        // ignore close errors
       }
-    });
+    }
+  };
 
-    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;
+  const doConnect = useEffectEvent(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
+          }
         }
-        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();
-      },
-    );
+  const doDisconnect = useEffectEvent(async () => {
+    setTools([]);
+    setAuthorizationUrl(null);
+    setConnectionState("disconnected");
+    await closeTransport();
+  });
 
-    // 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 doCompleteAuth = useEffectEvent(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 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,
-      ],
-    );
+  const tryAutoConnect = useEffectEvent(
+    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();
+    },
+  );
 
-    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,
+  // Auto-connect on mount when usable auth exists.
+  useEffect(() => {
+    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 = useMemo<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

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

package/src/resources/storage/McpLocalStorage.ts

@@ -22,56 +22,56 @@ function resolveStorage(opts: McpLocalStorageOptions): Storage | 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);
+export const McpLocalStorage = resource(function McpLocalStorage(
+  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 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 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
-      }
-    };
+  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));
-      },
-    };
-  },
-);
+  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

@@ -3,20 +3,22 @@ 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);
-    },
-  };
-});
+export const McpMemoryStorage = resource(
+  function McpMemoryStorage(): 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

@@ -1,4 +1,4 @@
-import type { ResourceElement } from "@assistant-ui/tap";
+import { type ResourceElement } from "@assistant-ui/tap";
 import type { MCPCustomServerRecord } from "../../mcp-scope";
 import type { MCPPersistedAuthState } from "../../auth/types";