@assistant-ui/react-mcp 0.1.20 → 0.1.22

package/src/resources/McpCustomServerPersistence.ts

@@ -0,0 +1,233 @@
+import {
+  useState,
+  useEffect,
+  useEffectEvent,
+  useRef,
+  useCallback,
+} from "react";
+import { resource } from "@assistant-ui/tap";
+import type { MCPCustomServerRecord } from "../mcp-scope";
+import type { MCPStorage } from "./storage/types";
+
+const reportCustomStorageFailure = (
+  operation: "load" | "save",
+  error: unknown,
+) => {
+  console.error(
+    `[assistant-ui/react-mcp] failed to ${operation} custom servers:`,
+    error,
+  );
+};
+
+const reportBlockedCustomServerPersistence = () => {
+  console.error(
+    "[assistant-ui/react-mcp] custom server changes remain in memory because loading the persisted list failed; remount the manager to retry",
+  );
+};
+
+const persistCustomServers = async (
+  storage: MCPStorage,
+  records: MCPCustomServerRecord[],
+) => {
+  try {
+    await storage.saveCustomServers(records);
+  } catch (error) {
+    reportCustomStorageFailure("save", error);
+  }
+};
+
+export type CustomServerPersistenceQueues = Map<string, Promise<void>>;
+
+const enqueueCustomServerTask = (
+  persistenceQueues: CustomServerPersistenceQueues,
+  scopeKey: string,
+  task: () => Promise<void>,
+) => {
+  const previous = persistenceQueues.get(scopeKey);
+  const next = (previous ?? Promise.resolve()).then(task);
+  persistenceQueues.set(scopeKey, next);
+  void next.then(() => {
+    if (persistenceQueues.get(scopeKey) === next) {
+      persistenceQueues.delete(scopeKey);
+    }
+  });
+};
+
+const enqueueCustomServerPersistence = (
+  persistenceQueues: CustomServerPersistenceQueues,
+  scopeKey: string,
+  storage: MCPStorage,
+  records: MCPCustomServerRecord[],
+) =>
+  enqueueCustomServerTask(persistenceQueues, scopeKey, () =>
+    persistCustomServers(storage, records),
+  );
+
+export const holdCustomServerPersistence = (
+  persistenceQueues: CustomServerPersistenceQueues,
+  scopeKey: string,
+) => {
+  let release = () => {};
+  const gate = new Promise<void>((resolve) => {
+    release = resolve;
+  });
+  enqueueCustomServerTask(persistenceQueues, scopeKey, () => gate);
+  return release;
+};
+
+const deduplicateCustomServers = (records: MCPCustomServerRecord[]) => {
+  const seen = new Set<string>();
+  return records.filter((record) => {
+    if (seen.has(record.id)) {
+      console.error(
+        `[assistant-ui/react-mcp] ignored duplicate custom server id "${record.id}" loaded from storage`,
+      );
+      return false;
+    }
+    seen.add(record.id);
+    return true;
+  });
+};
+
+type McpCustomServersResourceProps = {
+  storage: MCPStorage;
+  scopeKey: string;
+  persistenceQueues: CustomServerPersistenceQueues;
+};
+
+const useMcpCustomServersResource = ({
+  storage,
+  scopeKey,
+  persistenceQueues,
+}: McpCustomServersResourceProps) => {
+  const [customServers, setCustomServers] = useState<MCPCustomServerRecord[]>(
+    [],
+  );
+  const [isHydrated, setIsHydrated] = useState(false);
+
+  const customServersRef = useRef<MCPCustomServerRecord[]>([]);
+  const hydrationStateRef = useRef<"pending" | "succeeded" | "failed">(
+    "pending",
+  );
+  const hasPendingMutationRef = useRef(false);
+  const [removedBeforeHydration] = useState(() => new Set<string>());
+  const reportedBlockedPersistenceRef = useRef(false);
+
+  const hydrate = useEffectEvent(async (signal: { cancelled: boolean }) => {
+    // A revisited scope must not read behind writes still queued against it.
+    while (true) {
+      const pendingPersistence = persistenceQueues.get(scopeKey);
+      if (!pendingPersistence) break;
+      await pendingPersistence;
+      if (signal.cancelled) return;
+      if (persistenceQueues.get(scopeKey) === pendingPersistence) break;
+    }
+
+    let records: Awaited<ReturnType<typeof storage.loadCustomServers>>;
+    try {
+      const loadedRecords = await storage.loadCustomServers();
+      records = deduplicateCustomServers(loadedRecords);
+    } catch (error) {
+      if (!signal.cancelled) {
+        reportCustomStorageFailure("load", error);
+        hydrationStateRef.current = "failed";
+        if (hasPendingMutationRef.current) {
+          reportBlockedCustomServerPersistence();
+          reportedBlockedPersistenceRef.current = true;
+        }
+        setIsHydrated(true);
+      }
+      return;
+    }
+    // Merge rather than replace so any addCustomServer calls that
+    // happened before hydration resolved aren't silently overwritten.
+    // Persisted order wins; pre-hydration locals append.
+    const hadPendingMutation = hasPendingMutationRef.current;
+    const hydratedRecords = records.filter(
+      (record) => !removedBeforeHydration.has(record.id),
+    );
+    const mergedRecords = (() => {
+      const prev = customServersRef.current;
+      if (prev.length === 0) return hydratedRecords;
+      const persistedIds = new Set(hydratedRecords.map((r) => r.id));
+      return [
+        ...hydratedRecords,
+        ...prev.filter((r) => !persistedIds.has(r.id)),
+      ];
+    })();
+    customServersRef.current = mergedRecords;
+    hydrationStateRef.current = "succeeded";
+    hasPendingMutationRef.current = false;
+    if (hadPendingMutation) {
+      enqueueCustomServerPersistence(
+        persistenceQueues,
+        scopeKey,
+        storage,
+        mergedRecords,
+      );
+    }
+    if (signal.cancelled) return;
+    setCustomServers(mergedRecords);
+    setIsHydrated(true);
+  });
+
+  useEffect(() => {
+    const signal = { cancelled: false };
+    // Hydration reads persisted records asynchronously; there is no earlier
+    // point than mount at which to start it.
+    // eslint-disable-next-line react-hooks/set-state-in-effect
+    void hydrate(signal);
+    return () => {
+      signal.cancelled = true;
+    };
+  }, []);
+
+  const updateCustomServers = useCallback(
+    (
+      updater: (records: MCPCustomServerRecord[]) => MCPCustomServerRecord[],
+    ) => {
+      const next = updater(customServersRef.current);
+      customServersRef.current = next;
+      setCustomServers(next);
+
+      if (hydrationStateRef.current === "succeeded") {
+        enqueueCustomServerPersistence(
+          persistenceQueues,
+          scopeKey,
+          storage,
+          next,
+        );
+        return;
+      }
+
+      hasPendingMutationRef.current = true;
+      if (
+        hydrationStateRef.current === "failed" &&
+        !reportedBlockedPersistenceRef.current
+      ) {
+        reportBlockedCustomServerPersistence();
+        reportedBlockedPersistenceRef.current = true;
+      }
+    },
+    [persistenceQueues, scopeKey, storage],
+  );
+
+  const removeCustomServer = useCallback(
+    (id: string) => {
+      if (hydrationStateRef.current === "pending") {
+        removedBeforeHydration.add(id);
+      }
+      updateCustomServers((prev) => prev.filter((record) => record.id !== id));
+    },
+    [removedBeforeHydration, updateCustomServers],
+  );
+
+  return {
+    customServers,
+    isHydrated,
+    updateCustomServers,
+    removeCustomServer,
+  };
+};
+
+export const McpCustomServersResource = resource(useMcpCustomServersResource);

