@assistant-ui/react-mcp 0.0.16 → 0.0.18

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);
     },

package/src/utils/serverId.ts

@@ -17,3 +17,15 @@ export function assertValidServerId(id: string): void {
     );
   }
 }
+
+export function assertUniqueServerIds(ids: Iterable<string>): void {
+  const seen = new Set<string>();
+  for (const id of ids) {
+    if (seen.has(id)) {
+      throw new Error(
+        `McpManagerResource received duplicate MCP server id "${id}". Server ids must be unique because they are used for lookups, OAuth routing, and tool name prefixes.`,
+      );
+    }
+    seen.add(id);
+  }
+}