@assistant-ui/react-mcp 0.1.15 → 0.1.16

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

@@ -1,8 +1,10 @@
-import { createTapRoot, useResource } from "@assistant-ui/tap";
+import { createTapRoot, resource, useResource } from "@assistant-ui/tap";
+import { useState } from "react";
 import { auth, type FetchLike } from "@modelcontextprotocol/client";
 import { describe, expect, it } from "vitest";
 import { createOAuthProvider } from "../../auth/createOAuthProvider";
 
+import type { MCPStorage } from "./types";
 import {
   McpLocalStorage,
   normalizeCustomServerRecords,
@@ -125,10 +127,12 @@ describe("normalizePersistedAuthState", () => {
       normalizePersistedAuthState({
         token: "bearer-token",
         codeVerifier: "pkce-verifier",
+        state: "aui-mcp:ZG9jcw.nonce",
       }),
     ).toEqual({
       token: "bearer-token",
       codeVerifier: "pkce-verifier",
+      state: "aui-mcp:ZG9jcw.nonce",
     });
   });
 
@@ -394,6 +398,7 @@ describe("McpLocalStorage auth state", () => {
       JSON.parse(storage.getItem("test-mcp:auth:docs") ?? "null"),
     ).toMatchObject({
       codeVerifier: expect.any(String),
+      state: authorizationUrls[0]!.searchParams.get("state"),
       discoveryState: {
         authorizationServerUrl: "https://auth.example.com",
       },
@@ -418,3 +423,68 @@ describe("McpLocalStorage auth state", () => {
     });
   });
 });
+
+describe("McpLocalStorage instance identity", () => {
+  it("derives a scope from the prefix for the shared default backing", () => {
+    let storage!: MCPStorage;
+
+    createTapRoot(function McpStorageScopeRoot() {
+      storage = useResource(McpLocalStorage({ keyPrefix: "test-mcp" }));
+      return storage;
+    });
+
+    expect(storage.scopeId).toBe("local-storage:test-mcp");
+  });
+
+  it("declares no scope for a custom backing store unless one is named", () => {
+    const backing = createStorage();
+    let unnamed!: MCPStorage;
+    let named!: MCPStorage;
+
+    createTapRoot(function McpStorageCustomScopeRoot() {
+      unnamed = useResource(
+        McpLocalStorage({ keyPrefix: "test-mcp", storage: backing }),
+      );
+      return unnamed;
+    });
+    createTapRoot(function McpStorageNamedScopeRoot() {
+      named = useResource(
+        McpLocalStorage({
+          keyPrefix: "test-mcp",
+          storage: backing,
+          scopeId: "session:alpha",
+        }),
+      );
+      return named;
+    });
+
+    expect(unnamed.scopeId).toBeUndefined();
+    expect(named.scopeId).toBe("session:alpha");
+  });
+
+  it("returns the same instance across re-renders", () => {
+    const backing = createStorage();
+    const seen: MCPStorage[] = [];
+    let rerender!: () => void;
+
+    const useHost = () => {
+      const [, setTick] = useState(0);
+      rerender = () => setTick((n) => n + 1);
+      const storage = useResource(
+        McpLocalStorage({ keyPrefix: "test-mcp", storage: backing }),
+      );
+      seen.push(storage);
+      return storage;
+    };
+    const Host = resource(useHost);
+
+    createTapRoot(function McpStorageIdentityRoot() {
+      return useResource(Host());
+    });
+    rerender();
+    rerender();
+
+    expect(seen.length).toBeGreaterThan(1);
+    expect(new Set(seen).size).toBe(1);
+  });
+});

package/src/resources/storage/McpLocalStorage.ts

@@ -1,4 +1,5 @@
 import { resource } from "@assistant-ui/tap";
