@assistant-ui/react-mcp 0.0.22 → 0.1.0

package/src/resources/McpServerResource.test.ts

@@ -30,6 +30,14 @@ const mocks = vi.hoisted(() => {
     this.callTool = vi.fn();
     this.listResources = vi.fn(() => Promise.resolve({ resources: [] }));
     this.readResource = vi.fn();
+    this.requestHandlers = new Map();
+    this.notificationHandlers = new Map();
+    this.setRequestHandler = vi.fn((method, handler) => {
+      this.requestHandlers.set(method, handler);
+    });
+    this.setNotificationHandler = vi.fn((method, handler) => {
+      this.notificationHandlers.set(method, handler);
+    });
     clients.push(this);
   });
 
@@ -56,11 +64,9 @@ const mocks = vi.hoisted(() => {
   };
 });
 
-vi.mock("@modelcontextprotocol/sdk/client/index.js", () => ({
+vi.mock("@modelcontextprotocol/client", async (importOriginal) => ({
+  ...(await importOriginal()),
   Client: mocks.Client,
-}));
-
-vi.mock("@modelcontextprotocol/sdk/client/streamableHttp.js", () => ({
   StreamableHTTPClientTransport: mocks.StreamableHTTPClientTransport,
 }));
 
@@ -92,6 +98,33 @@ const waitFor = async (predicate: () => boolean) => {
   expect(predicate()).toBe(true);
 };
 
+const waitForResourceUpdate = async (predicate: () => boolean) => {
+  for (let i = 0; i < 20; i++) {
+    if (predicate()) return;
+    await flushMacrotask();
+  }
+  expect(predicate()).toBe(true);
+};
+
+const requestElicitation = (
+  client: any,
+  message: string,
+  requestedSchema: unknown,
+  context: { signal: AbortSignal } = {
+    signal: new AbortController().signal,
+  },
+) => {
+  const handler = client.requestHandlers.get("elicitation/create");
+  if (!handler) throw new Error("elicitation/create handler not registered");
+  return handler(
+    {
+      method: "elicitation/create",
+      params: { message, requestedSchema },
+    },
+    context,
+  );
+};
+
 const createStorage = (): MCPStorage => ({
   loadCustomServers: vi.fn(async () => []),
   saveCustomServers: vi.fn(async () => {}),
@@ -115,6 +148,8 @@ const mount = (
   props?: {
     auth?: MCPAuthConfig | undefined;
     connectionTimeout?: number | undefined;
+    cache?: { readonly defaultTtlMs?: number } | undefined;
+    elicitation?: boolean | undefined;
   },
   onMount?: (server: ClientOutput<"mcpServer">) => void,
 ) => {
@@ -133,6 +168,10 @@ const mount = (
         redirectUri: "https://example.com/callback",
         autoConnect: false,
         connectionTimeout,
+        cache: props?.cache,
+        ...(props?.elicitation !== undefined
+          ? { elicitation: props.elicitation }
+          : {}),
         onRemove: vi.fn(async () => {}),
       }),
     );
@@ -469,6 +508,698 @@ describe("McpServerResource completeAuth", () => {
   });
 });
 
+describe("McpServerResource elicitation", () => {
+  beforeEach(resetMocks);
+
+  it("surfaces pending elicitation requests in server state", async () => {
+    const root = mount();
+    const requestedSchema = {
+      type: "object",
+      properties: {
+        answer: { type: "string" },
+      },
+    };
+
+    try {
+      await root.getValue().connect();
+      const response = requestElicitation(
+        mocks.clients[0],
+        "Provide an answer",
+        requestedSchema,
+      );
+      await waitForResourceUpdate(
+        () => root.getValue().getState().pendingElicitations.length === 1,
+      );
+
+      const [elicitation] = root.getValue().getState().pendingElicitations;
+      expect(elicitation).toEqual({
+        id: expect.any(String),
+        message: "Provide an answer",
+        requestedSchema,
+      });
+
+      root.getValue().answerElicitation(elicitation!.id, {
+        action: "cancel",
+      });
+      await expect(response).resolves.toEqual({ action: "cancel" });
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("answers pending elicitations with accepted content", async () => {
+    const root = mount();
+
+    try {
+      await root.getValue().connect();
+      const response = requestElicitation(mocks.clients[0], "Choose a color", {
+        type: "object",
+        properties: {
+          color: { type: "string" },
+        },
+      });
+      await waitForResourceUpdate(
+        () => root.getValue().getState().pendingElicitations.length === 1,
+      );
+      const [elicitation] = root.getValue().getState().pendingElicitations;
+      const content = { color: "blue" };
+
+      expect(
+        root
+          .getValue()
+          .answerElicitation(elicitation!.id, { action: "accept", content }),
+      ).toBeUndefined();
+
+      await expect(response).resolves.toEqual({
+        action: "accept",
+        content,
+      });
+      await waitForResourceUpdate(
+        () => root.getValue().getState().pendingElicitations.length === 0,
+      );
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("keeps invalid accepted content pending until a valid response arrives", async () => {
+    const root = mount();
+
+    try {
+      await root.getValue().connect();
+      const response = requestElicitation(mocks.clients[0], "Choose a color", {
+        type: "object",
+        required: ["color"],
+        properties: {
+          color: { type: "string" },
+        },
+      });
+      let resolved = false;
+      void response.then(() => {
+        resolved = true;
+      });
+      await waitForResourceUpdate(
+        () => root.getValue().getState().pendingElicitations.length === 1,
+      );
+      const [elicitation] = root.getValue().getState().pendingElicitations;
+
+      expect(
+        root.getValue().answerElicitation(elicitation!.id, {
+          action: "accept",
+          content: { color: 1 },
+        }),
+      ).toEqual([{ property: "color", message: "Expected a string." }]);
+
+      await waitForResourceUpdate(
+        () =>
+          root.getValue().getState().pendingElicitations[0]?.error?.message ===
+          "Invalid elicitation content: color.",
+      );
+      expect(root.getValue().getState().pendingElicitations).toEqual([
+        expect.objectContaining({
+          id: elicitation!.id,
+          error: {
+            message: "Invalid elicitation content: color.",
+            properties: ["color"],
+          },
+        }),
+      ]);
+      await tick();
+      expect(resolved).toBe(false);
+
+      expect(
+        root.getValue().answerElicitation(elicitation!.id, {
+          action: "accept",
+          content: { color: "blue" },
+        }),
+      ).toBeUndefined();
+
+      await expect(response).resolves.toEqual({
+        action: "accept",
+        content: { color: "blue" },
+      });
+      await waitForResourceUpdate(
+        () => root.getValue().getState().pendingElicitations.length === 0,
+      );
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("declines an elicitation that carries a validation error", async () => {
+    const root = mount();
+
+    try {
+      await root.getValue().connect();
+      const response = requestElicitation(mocks.clients[0], "Choose a color", {
+        type: "object",
+        required: ["color"],
+        properties: {
+          color: { type: "string" },
+        },
+      });
+      await waitForResourceUpdate(
+        () => root.getValue().getState().pendingElicitations.length === 1,
+      );
+      const [elicitation] = root.getValue().getState().pendingElicitations;
+
+      root.getValue().answerElicitation(elicitation!.id, {
+        action: "accept",
+        content: {},
+      });
+      await waitForResourceUpdate(
+        () =>
+          root.getValue().getState().pendingElicitations[0]?.error !==
+          undefined,
+      );
+
+      root.getValue().answerElicitation(elicitation!.id, { action: "decline" });
+
+      await expect(response).resolves.toEqual({ action: "decline" });
+      await waitForResourceUpdate(
+        () => root.getValue().getState().pendingElicitations.length === 0,
+      );
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("declines pending elicitations", async () => {
+    const root = mount();
+
+    try {
+      await root.getValue().connect();
+      const response = requestElicitation(mocks.clients[0], "Confirm access", {
+        type: "object",
+        properties: {},
+      });
+      await waitForResourceUpdate(
+        () => root.getValue().getState().pendingElicitations.length === 1,
+      );
+      const [elicitation] = root.getValue().getState().pendingElicitations;
+
+      root.getValue().answerElicitation(elicitation!.id, {
+        action: "decline",
+      });
+
+      await expect(response).resolves.toEqual({ action: "decline" });
+      await waitForResourceUpdate(
+        () => root.getValue().getState().pendingElicitations.length === 0,
+      );
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("cancels pending elicitations", async () => {
+    const root = mount();
+
+    try {
+      await root.getValue().connect();
+      const response = requestElicitation(mocks.clients[0], "Confirm access", {
+        type: "object",
+        properties: {},
+      });
+      await waitForResourceUpdate(
+        () => root.getValue().getState().pendingElicitations.length === 1,
+      );
+      const [elicitation] = root.getValue().getState().pendingElicitations;
+
+      root.getValue().answerElicitation(elicitation!.id, {
+        action: "cancel",
+      });
+
+      await expect(response).resolves.toEqual({ action: "cancel" });
+      await waitForResourceUpdate(
+        () => root.getValue().getState().pendingElicitations.length === 0,
+      );
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("ignores answers for unknown elicitations", () => {
+    const root = mount();
+
+    try {
+      expect(
+        root
+          .getValue()
+          .answerElicitation("missing-elicitation", { action: "cancel" }),
+      ).toBeUndefined();
+      expect(root.getValue().getState().pendingElicitations).toEqual([]);
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("returns validation errors for accepted elicitation content that is not an object", async () => {
+    const root = mount();
+
+    try {
+      await root.getValue().connect();
+      const response = requestElicitation(mocks.clients[0], "Confirm access", {
+        type: "object",
+        properties: {},
+      });
+      await waitForResourceUpdate(
+        () => root.getValue().getState().pendingElicitations.length === 1,
+      );
+      const [elicitation] = root.getValue().getState().pendingElicitations;
+
+      expect(
+        root.getValue().answerElicitation(elicitation!.id, {
+          action: "accept",
+          content: null as never,
+        }),
+      ).toEqual([
+        { property: "content", message: "Response content must be an object." },
+      ]);
+      await waitForResourceUpdate(
+        () =>
+          root.getValue().getState().pendingElicitations[0]?.error?.message ===
+          "Invalid elicitation content: content.",
+      );
+      expect(root.getValue().getState().pendingElicitations).toEqual([
+        expect.objectContaining({
+          id: elicitation!.id,
+          error: {
+            message: "Invalid elicitation content: content.",
+            properties: ["content"],
+          },
+        }),
+      ]);
+
+      await root.getValue().disconnect();
+      await expect(response).resolves.toEqual({ action: "cancel" });
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("cancels pending elicitations on disconnect", async () => {
+    const root = mount();
+
+    try {
+      await root.getValue().connect();
+      const response = requestElicitation(mocks.clients[0], "Confirm access", {
+        type: "object",
+        properties: {},
+      });
+      await waitForResourceUpdate(
+        () => root.getValue().getState().pendingElicitations.length === 1,
+      );
+
+      await root.getValue().disconnect();
+
+      await expect(response).resolves.toEqual({ action: "cancel" });
+      await waitForResourceUpdate(
+        () => root.getValue().getState().pendingElicitations.length === 0,
+      );
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("cancels pending elicitations when the server aborts the request", async () => {
+    const root = mount();
+    const controller = new AbortController();
+
+    try {
+      await root.getValue().connect();
+      const response = requestElicitation(
+        mocks.clients[0],
+        "Confirm access",
+        {
+          type: "object",
+          properties: {},
+        },
+        { signal: controller.signal },
+      );
+      await waitForResourceUpdate(
+        () => root.getValue().getState().pendingElicitations.length === 1,
+      );
+
+      controller.abort();
+
+      await expect(response).resolves.toEqual({ action: "cancel" });
+      await waitForResourceUpdate(
+        () => root.getValue().getState().pendingElicitations.length === 0,
+      );
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("cancels pending elicitations on unmount", async () => {
+    const root = mount();
+    let didUnmount = false;
+
+    try {
+      await root.getValue().connect();
+      const response = requestElicitation(mocks.clients[0], "Confirm access", {
+        type: "object",
+        properties: {},
+      });
+      await waitForResourceUpdate(
+        () => root.getValue().getState().pendingElicitations.length === 1,
+      );
+
+      root.unmount();
+      didUnmount = true;
+      await flushMacrotask();
+
+      await expect(response).resolves.toEqual({ action: "cancel" });
+    } finally {
+      if (!didUnmount) root.unmount();
+    }
+  });
+
+  it("cancels requests from stale connections without surfacing them", async () => {
+    const root = mount();
+
+    try {
+      await root.getValue().connect();
+      await root.getValue().connect();
+
+      await expect(
+        requestElicitation(mocks.clients[0], "Confirm access", {
+          type: "object",
+          properties: {},
+        }),
+      ).resolves.toEqual({ action: "cancel" });
+      expect(root.getValue().getState().pendingElicitations).toEqual([]);
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("advertises form elicitation capability and registers its handler by default", async () => {
+    const root = mount();
+
+    try {
+      await root.getValue().connect();
+
+      expect(mocks.Client).toHaveBeenCalledWith(
+        {
+          name: "assistant-ui-mcp",
+          version: "0.0.0",
+        },
+        {
+          capabilities: {
+            elicitation: {},
+          },
+          listChanged: {
+            tools: {
+              autoRefresh: true,
+              debounceMs: 300,
+              onChanged: expect.any(Function),
+            },
+          },
+        },
+      );
+      expect(mocks.clients[0].setRequestHandler).toHaveBeenCalledWith(
+        "elicitation/create",
+        expect.any(Function),
+      );
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("does not advertise or handle elicitation when opted out", async () => {
+    const root = mount({ elicitation: false });
+
+    try {
+      await root.getValue().connect();
+
+      const clientOptions = mocks.Client.mock.calls[0]?.[1];
+      expect(clientOptions).not.toHaveProperty("capabilities.elicitation");
+      expect(mocks.clients[0].setRequestHandler).not.toHaveBeenCalledWith(
+        "elicitation/create",
+        expect.any(Function),
+      );
+    } finally {
+      root.unmount();
+    }
+  });
+});
+
+describe("McpServerResource tools listChanged", () => {
+  beforeEach(resetMocks);
+
+  it("opts into automatic listChanged tool refreshes", async () => {
+    const root = mount();
+
+    try {
+      await root.getValue().connect();
+
+      expect(mocks.Client).toHaveBeenCalledWith(
+        {
+          name: "assistant-ui-mcp",
+          version: "0.0.0",
+        },
+        {
+          capabilities: {
+            elicitation: {},
+          },
+          listChanged: {
+            tools: {
+              autoRefresh: true,
+              debounceMs: 300,
+              onChanged: expect.any(Function),
+            },
+          },
+        },
+      );
+      expect(mocks.clients[0].setNotificationHandler).not.toHaveBeenCalled();
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("applies refreshed tools from the listChanged callback", async () => {
+    mocks.listToolsResults.push(async () => ({
+      tools: [
+        {
+          name: "search",
+          description: "Search docs",
+          inputSchema: { type: "object" },
+        },
+      ],
+    }));
+    const root = mount();
+
+    try {
+      await root.getValue().connect();
+      await waitForResourceUpdate(
+        () => root.getValue().getState().tools.length === 1,
+      );
+      expect(root.getValue().getState().tools).toEqual([
+        {
+          name: "search",
+          description: "Search docs",
+          inputSchema: { type: "object" },
+        },
+      ]);
+
+      const onChanged =
+        mocks.Client.mock.calls[0]?.[1]?.listChanged?.tools?.onChanged;
+      if (!onChanged) throw new Error("Expected tools listChanged callback");
+      onChanged(null, [
+        {
+          name: "search",
+          description: "Search docs",
+          inputSchema: { type: "object" },
+        },
+        {
+          name: "summarize",
+          inputSchema: { type: "object", properties: {} },
+        },
+      ]);
+      await waitForResourceUpdate(
+        () => root.getValue().getState().tools.length === 2,
+      );
+
+      expect(mocks.clients[0].listTools).toHaveBeenCalledTimes(1);
+      expect(root.getValue().getState().tools).toEqual([
+        {
+          name: "search",
+          description: "Search docs",
+          inputSchema: { type: "object" },
+        },
+        {
+          name: "summarize",
+          inputSchema: { type: "object", properties: {} },
+        },
+      ]);
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("stores listChanged errors without replacing the current tools", async () => {
+    mocks.listToolsResults.push(async () => ({
+      tools: [
+        {
+          name: "search",
+          description: "Search docs",
+          inputSchema: { type: "object" },
+        },
+      ],
+    }));
+    const root = mount();
+
+    try {
+      await root.getValue().connect();
+      await waitForResourceUpdate(
+        () => root.getValue().getState().tools.length === 1,
+      );
+      const onChanged =
+        mocks.Client.mock.calls[0]?.[1]?.listChanged?.tools?.onChanged;
+      if (!onChanged) throw new Error("Expected tools listChanged callback");
+
+      onChanged(new Error("tool update failed"), null);
+      await waitForResourceUpdate(
+        () =>
+          root.getValue().getState().lastError?.message ===
+          "tool update failed",
+      );
+
+      expect(root.getValue().getState()).toMatchObject({
+        connectionState: "connected",
+        lastError: { message: "tool update failed" },
+        tools: [
+          {
+            name: "search",
+            description: "Search docs",
+            inputSchema: { type: "object" },
+          },
+        ],
+      });
+
+      onChanged(null, [
+        {
+          name: "summarize",
+          inputSchema: { type: "object", properties: {} },
+        },
+      ]);
+      await waitForResourceUpdate(
+        () => root.getValue().getState().lastError === null,
+      );
+
+      expect(root.getValue().getState().tools).toEqual([
+        {
+          name: "summarize",
+          inputSchema: { type: "object", properties: {} },
+        },
+      ]);
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("ignores listChanged callbacks from a stale connection generation", async () => {
+    mocks.listToolsResults.push(async () => ({
+      tools: [
+        {
+          name: "first",
+          inputSchema: { type: "object" },
+        },
+      ],
+    }));
+    mocks.listToolsResults.push(async () => ({
+      tools: [
+        {
+          name: "second",
+          inputSchema: { type: "object" },
+        },
+      ],
+    }));
+    const root = mount();
+
+    try {
+      await root.getValue().connect();
+      const staleOnChanged =
+        mocks.Client.mock.calls[0]?.[1]?.listChanged?.tools?.onChanged;
+      if (!staleOnChanged)
+        throw new Error("Expected tools listChanged callback");
+      await root.getValue().connect();
+      await waitForResourceUpdate(
+        () => root.getValue().getState().tools[0]?.name === "second",
+      );
+
+      expect(root.getValue().getState().tools).toEqual([
+        {
+          name: "second",
+          inputSchema: { type: "object" },
+        },
+      ]);
+
+      staleOnChanged(null, [
+        {
+          name: "stale-tool",
+          inputSchema: { type: "object" },
+        },
+      ]);
+      await flushMacrotask();
+
+      expect(root.getValue().getState().tools).toEqual([
+        {
+          name: "second",
+          inputSchema: { type: "object" },
+        },
+      ]);
+      expect(mocks.clients[0].listTools).toHaveBeenCalledTimes(1);
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("passes defaultCacheTtlMs when cache.defaultTtlMs is configured", async () => {
+    const root = mount({ cache: { defaultTtlMs: 5_000 } });
+
+    try {
+      await root.getValue().connect();
+
+      expect(mocks.Client).toHaveBeenCalledWith(
+        {
+          name: "assistant-ui-mcp",
+          version: "0.0.0",
+        },
+        {
+          capabilities: {
+            elicitation: {},
+          },
+          listChanged: {
+            tools: {
+              autoRefresh: true,
+              debounceMs: 300,
+              onChanged: expect.any(Function),
+            },
+          },
+          defaultCacheTtlMs: 5_000,
+        },
+      );
+    } finally {
+      root.unmount();
+    }
+  });
+
+  it("omits defaultCacheTtlMs when cache is not configured", async () => {
+    const root = mount();
+
+    try {
+      await root.getValue().connect();
+
+      const options = mocks.Client.mock.calls[0]?.[1];
+      expect(options).not.toHaveProperty("defaultCacheTtlMs");
+    } finally {
+      root.unmount();
+    }
+  });
+});
+
 describe("McpServerResource resource methods", () => {
   beforeEach(() => {
     mocks.clients.length = 0;