@assistant-ui/react-mcp 0.0.17 → 0.0.18

package/src/hooks/useMcpOAuthCallback.tsx

@@ -2,6 +2,20 @@ import { type FC, type ReactNode, useEffect, useRef, useState } from "react";
 import { useAui } from "@assistant-ui/store";
 import { decodeServerIdFromState } from "../auth/createOAuthProvider";
 
+export const createMcpOAuthCallbackError = (
+  err: unknown,
+  serverId: string | null,
+): Error => {
+  const message = err instanceof Error ? err.message : String(err);
+  if (serverId) {
+    return new Error(
+      `MCP OAuth callback for server "${serverId}" failed: ${message}`,
+      { cause: err },
+    );
+  }
+  return new Error(`MCP OAuth callback failed: ${message}`, { cause: err });
+};
+
 export type UseMcpOAuthCallbackOptions = {
   /** Defaults to `window.location.href`. */
   url?: string;
@@ -39,27 +53,28 @@ export function useMcpOAuthCallback(
     startedRef.current = url;
 
     (async () => {
+      let serverId: string | null = null;
       try {
         const parsed = new URL(url);
         const state = parsed.searchParams.get("state");
+        if (state) serverId = decodeServerIdFromState(state);
         const error = parsed.searchParams.get("error");
         if (error) {
           throw new Error(
             parsed.searchParams.get("error_description") ?? error,
           );
         }
-        if (!state) throw new Error("missing state parameter in callback URL");
-        const serverId = decodeServerIdFromState(state);
+        if (!state) throw new Error('missing "state" parameter');
         if (!serverId) {
-          throw new Error("callback state does not match an MCP server");
+          throw new Error("state was not created by assistant-ui MCP");
         }
         setResult({ status: "running", serverId, error: null });
         await aui.mcp().server({ id: serverId }).completeAuth(url);
         setResult({ status: "done", serverId, error: null });
         optsRef.current.onComplete?.(serverId);
       } catch (err) {
-        const e = err instanceof Error ? err : new Error(String(err));
-        setResult((prev) => ({ ...prev, status: "error", error: e }));
+        const e = createMcpOAuthCallbackError(err, serverId);
+        setResult({ status: "error", serverId, error: e });
         optsRef.current.onError?.(e);
       }
     })();

package/src/mcp-scope.ts

@@ -72,6 +72,8 @@ export type MCPServerMethods = {
   disconnect: () => Promise<void>;
   remove: () => Promise<void>;
   callTool: (name: string, args: unknown) => Promise<unknown>;
+  /** List resources exposed by the server. Returns the raw MCP `ListResourcesResult`. */
+  listResources: (params?: { cursor?: string | undefined }) => Promise<unknown>;
   /** Read a resource by URI. Returns the raw MCP `ReadResourceResult`. */
   readResource: (uri: string) => Promise<unknown>;
   /** OAuth only: pass full callback URL (e.g. window.location.href) */

package/src/resources/McpServerResource.test.ts

@@ -1,5 +1,6 @@
 import { createTapRoot, useResource } from "@assistant-ui/tap";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import type { MCPAuthConfig } from "../mcp-scope";
 import type { MCPStorage } from "./storage/types";
 
 const mocks = vi.hoisted(() => {
@@ -15,6 +16,7 @@ const mocks = vi.hoisted(() => {
       }>;
     }>
   > = [];
+  const finishAuthResults: Array<() => Promise<void>> = [];
 
   const Client = vi.fn().mockImplementation(function Client(this: any) {
     const index = clients.length;
@@ -23,6 +25,7 @@ const mocks = vi.hoisted(() => {
       () => listToolsResults[index]?.() ?? Promise.resolve({ tools: [] }),
     );
     this.callTool = vi.fn();
+    this.listResources = vi.fn(() => Promise.resolve({ resources: [] }));
     this.readResource = vi.fn();
     clients.push(this);
   });
@@ -30,8 +33,11 @@ const mocks = vi.hoisted(() => {
   const StreamableHTTPClientTransport = vi
     .fn()
     .mockImplementation(function StreamableHTTPClientTransport(this: any) {
+      const index = transports.length;
       this.close = vi.fn(() => Promise.resolve());
-      this.finishAuth = vi.fn(() => Promise.resolve());
+      this.finishAuth = vi.fn(
+        () => finishAuthResults[index]?.() ?? Promise.resolve(),
+      );
       transports.push(this);
     });
 
@@ -42,6 +48,7 @@ const mocks = vi.hoisted(() => {
     transports,
     connectResults,
     listToolsResults,
+    finishAuthResults,
   };
 });
 
@@ -61,6 +68,18 @@ const tick = async () => {
   await Promise.resolve();
 };
 
+const flushMacrotask = async () => {
+  await new Promise<void>((resolve) => {
+    const channel = new MessageChannel();
+    channel.port1.onmessage = () => {
+      channel.port1.close();
+      channel.port2.close();
+      resolve();
+    };
+    channel.port2.postMessage(null);
+  });
+};
+
 const waitFor = async (predicate: () => boolean) => {
   for (let i = 0; i < 20; i++) {
     if (predicate()) return;
@@ -77,7 +96,20 @@ const createStorage = (): MCPStorage => ({
   clearAuthState: vi.fn(async () => {}),
 });
 
-const mount = (props?: { connectionTimeout?: number | undefined }) => {
+const resetMocks = () => {
+  mocks.clients.length = 0;
+  mocks.transports.length = 0;
+  mocks.connectResults.length = 0;
+  mocks.listToolsResults.length = 0;
+  mocks.finishAuthResults.length = 0;
+  mocks.Client.mockClear();
+  mocks.StreamableHTTPClientTransport.mockClear();
+};
+
+const mount = (props?: {
+  auth?: MCPAuthConfig | undefined;
+  connectionTimeout?: number | undefined;
+}) => {
   const connectionTimeout =
     props && "connectionTimeout" in props ? props.connectionTimeout : 10_000;
 
@@ -88,7 +120,7 @@ const mount = (props?: { connectionTimeout?: number | undefined }) => {
         kind: "connector",
         name: "Docs",
         url: "https://example.com/mcp",
-        auth: { type: "none" },
+        auth: props?.auth ?? { type: "none" },
         storage: createStorage(),
         redirectUri: "https://example.com/callback",
         autoConnect: false,
@@ -102,12 +134,7 @@ const mount = (props?: { connectionTimeout?: number | undefined }) => {
 describe("McpServerResource connectionTimeout", () => {
   beforeEach(() => {
     vi.useFakeTimers();
-    mocks.clients.length = 0;
-    mocks.transports.length = 0;
-    mocks.connectResults.length = 0;
-    mocks.listToolsResults.length = 0;
-    mocks.Client.mockClear();
-    mocks.StreamableHTTPClientTransport.mockClear();
+    resetMocks();
   });
 
   afterEach(() => {
@@ -219,3 +246,119 @@ describe("McpServerResource connectionTimeout", () => {
     }
   });
 });
+
+describe("McpServerResource completeAuth", () => {
+  beforeEach(resetMocks);
+
+  it("rejects when the callback URL has no authorization code", async () => {
+    const root = mount();
+
+    try {
+      await expect(
+        root.getValue().completeAuth("https://example.com/callback?state=abc"),
+      ).rejects.toThrow("missing authorization code in callback URL");
+      await flushMacrotask();
+
+      expect(root.getValue().getState()).toMatchObject({
+        connectionState: "error",
+        lastError: {
+          message: "missing authorization code in callback URL",
+        },
+      });
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("rejects after storing finishAuth failures on the server state", async () => {
+    mocks.finishAuthResults.push(() =>
+      Promise.reject(new Error("invalid_grant")),
+    );
+    const root = mount({ auth: { type: "oauth" } });
+
+    try {
+      await expect(
+        root.getValue().completeAuth("https://example.com/callback?code=abc"),
+      ).rejects.toThrow("invalid_grant");
+      await flushMacrotask();
+
+      expect(root.getValue().getState()).toMatchObject({
+        connectionState: "error",
+        lastError: {
+          message: "invalid_grant",
+        },
+      });
+      expect(mocks.transports[0].finishAuth).toHaveBeenCalledWith("abc");
+      expect(mocks.transports[0].close).toHaveBeenCalledTimes(1);
+    } finally {
+      root.unmount();
+    }
+  });
+});
+
+describe("McpServerResource resource methods", () => {
+  beforeEach(() => {
+    mocks.clients.length = 0;
+    mocks.transports.length = 0;
+    mocks.connectResults.length = 0;
+    mocks.listToolsResults.length = 0;
+    mocks.Client.mockClear();
+    mocks.StreamableHTTPClientTransport.mockClear();
+  });
+
+  it("lists resources from a connected server", async () => {
+    const result = {
+      resources: [
+        {
+          uri: "docs://intro",
+          name: "Intro",
+          mimeType: "text/markdown",
+        },
+      ],
+    };
+    const root = mount();
+
+    try {
+      await root.getValue().connect();
+      mocks.clients[0].listResources.mockResolvedValueOnce(result);
+
+      await expect(root.getValue().listResources()).resolves.toBe(result);
+      expect(mocks.clients[0].listResources).toHaveBeenCalledTimes(1);
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("forwards the resource pagination cursor", async () => {
+    const result = {
+      resources: [{ uri: "docs://page-two", name: "Page two" }],
+    };
+    const root = mount();
+
+    try {
+      await root.getValue().connect();
+      mocks.clients[0].listResources.mockResolvedValueOnce(result);
+
+      await expect(
+        root.getValue().listResources({ cursor: "next-page" }),
+      ).resolves.toBe(result);
+      expect(mocks.clients[0].listResources).toHaveBeenCalledWith({
+        cursor: "next-page",
+      });
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("rejects listResources when the server is disconnected", async () => {
+    const root = mount();
+
+    try {
+      await expect(root.getValue().listResources()).rejects.toThrow(
+        'MCP server "docs" is not connected',
+      );
+    } finally {
+      root.unmount();
+    }
+  });
+});

package/src/resources/McpServerResource.ts

@@ -221,10 +221,12 @@ const useMcpServerResource = (
       await finalizeConnect(transport);
     } catch (err) {
       await closeTransport();
+      const error = err instanceof Error ? err : new Error(String(err));
       setLastError({
-        message: err instanceof Error ? err.message : String(err),
+        message: error.message,
       });
       setConnectionState("error");
+      throw error;
     }
   });
 
@@ -308,6 +310,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,184 @@
+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("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 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,111 @@ 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 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)
+  );
+};
+
+export const normalizeCustomServerRecords = (
+  value: unknown,
+): MCPCustomServerRecord[] => {
+  if (!Array.isArray(value)) return [];
+  return value.filter(isCustomServerRecord);
+};
+
+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 +169,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);
     },