+import { useMemo } from "react";
 import {
   OAuthMetadataSchema,
   OAuthClientInformationFullSchema,
@@ -15,6 +16,13 @@ export type McpLocalStorageOptions = {
   keyPrefix?: string;
   /** Override the underlying Storage. Defaults to globalThis.localStorage. */
   storage?: Storage;
+  /**
+   * Stable identity for the backing data, used to key server reconnects.
+   * Required to get reconnect-on-swap behavior when `storage` is overridden;
+   * without it a custom backing store declares no scope, since a prefix
+   * alone cannot distinguish two different stores.
+   */
+  scopeId?: string;
 };
 
 function resolveStorage(opts: McpLocalStorageOptions): Storage | null {
@@ -193,6 +201,7 @@ export const normalizePersistedAuthState = (
   if (isNonEmptyString(value.codeVerifier)) {
     state.codeVerifier = value.codeVerifier;
   }
+  if (isNonEmptyString(value.state)) state.state = value.state;
 
   const tokens = normalizeOAuthTokens(value.tokens);
   if (tokens) state.tokens = tokens;
@@ -208,54 +217,67 @@ export const normalizePersistedAuthState = (
 
 const useMcpLocalStorage = (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 () =>
-      normalizeCustomServerRecords(read<unknown>(customServersKey, [])),
-    saveCustomServers: async (records) => {
-      write(customServersKey, records);
-    },
-    loadAuthState: async (id) =>
-      normalizePersistedAuthState(read<unknown>(authKey(id), null)),
-    saveAuthState: async (id, state) => {
-      write(authKey(id), state);
-    },
-    clearAuthState: async (id) => {
-      remove(authKey(id));
-    },
-  };
+  // Deriving a scope from the prefix is only honest for the shared
+  // globalThis.localStorage; two custom backing stores under one prefix hold
+  // different data, so an overridden backing declares no scope unless the
+  // caller names one.
+  const scopeId =
+    opts.scopeId ??
+    (opts.storage === undefined ? `local-storage:${prefix}` : undefined);
+
+  // Callers key per-server coordination state on this instance, so it has to
+  // stay referentially stable for as long as the underlying store does.
+  return useMemo(() => {
+    const customServersKey = `${prefix}:custom-servers`;
+    const authKey = (id: string) => `${prefix}:auth:${id}`;
+
+    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 {
+      ...(scopeId !== undefined ? { scopeId } : {}),
+      loadCustomServers: async () =>
+        normalizeCustomServerRecords(read<unknown>(customServersKey, [])),
+      saveCustomServers: async (records) => {
+        write(customServersKey, records);
+      },
+      loadAuthState: async (id) =>
+        normalizePersistedAuthState(read<unknown>(authKey(id), null)),
+      saveAuthState: async (id, state) => {
+        write(authKey(id), state);
+      },
+      clearAuthState: async (id) => {
+        remove(authKey(id));
+      },
+    };
+  }, [prefix, storage, scopeId]);
 };
 
 export const McpLocalStorage = resource(useMcpLocalStorage);

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

@@ -0,0 +1,99 @@
+import { createTapRoot, resource, useResource } from "@assistant-ui/tap";
+import { describe, expect, it } from "vitest";
+import { useState } from "react";
+import type { MCPStorage } from "./types";
+import { McpMemoryStorage } from "./McpMemoryStorage";
+
+const mountWithRerender = () => {
+  const seen: MCPStorage[] = [];
+  let setTick!: (update: (value: number) => number) => void;
+
+  const useHost = () => {
+    const [, setValue] = useState(0);
+    setTick = setValue;
+    const storage = useResource(McpMemoryStorage());
+    seen.push(storage);
+    return storage;
+  };
+  const Host = resource(useHost);
+
+  createTapRoot(function MemoryStorageRoot() {
+    return useResource(Host());
+  });
+
+  return {
+    seen,
+    latest: () => seen[seen.length - 1]!,
+    rerender: () => setTick((value) => value + 1),
+  };
+};
+
+describe("McpMemoryStorage", () => {
+  it("returns the same instance across re-renders", () => {
+    const { seen, rerender } = mountWithRerender();
+
+    rerender();
+    rerender();
+
+    expect(seen.length).toBeGreaterThan(1);
+    expect(new Set(seen).size).toBe(1);
+  });
+
+  it("keeps persisted auth state across re-renders", async () => {
+    const { seen, latest, rerender } = mountWithRerender();
+    await seen[0]!.saveAuthState("docs", { codeVerifier: "pkce-verifier" });
+
+    rerender();
+
+    await expect(latest().loadAuthState("docs")).resolves.toEqual({
+      codeVerifier: "pkce-verifier",
+    });
+  });
+
+  it("keeps custom servers across re-renders", async () => {
+    const { seen, latest, rerender } = mountWithRerender();
+    await seen[0]!.saveCustomServers([
+      {
+        id: "docs",
+        name: "Docs",
+        url: "https://docs.example.com/mcp",
+        auth: { type: "none" },
+        createdAt: 1,
+      },
+    ]);
+
+    rerender();
+
+    await expect(latest().loadCustomServers()).resolves.toHaveLength(1);
+  });
+});
+
+describe("McpMemoryStorage scope identity", () => {
+  it("scopes each instance uniquely and keeps it stable per instance", () => {
+    let first!: MCPStorage;
+    let second!: MCPStorage;
+    const seen: (string | undefined)[] = [];
+    let rerender!: () => void;
+
+    const Host = resource(function useHost() {
+      const [, setTick] = useState(0);
+      rerender = () => setTick((n) => n + 1);
+      first = useResource(McpMemoryStorage());
+      seen.push(first.scopeId);
+      return first;
+    });
+    createTapRoot(function MemoryScopeRootA() {
+      return useResource(Host());
+    });
+    createTapRoot(function MemoryScopeRootB() {
+      second = useResource(McpMemoryStorage());
+      return second;
+    });
+    rerender();
+
+    expect(first.scopeId).toMatch(/^memory:/);
+    expect(second.scopeId).toMatch(/^memory:/);
+    expect(first.scopeId).not.toBe(second.scopeId);
+    expect(new Set(seen).size).toBe(1);
+  });
+});

package/src/resources/storage/McpMemoryStorage.ts

@@ -1,24 +1,30 @@
 import { resource } from "@assistant-ui/tap";
+import { useMemo } from "react";
+import { generateId } from "@assistant-ui/core";
 import type { MCPCustomServerRecord } from "../../mcp-scope";
 import type { MCPPersistedAuthState } from "../../auth/types";
 import type { MCPStorage } from "./types";
 
-const useMcpMemoryStorage = (): 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);
-    },
-  };
-};
+const useMcpMemoryStorage = (): MCPStorage =>
+  useMemo(() => {
+    let servers: MCPCustomServerRecord[] = [];
+    const auth = new Map<string, MCPPersistedAuthState>();
+    return {
+      // Each memory store is its own private data: a distinct instance is a
+      // distinct scope, so replacing one keys a reconnect.
+      scopeId: `memory:${generateId()}`,
+      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(useMcpMemoryStorage);

package/src/resources/storage/types.ts

@@ -3,6 +3,18 @@ import type { MCPCustomServerRecord } from "../../mcp-scope";
 import type { MCPPersistedAuthState } from "../../auth/types";
 
 export type MCPStorage = {
+  /**
+   * Stable identity of the backing store. Two storages with the same scopeId
+   * must read and write the same persisted data. When present, server
+   * connections and the OAuth write fence key on it, so swapping to a
+   * differently-scoped storage reconnects instead of leaving a live OAuth flow
+   * on the replaced store, and clearing through a same-scoped replacement still
+   * waits for writes queued against the storage it replaced. When absent, the
+   * fence falls back to object identity while connections never re-key at all,
+   * so a storage rebuilt on every render has to declare a scopeId; without one
+   * a clear runs unfenced against the writes queued by the object it replaced.
+   */
+  scopeId?: string;
   loadCustomServers: () => Promise<MCPCustomServerRecord[]>;
   saveCustomServers: (records: MCPCustomServerRecord[]) => Promise<void>;
   loadAuthState: (serverId: string) => Promise<MCPPersistedAuthState | null>;