diff --git a/packages/api/src/mcp/__tests__/MCPConnectionAgentLifecycle.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionAgentLifecycle.test.ts index 302013824d..19c972f0c6 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionAgentLifecycle.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionAgentLifecycle.test.ts @@ -44,7 +44,13 @@ jest.mock('~/auth', () => ({ })); jest.mock('~/mcp/mcpConfig', () => ({ - mcpConfig: { CONNECTION_CHECK_TTL: 0, TOOLS_LIST_MAX_PAGES: 50 }, + mcpConfig: { + CONNECTION_CHECK_TTL: 0, + TOOLS_LIST_MAX_PAGES: 50, + TOOLS_LIST_MAX_TOOLS: 1000, + TOOLS_LIST_MAX_BYTES: 5 * 1024 * 1024, + TOOLS_LIST_TIMEOUT_MS: 30000, + }, })); const mockLogger = logger as jest.Mocked; diff --git a/packages/api/src/mcp/__tests__/MCPConnectionFactory.oauthSdk.integration.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionFactory.oauthSdk.integration.test.ts index ea5b3ec788..3ef669bb94 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionFactory.oauthSdk.integration.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionFactory.oauthSdk.integration.test.ts @@ -44,6 +44,9 @@ jest.mock('~/mcp/mcpConfig', () => ({ CONNECTION_CHECK_TTL: 0, USER_CONNECTION_IDLE_TIMEOUT: 30 * 60 * 1000, TOOLS_LIST_MAX_PAGES: 50, + TOOLS_LIST_MAX_TOOLS: 1000, + TOOLS_LIST_MAX_BYTES: 5 * 1024 * 1024, + TOOLS_LIST_TIMEOUT_MS: 30000, }, })); diff --git a/packages/api/src/mcp/__tests__/MCPConnectionFetchTools.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionFetchTools.test.ts index 095c2f36ea..d0aec66fc9 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionFetchTools.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionFetchTools.test.ts @@ -10,6 +10,7 @@ import { logger } from '@librechat/data-schemas'; import { MCPConnection } from '~/mcp/connection'; +import { mcpConfig } from '~/mcp/mcpConfig'; jest.mock('@librechat/data-schemas', () => ({ logger: { @@ -29,7 +30,13 @@ jest.mock('~/auth', () => ({ /** Pin the page cap to a small value so the cap path is cheap to exercise. */ jest.mock('~/mcp/mcpConfig', () => ({ - mcpConfig: { TOOLS_LIST_MAX_PAGES: 3, CONNECTION_CHECK_TTL: 0 }, + mcpConfig: { + TOOLS_LIST_MAX_PAGES: 3, + TOOLS_LIST_MAX_TOOLS: 1000, + TOOLS_LIST_MAX_BYTES: 5 * 1024 * 1024, + TOOLS_LIST_TIMEOUT_MS: 30000, + CONNECTION_CHECK_TTL: 0, + }, })); const mockLogger = logger as jest.Mocked; @@ -51,9 +58,27 @@ function createConnectionWithListTools(listTools: jest.Mock): MCPConnection { return conn; } +function expectListToolsCall( + listTools: jest.Mock, + callNumber: number, + params: { cursor?: string } | undefined, +): void { + expect(listTools).toHaveBeenNthCalledWith( + callNumber, + params, + expect.objectContaining({ + timeout: expect.any(Number), + maxTotalTimeout: expect.any(Number), + }), + ); +} + describe('MCPConnection.fetchTools pagination', () => { beforeEach(() => { jest.clearAllMocks(); + mcpConfig.TOOLS_LIST_MAX_TOOLS = 1000; + mcpConfig.TOOLS_LIST_MAX_BYTES = 5 * 1024 * 1024; + mcpConfig.TOOLS_LIST_TIMEOUT_MS = 30000; }); it('returns the tools from a single page and makes one request when there is no nextCursor', async () => { @@ -64,7 +89,7 @@ describe('MCPConnection.fetchTools pagination', () => { expect(tools.map((t) => t.name)).toEqual(['a', 'b']); expect(listTools).toHaveBeenCalledTimes(1); - expect(listTools).toHaveBeenNthCalledWith(1, undefined); + expectListToolsCall(listTools, 1, undefined); expect(mockLogger.warn).not.toHaveBeenCalled(); }); @@ -87,9 +112,9 @@ describe('MCPConnection.fetchTools pagination', () => { expect(tools.map((t) => t.name)).toEqual(['a', 'b', 'c', 'd', 'e']); expect(listTools).toHaveBeenCalledTimes(3); - expect(listTools).toHaveBeenNthCalledWith(1, undefined); - expect(listTools).toHaveBeenNthCalledWith(2, { cursor: 'c1' }); - expect(listTools).toHaveBeenNthCalledWith(3, { cursor: 'c2' }); + expectListToolsCall(listTools, 1, undefined); + expectListToolsCall(listTools, 2, { cursor: 'c1' }); + expectListToolsCall(listTools, 3, { cursor: 'c2' }); expect(mockLogger.warn).not.toHaveBeenCalled(); }); @@ -109,6 +134,96 @@ describe('MCPConnection.fetchTools pagination', () => { expect(mockLogger.warn).toHaveBeenCalledWith(expect.stringContaining('pagination limit')); }); + it('stops at the aggregate tool-count budget and warns', async () => { + mcpConfig.TOOLS_LIST_MAX_TOOLS = 3; + const listTools = jest.fn(async (params?: { cursor?: string }) => { + if (params?.cursor == null) { + return { tools: [makeTool('a'), makeTool('b')], nextCursor: 'c1' }; + } + return { tools: [makeTool('c'), makeTool('d')], nextCursor: 'c2' }; + }); + const conn = createConnectionWithListTools(listTools); + + const tools = await conn.fetchTools(); + + expect(tools.map((t) => t.name)).toEqual(['a', 'b', 'c']); + expect(listTools).toHaveBeenCalledTimes(2); + expect(mockLogger.warn).toHaveBeenCalledWith(expect.stringContaining('tool count budget')); + }); + + it('does not request another page when the tool-count budget is exactly full', async () => { + mcpConfig.TOOLS_LIST_MAX_TOOLS = 2; + const listTools = jest.fn(async (params?: { cursor?: string }) => { + if (params?.cursor == null) { + return { tools: [makeTool('a'), makeTool('b')], nextCursor: 'c1' }; + } + return { tools: [makeTool('c')] }; + }); + const conn = createConnectionWithListTools(listTools); + + const tools = await conn.fetchTools(); + + expect(tools.map((t) => t.name)).toEqual(['a', 'b']); + expect(listTools).toHaveBeenCalledTimes(1); + expect(mockLogger.warn).toHaveBeenCalledWith(expect.stringContaining('tool count budget')); + }); + + it('stops at the aggregate byte budget and warns', async () => { + mcpConfig.TOOLS_LIST_MAX_BYTES = 170; + const listTools = jest.fn(async () => ({ + tools: [makeTool('a'), makeTool('b')], + nextCursor: 'c1', + })); + const conn = createConnectionWithListTools(listTools); + + const tools = await conn.fetchTools(); + + expect(tools.map((t) => t.name)).toEqual(['a']); + expect(listTools).toHaveBeenCalledTimes(1); + expect(mockLogger.warn).toHaveBeenCalledWith(expect.stringContaining('size budget')); + }); + + it('stops at the elapsed-time budget before requesting another page', async () => { + mcpConfig.TOOLS_LIST_TIMEOUT_MS = 1; + const listTools = jest.fn(async () => ({ tools: [makeTool('a')], nextCursor: 'c1' })); + const conn = createConnectionWithListTools(listTools); + const dateNow = jest + .spyOn(Date, 'now') + .mockReturnValueOnce(1000) + .mockReturnValueOnce(1000) + .mockReturnValueOnce(1001); + + const tools = await conn.fetchTools(); + + expect(tools.map((t) => t.name)).toEqual(['a']); + expect(listTools).toHaveBeenCalledTimes(1); + expect(mockLogger.warn).toHaveBeenCalledWith(expect.stringContaining('time budget')); + dateNow.mockRestore(); + }); + + it('passes the elapsed-time budget to the SDK request timeout', async () => { + mcpConfig.TOOLS_LIST_TIMEOUT_MS = 25; + const listTools = jest.fn( + async ( + _params?: { cursor?: string }, + _options?: { timeout: number; maxTotalTimeout: number }, + ) => { + throw new Error('Request timed out'); + }, + ); + const conn = createConnectionWithListTools(listTools); + + const tools = await conn.fetchTools(); + + expect(tools).toEqual([]); + expect(listTools).toHaveBeenCalledTimes(1); + const options = listTools.mock.calls[0][1]!; + expect(options.timeout).toBeGreaterThan(0); + expect(options.timeout).toBeLessThanOrEqual(25); + expect(options.maxTotalTimeout).toBe(options.timeout); + expect(mockLogger.error).toHaveBeenCalledWith(expect.stringContaining('Request timed out')); + }); + it('stops and warns when the server repeats a cursor instead of looping forever', async () => { const listTools = jest.fn().mockResolvedValue({ tools: [makeTool('x')], nextCursor: 'same' }); const conn = createConnectionWithListTools(listTools); @@ -151,7 +266,7 @@ describe('MCPConnection.fetchTools pagination', () => { expect(tools.map((t) => t.name)).toEqual(['a', 'b']); expect(listTools).toHaveBeenCalledTimes(2); - expect(listTools).toHaveBeenNthCalledWith(2, { cursor: '' }); + expectListToolsCall(listTools, 2, { cursor: '' }); }); it('returns the pages already fetched when a later page fails, without throwing', async () => { diff --git a/packages/api/src/mcp/__tests__/MCPConnectionSSRF.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionSSRF.test.ts index c6aec18f67..702a073721 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionSSRF.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionSSRF.test.ts @@ -71,7 +71,13 @@ jest.mock('~/auth', () => ({ })); jest.mock('~/mcp/mcpConfig', () => ({ - mcpConfig: { CONNECTION_CHECK_TTL: 0, TOOLS_LIST_MAX_PAGES: 50 }, + mcpConfig: { + CONNECTION_CHECK_TTL: 0, + TOOLS_LIST_MAX_PAGES: 50, + TOOLS_LIST_MAX_TOOLS: 1000, + TOOLS_LIST_MAX_BYTES: 5 * 1024 * 1024, + TOOLS_LIST_TIMEOUT_MS: 30000, + }, })); const mockedResolveHostnameSSRF = resolveHostnameSSRF as jest.MockedFunction< diff --git a/packages/api/src/mcp/__tests__/dbSourced.integration.test.ts b/packages/api/src/mcp/__tests__/dbSourced.integration.test.ts index 4ee078795c..981f537a57 100644 --- a/packages/api/src/mcp/__tests__/dbSourced.integration.test.ts +++ b/packages/api/src/mcp/__tests__/dbSourced.integration.test.ts @@ -43,7 +43,13 @@ jest.mock('~/auth', () => ({ })); jest.mock('~/mcp/mcpConfig', () => ({ - mcpConfig: { CONNECTION_CHECK_TTL: 0, TOOLS_LIST_MAX_PAGES: 50 }, + mcpConfig: { + CONNECTION_CHECK_TTL: 0, + TOOLS_LIST_MAX_PAGES: 50, + TOOLS_LIST_MAX_TOOLS: 1000, + TOOLS_LIST_MAX_BYTES: 5 * 1024 * 1024, + TOOLS_LIST_TIMEOUT_MS: 30000, + }, })); /** Track all Agents for cleanup */ diff --git a/packages/api/src/mcp/connection.ts b/packages/api/src/mcp/connection.ts index bd9884ef64..87116b6a6e 100644 --- a/packages/api/src/mcp/connection.ts +++ b/packages/api/src/mcp/connection.ts @@ -30,6 +30,7 @@ import { mcpConfig } from './mcpConfig'; type FetchLike = (url: string | URL, init?: RequestInit) => Promise; type ManagedDispatcher = Agent | ProxyAgent; type ParsedIP = { version: 4 | 6; bits: 32 | 128; value: bigint }; +type MCPTool = MCPListToolsResult['tools'][number]; const BIGINT_ZERO = BigInt(0); const BIGINT_ONE = BigInt(1); @@ -37,6 +38,29 @@ const BIGINT_EIGHT = BigInt(8); const BIGINT_SIXTEEN = BigInt(16); const UINT16_MASK = BigInt(0xffff); +function getApproximateToolBytes(tool: MCPTool): number { + try { + return Buffer.byteLength(JSON.stringify(tool), 'utf8'); + } catch { + return Number.POSITIVE_INFINITY; + } +} + +function getToolsListBudgetExceededReason( + toolCount: number, + totalBytes: number, + maxTools: number, + maxBytes: number, +): string | null { + if (toolCount >= maxTools) { + return 'tool count'; + } + if (totalBytes >= maxBytes) { + return 'size'; + } + return null; +} + type MCPProxyConfig = | { type: 'explicit'; @@ -2202,29 +2226,77 @@ export class MCPConnection extends EventEmitter { * server that spans multiple pages (e.g. an aggregating gateway exposing many * tools) is loaded in full instead of being truncated to the first page. * - * Pagination is bounded by {@link mcpConfig.TOOLS_LIST_MAX_PAGES} and a - * repeated-cursor guard. On error, the tools already fetched are returned rather - * than discarded, and the method never throws. + * Pagination is bounded by {@link mcpConfig.TOOLS_LIST_MAX_PAGES}, aggregate + * tool count, approximate serialized size, elapsed time, and a repeated-cursor + * guard. On error, the tools already fetched are returned rather than discarded, + * and the method never throws. */ async fetchTools(): Promise { const maxPages = mcpConfig.TOOLS_LIST_MAX_PAGES; + const maxTools = mcpConfig.TOOLS_LIST_MAX_TOOLS; + const maxBytes = mcpConfig.TOOLS_LIST_MAX_BYTES; + const deadline = Date.now() + mcpConfig.TOOLS_LIST_TIMEOUT_MS; const allTools: MCPListToolsResult['tools'] = []; const seenCursors = new Set(); let cursor: string | undefined; + let totalBytes = 0; for (let page = 1; page <= maxPages; page++) { - const result = await this.listToolsPage(cursor); + const exhaustedBudget = getToolsListBudgetExceededReason( + allTools.length, + totalBytes, + maxTools, + maxBytes, + ); + if (exhaustedBudget != null) { + this.warnToolsListBudgetExceeded(exhaustedBudget, allTools.length); + return allTools; + } + + const remainingMs = deadline - Date.now(); + if (remainingMs <= 0) { + this.warnToolsListBudgetExceeded('time', allTools.length); + return allTools; + } + + const result = await this.listToolsPage(cursor, remainingMs); if (result == null) { /** Request failed mid-pagination: return the pages already fetched instead of discarding them. */ return allTools; } - allTools.push(...result.tools); + for (const tool of result.tools) { + if (allTools.length >= maxTools) { + this.warnToolsListBudgetExceeded('tool count', allTools.length); + return allTools; + } + + const toolBytes = getApproximateToolBytes(tool); + if (totalBytes + toolBytes > maxBytes) { + this.warnToolsListBudgetExceeded('size', allTools.length); + return allTools; + } + + allTools.push(tool); + totalBytes += toolBytes; + } const { nextCursor } = result; if (nextCursor == null) { return allTools; } + + const nextPageBudget = getToolsListBudgetExceededReason( + allTools.length, + totalBytes, + maxTools, + maxBytes, + ); + if (nextPageBudget != null) { + this.warnToolsListBudgetExceeded(nextPageBudget, allTools.length); + return allTools; + } + if (seenCursors.has(nextCursor)) { logger.warn( `${this.getLogPrefix()} MCP server returned a repeated tools/list cursor; stopping pagination after ${page} page(s).`, @@ -2242,10 +2314,22 @@ export class MCPConnection extends EventEmitter { return allTools; } + private warnToolsListBudgetExceeded(reason: string, toolCount: number): void { + logger.warn( + `${this.getLogPrefix()} Stopping tools/list pagination because the ${reason} budget was reached after ${toolCount} tool(s).`, + ); + } + /** Fetches a single `tools/list` page, returning null (and logging) on failure so pagination can stop gracefully. */ - private async listToolsPage(cursor: string | undefined): Promise { + private async listToolsPage( + cursor: string | undefined, + timeoutMs: number, + ): Promise { try { - return await this.client.listTools(cursor != null ? { cursor } : undefined); + return await this.client.listTools(cursor != null ? { cursor } : undefined, { + timeout: timeoutMs, + maxTotalTimeout: timeoutMs, + }); } catch (error) { this.emitError(error, 'Failed to fetch tools'); return null; diff --git a/packages/api/src/mcp/mcpConfig.ts b/packages/api/src/mcp/mcpConfig.ts index ea75220958..68ae3cdfe3 100644 --- a/packages/api/src/mcp/mcpConfig.ts +++ b/packages/api/src/mcp/mcpConfig.ts @@ -26,6 +26,12 @@ export const mcpConfig: { /** Max number of `tools/list` pages to request when an MCP server paginates its tool list. * Bounds the pagination loop so a misbehaving server cannot stall tool discovery. Default: 50 */ TOOLS_LIST_MAX_PAGES: number; + /** Max total tools to retain from paginated `tools/list` responses. Default: 1000 */ + TOOLS_LIST_MAX_TOOLS: number; + /** Max approximate JSON bytes to retain from paginated `tools/list` responses. Default: 5 MiB */ + TOOLS_LIST_MAX_BYTES: number; + /** Max elapsed time (ms) for paginated `tools/list` discovery. Default: 30000 */ + TOOLS_LIST_TIMEOUT_MS: number; /** Idle timeout (ms) after which user connections are disconnected. Default: 15 minutes */ USER_CONNECTION_IDLE_TIMEOUT: number; /** Max connect/disconnect cycles before the circuit breaker trips. Default: 7 */ @@ -52,6 +58,12 @@ export const mcpConfig: { CONNECTION_CHECK_TTL: math(process.env.MCP_CONNECTION_CHECK_TTL ?? 60000), /** Max number of `tools/list` pages to request when an MCP server paginates its tool list. Clamped to >= 1. Default: 50 */ TOOLS_LIST_MAX_PAGES: Math.max(1, math(process.env.MCP_TOOLS_LIST_MAX_PAGES ?? 50)), + /** Max total tools to retain from paginated `tools/list` responses. Clamped to >= 1. Default: 1000 */ + TOOLS_LIST_MAX_TOOLS: Math.max(1, math(process.env.MCP_TOOLS_LIST_MAX_TOOLS ?? 1000)), + /** Max approximate JSON bytes to retain from paginated `tools/list` responses. Clamped to >= 1. Default: 5 MiB */ + TOOLS_LIST_MAX_BYTES: Math.max(1, math(process.env.MCP_TOOLS_LIST_MAX_BYTES ?? 5 * 1024 * 1024)), + /** Max elapsed time (ms) for paginated `tools/list` discovery. Clamped to >= 1. Default: 30000 */ + TOOLS_LIST_TIMEOUT_MS: Math.max(1, math(process.env.MCP_TOOLS_LIST_TIMEOUT_MS ?? 30_000)), /** Idle timeout (ms) after which user connections are disconnected. Default: 15 minutes */ USER_CONNECTION_IDLE_TIMEOUT: math( process.env.MCP_USER_CONNECTION_IDLE_TIMEOUT ?? 15 * 60 * 1000, diff --git a/packages/api/src/mcp/registry/__tests__/MCPReinitRecovery.integration.test.ts b/packages/api/src/mcp/registry/__tests__/MCPReinitRecovery.integration.test.ts index dd73e89d43..cd21502432 100644 --- a/packages/api/src/mcp/registry/__tests__/MCPReinitRecovery.integration.test.ts +++ b/packages/api/src/mcp/registry/__tests__/MCPReinitRecovery.integration.test.ts @@ -56,7 +56,13 @@ jest.mock('~/cluster', () => ({ })); jest.mock('~/mcp/mcpConfig', () => ({ - mcpConfig: { CONNECTION_CHECK_TTL: 0, TOOLS_LIST_MAX_PAGES: 50 }, + mcpConfig: { + CONNECTION_CHECK_TTL: 0, + TOOLS_LIST_MAX_PAGES: 50, + TOOLS_LIST_MAX_TOOLS: 1000, + TOOLS_LIST_MAX_BYTES: 5 * 1024 * 1024, + TOOLS_LIST_TIMEOUT_MS: 30000, + }, })); jest.mock('~/mcp/registry/db/ServerConfigsDB', () => ({