|
5 | 5 | LATEST_PROTOCOL_VERSION, |
6 | 6 | type ListToolsResult, |
7 | 7 | SUPPORTED_PROTOCOL_VERSIONS, |
8 | | - type Tool, |
9 | 8 | ToolListChangedNotificationSchema, |
10 | 9 | } from '@modelcontextprotocol/sdk/types.js' |
11 | 10 | import { createLogger } from '@sim/logger' |
@@ -275,43 +274,101 @@ export class McpClient { |
275 | 274 | const maxTotalTimeoutMs = MCP_CLIENT_CONSTANTS.LIST_TOOLS_MAX_TOTAL_TIMEOUT_MS |
276 | 275 | const startedAt = Date.now() |
277 | 276 |
|
| 277 | + // The SDK's `listTools()` returns a single page; a server that paginates via |
| 278 | + // `nextCursor` would otherwise be silently truncated to page one. Follow the |
| 279 | + // cursor, bounded by four independent budgets — pages, tool count, byte size, |
| 280 | + // and aggregate wall-clock — plus a repeated-cursor guard, since a page cap |
| 281 | + // alone can't stop a server that returns a fresh cursor with no new tools. |
| 282 | + const deadline = startedAt + maxTotalTimeoutMs |
| 283 | + const tools: McpTool[] = [] |
| 284 | + const seenCursors = new Set<string>() |
| 285 | + let cursor: string | undefined |
| 286 | + let bytes = 0 |
| 287 | + let truncated: string | undefined |
| 288 | + |
278 | 289 | try { |
279 | | - const result: ListToolsResult = await this.client.listTools(undefined, { |
280 | | - // resetTimeoutOnProgress only takes effect when onprogress is supplied. |
281 | | - timeout: idleTimeoutMs, |
282 | | - maxTotalTimeout: maxTotalTimeoutMs, |
283 | | - resetTimeoutOnProgress: true, |
284 | | - onprogress: (progress) => { |
285 | | - logger.debug(`Tool discovery progress from ${this.config.name}`, { |
| 290 | + for (let page = 0; page < MCP_CLIENT_CONSTANTS.LIST_TOOLS_MAX_PAGES; page++) { |
| 291 | + const remainingMs = deadline - Date.now() |
| 292 | + if (remainingMs <= 0) { |
| 293 | + truncated = 'aggregate timeout' |
| 294 | + break |
| 295 | + } |
| 296 | + const result: ListToolsResult = await this.client.listTools( |
| 297 | + cursor ? { cursor } : undefined, |
| 298 | + { |
| 299 | + // resetTimeoutOnProgress only takes effect when onprogress is supplied. |
| 300 | + timeout: Math.min(idleTimeoutMs, remainingMs), |
| 301 | + maxTotalTimeout: remainingMs, |
| 302 | + resetTimeoutOnProgress: true, |
| 303 | + onprogress: (progress) => { |
| 304 | + logger.debug(`Tool discovery progress from ${this.config.name}`, { |
| 305 | + serverId: this.config.id, |
| 306 | + progress: progress.progress, |
| 307 | + total: progress.total, |
| 308 | + }) |
| 309 | + }, |
| 310 | + } |
| 311 | + ) |
| 312 | + |
| 313 | + if (!result.tools || !Array.isArray(result.tools)) { |
| 314 | + logger.warn(`Invalid tools response from server ${this.config.name}:`, result) |
| 315 | + break |
| 316 | + } |
| 317 | + |
| 318 | + for (const tool of result.tools) { |
| 319 | + if (tools.length >= MCP_CLIENT_CONSTANTS.LIST_TOOLS_MAX_TOOLS) { |
| 320 | + truncated = 'tool count' |
| 321 | + break |
| 322 | + } |
| 323 | + bytes += JSON.stringify(tool).length |
| 324 | + if (bytes > MCP_CLIENT_CONSTANTS.LIST_TOOLS_MAX_BYTES) { |
| 325 | + truncated = 'byte size' |
| 326 | + break |
| 327 | + } |
| 328 | + tools.push({ |
| 329 | + name: tool.name, |
| 330 | + description: tool.description, |
| 331 | + inputSchema: tool.inputSchema as McpTool['inputSchema'], |
286 | 332 | serverId: this.config.id, |
287 | | - progress: progress.progress, |
288 | | - total: progress.total, |
| 333 | + serverName: this.config.name, |
289 | 334 | }) |
290 | | - }, |
291 | | - }) |
| 335 | + } |
| 336 | + if (truncated) break |
| 337 | + |
| 338 | + const next = result.nextCursor |
| 339 | + if (!next) break // missing/empty cursor = end of results (spec) |
| 340 | + if (seenCursors.has(next)) { |
| 341 | + truncated = 'repeated cursor' |
| 342 | + break |
| 343 | + } |
| 344 | + seenCursors.add(next) |
| 345 | + cursor = next |
| 346 | + } |
292 | 347 |
|
293 | | - if (!result.tools || !Array.isArray(result.tools)) { |
294 | | - logger.warn(`Invalid tools response from server ${this.config.name}:`, result) |
295 | | - return [] |
| 348 | + if (truncated || seenCursors.size >= MCP_CLIENT_CONSTANTS.LIST_TOOLS_MAX_PAGES - 1) { |
| 349 | + logger.warn(`Tool discovery truncated for server ${this.config.name}`, { |
| 350 | + serverId: this.config.id, |
| 351 | + reason: truncated ?? 'page cap', |
| 352 | + toolsCollected: tools.length, |
| 353 | + pagesFetched: seenCursors.size + 1, |
| 354 | + }) |
296 | 355 | } |
297 | 356 |
|
298 | | - return result.tools.map((tool: Tool) => ({ |
299 | | - name: tool.name, |
300 | | - description: tool.description, |
301 | | - inputSchema: tool.inputSchema as McpTool['inputSchema'], |
302 | | - serverId: this.config.id, |
303 | | - serverName: this.config.name, |
304 | | - })) |
| 357 | + return tools |
305 | 358 | } catch (error) { |
306 | 359 | logger.error(`Failed to list tools from server ${this.config.name}`, { |
307 | 360 | serverId: this.config.id, |
308 | 361 | phase: 'tools/list', |
309 | 362 | durationMs: Date.now() - startedAt, |
310 | 363 | idleTimeoutMs, |
311 | 364 | maxTotalTimeoutMs, |
| 365 | + pagesFetched: seenCursors.size + 1, |
| 366 | + toolsCollected: tools.length, |
312 | 367 | sessionIdPresent: Boolean(this.transport.sessionId), |
313 | 368 | error: getMcpSafeErrorDiagnostics(error), |
314 | 369 | }) |
| 370 | + // Partial results from earlier pages are still useful; only fail if page one failed. |
| 371 | + if (tools.length > 0) return tools |
315 | 372 | throw error |
316 | 373 | } |
317 | 374 | } |
|
0 commit comments