@assistant-ui/react-mcp 0.0.17 → 0.0.19

package/src/resources/McpServerResource.ts

@@ -46,6 +46,64 @@ const useMcpServerResource = (
 
   const clientRef = useRef<Client | null>(null);
   const transportRef = useRef<StreamableHTTPClientTransport | null>(null);
+  const pendingTransportRef = useRef<StreamableHTTPClientTransport | null>(
+    null,
+  );
+  const transportCloseQueueRef = useRef(Promise.resolve());
+  const connectionGenerationRef = useRef(0);
+  const pendingDisposalRef = useRef<{ cancelled: boolean } | null>(null);
+  const mountedRef = useRef(true);
+
+  const closeTransportSafely = async (
+    transport: StreamableHTTPClientTransport,
+  ): Promise<void> => {
+    try {
+      await transport.close();
+    } catch {
+      // ignore close errors
+    }
+  };
+
+  const closeQueuedTransports = (
+    transports: StreamableHTTPClientTransport[],
+  ): Promise<void> => {
+    const task = transportCloseQueueRef.current.then(async () => {
+      await Promise.all(transports.map(closeTransportSafely));
+    });
+    transportCloseQueueRef.current = task;
+    return task;
+  };
+
+  const closePendingTransport = async () => {
+    const transport = pendingTransportRef.current;
+    pendingTransportRef.current = null;
+    await closeQueuedTransports(transport ? [transport] : []);
+  };
+
+  const closeTransports = async () => {
+    const pendingTransport = pendingTransportRef.current;
+    const activeTransport = transportRef.current;
+    pendingTransportRef.current = null;
+    transportRef.current = null;
+    clientRef.current = null;
+
+    const transports = new Set(
+      [pendingTransport, activeTransport].filter(
+        (transport): transport is StreamableHTTPClientTransport =>
+          transport !== null,
+      ),
+    );
+    await closeQueuedTransports([...transports]);
+  };
+
+  const isCurrentConnection = (generation: number) =>
+    mountedRef.current && generation === connectionGenerationRef.current;
+
+  const createInterruptedAuthError = (cause?: unknown) =>
+    new Error(
+      `MCP server "${props.id}" authorization was interrupted before completion.`,
+      cause === undefined ? undefined : { cause },
+    );
 
   const withConnectionTimeout = useEffectEvent(
     async <T>(
@@ -105,7 +163,10 @@ const useMcpServerResource = (
   );
 
   const finalizeConnect = useEffectEvent(
-    async (transport: StreamableHTTPClientTransport) => {
+    async (
+      transport: StreamableHTTPClientTransport,
+      generation: number,
+    ): Promise<boolean> => {
       const client = new Client({
         name: "assistant-ui-mcp",
         version: "0.0.0",
@@ -120,6 +181,7 @@ const useMcpServerResource = (
         "connecting",
         startedAt,
       );
+      if (!isCurrentConnection(generation)) return false;
       // 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
@@ -129,6 +191,9 @@ const useMcpServerResource = (
         "listing tools",
         startedAt,
       );
+      if (!isCurrentConnection(generation)) return false;
+
+      pendingTransportRef.current = null;
       clientRef.current = client;
       transportRef.current = transport;
       setTools(
@@ -142,25 +207,16 @@ const useMcpServerResource = (
         }),
       );
       setConnectionState("connected");
+      return true;
     },
   );
 
-  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 = useEffectEvent(async () => {
+    const generation = ++connectionGenerationRef.current;
     // Close any prior transport/client so a re-connect doesn't leak.
-    await closeTransport();
+    await closeTransports();
+    if (!isCurrentConnection(generation)) return;
+
     setConnectionState("connecting");
     setLastError(null);
     setAuthorizationUrl(null);
@@ -170,24 +226,29 @@ const useMcpServerResource = (
     let transport: StreamableHTTPClientTransport | null = null;
     try {
       transport = await buildTransport();
+      if (!isCurrentConnection(generation)) {
+        await closeQueuedTransports([transport]);
+        return;
+      }
+      pendingTransportRef.current = transport;
       // 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);
+      await finalizeConnect(transport, generation);
     } catch (err) {
+      if (!isCurrentConnection(generation)) return;
+
       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.
+        pendingTransportRef.current = null;
         transportRef.current = transport;
         setConnectionState("authRequired");
       } else {
         if (transport) {
-          try {
-            await transport.close();
-          } catch {
-            // ignore close errors
-          }
+          pendingTransportRef.current = null;
+          await closeQueuedTransports([transport]);
         }
         setLastError({
           message: err instanceof Error ? err.message : String(err),
@@ -198,13 +259,18 @@ const useMcpServerResource = (
   });
 
   const doDisconnect = useEffectEvent(async () => {
+    connectionGenerationRef.current += 1;
     setTools([]);
     setAuthorizationUrl(null);
     setConnectionState("disconnected");
-    await closeTransport();
+    await closeTransports();
   });
 
   const doCompleteAuth = useEffectEvent(async (callbackUrl: string) => {
+    const generation = ++connectionGenerationRef.current;
+    await closePendingTransport();
+    if (!isCurrentConnection(generation)) throw createInterruptedAuthError();
+
     setConnectionState("authPending");
     setLastError(null);
     try {
@@ -214,17 +280,30 @@ const useMcpServerResource = (
       let transport = transportRef.current;
       if (!transport) {
         transport = await buildTransport();
-        transportRef.current = transport;
+        if (!isCurrentConnection(generation)) {
+          await closeQueuedTransports([transport]);
+          throw createInterruptedAuthError();
+        }
       }
+      transportRef.current = null;
+      clientRef.current = null;
+      pendingTransportRef.current = transport;
       await transport.finishAuth(code);
+      if (!isCurrentConnection(generation)) throw createInterruptedAuthError();
       setAuthorizationUrl(null);
-      await finalizeConnect(transport);
+      const connected = await finalizeConnect(transport, generation);
+      if (!connected) throw createInterruptedAuthError();
     } catch (err) {
-      await closeTransport();
+      const error = err instanceof Error ? err : new Error(String(err));
+      if (!isCurrentConnection(generation))
+        throw createInterruptedAuthError(error);
+
+      await closeTransports();
       setLastError({
-        message: err instanceof Error ? err.message : String(err),
+        message: error.message,
       });
       setConnectionState("error");
+      throw error;
     }
   });
 
@@ -253,14 +332,22 @@ const useMcpServerResource = (
 
   // Auto-connect on mount when usable auth exists.
   useEffect(() => {
+    const previousDisposal = pendingDisposalRef.current;
+    if (previousDisposal) previousDisposal.cancelled = true;
+    const pendingDisposal = { cancelled: false };
+    pendingDisposalRef.current = pendingDisposal;
+    mountedRef.current = true;
     const signal = { cancelled: false };
     void tryAutoConnect(signal);
     return () => {
+      mountedRef.current = false;
       signal.cancelled = true;
-      const t = transportRef.current;
-      transportRef.current = null;
-      clientRef.current = null;
-      if (t) t.close().catch(() => {});
+      // Defer disposal so StrictMode can replay setup before closing the transport.
+      queueMicrotask(() => {
+        if (pendingDisposal.cancelled) return;
+        connectionGenerationRef.current += 1;
+        void closeTransports();
+      });
     };
   }, []);
 
@@ -308,6 +395,13 @@ const useMcpServerResource = (
         arguments: args as Record<string, unknown> | undefined,
       });
     },
+    listResources: async (params) => {
+      const client = clientRef.current;
+      if (!client) {
+        throw new Error(`MCP server "${props.id}" is not connected`);
+      }
+      return await client.listResources(params);
+    },
     readResource: async (uri) => {
       const client = clientRef.current;
       if (!client) {

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

@@ -0,0 +1,240 @@
+import { createTapRoot, useResource } from "@assistant-ui/tap";
+import { describe, expect, it } from "vitest";
+
+import {
+  McpLocalStorage,
+  normalizeCustomServerRecords,
+  normalizePersistedAuthState,
+} from "./McpLocalStorage";
+
+const validRecord = {
+  id: "docs",
+  name: "Docs",
+  url: "https://docs.example.com/mcp",
+  auth: { type: "none" },
+  createdAt: 1,
+} as const;
+
+describe("normalizeCustomServerRecords", () => {
+  it("returns an empty list when the persisted value is not an array", () => {
+    expect(normalizeCustomServerRecords({ bad: true })).toEqual([]);
+    expect(normalizeCustomServerRecords(null)).toEqual([]);
+  });
+
+  it("filters malformed custom server entries", () => {
+    expect(
+      normalizeCustomServerRecords([
+        validRecord,
+        null,
+        { ...validRecord, id: "" },
+        { ...validRecord, id: "docs__search" },
+        { ...validRecord, name: "" },
+        { ...validRecord, url: " " },
+        { ...validRecord, name: 123 },
+        { ...validRecord, url: null },
+        { ...validRecord, createdAt: Number.NaN },
+        { ...validRecord, auth: { type: "bearer", token: "" } },
+        { ...validRecord, auth: { type: "bearer", token: 123 } },
+      ]),
+    ).toEqual([validRecord]);
+  });
+
+  it("strips malformed connection timeouts without dropping servers", () => {
+    expect(
+      normalizeCustomServerRecords([
+        { ...validRecord, id: "string-timeout", connectionTimeout: "10000" },
+        { ...validRecord, id: "null-timeout", connectionTimeout: null },
+        { ...validRecord, id: "negative-timeout", connectionTimeout: -1 },
+        { ...validRecord, id: "nan-timeout", connectionTimeout: Number.NaN },
+        {
+          ...validRecord,
+          id: "infinite-timeout",
+          connectionTimeout: Number.POSITIVE_INFINITY,
+        },
+      ]),
+    ).toEqual([
+      { ...validRecord, id: "string-timeout" },
+      { ...validRecord, id: "null-timeout" },
+      { ...validRecord, id: "negative-timeout" },
+      { ...validRecord, id: "nan-timeout" },
+      { ...validRecord, id: "infinite-timeout" },
+    ]);
+  });
+
+  it("accepts finite non-negative connection timeouts", () => {
+    expect(
+      normalizeCustomServerRecords([
+        { ...validRecord, connectionTimeout: 0 },
+        { ...validRecord, connectionTimeout: 10_000 },
+      ]),
+    ).toEqual([
+      { ...validRecord, connectionTimeout: 0 },
+      { ...validRecord, connectionTimeout: 10_000 },
+    ]);
+  });
+
+  it("accepts persisted bearer and oauth auth configs", () => {
+    expect(
+      normalizeCustomServerRecords([
+        {
+          ...validRecord,
+          id: "private-docs",
+          auth: { type: "bearer", token: "token" },
+        },
+        {
+          ...validRecord,
+          id: "oauth-docs",
+          auth: {
+            type: "oauth",
+            scopes: ["docs.read"],
+            authorizationEndpoint: "https://docs.example.com/oauth/authorize",
+            tokenEndpoint: "https://docs.example.com/oauth/token",
+          },
+        },
+      ]),
+    ).toHaveLength(2);
+  });
+});
+
+describe("normalizePersistedAuthState", () => {
+  it("returns null when the persisted value is not an object", () => {
+    expect(normalizePersistedAuthState(null)).toBeNull();
+    expect(normalizePersistedAuthState("bad")).toBeNull();
+    expect(normalizePersistedAuthState([])).toBeNull();
+  });
+
+  it("drops malformed auth state fields", () => {
+    expect(
+      normalizePersistedAuthState({
+        token: "",
+        codeVerifier: 123,
+        tokens: {
+          access_token: "access-token",
+          token_type: "Bearer",
+          expires_in: Number.POSITIVE_INFINITY,
+        },
+        clientInformation: { client_id: "client-id" },
+      }),
+    ).toBeNull();
+  });
+
+  it("keeps valid bearer and OAuth callback state", () => {
+    expect(
+      normalizePersistedAuthState({
+        token: "bearer-token",
+        codeVerifier: "pkce-verifier",
+      }),
+    ).toEqual({
+      token: "bearer-token",
+      codeVerifier: "pkce-verifier",
+    });
+  });
+
+  it("keeps valid OAuth tokens and client information", () => {
+    const tokens = {
+      access_token: "access-token",
+      token_type: "Bearer",
+      refresh_token: "refresh-token",
+      expires_in: 3600,
+      scope: "docs.read",
+    };
+    const clientInformation = {
+      client_id: "client-id",
+      client_secret: "client-secret",
+      redirect_uris: ["http://localhost/callback"],
+    };
+
+    expect(
+      normalizePersistedAuthState({
+        tokens,
+        clientInformation,
+      }),
+    ).toEqual({
+      tokens,
+      clientInformation,
+    });
+  });
+
+  it("keeps valid fields when neighboring fields are malformed", () => {
+    expect(
+      normalizePersistedAuthState({
+        token: "bearer-token",
+        codeVerifier: 123,
+        tokens: { access_token: "access-token" },
+        clientInformation: "not-client-info",
+      }),
+    ).toEqual({
+      token: "bearer-token",
+    });
+  });
+});
+
+const createStorage = (): Storage => {
+  const data = new Map<string, string>();
+  return {
+    get length() {
+      return data.size;
+    },
+    clear: () => {
+      data.clear();
+    },
+    getItem: (key) => data.get(key) ?? null,
+    key: (index) => Array.from(data.keys())[index] ?? null,
+    removeItem: (key) => {
+      data.delete(key);
+    },
+    setItem: (key, value) => {
+      data.set(key, value);
+    },
+  };
+};
+
+const loadStorage = (storage: Storage) =>
+  createTapRoot(function McpStorageRoot() {
+    return useResource(
+      McpLocalStorage({
+        keyPrefix: "test-mcp",
+        storage,
+      }),
+    );
+  }).getValue();
+
+describe("McpLocalStorage custom servers", () => {
+  it("strips malformed connection timeouts when loading", async () => {
+    const storage = createStorage();
+    storage.setItem(
+      "test-mcp:custom-servers",
+      JSON.stringify([
+        validRecord,
+        {
+          ...validRecord,
+          id: "bad-timeout",
+          connectionTimeout: "immediately",
+        },
+      ]),
+    );
+
+    await expect(loadStorage(storage).loadCustomServers()).resolves.toEqual([
+      validRecord,
+      { ...validRecord, id: "bad-timeout" },
+    ]);
+  });
+});
+
+describe("McpLocalStorage auth state", () => {
+  it("normalizes loaded auth state from localStorage", async () => {
+    const storage = createStorage();
+    storage.setItem(
+      "test-mcp:auth:docs",
+      JSON.stringify({
+        token: "bearer-token",
+        codeVerifier: 123,
+        tokens: "not-tokens",
+      }),
+    );
+
+    await expect(loadStorage(storage).loadAuthState("docs")).resolves.toEqual({
+      token: "bearer-token",
+    });
+  });
+});

package/src/resources/storage/McpLocalStorage.ts

@@ -1,6 +1,11 @@
 import { resource } from "@assistant-ui/tap";
-import type { MCPCustomServerRecord } from "../../mcp-scope";
+import {
+  OAuthClientInformationFullSchema,
+  OAuthTokensSchema,
+} from "@modelcontextprotocol/sdk/shared/auth.js";
+import type { MCPAuthConfig, MCPCustomServerRecord } from "../../mcp-scope";
 import type { MCPPersistedAuthState } from "../../auth/types";
+import { assertValidServerId } from "../../utils/serverId";
 import type { MCPStorage } from "./types";
 
 export type McpLocalStorageOptions = {
@@ -22,6 +27,132 @@ function resolveStorage(opts: McpLocalStorageOptions): Storage | null {
   return null;
 }
 
+const isRecord = (value: unknown): value is Record<string, unknown> =>
+  typeof value === "object" && value !== null && !Array.isArray(value);
+
+const isOptionalString = (value: unknown): value is string | undefined =>
+  value === undefined || typeof value === "string";
+
+const isNonEmptyString = (value: unknown): value is string =>
+  typeof value === "string" && value.trim().length > 0;
+
+const isOptionalNonEmptyString = (
+  value: unknown,
+): value is string | undefined =>
+  value === undefined || isNonEmptyString(value);
+
+const isOptionalStringArray = (value: unknown): value is string[] | undefined =>
+  value === undefined ||
+  (Array.isArray(value) && value.every((item) => typeof item === "string"));
+
+const isOptionalConnectionTimeout = (
+  value: unknown,
+): value is number | undefined =>
+  value === undefined ||
+  (typeof value === "number" && Number.isFinite(value) && value >= 0);
+
+const isValidServerId = (id: string): boolean => {
+  try {
+    assertValidServerId(id);
+    return true;
+  } catch {
+    return false;
+  }
+};
+
+const isMCPAuthConfig = (auth: unknown): auth is MCPAuthConfig => {
+  if (!isRecord(auth)) return false;
+
+  switch (auth.type) {
+    case "none":
+      return true;
+    case "bearer":
+      return isOptionalNonEmptyString(auth.token);
+    case "oauth":
+      return (
+        isOptionalStringArray(auth.scopes) &&
+        isOptionalString(auth.authorizationEndpoint) &&
+        isOptionalString(auth.tokenEndpoint) &&
+        isOptionalString(auth.registrationEndpoint) &&
+        isOptionalString(auth.clientId) &&
+        isOptionalString(auth.clientSecret)
+      );
+    default:
+      return false;
+  }
+};
+
+const isCustomServerRecord = (
+  value: unknown,
+): value is MCPCustomServerRecord => {
+  if (!isRecord(value)) return false;
+  if (typeof value.id !== "string" || !isValidServerId(value.id)) {
+    return false;
+  }
+  return (
+    isNonEmptyString(value.name) &&
+    isNonEmptyString(value.url) &&
+    Number.isFinite(value.createdAt) &&
+    isMCPAuthConfig(value.auth) &&
+    isOptionalConnectionTimeout(value.connectionTimeout)
+  );
+};
+
+const normalizeCustomServerRecord = (
+  value: unknown,
+): MCPCustomServerRecord | null => {
+  if (isCustomServerRecord(value)) return value;
+  if (!isRecord(value)) return null;
+
+  const record = { ...value };
+  delete record.connectionTimeout;
+  return isCustomServerRecord(record) ? record : null;
+};
+
+export const normalizeCustomServerRecords = (
+  value: unknown,
+): MCPCustomServerRecord[] => {
+  if (!Array.isArray(value)) return [];
+  return value.flatMap((item) => {
+    const record = normalizeCustomServerRecord(item);
+    return record === null ? [] : [record];
+  });
+};
+
+const normalizeOAuthTokens = (
+  value: unknown,
+): MCPPersistedAuthState["tokens"] | undefined => {
+  const result = OAuthTokensSchema.safeParse(value);
+  return result.success ? result.data : undefined;
+};
+
+const normalizeClientInformation = (
+  value: unknown,
+): MCPPersistedAuthState["clientInformation"] | undefined => {
+  const result = OAuthClientInformationFullSchema.safeParse(value);
+  return result.success ? result.data : undefined;
+};
+
+export const normalizePersistedAuthState = (
+  value: unknown,
+): MCPPersistedAuthState | null => {
+  if (!isRecord(value)) return null;
+
+  const state: MCPPersistedAuthState = {};
+  if (isNonEmptyString(value.token)) state.token = value.token;
+  if (isNonEmptyString(value.codeVerifier)) {
+    state.codeVerifier = value.codeVerifier;
+  }
+
+  const tokens = normalizeOAuthTokens(value.tokens);
+  if (tokens) state.tokens = tokens;
+
+  const clientInformation = normalizeClientInformation(value.clientInformation);
+  if (clientInformation) state.clientInformation = clientInformation;
+
+  return Object.keys(state).length > 0 ? state : null;
+};
+
 const useMcpLocalStorage = (opts: McpLocalStorageOptions = {}): MCPStorage => {
   const prefix = opts.keyPrefix ?? "aui-mcp";
   const customServersKey = `${prefix}:custom-servers`;
@@ -59,12 +190,12 @@ const useMcpLocalStorage = (opts: McpLocalStorageOptions = {}): MCPStorage => {
 
   return {
     loadCustomServers: async () =>
-      read<MCPCustomServerRecord[]>(customServersKey, []),
+      normalizeCustomServerRecords(read<unknown>(customServersKey, [])),
     saveCustomServers: async (records) => {
       write(customServersKey, records);
     },
     loadAuthState: async (id) =>
-      read<MCPPersistedAuthState | null>(authKey(id), null),
+      normalizePersistedAuthState(read<unknown>(authKey(id), null)),
     saveAuthState: async (id, state) => {
       write(authKey(id), state);
     },