diff --git a/api/server/controllers/mcp.js b/api/server/controllers/mcp.js index 1d1f6eed7c..447c3f1abb 100644 --- a/api/server/controllers/mcp.js +++ b/api/server/controllers/mcp.js @@ -13,6 +13,7 @@ const { createAuthIdentityContext, MCPConnection, MCPErrorCodes, + MCPCatalogCapacityError, splitMCPToolKey, normalizeServerName, findShadowedServerNames, @@ -201,21 +202,35 @@ const getMCPTools = async (req, res) => { user: req.user, tenantId: getTenantId(), }); - const { serverTools: serverToolsMap, serversWithoutTools } = await loadMCPServerCatalogs({ - user: req.user, - servers: configuredServers.map((serverName) => ({ - serverName, - serverConfig: mcpConfig[serverName], - })), - upstreamTokenProvider: createOpenIDSessionTokenProvider({ - req, - res, + const catalogAbortController = new AbortController(); + const abortCatalogLoad = () => { + if (!res.writableEnded) { + catalogAbortController.abort(); + } + }; + res.once('close', abortCatalogLoad); + let catalogResult; + try { + catalogResult = await loadMCPServerCatalogs({ user: req.user, - identityContext: oboIdentityContext, - tokenPreference: 'access_token', - }), - oboIdentityContext, - }); + servers: configuredServers.map((serverName) => ({ + serverName, + serverConfig: mcpConfig[serverName], + })), + upstreamTokenProvider: createOpenIDSessionTokenProvider({ + req, + res, + user: req.user, + identityContext: oboIdentityContext, + tokenPreference: 'access_token', + }), + oboIdentityContext, + signal: catalogAbortController.signal, + }); + } finally { + res.off('close', abortCatalogLoad); + } + const { serverTools: serverToolsMap, serversWithoutTools } = catalogResult; if (serversWithoutTools.length > 0) { logger.debug( `[getMCPTools] No tools (${serversWithoutTools.length}): ${serversWithoutTools.join(', ')}`, @@ -286,7 +301,11 @@ const getMCPTools = async (req, res) => { res.status(200).json({ servers: mcpServers }); } catch (error) { logger.error('[getMCPTools]', error); - res.status(500).json({ message: error.message }); + if (res.destroyed || res.headersSent) { + return; + } + const status = error instanceof MCPCatalogCapacityError ? 503 : 500; + res.status(status).json({ message: error.message }); } }; /** diff --git a/api/server/routes/__tests__/mcp.spec.js b/api/server/routes/__tests__/mcp.spec.js index fe5ca60df5..aa5cb4c2e5 100644 --- a/api/server/routes/__tests__/mcp.spec.js +++ b/api/server/routes/__tests__/mcp.spec.js @@ -3324,6 +3324,7 @@ describe('MCP Routes', () => { servers: [{ serverName: 'user-server', serverConfig }], upstreamTokenProvider: expect.any(Function), oboIdentityContext: expect.any(Object), + signal: expect.any(AbortSignal), }); }); @@ -3376,6 +3377,7 @@ describe('MCP Routes', () => { servers: [{ serverName: 'connected-server', serverConfig }], upstreamTokenProvider: expect.any(Function), oboIdentityContext: expect.any(Object), + signal: expect.any(AbortSignal), }); }); @@ -3474,8 +3476,22 @@ describe('MCP Routes', () => { })), upstreamTokenProvider: expect.any(Function), oboIdentityContext: expect.any(Object), + signal: expect.any(AbortSignal), }); }); + + it('returns a retryable failure when process-wide catalog capacity is saturated', async () => { + const { MCPCatalogCapacityError } = require('@librechat/api'); + mockResolveAllMcpConfigs.mockResolvedValueOnce({ + saturated: { type: 'sse', url: 'https://saturated.example.com/sse' }, + }); + mockLoadMCPServerCatalogs.mockRejectedValueOnce(new MCPCatalogCapacityError()); + + const response = await request(app).get('/api/mcp/tools'); + + expect(response.status).toBe(503); + expect(response.body).toEqual({ message: 'MCP catalog capacity is temporarily exhausted' }); + }); }); describe('GET /servers', () => { diff --git a/api/server/services/MCP.js b/api/server/services/MCP.js index 3f636d7591..2e24ad162b 100644 --- a/api/server/services/MCP.js +++ b/api/server/services/MCP.js @@ -399,11 +399,12 @@ async function getAssistantToolDefinitions({ req, res, tools }) { getAllServerConfigs: (userId, configServers, role) => registry.getAllServerConfigs(userId, configServers, role), getMCPServerTools, - getServerToolFunctionsSnapshot: async (userId, serverName, serverConfig) => + getServerToolFunctionsSnapshot: async (userId, serverName, serverConfig, options) => (await getMCPManager()?.getServerToolFunctionsSnapshot( userId, serverName, serverConfig, + options, )) ?? { tools: null, }, diff --git a/api/server/services/Tools/mcp.js b/api/server/services/Tools/mcp.js index b1eb82f303..25906f6377 100644 --- a/api/server/services/Tools/mcp.js +++ b/api/server/services/Tools/mcp.js @@ -40,13 +40,20 @@ const MCP_REINITIALIZE_FAILURE_REASONS = { * @param {Array<{ serverName: string, serverConfig: object }>} params.servers * @param {import('@librechat/api').UpstreamTokenProvider} [params.upstreamTokenProvider] - Live upstream-token closure for OBO discovery, built at the request boundary so this layer never receives the raw Express request. * @param {import('@librechat/api').AuthIdentityContext} [params.oboIdentityContext] - Non-template-visible OBO identity context built from the real request user. + * @param {AbortSignal} [params.signal] - Cancels queued and in-flight catalog reads when the request ends. */ -async function loadMCPServerCatalogs({ user, servers, upstreamTokenProvider, oboIdentityContext }) { +async function loadMCPServerCatalogs({ + user, + servers, + upstreamTokenProvider, + oboIdentityContext, + signal, +}) { const flowManager = getFlowStateManager(getLogStores(CacheKeys.FLOWS)); const tokenMethods = { findToken, updateToken, createToken, deleteTokens }; const mcpManager = getMCPManager(); return loadCatalogs( - { user, servers }, + { user, servers, signal }, { loadUserMCPAuthMap: (userId, serverNames) => getUserMCPAuthMap({ @@ -67,8 +74,8 @@ async function loadMCPServerCatalogs({ user, servers, upstreamTokenProvider, obo }), formatServerTools: formatMCPServerTools, getCachedServerTools: getMCPServerTools, - getServerToolFunctionsSnapshot: (userId, serverName, serverConfig) => - mcpManager.getServerToolFunctionsSnapshot(userId, serverName, serverConfig), + getServerToolFunctionsSnapshot: (userId, serverName, serverConfig, options) => + mcpManager.getServerToolFunctionsSnapshot(userId, serverName, serverConfig, options), cacheServerTools: cacheMCPServerTools, }, ); diff --git a/api/server/services/Tools/mcp.spec.js b/api/server/services/Tools/mcp.spec.js index 8eb8bac327..89e1ad5c82 100644 --- a/api/server/services/Tools/mcp.spec.js +++ b/api/server/services/Tools/mcp.spec.js @@ -85,7 +85,9 @@ describe('loadMCPServerCatalogs', () => { }); deps.formatServerTools('config-only', []); await deps.getCachedServerTools(user.id, 'config-only', servers[0].serverConfig); - await deps.getServerToolFunctionsSnapshot(user.id, 'config-only', servers[0].serverConfig); + await deps.getServerToolFunctionsSnapshot(user.id, 'config-only', servers[0].serverConfig, { + deadlineMs: 123, + }); await deps.cacheServerTools({ serverName: 'config-only' }); return { serverTools: new Map([['config-only', {}]]), serversWithoutTools: [] }; }); @@ -126,6 +128,7 @@ describe('loadMCPServerCatalogs', () => { user.id, 'config-only', servers[0].serverConfig, + { deadlineMs: 123 }, ); expect(mockCacheMCPServerTools).toHaveBeenCalledWith({ serverName: 'config-only' }); expect(result).toEqual({ diff --git a/packages/api/src/mcp/MCPManager.ts b/packages/api/src/mcp/MCPManager.ts index c43d4d747d..425aeda004 100644 --- a/packages/api/src/mcp/MCPManager.ts +++ b/packages/api/src/mcp/MCPManager.ts @@ -450,12 +450,18 @@ export class MCPManager extends UserConnectionManager { userId: string, serverName: string, serverConfig?: t.ParsedServerConfig, + options?: { deadlineMs?: number; signal?: AbortSignal }, ): Promise<{ tools: t.LCAvailableTools | null; publicationGeneration?: string; publicationRevision?: string; }> { try { + const signal = createDeadlineAbortSignal(options?.deadlineMs, options?.signal); + const readToolCatalog = (connection: MCPConnection) => + options == null + ? MCPServerInspector.getToolCatalog(serverName, connection) + : MCPServerInspector.getToolCatalog(serverName, connection, options.deadlineMs, signal); const registry = MCPServersRegistry.getInstance(); const effectiveConfig = serverConfig ?? (await registry.getServerConfig(serverName, userId)); const useAppConnection = @@ -466,7 +472,7 @@ export class MCPManager extends UserConnectionManager { ? await this.appConnections?.get(serverName) : null; if (existingAppConnection != null) { - return MCPServerInspector.getToolCatalog(serverName, existingAppConnection); + return readToolCatalog(existingAppConnection); } let awaitedRecovery: Promise | undefined; @@ -507,12 +513,12 @@ export class MCPManager extends UserConnectionManager { if (recovery && recovery !== awaitedRecovery) { awaitedRecovery = recovery; await this.releaseConnection(connection); - await this.waitForConnectionRecovery(recovery); + await this.waitForConnectionRecovery(recovery, signal); continue; } try { - const { tools } = await MCPServerInspector.getToolCatalog(serverName, connection); + const { tools } = await readToolCatalog(connection); const generationAfterFetch = await getMCPToolsChangedGeneration({ userId, serverName }); if ( publicationGeneration != null && diff --git a/packages/api/src/mcp/__tests__/MCPManager.test.ts b/packages/api/src/mcp/__tests__/MCPManager.test.ts index eb27129e98..b7a785d0f1 100644 --- a/packages/api/src/mcp/__tests__/MCPManager.test.ts +++ b/packages/api/src/mcp/__tests__/MCPManager.test.ts @@ -3878,13 +3878,27 @@ describe('MCPManager', () => { )?.[1]; const tools: t.MCPTool[] = [{ name: 'current', inputSchema: { type: 'object' } }]; listener(tools); - const snapshot = await manager.getServerToolFunctionsSnapshot(userId, serverName); + const deadlineMs = Date.now() + 3000; + const snapshot = await manager.getServerToolFunctionsSnapshot( + userId, + serverName, + undefined, + { + deadlineMs, + }, + ); expect(generationSpy).toHaveBeenCalledWith({ userId, serverName }); expect(snapshot).toEqual({ tools: expectedToolFunctions, publicationGeneration: 'generation-a', }); + expect(MCPServerInspector.getToolCatalog).toHaveBeenCalledWith( + serverName, + mockConnection, + deadlineMs, + expect.any(AbortSignal), + ); expect(notifySpy).toHaveBeenCalledWith({ tools, userId, diff --git a/packages/api/src/mcp/catalog/recovery.spec.ts b/packages/api/src/mcp/catalog/recovery.spec.ts index 053230268c..14aae00773 100644 --- a/packages/api/src/mcp/catalog/recovery.spec.ts +++ b/packages/api/src/mcp/catalog/recovery.spec.ts @@ -3,7 +3,11 @@ import { Constants } from 'librechat-data-provider'; import { ErrorCode, McpError } from '@modelcontextprotocol/sdk/types.js'; import type { IUser } from '@librechat/data-schemas'; import type { LCAvailableTools, ParsedServerConfig, ToolDiscoveryOptions } from '../types'; -import { loadMCPServerCatalogs, recoverMCPServerCatalogs } from './recovery'; +import { + MCPCatalogCapacityError, + loadMCPServerCatalogs, + recoverMCPServerCatalogs, +} from './recovery'; jest.mock('@librechat/data-schemas', () => ({ logger: { info: jest.fn(), warn: jest.fn(), error: jest.fn(), debug: jest.fn() }, @@ -71,7 +75,7 @@ describe('recoverMCPServerCatalogs', () => { expect(result.size).toBe(2); }); - it('limits passive discovery to three concurrent servers', async () => { + it('limits passive discovery to three concurrent servers across simultaneous catalog loads', async () => { const servers = Array.from({ length: 7 }, (_, index) => ({ serverName: `server-${index}`, serverConfig: serverConfig(`server-${index}`), @@ -85,19 +89,39 @@ describe('recoverMCPServerCatalogs', () => { active -= 1; return { tools: [] }; }); + const getServerToolFunctionsSnapshot = jest.fn(async () => { + active += 1; + maxActive = Math.max(maxActive, active); + await new Promise((resolve) => setTimeout(resolve, 5)); + active -= 1; + return { tools: null }; + }); - const result = await recoverMCPServerCatalogs( - { user, servers }, - { - loadUserMCPAuthMap: jest.fn().mockResolvedValue({}), - discoverServerTools, - formatServerTools: jest.fn().mockReturnValue({}), - }, + const deps = { + getCachedServerTools: jest.fn().mockResolvedValue(null), + getServerToolFunctionsSnapshot, + cacheServerTools: jest.fn(), + loadUserMCPAuthMap: jest.fn().mockResolvedValue({}), + discoverServerTools, + formatServerTools: jest.fn().mockReturnValue({}), + }; + const results = await Promise.allSettled( + Array.from({ length: 4 }, () => loadMCPServerCatalogs({ user, servers }, deps)), ); - expect(discoverServerTools).toHaveBeenCalledTimes(7); + expect(getServerToolFunctionsSnapshot).toHaveBeenCalledTimes(21); + expect(discoverServerTools).toHaveBeenCalledTimes(21); expect(maxActive).toBe(3); - expect(result.size).toBe(7); + expect(results).toHaveLength(4); + expect(results.filter(({ status }) => status === 'fulfilled')).toHaveLength(3); + expect(results.filter(({ status }) => status === 'rejected')).toEqual([ + { status: 'rejected', reason: expect.any(MCPCatalogCapacityError) }, + ]); + expect( + results.flatMap((result) => + result.status === 'fulfilled' ? [...result.value.serverTools.keys()] : [], + ), + ).toHaveLength(21); }); it('logs configuration-impossible discovery failures at debug, keeping error for the unexpected', async () => { @@ -500,22 +524,31 @@ describe('recoverMCPServerCatalogs — bounded, skippable discovery', () => { it('bounds snapshot refreshes, which each issue a real tools/list', async () => { let active = 0; let maxActive = 0; + const before = Date.now(); const servers = Array.from({ length: 9 }, (_, index) => ({ serverName: `server-${index}`, serverConfig: serverConfig(`server-${index}`), })); + const getServerToolFunctionsSnapshot = jest.fn( + async ( + _userId: string, + _serverName: string, + _serverConfig: ParsedServerConfig, + _options?: { deadlineMs?: number; signal?: AbortSignal }, + ) => { + active += 1; + maxActive = Math.max(maxActive, active); + await new Promise((resolve) => setTimeout(resolve, 20)); + active -= 1; + return { tools: null }; + }, + ); await loadMCPServerCatalogs( { user, servers }, { getCachedServerTools: jest.fn().mockResolvedValue(null), - getServerToolFunctionsSnapshot: jest.fn(async () => { - active += 1; - maxActive = Math.max(maxActive, active); - await new Promise((resolve) => setTimeout(resolve, 20)); - active -= 1; - return { tools: null }; - }), + getServerToolFunctionsSnapshot, cacheServerTools: jest.fn(), loadUserMCPAuthMap: jest.fn().mockResolvedValue({}), discoverServerTools: jest.fn().mockResolvedValue({ tools: null }), @@ -524,5 +557,49 @@ describe('recoverMCPServerCatalogs — bounded, skippable discovery', () => { ); expect(maxActive).toBe(3); + const options = getServerToolFunctionsSnapshot.mock.calls[0]?.[3]; + expect(options?.deadlineMs).toBeGreaterThanOrEqual(before + 30_000); + expect(options?.deadlineMs).toBeLessThanOrEqual(Date.now() + 30_000); + }); + + it('cancels queued catalog work when the originating request ends', async () => { + const controller = new AbortController(); + const releases: Array<() => void> = []; + const servers = Array.from({ length: 9 }, (_, index) => ({ + serverName: `server-${index}`, + serverConfig: serverConfig(`server-${index}`), + })); + const getServerToolFunctionsSnapshot = jest.fn( + () => + new Promise<{ tools: null }>((resolve) => { + releases.push(() => resolve({ tools: null })); + }), + ); + + const loading = loadMCPServerCatalogs( + { user, servers, signal: controller.signal }, + { + getCachedServerTools: jest.fn().mockResolvedValue(null), + getServerToolFunctionsSnapshot, + cacheServerTools: jest.fn(), + loadUserMCPAuthMap: jest.fn().mockResolvedValue({}), + discoverServerTools: jest.fn().mockResolvedValue({ tools: null }), + formatServerTools: jest.fn().mockReturnValue({}), + }, + ); + await new Promise((resolve) => setImmediate(resolve)); + expect(getServerToolFunctionsSnapshot).toHaveBeenCalledTimes(3); + + controller.abort(); + releases.splice(0).forEach((release) => release()); + await loading; + + expect(getServerToolFunctionsSnapshot).toHaveBeenCalledTimes(3); + expect(getServerToolFunctionsSnapshot).toHaveBeenCalledWith( + user.id, + 'server-0', + servers[0].serverConfig, + expect.objectContaining({ signal: controller.signal }), + ); }); }); diff --git a/packages/api/src/mcp/catalog/recovery.ts b/packages/api/src/mcp/catalog/recovery.ts index f4e9f7ab95..ee7a3112f3 100644 --- a/packages/api/src/mcp/catalog/recovery.ts +++ b/packages/api/src/mcp/catalog/recovery.ts @@ -4,12 +4,14 @@ import type { Tool } from '@modelcontextprotocol/sdk/types.js'; import type { IUser } from '@librechat/data-schemas'; import type { LCAvailableTools, ParsedServerConfig, ToolDiscoveryOptions } from '../types'; import { hasCustomUserVars, getMissingCustomUserVars } from '../utils'; -import { createConcurrencyLimiter } from '~/utils/promise'; import { getServerCustomUserVars } from '../auth'; +import { mcpConfig } from '../mcpConfig'; -/** Bounds outbound MCP fan-out from one catalog request: snapshot refreshes issue a real - * `tools/list`, so they burst as readily as passive discovery does. */ +/** Bounds all outbound catalog reads across this runtime. */ const CATALOG_FANOUT_CONCURRENCY = 3; +const activeCatalogLanes = new Set(); +const pendingCatalogLanes: CatalogWorkLane[] = []; +let activeCatalogWork = 0; /** * Bounds one server's discovery end to end — connect, `tools/list` pagination, and the * unauthenticated fallback all draw down this single budget, so a slot is held for at most this @@ -48,6 +50,7 @@ export interface MCPServerCatalogLoaderDeps extends MCPServerCatalogRecoveryDeps userId: string, serverName: string, serverConfig: ParsedServerConfig, + options?: { deadlineMs?: number; signal?: AbortSignal }, ) => Promise; cacheServerTools: (params: { userId: string; @@ -64,10 +67,34 @@ export interface MCPServerCatalogLoaderResult { serversWithoutTools: string[]; } +interface MCPServerCatalogEntry extends MCPServerCatalogSnapshot { + serverName: string; + serverConfig: ParsedServerConfig; + source: 'cache' | 'snapshot'; +} + interface RecoveryCandidate extends MCPServerCatalogRecoveryInput { customUserVars?: Record; } +interface CatalogWorkLane { + tasks: ReadonlyArray<() => Promise>; + resolve: (admitted: boolean) => void; + nextIndex: number; + inFlight: number; + signal?: AbortSignal; + onAbort?: () => void; +} + +export class MCPCatalogCapacityError extends Error { + readonly code = 'MCP_CATALOG_CAPACITY'; + + constructor() { + super('MCP catalog capacity is temporarily exhausted'); + this.name = 'MCPCatalogCapacityError'; + } +} + /** Bounds one server's discovery, honouring a shorter operator `initTimeout`. */ function resolveBudget(serverConfig: ParsedServerConfig): number { const { initTimeout } = serverConfig; @@ -81,6 +108,7 @@ async function discoverCandidate( user: IUser, { serverName, serverConfig, customUserVars }: RecoveryCandidate, deps: MCPServerCatalogRecoveryDeps, + signal?: AbortSignal, ): Promise<[string, LCAvailableTools | null]> { try { const result = await deps.discoverServerTools({ @@ -89,6 +117,7 @@ async function discoverCandidate( configServers: { [serverName]: serverConfig }, customUserVars, deadlineMs: Date.now() + resolveBudget(serverConfig), + signal, }); return [ serverName, @@ -110,6 +139,82 @@ async function discoverCandidate( } } +function finishCatalogLane(lane: CatalogWorkLane): void { + if (lane.nextIndex < lane.tasks.length || lane.inFlight > 0) { + return; + } + activeCatalogLanes.delete(lane); + if (lane.signal != null && lane.onAbort != null) { + lane.signal.removeEventListener('abort', lane.onAbort); + } + lane.resolve(true); +} + +/** Shares the fixed process budget round-robin across at most three request lanes. */ +function scheduleCatalogWork(): void { + while (activeCatalogWork < CATALOG_FANOUT_CONCURRENCY && pendingCatalogLanes.length > 0) { + const lane = pendingCatalogLanes.shift(); + if (lane == null) { + return; + } + const task = lane.tasks[lane.nextIndex]; + lane.nextIndex += 1; + lane.inFlight += 1; + activeCatalogWork += 1; + if (lane.nextIndex < lane.tasks.length) { + pendingCatalogLanes.push(lane); + } + void task() + .catch((error) => { + logger.error('[MCP catalog] Scheduled outbound work failed:', error); + }) + .finally(() => { + lane.inFlight -= 1; + activeCatalogWork -= 1; + finishCatalogLane(lane); + scheduleCatalogWork(); + }); + } +} + +function runCatalogWork( + tasks: ReadonlyArray<() => Promise>, + signal?: AbortSignal, +): Promise { + if (tasks.length === 0) { + return Promise.resolve(true); + } + if (signal?.aborted) { + return Promise.resolve(true); + } + if (activeCatalogLanes.size >= CATALOG_FANOUT_CONCURRENCY) { + logger.debug('[MCP catalog] Skipping request: outbound capacity reached'); + return Promise.resolve(false); + } + return new Promise((resolve) => { + const lane: CatalogWorkLane = { + tasks, + resolve, + nextIndex: 0, + inFlight: 0, + signal, + }; + lane.onAbort = () => { + lane.nextIndex = lane.tasks.length; + for (let index = pendingCatalogLanes.length - 1; index >= 0; index -= 1) { + if (pendingCatalogLanes[index] === lane) { + pendingCatalogLanes.splice(index, 1); + } + } + finishCatalogLane(lane); + }; + signal?.addEventListener('abort', lane.onAbort, { once: true }); + activeCatalogLanes.add(lane); + pendingCatalogLanes.push(lane); + scheduleCatalogWork(); + }); +} + /** * Passively discovers cold MCP catalogs for one request. * @@ -120,10 +225,14 @@ async function discoverCandidate( * proves pointless, and bounds everything else by a per-server budget. */ export async function recoverMCPServerCatalogs( - params: { user: IUser; servers: readonly MCPServerCatalogRecoveryInput[] }, + params: { + user: IUser; + servers: readonly MCPServerCatalogRecoveryInput[]; + signal?: AbortSignal; + }, deps: MCPServerCatalogRecoveryDeps, ): Promise> { - const { user, servers } = params; + const { user, servers, signal } = params; /** Only the config tier retries a failed stub on its own clock. A `yaml`- or `user`-sourced * stub has no such timer, so skipping it unconditionally would hide the server for good — * exactly the state this recovery exists to escape. */ @@ -168,21 +277,35 @@ export async function recoverMCPServerCatalogs( return new Map(); } - const recover = createConcurrencyLimiter(CATALOG_FANOUT_CONCURRENCY); - const results = await Promise.all( - authorized.map((candidate) => recover(() => discoverCandidate(user, candidate, deps))), + const results: Array<[string, LCAvailableTools | null] | undefined> = []; + const admitted = await runCatalogWork( + authorized.map((candidate, index) => async () => { + results[index] = await discoverCandidate(user, candidate, deps, signal); + }), + signal, ); + if (!admitted) { + throw new MCPCatalogCapacityError(); + } - return new Map(results.filter((entry): entry is [string, LCAvailableTools] => entry[1] != null)); + return new Map( + results.filter( + (entry): entry is [string, LCAvailableTools] => entry != null && entry[1] != null, + ), + ); } /** Loads cached, connected, then passive MCP catalogs for a marketplace-style list request. */ export async function loadMCPServerCatalogs( - params: { user: IUser; servers: readonly MCPServerCatalogRecoveryInput[] }, + params: { + user: IUser; + servers: readonly MCPServerCatalogRecoveryInput[]; + signal?: AbortSignal; + }, deps: MCPServerCatalogLoaderDeps, ): Promise { - const { user, servers } = params; - const cached = await Promise.all( + const { user, servers, signal } = params; + const cached: MCPServerCatalogEntry[] = await Promise.all( servers.map(async ({ serverName, serverConfig }) => { try { const tools = await deps.getCachedServerTools(user.id, serverName, serverConfig); @@ -194,32 +317,36 @@ export async function loadMCPServerCatalogs( }), ); - /** A snapshot is not a local read — both connection paths issue a fresh `tools/list` — so a - * cache reset across many servers would otherwise burst unbounded outbound requests. */ - const refresh = createConcurrencyLimiter(CATALOG_FANOUT_CONCURRENCY); - const snapshots = await Promise.all( - cached.map((entry) => { - if (entry.tools != null) { - return entry; - } - return refresh(async () => { - try { - const snapshot = await deps.getServerToolFunctionsSnapshot( - user.id, - entry.serverName, - entry.serverConfig, - ); - return { ...entry, ...snapshot, source: 'snapshot' as const }; - } catch (error) { - logger.error( - `[MCP catalog loader] Failed to read connected tools for ${entry.serverName}:`, - error, - ); - return { ...entry, tools: null, source: 'snapshot' as const }; - } - }); - }), + const snapshots = [...cached]; + const snapshotsAdmitted = await runCatalogWork( + cached.flatMap((entry, index) => + entry.tools != null + ? [] + : [ + async () => { + try { + const snapshot = await deps.getServerToolFunctionsSnapshot( + user.id, + entry.serverName, + entry.serverConfig, + { deadlineMs: Date.now() + mcpConfig.TOOLS_LIST_TIMEOUT_MS, signal }, + ); + snapshots[index] = { ...entry, ...snapshot, source: 'snapshot' as const }; + } catch (error) { + logger.error( + `[MCP catalog loader] Failed to read connected tools for ${entry.serverName}:`, + error, + ); + snapshots[index] = { ...entry, tools: null, source: 'snapshot' as const }; + } + }, + ], + ), + signal, ); + if (!snapshotsAdmitted) { + throw new MCPCatalogCapacityError(); + } const coldServers = snapshots .filter(({ tools }) => tools == null) @@ -227,8 +354,11 @@ export async function loadMCPServerCatalogs( let recovered = new Map(); if (coldServers.length > 0) { try { - recovered = await recoverMCPServerCatalogs({ user, servers: coldServers }, deps); + recovered = await recoverMCPServerCatalogs({ user, servers: coldServers, signal }, deps); } catch (error) { + if (error instanceof MCPCatalogCapacityError) { + throw error; + } logger.error('[MCP catalog loader] Failed to recover cold server catalogs:', error); } } diff --git a/packages/api/src/mcp/registry/MCPServerInspector.ts b/packages/api/src/mcp/registry/MCPServerInspector.ts index fe05d41e6c..76b5155a61 100644 --- a/packages/api/src/mcp/registry/MCPServerInspector.ts +++ b/packages/api/src/mcp/registry/MCPServerInspector.ts @@ -177,8 +177,10 @@ export class MCPServerInspector { public static async getToolCatalog( serverName: string, connection: MCPConnection, + deadlineMs?: number, + signal?: AbortSignal, ): Promise<{ tools: t.LCAvailableTools; publicationRevision?: string }> { - const snapshot = await connection.fetchOrderedToolsSnapshot(); + const snapshot = await connection.fetchOrderedToolsSnapshot(deadlineMs, signal); if (!snapshot.complete) { throw new Error(`Incomplete tools/list snapshot for MCP server ${serverName}`); }