Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 34 additions & 12 deletions src/transport/probe.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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).
Expand Down
75 changes: 75 additions & 0 deletions test/probe.test.ts
Original file line number Diff line number Diff line change
@@ -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" });
});
});
Loading