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
78 changes: 76 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,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: {
Expand Down
113 changes: 91 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,113 @@ 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<string>()
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,
})
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,
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
}
}
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