package/src/resources/McpElicitationLifecycle.test.ts

@@ -0,0 +1,83 @@
+import { createTapRoot } from "@assistant-ui/tap";
+import { describe, expect, it, vi } from "vitest";
+import { useMcpElicitationLifecycle } from "./McpElicitationLifecycle";
+
+describe("useMcpElicitationLifecycle", () => {
+  it("validates accepted content before resolving the request", async () => {
+    const root = createTapRoot(function Root() {
+      return useMcpElicitationLifecycle();
+    });
+    const requestedSchema = {
+      type: "object" as const,
+      properties: { answer: { type: "string" as const } },
+    };
+
+    try {
+      const response = root.getValue().requestElicitation(
+        {
+          method: "elicitation/create",
+          params: { message: "Answer", requestedSchema },
+        },
+        new AbortController().signal,
+        () => true,
+      );
+      await vi.waitFor(() =>
+        expect(root.getValue().pendingElicitations).toHaveLength(1),
+      );
+      const id = root.getValue().pendingElicitations[0]!.id;
+
+      expect(
+        root.getValue().answerElicitation(id, {
+          action: "accept",
+          content: { answer: 42 },
+        }),
+      ).toEqual([{ property: "answer", message: "Expected a string." }]);
+      await vi.waitFor(() =>
+        expect(
+          root.getValue().pendingElicitations[0]?.error?.properties,
+        ).toEqual(["answer"]),
+      );
+
+      expect(
+        root.getValue().answerElicitation(id, {
+          action: "accept",
+          content: { answer: "yes" },
+        }),
+      ).toBeUndefined();
+      await expect(response).resolves.toEqual({
+        action: "accept",
+        content: { answer: "yes" },
+      });
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("cancels a request when its signal aborts", async () => {
+    const root = createTapRoot(function Root() {
+      return useMcpElicitationLifecycle();
+    });
+    const controller = new AbortController();
+
+    try {
+      const response = root.getValue().requestElicitation(
+        {
+          method: "elicitation/create",
+          params: {
+            message: "Answer",
+            requestedSchema: { type: "object", properties: {} },
+          },
+        },
+        controller.signal,
+        () => true,
+      );
+      controller.abort();
+      await expect(response).resolves.toEqual({ action: "cancel" });
+      await vi.waitFor(() =>
+        expect(root.getValue().pendingElicitations).toHaveLength(0),
+      );
+    } finally {
+      root.unmount();
+    }
+  });
+});

package/src/resources/McpElicitationLifecycle.ts

@@ -0,0 +1,156 @@
+import { useRef, useState } from "react";
+import type { ElicitRequest, ElicitResult } from "@modelcontextprotocol/client";
+import type { MCPElicitation, MCPElicitationResponse } from "../mcp-scope";
+import { createMcpId } from "../utils/createMcpId";
+import { validateElicitationContent } from "./validateElicitationContent";
+
+export const useMcpElicitationLifecycle = () => {
+  const [pendingElicitations, setPendingElicitations] = useState<
+    MCPElicitation[]
+  >([]);
+  const elicitationResolversRef = useRef(
+    new Map<
+      string,
+      {
+        resolve: (result: ElicitResult) => void;
+        signal: AbortSignal;
+        onAbort: () => void;
+        requestedSchema: unknown;
+      }
+    >(),
+  );
+  const resolvePendingElicitation = (id: string, result: ElicitResult) => {
+    const entry = elicitationResolversRef.current.get(id);
+    if (!entry) return false;
+    elicitationResolversRef.current.delete(id);
+    entry.signal.removeEventListener("abort", entry.onAbort);
+    setPendingElicitations((current) =>
+      current.filter((elicitation) => elicitation.id !== id),
+    );
+    entry.resolve(result);
+    return true;
+  };
+
+  const setPendingElicitationError = (
+    id: string,
+    error: NonNullable<MCPElicitation["error"]>,
+  ) => {
+    if (!elicitationResolversRef.current.has(id)) return false;
+    setPendingElicitations((current) =>
+      current.map((elicitation) =>
+        elicitation.id === id ? { ...elicitation, error } : elicitation,
+      ),
+    );
+    return true;
+  };
+
+  const cancelPendingElicitations = () => {
+    for (const [id] of elicitationResolversRef.current) {
+      resolvePendingElicitation(id, { action: "cancel" });
+    }
+  };
+
+  const requestElicitation = (
+    request: ElicitRequest,
+    signal: AbortSignal,
+    isCurrentConnection: () => boolean,
+  ): Promise<ElicitResult> => {
+    if (!isCurrentConnection()) {
+      return Promise.resolve({ action: "cancel" });
+    }
+    if (!("requestedSchema" in request.params)) {
+      return Promise.resolve({ action: "cancel" });
+    }
+    const { message, requestedSchema } = request.params;
+
+    const id = createMcpId();
+    const promise = new Promise<ElicitResult>((resolve) => {
+      const onAbort = () => {
+        resolvePendingElicitation(id, { action: "cancel" });
+      };
+      elicitationResolversRef.current.set(id, {
+        resolve,
+        signal,
+        onAbort,
+        requestedSchema,
+      });
+    });
+    setPendingElicitations((current) => [
+      ...current,
+      {
+        id,
+        message,
+        requestedSchema,
+      },
+    ]);
+    const entry = elicitationResolversRef.current.get(id);
+    if (entry) {
+      if (signal.aborted) {
+        entry.onAbort();
+      } else {
+        signal.addEventListener("abort", entry.onAbort, {
+          once: true,
+        });
+      }
+    }
+    return promise;
+  };
+
+  return {
+    pendingElicitations,
+    cancelPendingElicitations,
+    requestElicitation,
+    answerElicitation: (
+      id: string,
+      response: MCPElicitationResponse,
+    ): readonly { property: string; message: string }[] | undefined => {
+      if (response.action === "accept") {
+        const entry = elicitationResolversRef.current.get(id);
+        if (!entry) return undefined;
+
+        if (
+          typeof response.content !== "object" ||
+          response.content === null ||
+          Array.isArray(response.content)
+        ) {
+          const errors = [
+            {
+              property: "content",
+              message: "Response content must be an object.",
+            },
+          ];
+          setPendingElicitationError(id, {
+            message: "Invalid elicitation content: content.",
+            properties: ["content"],
+          });
+          return errors;
+        }
+
+        const errors = validateElicitationContent(
+          entry.requestedSchema,
+          response.content,
+        );
+        if (errors.length > 0) {
+          const properties = [
+            ...new Set(errors.map((error) => error.property)),
+          ];
+          setPendingElicitationError(id, {
+            message: `Invalid elicitation content: ${properties.join(", ")}.`,
+            properties,
+          });
+          return errors;
+        }
+
+        const result: ElicitResult = {
+          action: "accept",
+          content: response.content as ElicitResult["content"],
+        };
+        resolvePendingElicitation(id, result);
+        return undefined;
+      }
+
+      resolvePendingElicitation(id, { action: response.action });
+      return undefined;
+    },
+  };
+};

package/src/resources/McpManagerResource.test.ts

@@ -416,6 +416,42 @@ describe("McpManagerResource storage failures", () => {
       consoleError.mockRestore();
     }
   });
+
+  it("preserves a custom server when its removal fails", async () => {
+    const docsServer: MCPCustomServerRecord = {
+      id: "docs",
+      name: "Docs",
+      url: "https://example.com/docs/mcp",
+      auth: { type: "none" },
+      createdAt: 1,
+    };
+    const error = new Error("storage unavailable");
+    const clearAuthState = vi.fn(async () => {});
+    clearAuthState.mockRejectedValueOnce(error);
+    const root = mount(
+      [],
+      McpCustomStorage({
+        loadCustomServers: vi.fn(async () => [docsServer]),
+        saveCustomServers: vi.fn(async () => {}),
+        loadAuthState: vi.fn(async () => null),
+        saveAuthState: vi.fn(async () => {}),
+        clearAuthState,
+      }),
+    );
+
+    try {
+      await vi.waitFor(() =>
+        expect(root.getValue().getState().customServers[0]?.id).toBe("docs"),
+      );
+
+      await expect(root.getValue().removeServer("docs")).rejects.toBe(error);
+      expect(root.getValue().getState().customServers).toMatchObject([
+        { id: "docs", name: "Docs" },
+      ]);
+    } finally {
+      root.unmount();
+    }
+  });
 });
 
 describe("McpManagerResource storage ordering", () => {