Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
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
52 changes: 50 additions & 2 deletions apps/sim/lib/mcp/client.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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),
})
Expand All @@ -196,10 +196,58 @@ 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('logs connection diagnostics without header values', async () => {
const client = new McpClient({
config: {
Expand Down
101 changes: 79 additions & 22 deletions apps/sim/lib/mcp/client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down Expand Up @@ -275,43 +274,101 @@ 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 tools: McpTool[] = []
const seenCursors = new Set<string>()
let cursor: string | undefined
let bytes = 0
let truncated: string | undefined

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 < MCP_CLIENT_CONSTANTS.LIST_TOOLS_MAX_PAGES; 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,
})
},
}
)

if (!result.tools || !Array.isArray(result.tools)) {
logger.warn(`Invalid tools response from server ${this.config.name}:`, result)
break
}

for (const tool of result.tools) {
if (tools.length >= MCP_CLIENT_CONSTANTS.LIST_TOOLS_MAX_TOOLS) {
truncated = 'tool count'
break
}
bytes += JSON.stringify(tool).length
Comment thread
waleedlatif1 marked this conversation as resolved.
Outdated
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) 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 []
if (truncated || seenCursors.size >= MCP_CLIENT_CONSTANTS.LIST_TOOLS_MAX_PAGES - 1) {
logger.warn(`Tool discovery truncated for server ${this.config.name}`, {
serverId: this.config.id,
reason: truncated ?? 'page cap',
toolsCollected: tools.length,
pagesFetched: seenCursors.size + 1,
})
Comment thread
waleedlatif1 marked this conversation as resolved.
}

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,
phase: 'tools/list',
durationMs: Date.now() - startedAt,
idleTimeoutMs,
maxTotalTimeoutMs,
pagesFetched: seenCursors.size + 1,
toolsCollected: tools.length,
sessionIdPresent: Boolean(this.transport.sessionId),
error: getMcpSafeErrorDiagnostics(error),
})
// Partial results from earlier pages are still useful; only fail if page one failed.
if (tools.length > 0) return tools
Comment thread
waleedlatif1 marked this conversation as resolved.
Outdated
throw error
}
}
Expand Down
6 changes: 6 additions & 0 deletions apps/sim/lib/mcp/utils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
Loading