diff --git a/packages/api/src/mcp/MCPConnectionFactory.ts b/packages/api/src/mcp/MCPConnectionFactory.ts index e545c1b102..4fe087b3ce 100644 --- a/packages/api/src/mcp/MCPConnectionFactory.ts +++ b/packages/api/src/mcp/MCPConnectionFactory.ts @@ -73,6 +73,7 @@ export class MCPConnectionFactory { protected oauthEnd?: () => Promise; protected returnOnOAuth?: boolean; protected readonly connectionTimeout?: number; + protected readonly deadlineMs?: number; protected readonly oboTokenResolver?: OboTokenResolver; protected readonly oboTrustChecker?: OboTrustChecker; protected readonly upstreamTokenProvider?: UpstreamTokenProvider; @@ -219,7 +220,7 @@ export class MCPConnectionFactory { connection.once('oauthRequired', oauthHandler); try { - const connectTimeout = this.connectionTimeout ?? this.serverConfig.initTimeout ?? 30000; + const connectTimeout = this.resolveConnectTimeout(30000); await withTimeout( connection.connect(), connectTimeout, @@ -227,7 +228,7 @@ export class MCPConnectionFactory { ); if (await connection.isConnected()) { - const snapshot = await connection.fetchOrderedToolsSnapshot(); + const snapshot = await connection.fetchOrderedToolsSnapshot(this.deadlineMs); connection.removeListener('oauthRequired', oauthHandler); return { tools: snapshot.complete ? snapshot.tools : null, @@ -242,24 +243,29 @@ export class MCPConnectionFactory { `${this.logPrefix} [Discovery] Connection failed, attempting unauthenticated tool listing`, ); } + + /** The authenticated attempt is done with, but `withTimeout` does not cancel the `connect()` + * it gave up on — that socket stays open. Dispose before the fallback opens a second one so + * a single discovery never holds two concurrent connects to the same server. */ + connection.removeListener('oauthRequired', oauthHandler); + await this.disposeQuietly(connection); + connection = null; + oauthHandler = null; + } + + if (this.isPastDeadline()) { + logger.debug( + `${this.logPrefix} [Discovery] Budget exhausted; skipping unauthenticated tool listing`, + ); + return { tools: null, connection: null, oauthRequired, oauthUrl }; } try { const tools = await this.attemptUnauthenticatedToolListing(); - if (connection && oauthHandler) { - connection.removeListener('oauthRequired', oauthHandler); - } if (tools && tools.length > 0) { logger.info( `${this.logPrefix} [Discovery] Successfully discovered ${tools.length} tools without auth`, ); - if (connection) { - try { - await connection.dispose(); - } catch { - // Ignore cleanup errors - } - } return { tools, connection: null, oauthRequired, oauthUrl }; } MCPConnection.decrementCycleCount(this.serverName); @@ -268,21 +274,30 @@ export class MCPConnectionFactory { logger.debug(`${this.logPrefix} [Discovery] Unauthenticated tool listing failed`); } - if (connection && oauthHandler) { - connection.removeListener('oauthRequired', oauthHandler); - } - - if (connection) { - try { - await connection.dispose(); - } catch { - // Ignore cleanup errors - } - } - return { tools: null, connection: null, oauthRequired, oauthUrl }; } + /** Clamps a single `connect()` to whatever remains of the caller's overall discovery budget. */ + private resolveConnectTimeout(fallback: number): number { + const configured = this.connectionTimeout ?? this.serverConfig.initTimeout ?? fallback; + if (this.deadlineMs == null) { + return configured; + } + return Math.max(1, Math.min(configured, this.deadlineMs - Date.now())); + } + + private isPastDeadline(): boolean { + return this.deadlineMs != null && Date.now() >= this.deadlineMs; + } + + private async disposeQuietly(connection: MCPConnection): Promise { + try { + await connection.dispose(); + } catch { + // Ignore cleanup errors + } + } + protected async attemptUnauthenticatedToolListing(): Promise { const unauthConnection = new MCPConnection({ serverName: this.serverName, @@ -305,23 +320,19 @@ export class MCPConnectionFactory { }); try { - const connectTimeout = this.connectionTimeout ?? this.serverConfig.initTimeout ?? 15000; + const connectTimeout = this.resolveConnectTimeout(15000); await withTimeout(unauthConnection.connect(), connectTimeout, `Unauth connection timeout`); if (await unauthConnection.isConnected()) { - const snapshot = await unauthConnection.fetchOrderedToolsSnapshot(); - await unauthConnection.dispose(); + const snapshot = await unauthConnection.fetchOrderedToolsSnapshot(this.deadlineMs); + await this.disposeQuietly(unauthConnection); return snapshot.complete ? snapshot.tools : null; } } catch { logger.debug(`${this.logPrefix} [Discovery] Unauthenticated connection attempt failed`); } - try { - await unauthConnection.dispose(); - } catch { - // Ignore cleanup errors - } + await this.disposeQuietly(unauthConnection); return null; } @@ -345,6 +356,7 @@ export class MCPConnectionFactory { this.allowedAddresses = basic.allowedAddresses; this.ephemeralConnection = basic.ephemeralConnection === true; this.connectionTimeout = options?.connectionTimeout; + this.deadlineMs = options?.deadlineMs; this.tenantContext = tenantStorage?.getStore?.(); this.tenantId = this.tenantContext?.tenantId ?? getTenantId(); this.logPrefix = options?.user ? `[MCP][User: ${options.user.id}]` : '[MCP]'; diff --git a/packages/api/src/mcp/MCPManager.ts b/packages/api/src/mcp/MCPManager.ts index 8d658bdecc..3ce946a875 100644 --- a/packages/api/src/mcp/MCPManager.ts +++ b/packages/api/src/mcp/MCPManager.ts @@ -280,7 +280,7 @@ export class MCPManager extends UserConnectionManager { ? await this.appConnections?.get(serverName) : null; if (existingAppConnection && (await existingAppConnection.isConnected())) { - const snapshot = await existingAppConnection.fetchOrderedToolsSnapshot(); + const snapshot = await existingAppConnection.fetchOrderedToolsSnapshot(args.deadlineMs); return { tools: snapshot.complete ? snapshot.tools : null, oauthRequired: false, @@ -350,6 +350,7 @@ export class MCPManager extends UserConnectionManager { requestBody: args.requestBody, graphTokenResolver: args.graphTokenResolver, connectionTimeout: args.connectionTimeout, + deadlineMs: args.deadlineMs, }); return finalizeDiscoveryResult(result); } @@ -370,6 +371,7 @@ export class MCPManager extends UserConnectionManager { requestBody: args.requestBody, graphTokenResolver: args.graphTokenResolver, connectionTimeout: args.connectionTimeout, + deadlineMs: args.deadlineMs, oboTokenResolver: args.oboTokenResolver, oboTrustChecker: args.oboTrustChecker, upstreamTokenProvider: args.upstreamTokenProvider, diff --git a/packages/api/src/mcp/__tests__/MCPConnectionDisposeRace.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionDisposeRace.test.ts new file mode 100644 index 0000000000..8851bff659 --- /dev/null +++ b/packages/api/src/mcp/__tests__/MCPConnectionDisposeRace.test.ts @@ -0,0 +1,88 @@ +/** + * Real-SDK coverage for disposal landing mid-connect. + * + * A discovery caller that gives up on a slow `connect()` disposes immediately, so it can dispose + * while `constructTransport()` is still pending. `dispose()` then finds no transport to close, and + * without a post-await check the abandoned attempt goes on to connect, leaving a live session + * nobody owns. Only transport construction is delayed here; the client, transport, and server are + * real SDK objects, so the assertions are about a genuinely open or closed session rather than + * about which mock was called. + */ +import { Server } from '@modelcontextprotocol/sdk/server/index.js'; +import { InMemoryTransport } from '@modelcontextprotocol/sdk/inMemory.js'; +import { ListToolsRequestSchema } from '@modelcontextprotocol/sdk/types.js'; +import { MCPConnection } from '~/mcp/connection'; + +jest.setTimeout(10_000); + +jest.mock('@librechat/data-schemas', () => ({ + logger: { + info: jest.fn(), + warn: jest.fn(), + error: jest.fn(), + debug: jest.fn(), + }, +})); + +jest.mock('~/auth', () => ({ + createSSRFSafeUndiciConnect: jest.fn(() => undefined), + isOAuthUrlAllowed: jest.fn(() => false), + isSSRFTarget: jest.fn(() => false), + resolveHostnameSSRF: jest.fn(async () => false), +})); + +describe('MCPConnection disposal during connect', () => { + let server: Server | undefined; + + afterEach(async () => { + await server?.close().catch(() => undefined); + server = undefined; + }); + + it('leaves no live session when disposal lands while the transport is being constructed', async () => { + server = new Server( + { name: 'dispose-race-server', version: '1.0.0' }, + { capabilities: { tools: {} } }, + ); + server.setRequestHandler(ListToolsRequestSchema, async () => ({ tools: [] })); + + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair(); + await server.connect(serverTransport); + + const connection = new MCPConnection({ + serverName: 'dispose-race', + serverConfig: { type: 'streamable-http', url: 'http://localhost/mcp' }, + useSSRFProtection: false, + }); + + let serverSawClose = false; + serverTransport.onclose = () => { + serverSawClose = true; + }; + + /** The only stub: hold construction open so disposal can land inside this window. */ + let releaseTransport: (() => void) | undefined; + const transportPending = new Promise((resolve) => { + releaseTransport = resolve; + }); + jest + .spyOn( + connection as unknown as { constructTransport: () => Promise }, + 'constructTransport', + ) + .mockImplementation(async () => { + await transportPending; + return clientTransport; + }); + + const connecting = connection.connectClient(); + await connection.dispose(); + releaseTransport?.(); + await connecting; + + /** A real `tools/list` is the honest probe: it succeeds over any session left open. */ + await expect(connection.client.listTools()).rejects.toThrow(); + expect(await connection.isConnected()).toBe(false); + expect(serverSawClose).toBe(true); + }); +}); diff --git a/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts index 6e2b5e3477..d01b4f0677 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts @@ -4190,6 +4190,100 @@ describe('MCPConnectionFactory', () => { expect(result.tools).toBeNull(); expect(mockLogger.debug).toHaveBeenCalled(); }); + + describe('deadline bounding', () => { + /** Builds distinct per-construction connection mocks so the authenticated attempt and the + * unauthenticated fallback can be told apart, recording lifecycle order across both. */ + const trackConnections = (events: string[]) => { + let index = 0; + mockMCPConnection.mockImplementation(() => { + const label = index++ === 0 ? 'auth' : 'unauth'; + const instance = { + connect: jest.fn(async () => { + events.push(`${label}:connect`); + throw new Error('Connection failed'); + }), + isConnected: jest.fn().mockResolvedValue(false), + setOAuthTokens: jest.fn(), + on: jest.fn(), + once: jest.fn(), + off: jest.fn(), + removeListener: jest.fn(), + emit: jest.fn(), + dispose: jest.fn(async () => { + events.push(`${label}:dispose`); + }), + } as unknown as jest.Mocked; + return instance; + }); + }; + + it('disposes the authenticated attempt before the fallback opens a second connection', async () => { + const events: string[] = []; + trackConnections(events); + + const result = await MCPConnectionFactory.discoverTools({ + serverName: 'test-server', + serverConfig: mockServerConfig, + }); + + expect(result.tools).toBeNull(); + expect(events).toEqual([ + 'auth:connect', + 'auth:dispose', + 'unauth:connect', + 'unauth:dispose', + ]); + }); + + it('skips the unauthenticated fallback once the deadline has passed', async () => { + const events: string[] = []; + trackConnections(events); + + const result = await MCPConnectionFactory.discoverTools( + { serverName: 'test-server', serverConfig: mockServerConfig }, + { deadlineMs: Date.now() - 1 }, + ); + + expect(result.tools).toBeNull(); + expect(events).toEqual(['auth:connect', 'auth:dispose']); + expect(mockMCPConnection).toHaveBeenCalledTimes(1); + }); + + it('forwards the deadline to the tools/list snapshot', async () => { + const deadlineMs = Date.now() + 5000; + mockConnectionInstance.connect.mockResolvedValue(undefined); + mockConnectionInstance.isConnected.mockResolvedValue(true); + mockConnectionInstance.fetchOrderedToolsSnapshot = jest + .fn() + .mockResolvedValue({ tools: [], complete: true }); + + await MCPConnectionFactory.discoverTools( + { serverName: 'test-server', serverConfig: mockServerConfig }, + { deadlineMs }, + ); + + expect(mockConnectionInstance.fetchOrderedToolsSnapshot).toHaveBeenCalledWith(deadlineMs); + }); + + it('clamps a long initTimeout to the remaining budget instead of waiting it out', async () => { + mockProcessMCPEnv.mockReturnValue({ + ...mockServerConfig, + initTimeout: 30000, + } as t.MCPOptions); + mockConnectionInstance.connect.mockImplementation(() => new Promise(() => {})); + mockConnectionInstance.isConnected.mockResolvedValue(false); + + const start = Date.now(); + const result = await MCPConnectionFactory.discoverTools( + { serverName: 'test-server', serverConfig: mockServerConfig }, + { deadlineMs: Date.now() + 50 }, + ); + + expect(result.tools).toBeNull(); + expect(Date.now() - start).toBeLessThan(2000); + }); + }); }); describe('proactive OAuth flow', () => { diff --git a/packages/api/src/mcp/__tests__/MCPConnectionFetchTools.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionFetchTools.test.ts index e0cc7a34d6..8805b887b8 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionFetchTools.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionFetchTools.test.ts @@ -263,6 +263,84 @@ describe('MCPConnection.fetchTools pagination', () => { expect(JSON.stringify(mockLogger.error.mock.calls)).not.toContain('Request timed out'); }); + it('caps the request timeout by a caller deadline shorter than the global budget', async () => { + mcpConfig.TOOLS_LIST_TIMEOUT_MS = 30000; + const listTools = jest.fn().mockResolvedValue({ tools: [makeTool('a')] }); + const conn = createConnectionWithListTools(listTools); + + const snapshot = await conn.fetchToolsSnapshot(Date.now() + 40); + + expect(snapshot.complete).toBe(true); + const options = listTools.mock.calls[0][1]!; + expect(options.timeout).toBeGreaterThan(0); + expect(options.timeout).toBeLessThanOrEqual(40); + }); + + it('keeps the global budget when the caller deadline is further out', async () => { + mcpConfig.TOOLS_LIST_TIMEOUT_MS = 25; + const listTools = jest.fn().mockResolvedValue({ tools: [makeTool('a')] }); + const conn = createConnectionWithListTools(listTools); + + await conn.fetchToolsSnapshot(Date.now() + 10_000); + + const options = listTools.mock.calls[0][1]!; + expect(options.timeout).toBeLessThanOrEqual(25); + }); + + it('stops paginating and reports incomplete when the caller deadline expires mid-walk', async () => { + mcpConfig.TOOLS_LIST_TIMEOUT_MS = 30000; + const listTools = jest.fn(async () => ({ tools: [makeTool('a')], nextCursor: 'c1' })); + const conn = createConnectionWithListTools(listTools); + const deadline = Date.now() + 20; + const dateNow = jest + .spyOn(Date, 'now') + .mockReturnValueOnce(deadline - 20) + .mockReturnValueOnce(deadline - 20) + .mockReturnValue(deadline + 1); + + const snapshot = await conn.fetchToolsSnapshot(deadline); + + expect(snapshot.tools.map((t) => t.name)).toEqual(['a']); + expect(snapshot.complete).toBe(false); + expect(listTools).toHaveBeenCalledTimes(1); + dateNow.mockRestore(); + }); + + it('returns an incomplete empty snapshot without a request when the deadline has passed', async () => { + const listTools = jest.fn(); + const conn = createConnectionWithListTools(listTools); + + const snapshot = await conn.fetchToolsSnapshot(Date.now() - 1); + + expect(snapshot.tools).toEqual([]); + expect(snapshot.complete).toBe(false); + expect(listTools).not.toHaveBeenCalled(); + }); + + it('stops waiting on an in-flight refresh that outlasts the caller deadline', async () => { + mcpConfig.TOOLS_LIST_TIMEOUT_MS = 30000; + const conn = createConnectionWithListTools(jest.fn()); + const mutable = conn as unknown as { + toolListChangeGeneration: number; + toolListRefreshPromise: Promise | null; + }; + /** A `list_changed` lands mid-fetch, so the ordered read must wait on a refresh... */ + const listTools = jest.fn(async () => { + mutable.toolListChangeGeneration = 1; + return { tools: [makeTool('a')] }; + }); + conn.client.listTools = listTools; + /** ...and that refresh never settles, standing in for one on the connection's own budget. */ + mutable.toolListRefreshPromise = new Promise(() => {}); + jest.spyOn(conn.client, 'getServerCapabilities').mockReturnValue({ tools: {} }); + + const start = Date.now(); + const snapshot = await conn.fetchOrderedToolsSnapshot(Date.now() + 30); + + expect(snapshot.complete).toBe(false); + expect(Date.now() - start).toBeLessThan(2000); + }); + 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); diff --git a/packages/api/src/mcp/catalog/recovery.spec.ts b/packages/api/src/mcp/catalog/recovery.spec.ts index 255b606e51..d9b374cd79 100644 --- a/packages/api/src/mcp/catalog/recovery.spec.ts +++ b/packages/api/src/mcp/catalog/recovery.spec.ts @@ -234,31 +234,35 @@ describe('recoverMCPServerCatalogs — bounded, skippable discovery', () => { formatServerTools: jest.fn().mockReturnValue({}), }); - it('bounds each connection attempt the factory makes', async () => { + it('bounds each server discovery end to end rather than per attempt', async () => { const discoverServerTools = jest.fn().mockResolvedValue({ tools: [] }); + const before = Date.now(); await recoverMCPServerCatalogs( { user, servers: [{ serverName: 'slow', serverConfig: serverConfig('slow') }] }, recoveryDeps(discoverServerTools), ); - expect(discoverServerTools).toHaveBeenCalledWith( - expect.objectContaining({ serverName: 'slow', connectionTimeout: 1500 }), - ); + const [options] = discoverServerTools.mock.calls[0]; + expect(options.serverName).toBe('slow'); + expect(options.connectionTimeout).toBeUndefined(); + expect(options.deadlineMs).toBeGreaterThanOrEqual(before + 3000); + expect(options.deadlineMs).toBeLessThanOrEqual(Date.now() + 3000); }); it('keeps a shorter configured initTimeout instead of raising it to the cap', async () => { const discoverServerTools = jest.fn().mockResolvedValue({ tools: [] }); const impatient = { ...serverConfig('impatient'), initTimeout: 900 } as ParsedServerConfig; + const before = Date.now(); await recoverMCPServerCatalogs( { user, servers: [{ serverName: 'impatient', serverConfig: impatient }] }, recoveryDeps(discoverServerTools), ); - expect(discoverServerTools).toHaveBeenCalledWith( - expect.objectContaining({ connectionTimeout: 900 }), - ); + const [options] = discoverServerTools.mock.calls[0]; + expect(options.deadlineMs).toBeGreaterThanOrEqual(before + 900); + expect(options.deadlineMs).toBeLessThanOrEqual(Date.now() + 900); }); it('leaves a server the config tier marked unreachable to that tier’s retry window', async () => { diff --git a/packages/api/src/mcp/catalog/recovery.ts b/packages/api/src/mcp/catalog/recovery.ts index 06d61b8b0f..6d17b6dd02 100644 --- a/packages/api/src/mcp/catalog/recovery.ts +++ b/packages/api/src/mcp/catalog/recovery.ts @@ -10,18 +10,12 @@ import { getServerCustomUserVars } from '../auth'; * `tools/list`, so they burst as readily as passive discovery does. */ const CATALOG_FANOUT_CONCURRENCY = 3; /** - * Bounds `connection.connect()` only — it is the sole segment of discovery a caller can bound - * today. It does NOT bound the whole operation: `discoverToolsInternal` spends this value once - * per connection attempt (authenticated, then unauthenticated), and `fetchToolsSnapshot` then - * applies its own `TOOLS_LIST_TIMEOUT_MS` (30s) to `tools/list`. A server that connects quickly - * and stalls while listing therefore still holds its slot for that longer window. - * - * Bounding discovery end to end needs a deadline threaded through `MCPConnectionFactory` into - * both `connect()` and `fetchToolsSnapshot()`; until that exists, this keeps the common - * unreachable case cheap, because recovery targets a server that is reachable and authorized but - * whose catalog cache expired, and such a server connects well inside this window. + * 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 + * long regardless of where the server stalls. Recovery targets a server that is reachable and + * authorized but whose catalog cache expired, and such a server answers well inside this window. */ -const RECOVERY_ATTEMPT_TIMEOUT_MS = 1500; +const RECOVERY_BUDGET_MS = 3000; export interface MCPServerCatalogRecoveryInput { serverName: string; @@ -73,13 +67,13 @@ interface RecoveryCandidate extends MCPServerCatalogRecoveryInput { customUserVars?: Record; } -/** Bounds one connection attempt, honouring a shorter operator `initTimeout`. */ -function resolveAttemptTimeout(serverConfig: ParsedServerConfig): number { +/** Bounds one server's discovery, honouring a shorter operator `initTimeout`. */ +function resolveBudget(serverConfig: ParsedServerConfig): number { const { initTimeout } = serverConfig; if (typeof initTimeout === 'number') { - return Math.min(initTimeout, RECOVERY_ATTEMPT_TIMEOUT_MS); + return Math.min(initTimeout, RECOVERY_BUDGET_MS); } - return RECOVERY_ATTEMPT_TIMEOUT_MS; + return RECOVERY_BUDGET_MS; } async function discoverCandidate( @@ -93,7 +87,7 @@ async function discoverCandidate( serverName, configServers: { [serverName]: serverConfig }, customUserVars, - connectionTimeout: resolveAttemptTimeout(serverConfig), + deadlineMs: Date.now() + resolveBudget(serverConfig), }); return [ serverName, @@ -112,7 +106,7 @@ async function discoverCandidate( * and is disposed, and the tool cache refuses unfenced writes, so the result is served only to * the requesting user. Recovery therefore stays stateless and individually cheap rather than * scheduling around a result it is not allowed to keep — it skips only what configuration alone - * proves pointless, and fails fast on everything else. + * proves pointless, and bounds everything else by a per-server budget. */ export async function recoverMCPServerCatalogs( params: { user: IUser; servers: readonly MCPServerCatalogRecoveryInput[] }, diff --git a/packages/api/src/mcp/connection.ts b/packages/api/src/mcp/connection.ts index aaec74f097..a19514123b 100644 --- a/packages/api/src/mcp/connection.ts +++ b/packages/api/src/mcp/connection.ts @@ -1926,6 +1926,12 @@ export class MCPConnection extends EventEmitter { } this.transport = await runOutsideTracing(() => this.constructTransport(this.options)); + /** `dispose()` can land while the transport is still being constructed — it finds nothing + * to close and returns, so without this check the attempt would go on to connect and + * leave a live connection on a disposed object. Ownership of teardown is ours here. */ + if (await this.abandonIfDisposed()) { + return; + } this.patchTransportSend(); const connectTimeout = this.options.initTimeout ?? DEFAULT_INIT_TIMEOUT; @@ -1937,6 +1943,9 @@ export class MCPConnection extends EventEmitter { ), ); + if (await this.abandonIfDisposed()) { + return; + } this.setupTransportOnMessageHandler(); this.connectionState = 'connected'; this.emit('connectionChange', 'connected'); @@ -2059,6 +2068,33 @@ export class MCPConnection extends EventEmitter { return this.connectPromise; } + /** + * Tears down a transport this attempt created after the connection was already disposed. + * Returns whether the caller should abandon the rest of the connect sequence. + */ + private async abandonIfDisposed(): Promise { + if (!this.isDisposed) { + return false; + } + logger.debug(`${this.getLogPrefix()} Disposed mid-connect; discarding the transport it opened`); + const transport = this.transport; + this.transport = null; + /** Closing the client only closes a transport the client has already adopted, which it has + * not when disposal beat `client.connect()`. Close the transport itself first, or the + * session this attempt opened outlives the connection that owned it. */ + try { + await transport?.close(); + } catch { + // Ignore cleanup errors + } + try { + await this.client.close(); + } catch { + // Ignore cleanup errors + } + return true; + } + private patchTransportSend(): void { if (!this.transport) { return; @@ -2348,8 +2384,12 @@ export class MCPConnection extends EventEmitter { * Fetches a bounded tool snapshot while preserving whether every requested page succeeded. * Notification refreshes use `complete` to avoid replacing a known-good cache with an empty or * partial list after a transient `tools/list` failure. + * + * @param deadlineMs Absolute epoch-ms cap for a caller working to a fixed budget. Pagination + * stops at whichever comes first, this or `TOOLS_LIST_TIMEOUT_MS`, and the partial result is + * returned as incomplete so it is never published as an authoritative catalog. */ - public async fetchToolsSnapshot(): Promise { + public async fetchToolsSnapshot(deadlineMs?: number): Promise { const maxPages = mcpConfig.TOOLS_LIST_MAX_PAGES; const maxTools = mcpConfig.TOOLS_LIST_MAX_TOOLS; const maxBytes = mcpConfig.TOOLS_LIST_MAX_BYTES; @@ -2357,7 +2397,8 @@ export class MCPConnection extends EventEmitter { * from a `tools/list` that started later. Every app-level publisher reads its ordering off * the snapshot it received, which is the only way to know when the data was actually read. */ const ordering = await this.reserveToolsPublicationRevision(); - const deadline = Date.now() + mcpConfig.TOOLS_LIST_TIMEOUT_MS; + const budgetDeadline = Date.now() + mcpConfig.TOOLS_LIST_TIMEOUT_MS; + const deadline = deadlineMs != null ? Math.min(budgetDeadline, deadlineMs) : budgetDeadline; const allTools: MCPListToolsResult['tools'] = []; const seenCursors = new Set(); let cursor: string | undefined; @@ -2467,11 +2508,14 @@ export class MCPConnection extends EventEmitter { * Returns a complete snapshot that cannot precede a concurrent `list_changed` refresh. * If the notification refresh cannot complete, the caller receives an incomplete snapshot * instead of publishing a request result that may already be stale. + * + * @param deadlineMs Absolute epoch-ms cap; see `fetchToolsSnapshot`. It also bounds the wait + * for a concurrent refresh, so a caller on a budget is never held by another caller's fetch. */ - public async fetchOrderedToolsSnapshot(): Promise { + public async fetchOrderedToolsSnapshot(deadlineMs?: number): Promise { const startEpoch = this.toolListRefreshEpoch; const startGeneration = this.toolListChangeGeneration; - const snapshot = await this.fetchToolsSnapshot(); + const snapshot = await this.fetchToolsSnapshot(deadlineMs); if ( startEpoch === this.toolListRefreshEpoch && @@ -2484,12 +2528,22 @@ export class MCPConnection extends EventEmitter { startEpoch === this.toolListRefreshEpoch && this.handledToolListChangeGeneration < this.toolListChangeGeneration ) { + if (deadlineMs != null && Date.now() >= deadlineMs) { + break; + } this.startToolListRefresh(); const refresh = this.toolListRefreshPromise; if (!refresh) { break; } - await refresh; + /** The refresh runs on the connection's own budget, not the caller's, so a refresh already + * in flight can outlast this deadline. Stop waiting on it rather than adopting its budget; + * it keeps running for whoever else wants it and this caller reports an incomplete read. */ + if (deadlineMs == null) { + await refresh; + } else if (!(await this.settlesBefore(refresh, deadlineMs))) { + break; + } if (this.toolListRefreshRetryTimer) { break; } @@ -2513,6 +2567,20 @@ export class MCPConnection extends EventEmitter { return { tools: [], complete: false }; } + /** Waits for `promise` only until `deadlineMs`, reporting whether it settled in time. */ + private async settlesBefore(promise: Promise, deadlineMs: number): Promise { + let timer: NodeJS.Timeout | undefined; + const expiry = new Promise((resolve) => { + timer = setTimeout(() => resolve(false), Math.max(0, deadlineMs - Date.now())); + timer.unref?.(); + }); + try { + return await Promise.race([promise.then(() => true), expiry]); + } finally { + clearTimeout(timer); + } + } + 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).`, diff --git a/packages/api/src/mcp/types/index.ts b/packages/api/src/mcp/types/index.ts index 1dba82518c..07a13c91c7 100644 --- a/packages/api/src/mcp/types/index.ts +++ b/packages/api/src/mcp/types/index.ts @@ -227,6 +227,10 @@ export interface UserConnectionContext { requestScopedConnections?: RequestScopedMCPConnectionStore; graphTokenResolver?: GraphTokenResolver; connectionTimeout?: number; + /** Absolute epoch-ms bound on the whole connect-and-list operation. `connectionTimeout` bounds + * only a single `connect()`, so a caller that must return within a fixed budget sets this to + * cap every segment, including `tools/list` pagination and the unauthenticated fallback. */ + deadlineMs?: number; } export interface RequestScopedMCPConnectionStore { @@ -284,6 +288,8 @@ export interface ToolDiscoveryOptions { requestBody?: RequestBody; graphTokenResolver?: GraphTokenResolver; connectionTimeout?: number; + /** Absolute epoch-ms bound on the whole discovery operation; see `UserConnectionContext`. */ + deadlineMs?: number; /** Pre-resolved config-source servers for tenant-scoped lookup */ configServers?: Record; oboTokenResolver?: OboTokenResolver;