diff --git a/apps/sim/lib/mcp/client.test.ts b/apps/sim/lib/mcp/client.test.ts index f5e6b85298c..8a87eac5a5f 100644 --- a/apps/sim/lib/mcp/client.test.ts +++ b/apps/sim/lib/mcp/client.test.ts @@ -177,7 +177,7 @@ describe('McpClient notification handler', () => { undefined, expect.objectContaining({ timeout: 30_000, - maxTotalTimeout: 60_000, + maxTotalTimeout: expect.any(Number), resetTimeoutOnProgress: true, onprogress: expect.any(Function), }) @@ -196,10 +196,84 @@ describe('McpClient notification handler', () => { expect(mockSdkListTools).toHaveBeenCalledWith( undefined, - expect.objectContaining({ timeout: 60_000, maxTotalTimeout: 60_000 }) + expect.objectContaining({ timeout: 60_000, maxTotalTimeout: expect.any(Number) }) ) }) + it('follows nextCursor pagination and aggregates all pages', async () => { + mockSdkListTools + .mockResolvedValueOnce({ tools: [{ name: 'a' }, { name: 'b' }], nextCursor: 'c1' }) + .mockResolvedValueOnce({ tools: [{ name: 'c' }], nextCursor: 'c2' }) + .mockResolvedValueOnce({ tools: [{ name: 'd' }] }) + const client = new McpClient({ + config: createConfig(), + securityPolicy: { requireConsent: false, auditLevel: 'basic' }, + }) + + await client.connect() + const tools = await client.listTools() + + expect(tools.map((t) => t.name)).toEqual(['a', 'b', 'c', 'd']) + expect(mockSdkListTools).toHaveBeenNthCalledWith(2, { cursor: 'c1' }, expect.anything()) + expect(mockSdkListTools).toHaveBeenNthCalledWith(3, { cursor: 'c2' }, expect.anything()) + }) + + it('stops paginating when the server repeats a cursor (loop guard)', async () => { + mockSdkListTools.mockResolvedValue({ tools: [{ name: 'x' }], nextCursor: 'same' }) + const client = new McpClient({ + config: createConfig(), + securityPolicy: { requireConsent: false, auditLevel: 'basic' }, + }) + + await client.connect() + const tools = await client.listTools() + + // Page 1 sets cursor 'same'; page 2 returns 'same' again → guard stops. Not 50 pages. + expect(mockSdkListTools).toHaveBeenCalledTimes(2) + expect(tools).toHaveLength(2) + }) + + it('returns partial tools when a later page fails', async () => { + mockSdkListTools + .mockResolvedValueOnce({ tools: [{ name: 'a' }], nextCursor: 'c1' }) + .mockRejectedValueOnce(new Error('page 2 blew up')) + const client = new McpClient({ + config: createConfig(), + securityPolicy: { requireConsent: false, auditLevel: 'basic' }, + }) + + await client.connect() + const tools = await client.listTools() + + expect(tools.map((t) => t.name)).toEqual(['a']) + }) + + it('keeps an empty partial (does not throw) when page one succeeds but a later page fails', async () => { + // Page one is valid but empty with a cursor; page two fails. Page one succeeded, so + // discovery must not fail the server — it returns [] rather than throwing. + mockSdkListTools + .mockResolvedValueOnce({ tools: [], nextCursor: 'c1' }) + .mockRejectedValueOnce(new Error('page 2 blew up')) + const client = new McpClient({ + config: createConfig(), + securityPolicy: { requireConsent: false, auditLevel: 'basic' }, + }) + + await client.connect() + await expect(client.listTools()).resolves.toEqual([]) + }) + + it('throws when the very first page fails', async () => { + mockSdkListTools.mockRejectedValueOnce(new Error('page 1 blew up')) + const client = new McpClient({ + config: createConfig(), + securityPolicy: { requireConsent: false, auditLevel: 'basic' }, + }) + + await client.connect() + await expect(client.listTools()).rejects.toThrow() + }) + it('logs connection diagnostics without header values', async () => { const client = new McpClient({ config: { diff --git a/apps/sim/lib/mcp/client.ts b/apps/sim/lib/mcp/client.ts index 7e7590fc7a2..3e22087c8cf 100644 --- a/apps/sim/lib/mcp/client.ts +++ b/apps/sim/lib/mcp/client.ts @@ -5,7 +5,6 @@ import { LATEST_PROTOCOL_VERSION, type ListToolsResult, SUPPORTED_PROTOCOL_VERSIONS, - type Tool, ToolListChangedNotificationSchema, } from '@modelcontextprotocol/sdk/types.js' import { createLogger } from '@sim/logger' @@ -275,33 +274,98 @@ export class McpClient { const maxTotalTimeoutMs = MCP_CLIENT_CONSTANTS.LIST_TOOLS_MAX_TOTAL_TIMEOUT_MS const startedAt = Date.now() + // The SDK's `listTools()` returns a single page; a server that paginates via + // `nextCursor` would otherwise be silently truncated to page one. Follow the + // cursor, bounded by four independent budgets — pages, tool count, byte size, + // and aggregate wall-clock — plus a repeated-cursor guard, since a page cap + // alone can't stop a server that returns a fresh cursor with no new tools. + const deadline = startedAt + maxTotalTimeoutMs + const maxPages = MCP_CLIENT_CONSTANTS.LIST_TOOLS_MAX_PAGES + const tools: McpTool[] = [] + const seenCursors = new Set() + let cursor: string | undefined + let bytes = 0 + let pagesFetched = 0 + let truncated: string | undefined + let reachedEnd = false + try { - const result: ListToolsResult = await this.client.listTools(undefined, { - // resetTimeoutOnProgress only takes effect when onprogress is supplied. - timeout: idleTimeoutMs, - maxTotalTimeout: maxTotalTimeoutMs, - resetTimeoutOnProgress: true, - onprogress: (progress) => { - logger.debug(`Tool discovery progress from ${this.config.name}`, { + for (let page = 0; page < maxPages; page++) { + const remainingMs = deadline - Date.now() + if (remainingMs <= 0) { + truncated = 'aggregate timeout' + break + } + const result: ListToolsResult = await this.client.listTools( + cursor ? { cursor } : undefined, + { + // resetTimeoutOnProgress only takes effect when onprogress is supplied. + timeout: Math.min(idleTimeoutMs, remainingMs), + maxTotalTimeout: remainingMs, + resetTimeoutOnProgress: true, + onprogress: (progress) => { + logger.debug(`Tool discovery progress from ${this.config.name}`, { + serverId: this.config.id, + progress: progress.progress, + total: progress.total, + }) + }, + } + ) + pagesFetched++ + + if (!result.tools || !Array.isArray(result.tools)) { + logger.warn(`Invalid tools response from server ${this.config.name}:`, result) + reachedEnd = true + break + } + + for (const tool of result.tools) { + if (tools.length >= MCP_CLIENT_CONSTANTS.LIST_TOOLS_MAX_TOOLS) { + truncated = 'tool count' + break + } + bytes += Buffer.byteLength(JSON.stringify(tool), 'utf8') + if (bytes > MCP_CLIENT_CONSTANTS.LIST_TOOLS_MAX_BYTES) { + truncated = 'byte size' + break + } + tools.push({ + name: tool.name, + description: tool.description, + inputSchema: tool.inputSchema as McpTool['inputSchema'], serverId: this.config.id, - progress: progress.progress, - total: progress.total, + serverName: this.config.name, }) - }, - }) + } + if (truncated) break + + const next = result.nextCursor + if (!next) { + reachedEnd = true + break // missing/empty cursor = end of results (spec) + } + if (seenCursors.has(next)) { + truncated = 'repeated cursor' + break + } + seenCursors.add(next) + cursor = next + } - if (!result.tools || !Array.isArray(result.tools)) { - logger.warn(`Invalid tools response from server ${this.config.name}:`, result) - return [] + // The loop can exhaust `maxPages` with a cursor still pending — that's a page-cap + // truncation, distinct from a natural end (`reachedEnd`) or an explicit budget hit. + if (!truncated && !reachedEnd) truncated = 'page cap' + if (truncated) { + logger.warn(`Tool discovery truncated for server ${this.config.name}`, { + serverId: this.config.id, + reason: truncated, + toolsCollected: tools.length, + pagesFetched, + }) } - return result.tools.map((tool: Tool) => ({ - name: tool.name, - description: tool.description, - inputSchema: tool.inputSchema as McpTool['inputSchema'], - serverId: this.config.id, - serverName: this.config.name, - })) + return tools } catch (error) { logger.error(`Failed to list tools from server ${this.config.name}`, { serverId: this.config.id, @@ -309,9 +373,14 @@ export class McpClient { durationMs: Date.now() - startedAt, idleTimeoutMs, maxTotalTimeoutMs, + pagesFetched, + toolsCollected: tools.length, sessionIdPresent: Boolean(this.transport.sessionId), error: getMcpSafeErrorDiagnostics(error), }) + // At least one page succeeded → keep its (possibly empty) partial result rather than + // failing discovery and marking the server unhealthy; only a page-one failure throws. + if (pagesFetched > 0) return tools throw error } } diff --git a/apps/sim/lib/mcp/utils.ts b/apps/sim/lib/mcp/utils.ts index 2f04d3039c3..8e64ee476c0 100644 --- a/apps/sim/lib/mcp/utils.ts +++ b/apps/sim/lib/mcp/utils.ts @@ -55,6 +55,12 @@ export const MCP_CLIENT_CONSTANTS = { LIST_TOOLS_TIMEOUT_MS: 30_000, /** Hard ceiling for tools/list regardless of progress (SDK maxTotalTimeout safeguard). */ LIST_TOOLS_MAX_TOTAL_TIMEOUT_MS: 60_000, + /** Max `tools/list` pages followed via `nextCursor` before truncating (see fetch loop). */ + LIST_TOOLS_MAX_PAGES: 50, + /** Max tools aggregated across all pages before truncating. */ + LIST_TOOLS_MAX_TOOLS: 1000, + /** Max total tool-payload bytes aggregated across all pages before truncating. */ + LIST_TOOLS_MAX_BYTES: 5 * 1024 * 1024, FAILURE_CACHE_TTL_MS: 120_000, } as const