From 39963705141df07055e2ef3e8f959d365cb7aaf1 Mon Sep 17 00:00:00 2001
From: "Sharad." <sharadvc@users.noreply.github.com>
Date: Sat, 12 Sep 2026 14:22:51 +0000
Subject: [PATCH] Fix probe to follow listTools pagination cursor

Add a regression test with a mocked two-page listTools response, then aggregate all pages (and the same for resources/prompts) via nextCursor so audits cover the full server surface.

Fixes AgentPostmortem/MCP-audit#31
---
 src/transport/probe.ts | 46 +++++++++++++++++++-------
 test/probe.test.ts     | 75 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 109 insertions(+), 12 deletions(-)
 create mode 100644 test/probe.test.ts

diff --git a/src/transport/probe.ts b/src/transport/probe.ts
index 4f9bad7..ac91293 100644
--- a/src/transport/probe.ts
+++ b/src/transport/probe.ts
@@ -34,18 +34,15 @@ export async function probe(
     const version = client.getServerVersion();
     const instructions = client.getInstructions();
 
-    const tools = await safeList(async () => {
-      const res = await client.listTools();
-      return res.tools as ToolSpec[];
-    });
-    const resources = await safeList(async () => {
-      const res = await client.listResources();
-      return res.resources as ResourceSpec[];
-    });
-    const prompts = await safeList(async () => {
-      const res = await client.listPrompts();
-      return res.prompts as PromptSpec[];
-    });
+    const tools = await safeList(async () =>
+      (await listAllPages(client.listTools.bind(client))) as ToolSpec[],
+    );
+    const resources = await safeList(async () =>
+      (await listAllPages(client.listResources.bind(client))) as ResourceSpec[],
+    );
+    const prompts = await safeList(async () =>
+      (await listAllPages(client.listPrompts.bind(client))) as PromptSpec[],
+    );
 
     return {
       transport: meta.kind,
@@ -65,6 +62,31 @@ export async function probe(
   }
 }
 
+type PaginatedListResult<T> = {
+  tools?: T[];
+  resources?: T[];
+  prompts?: T[];
+  nextCursor?: string;
+};
+
+/**
+ * Follow {@link nextCursor} until all pages of a list capability are loaded.
+ */
+async function listAllPages<T>(
+  listPage: (params?: { cursor?: string }) => Promise<PaginatedListResult<T>>,
+): Promise<T[]> {
+  const items: T[] = [];
+  let cursor: string | undefined;
+  do {
+    const res = await listPage(cursor ? { cursor } : undefined);
+    const page =
+      res.tools ?? res.resources ?? res.prompts ?? ([] as T[]);
+    items.push(...page);
+    cursor = res.nextCursor;
+  } while (cursor);
+  return items;
+}
+
 /**
  * List a capability, returning an empty array when the server does not support
  * it (the SDK throws a "Method not found" error in that case).
diff --git a/test/probe.test.ts b/test/probe.test.ts
new file mode 100644
index 0000000..7f6d853
--- /dev/null
+++ b/test/probe.test.ts
@@ -0,0 +1,75 @@
+import { describe, it, expect, vi } from "vitest";
+import type { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";
+
+const { listToolsMock } = vi.hoisted(() => ({
+  listToolsMock: vi.fn(),
+}));
+
+vi.mock("@modelcontextprotocol/sdk/client/index.js", () => {
+  class Client {
+    async connect(_transport: Transport) {}
+    getServerVersion() {
+      return { name: "paginated-server", version: "1.0.0" };
+    }
+    getInstructions() {
+      return undefined;
+    }
+    async listTools(params?: { cursor?: string }) {
+      return listToolsMock(params);
+    }
+    async listResources() {
+      return { resources: [] };
+    }
+    async listPrompts() {
+      return { prompts: [] };
+    }
+    async close() {}
+  }
+  return { Client };
+});
+
+import { probe } from "../src/transport/probe.js";
+
+describe("probe", () => {
+  it("aggregates all pages from paginated listTools", async () => {
+    listToolsMock.mockImplementation((params?: { cursor?: string }) => {
+      if (!params?.cursor) {
+        return {
+          tools: [
+            {
+              name: "page-one-tool",
+              description: "first page",
+              inputSchema: { type: "object" },
+            },
+          ],
+          nextCursor: "page-2",
+        };
+      }
+      if (params.cursor === "page-2") {
+        return {
+          tools: [
+            {
+              name: "page-two-tool",
+              description: "second page",
+              inputSchema: { type: "object" },
+            },
+          ],
+        };
+      }
+      throw new Error(`unexpected cursor: ${params.cursor}`);
+    });
+
+    const target = await probe({} as Transport, {
+      kind: "stdio",
+      source: "test://paginated",
+    });
+
+    expect(target.tools.map((t) => t.name)).toEqual([
+      "page-one-tool",
+      "page-two-tool",
+    ]);
+    expect(listToolsMock).toHaveBeenCalledTimes(2);
+    expect(listToolsMock.mock.calls[0]?.[0]).toBeUndefined();
+    expect(listToolsMock.mock.calls[1]?.[0]).toEqual({ cursor: "page-2" });
+  });
+});