diff --git a/packages/api/src/mcp/MCPConnectionFactory.ts b/packages/api/src/mcp/MCPConnectionFactory.ts index a7808feaab..01f20cb1bf 100644 --- a/packages/api/src/mcp/MCPConnectionFactory.ts +++ b/packages/api/src/mcp/MCPConnectionFactory.ts @@ -433,6 +433,13 @@ export class MCPConnectionFactory { if (cleanupOAuthHandlers) { cleanupOAuthHandlers(); } + try { + await connection.dispose(); + } catch (disconnectError) { + logger.warn(`${this.logPrefix} Failed to clean up rejected MCP connection`, { + error: disconnectError, + }); + } throw error; } } @@ -1075,24 +1082,61 @@ export class MCPConnectionFactory { // The grace covers the reconnect after `oauthHandled` (retry backoff + transport connect), // which happens *after* the handling wait, so a user who authorizes near the deadline still // gets a connection instead of a timeout. + const oauthHandlingTimeout = Number.isFinite(mcpConfig.OAUTH_HANDLING_TIMEOUT) + ? mcpConfig.OAUTH_HANDLING_TIMEOUT + : 10 * 60 * 1000; const connectTimeout = this.useOAuth - ? Math.max(baseTimeout, mcpConfig.OAUTH_HANDLING_TIMEOUT + 60000) + ? Math.max(baseTimeout, oauthHandlingTimeout + 60000) : baseTimeout; - await withTimeout( - this.connectTo(connection), - connectTimeout, - `Connection timeout after ${connectTimeout}ms`, - ); + const retryController = new AbortController(); + let timeoutId: NodeJS.Timeout | undefined; + const timeout = new Promise((_, reject) => { + timeoutId = setTimeout(() => { + retryController.abort(); + reject(new Error(`Connection timeout after ${connectTimeout}ms`)); + }, connectTimeout); + }); + + try { + await Promise.race([this.connectTo(connection, retryController.signal), timeout]); + } finally { + if (timeoutId) { + clearTimeout(timeoutId); + } + retryController.abort(); + } if (await connection.isConnected()) return; logger.error(`${this.logPrefix} Failed to establish connection.`); } - private async connectTo(connection: MCPConnection): Promise { + private waitForRetry(delayMs: number, signal: AbortSignal): Promise { + return new Promise((resolve, reject) => { + if (signal.aborted) { + reject(new Error('Connection retry cancelled')); + return; + } + + const onAbort = () => { + clearTimeout(timeoutId); + reject(new Error('Connection retry cancelled')); + }; + const timeoutId = setTimeout(() => { + signal.removeEventListener('abort', onAbort); + resolve(); + }, delayMs); + signal.addEventListener('abort', onAbort, { once: true }); + }); + } + + private async connectTo(connection: MCPConnection, signal: AbortSignal): Promise { const maxAttempts = 3; let attempts = 0; while (attempts < maxAttempts) { + if (signal.aborted) { + throw new Error('Connection retry cancelled'); + } try { await connection.connect(); if (await connection.isConnected()) { @@ -1102,6 +1146,10 @@ export class MCPConnectionFactory { } catch (error) { attempts++; + if (signal.aborted) { + throw error; + } + if (this.useOAuth && this.isOAuthError(error)) { logger.info(`${this.logPrefix} OAuth required, stopping connection attempts`); throw error; @@ -1111,7 +1159,7 @@ export class MCPConnectionFactory { logger.error(`${this.logPrefix} Failed to connect after ${maxAttempts} attempts`, error); throw error; } - await new Promise((resolve) => setTimeout(resolve, 2000 * attempts)); + await this.waitForRetry(2000 * attempts, signal); } } } diff --git a/packages/api/src/mcp/MCPManager.ts b/packages/api/src/mcp/MCPManager.ts index 0269937e5b..6d39acfd68 100644 --- a/packages/api/src/mcp/MCPManager.ts +++ b/packages/api/src/mcp/MCPManager.ts @@ -385,8 +385,6 @@ Please follow these instructions when using tools from the respective MCP server const logPrefix = userId ? `[MCP][User: ${userId}][${serverName}]` : `[MCP][${serverName}]`; try { - if (userId && user) this.updateUserLastActivity(userId); - connection = await this.getConnection({ serverName, user, @@ -421,8 +419,8 @@ Please follow these instructions when using tools from the respective MCP server ); } const isDbSourced = isUserSourced(rawConfig); - disconnectAfterCall = - !!userId && requiresEphemeralUserConnection(rawConfig) && !requestScopedConnections; + const ephemeralConnection = !!userId && requiresEphemeralUserConnection(rawConfig); + disconnectAfterCall = ephemeralConnection && !requestScopedConnections; /** Pre-process Graph token placeholders (async) before the synchronous processMCPEnv pass */ const graphProcessedConfig = isDbSourced @@ -528,7 +526,9 @@ Please follow these instructions when using tools from the respective MCP server ...options, }, ); - if (userId) { + const hasPersistentUserConnections = + !!userId && (this.userConnections.get(userId)?.size ?? 0) > 0; + if (!ephemeralConnection && hasPersistentUserConnections) { this.updateUserLastActivity(userId); } this.checkIdleConnections(); diff --git a/packages/api/src/mcp/UserConnectionManager.ts b/packages/api/src/mcp/UserConnectionManager.ts index a995c4ab65..12902b2c85 100644 --- a/packages/api/src/mcp/UserConnectionManager.ts +++ b/packages/api/src/mcp/UserConnectionManager.ts @@ -104,7 +104,6 @@ export abstract class UserConnectionManager { requestScopedConnections.connections.delete(requestConnectionKey); } else if (await existing.isConnected()) { logger.debug(`[MCP][User: ${userId}][${serverName}] Reusing request-scoped connection`); - this.updateUserLastActivity(userId); return existing; } else { requestScopedConnections.connections.delete(requestConnectionKey); @@ -525,8 +524,9 @@ export abstract class UserConnectionManager { } logger.info(`[MCP][User: ${userId}][${serverName}] Connection successfully established`); - // Update timestamp on creation - this.updateUserLastActivity(userId); + if (!ephemeralConnection) { + this.updateUserLastActivity(userId); + } return connection; } catch (error) { logger.error(`[MCP][User: ${userId}][${serverName}] Failed to establish connection`, error); 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 3ef669bb94..e087ab9f20 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionFactory.oauthSdk.integration.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionFactory.oauthSdk.integration.test.ts @@ -42,6 +42,7 @@ jest.mock('~/auth', () => ({ jest.mock('~/mcp/mcpConfig', () => ({ mcpConfig: { CONNECTION_CHECK_TTL: 0, + OAUTH_HANDLING_TIMEOUT: 10 * 60 * 1000, USER_CONNECTION_IDLE_TIMEOUT: 30 * 60 * 1000, TOOLS_LIST_MAX_PAGES: 50, TOOLS_LIST_MAX_TOOLS: 1000, diff --git a/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts index 73048c7ab3..bd67e68995 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts @@ -56,6 +56,10 @@ class InspectableMCPConnectionFactory extends MCPConnectionFactory { return await this.createConnection(); } + public async attemptToConnectForTest(connection: MCPConnection): Promise { + await this.attemptToConnect(connection); + } + public getRequestScopedOAuthState(): { signal?: AbortSignal; oauthStart?: (authURL: string) => Promise; @@ -114,6 +118,7 @@ describe('MCPConnectionFactory', () => { off: jest.fn().mockReturnValue(mockConnectionInstance), removeListener: jest.fn().mockReturnValue(mockConnectionInstance), emit: jest.fn(), + dispose: jest.fn().mockResolvedValue(undefined), } as unknown as jest.Mocked; mockMCPConnection.mockImplementation(() => mockConnectionInstance); @@ -145,6 +150,40 @@ describe('MCPConnectionFactory', () => { (getTenantId as jest.Mock).mockReturnValue(undefined); }); + it('cancels pending retry backoff when the connection deadline expires', async () => { + jest.useFakeTimers(); + try { + const config = { + command: 'node', + args: ['server.js'], + initTimeout: 100, + } as t.MCPOptions; + mockProcessMCPEnv.mockReturnValue(config); + const factory = new InspectableMCPConnectionFactory({ + serverName: 'timed-out-server', + serverConfig: config, + }); + mockConnectionInstance.connect.mockRejectedValue(new Error('connection failed')); + mockConnectionInstance.isConnected.mockResolvedValue(false); + + const attempt = factory.attemptToConnectForTest(mockConnectionInstance); + const rejection = attempt.then( + () => null, + (error: Error) => error, + ); + + await jest.advanceTimersByTimeAsync(101); + await expect(rejection).resolves.toMatchObject({ + message: 'Connection timeout after 100ms', + }); + await jest.advanceTimersByTimeAsync(5000); + + expect(mockConnectionInstance.connect).toHaveBeenCalledTimes(1); + } finally { + jest.useRealTimers(); + } + }); + describe('static create method', () => { it('should create a basic connection without OAuth', async () => { const basicOptions = { diff --git a/packages/api/src/mcp/__tests__/MCPManager.test.ts b/packages/api/src/mcp/__tests__/MCPManager.test.ts index 865de69eb4..5b88acf1a9 100644 --- a/packages/api/src/mcp/__tests__/MCPManager.test.ts +++ b/packages/api/src/mcp/__tests__/MCPManager.test.ts @@ -445,6 +445,86 @@ describe('MCPManager', () => { }); }); + describe('callTool - Activity Tracking', () => { + const mockUser = { id: 'activity-user' } as IUser; + const mockFlowManager = {} as Parameters[0]['flowManager']; + const serverConfig: t.SSEOptions = { + type: 'sse', + url: 'https://api.example.com', + }; + + function createConnection(): MCPConnection { + return { + isConnected: jest.fn().mockResolvedValue(true), + setRequestHeaders: jest.fn(), + timeout: 30000, + client: { + request: jest.fn().mockResolvedValue({ + content: [{ type: 'text', text: 'Tool result' }], + isError: false, + }), + }, + } as unknown as MCPConnection; + } + + function getManagerInternals(manager: MCPManager): { + userConnections: Map>; + updateUserLastActivity: (trackedUserId: string) => void; + } { + return manager as unknown as { + userConnections: Map>; + updateUserLastActivity: (trackedUserId: string) => void; + }; + } + + beforeEach(() => { + (graphUtils.preProcessGraphTokens as jest.Mock).mockImplementation( + async (options) => options, + ); + }); + + it('updates activity when a cached connection is replaced during an in-flight call', async () => { + const manager = new MCPManager(); + const activeConnection = createConnection(); + const replacementConnection = createConnection(); + const internals = getManagerInternals(manager); + internals.userConnections.set(mockUser.id, new Map([[serverName, replacementConnection]])); + jest.spyOn(manager, 'getConnection').mockResolvedValue(activeConnection); + const updateActivity = jest.spyOn(internals, 'updateUserLastActivity'); + + await manager.callTool({ + user: mockUser, + serverName, + serverConfig, + toolName: 'test_tool', + provider: 'openai', + flowManager: mockFlowManager, + }); + + expect(updateActivity).toHaveBeenCalledWith(mockUser.id); + }); + + it('does not create activity entries for app-shared connections', async () => { + const manager = new MCPManager(); + const appConnection = createConnection(); + const internals = getManagerInternals(manager); + jest.spyOn(manager, 'getConnection').mockResolvedValue(appConnection); + const updateActivity = jest.spyOn(internals, 'updateUserLastActivity'); + + await manager.callTool({ + user: mockUser, + serverName, + serverConfig, + toolName: 'test_tool', + provider: 'openai', + flowManager: mockFlowManager, + }); + + expect(updateActivity).not.toHaveBeenCalled(); + expect(manager.getConnectionStats().activityEntries).toBe(0); + }); + }); + describe('callTool - Graph Token Integration', () => { const mockUser: Partial = { id: 'user-123', diff --git a/packages/api/src/mcp/__tests__/scope.integration.test.ts b/packages/api/src/mcp/__tests__/scope.integration.test.ts new file mode 100644 index 0000000000..9f49d161de --- /dev/null +++ b/packages/api/src/mcp/__tests__/scope.integration.test.ts @@ -0,0 +1,307 @@ +import { z } from 'zod'; +import * as http from 'http'; +import { Agent } from 'undici'; +import { randomUUID } from 'crypto'; +import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; +import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js'; +import type { IUser } from '@librechat/data-schemas'; +import type { Socket } from 'net'; +import type { ParsedServerConfig, RequestScopedMCPConnectionStore } from '~/mcp/types'; +import type { FlowStateManager } from '~/flow/manager'; +import type { MCPRequestContext } from '~/mcp/request'; +import type { MCPOAuthTokens } from '~/mcp/oauth'; +import { createMCPRequestContext, cleanupMCPRequestContext } from '~/mcp/request'; +import { MCPServersRegistry } from '~/mcp/registry/MCPServersRegistry'; +import { getFreePort } from '~/mcp/__tests__/helpers/oauthTestServer'; +import { ConnectionsRepository } from '~/mcp/ConnectionsRepository'; +import { MCPConnection } from '~/mcp/connection'; +import { MCPManager } from '~/mcp/MCPManager'; + +jest.mock('@librechat/data-schemas', () => ({ + logger: { + info: jest.fn(), + warn: jest.fn(), + error: jest.fn(), + debug: jest.fn(), + }, + getTenantId: jest.fn(() => undefined), + tenantStorage: { + getStore: jest.fn(() => undefined), + run: jest.fn((_context: object, fn: () => unknown) => fn()), + }, +})); + +jest.mock('~/auth', () => ({ + createSSRFSafeUndiciConnect: jest.fn(() => undefined), + isOAuthUrlAllowed: jest.fn(() => false), + isSSRFTarget: jest.fn(() => false), + resolveHostnameSSRF: jest.fn(async () => false), +})); + +jest.mock('~/auth/domain', () => ({ + isMCPDomainAllowed: jest.fn(async () => true), +})); + +interface RequestScopedTestServer { + url: string; + brokenUrl: string; + close: () => Promise; + deleteCount: () => number; + liveSessionCount: () => number; + sessionsCreated: () => number; + toolCallCount: () => number; + observedRunIds: () => string[]; +} + +function trackSockets(httpServer: http.Server): () => Promise { + const sockets = new Set(); + httpServer.on('connection', (socket: Socket) => { + sockets.add(socket); + socket.once('close', () => sockets.delete(socket)); + }); + return () => + new Promise((resolve) => { + for (const socket of sockets) { + socket.destroy(); + } + sockets.clear(); + httpServer.close(() => resolve()); + }); +} + +async function createRequestScopedTestServer(): Promise { + const sessions = new Map(); + const runIds: string[] = []; + let created = 0; + let deletes = 0; + let toolCalls = 0; + + const httpServer = http.createServer(async (req, res) => { + if (req.url === '/broken') { + res.writeHead(500).end('broken MCP'); + return; + } + + if (req.method === 'POST') { + const runId = req.headers['x-run-id']; + if (typeof runId === 'string') { + runIds.push(runId); + } + } else if (req.method === 'DELETE') { + deletes += 1; + } + + const sessionId = req.headers['mcp-session-id'] as string | undefined; + let transport = sessionId ? sessions.get(sessionId) : undefined; + const isNewTransport = !transport; + + if (!transport) { + transport = new StreamableHTTPServerTransport({ sessionIdGenerator: () => randomUUID() }); + const mcp = new McpServer({ name: 'request-scoped-test', version: '0.0.1' }); + mcp.tool('echo', 'Echo a value', { value: z.string() }, async ({ value }) => { + toolCalls += 1; + return { content: [{ type: 'text', text: value }] }; + }); + await mcp.connect(transport); + } + + await transport.handleRequest(req, res); + + if (isNewTransport && transport.sessionId) { + created += 1; + const registeredId = transport.sessionId; + sessions.set(registeredId, transport); + transport.onclose = () => sessions.delete(registeredId); + } + }); + + const destroySockets = trackSockets(httpServer); + const port = await getFreePort(); + await new Promise((resolve) => httpServer.listen(port, '127.0.0.1', resolve)); + + return { + url: `http://127.0.0.1:${port}/mcp`, + brokenUrl: `http://127.0.0.1:${port}/broken`, + deleteCount: () => deletes, + liveSessionCount: () => sessions.size, + sessionsCreated: () => created, + toolCallCount: () => toolCalls, + observedRunIds: () => [...runIds], + close: async () => { + const closing = [...sessions.values()].map((transport) => + transport.close().catch(() => undefined), + ); + sessions.clear(); + await Promise.all(closing); + await destroySockets(); + }, + }; +} + +const user = { id: 'scale-user', role: 'USER' } as unknown as IUser; +const flowManager = {} as FlowStateManager; + +function createManager(): MCPManager { + const manager = new MCPManager(); + manager.appConnections = { + has: jest.fn(async () => false), + getConnectionCount: jest.fn(() => 0), + } as unknown as ConnectionsRepository; + return manager; +} + +function createServerConfig(url: string): ParsedServerConfig { + return { + type: 'streamable-http', + url, + source: 'yaml', + requiresOAuth: false, + initTimeout: 500, + headers: { 'X-Run-Id': '{{LIBRECHAT_BODY_MESSAGEID}}' }, + }; +} + +function callEcho({ + manager, + config, + requestScopedConnections, + runId, + value, + serverName = 'request-scoped', +}: { + manager: MCPManager; + config: ParsedServerConfig; + requestScopedConnections: RequestScopedMCPConnectionStore; + runId: string; + value: string; + serverName?: string; +}) { + return manager.callTool({ + user, + serverName, + serverConfig: config, + toolName: 'echo', + provider: 'openai', + toolArguments: { value }, + requestBody: { + conversationId: 'conversation-1', + parentMessageId: 'parent-1', + messageId: runId, + }, + requestScopedConnections, + flowManager, + }); +} + +describe('request-scoped MCP lifecycle integration', () => { + let server: RequestScopedTestServer; + let manager: MCPManager; + const contexts = new Set(); + + beforeEach(async () => { + server = await createRequestScopedTestServer(); + manager = createManager(); + jest.spyOn(MCPServersRegistry, 'getInstance').mockReturnValue({ + resolveAllowlists: jest.fn(async () => ({ + allowedDomains: null, + allowedAddresses: null, + useSSRFProtection: false, + })), + } as unknown as MCPServersRegistry); + }); + + afterEach(async () => { + await Promise.all([...contexts].map((context) => cleanupMCPRequestContext(context))); + contexts.clear(); + MCPConnection.clearCooldown('request-scoped'); + MCPConnection.clearCooldown('broken-request-scoped'); + jest.restoreAllMocks(); + await server.close(); + }); + + function createContext(): MCPRequestContext { + const context = createMCPRequestContext(); + contexts.add(context); + return context; + } + + it('coalesces a concurrent burst, tears down the run, and isolates the next run', async () => { + const config = createServerConfig(server.url); + const firstRun = createContext(); + const burstSize = 20; + + await Promise.all( + Array.from({ length: burstSize }, (_, index) => + callEcho({ + manager, + config, + requestScopedConnections: firstRun, + runId: 'run-1', + value: `value-${index}`, + }), + ), + ); + + expect(server.sessionsCreated()).toBe(1); + expect(server.liveSessionCount()).toBe(1); + expect(server.toolCallCount()).toBe(burstSize); + expect(new Set(server.observedRunIds())).toEqual(new Set(['run-1'])); + expect(manager.getConnectionStats().activityEntries).toBe(0); + + await cleanupMCPRequestContext(firstRun); + contexts.delete(firstRun); + + expect(server.deleteCount()).toBe(1); + expect(server.liveSessionCount()).toBe(0); + + const secondRun = createContext(); + await callEcho({ + manager, + config, + requestScopedConnections: secondRun, + runId: 'run-2', + value: 'second-run', + }); + + expect(server.sessionsCreated()).toBe(2); + expect(server.liveSessionCount()).toBe(1); + expect(new Set(server.observedRunIds())).toEqual(new Set(['run-1', 'run-2'])); + }); + + it('clears a failed run so the same server can recover in a fresh run', async () => { + const brokenRun = createContext(); + const destroyAgent = jest.spyOn(Agent.prototype, 'destroy'); + + await expect( + callEcho({ + manager, + config: createServerConfig(server.brokenUrl), + requestScopedConnections: brokenRun, + runId: 'broken-run', + value: 'unreachable', + serverName: 'broken-request-scoped', + }), + ).rejects.toThrow(); + + expect(brokenRun.connections.size).toBe(0); + expect(brokenRun.pending.size).toBe(0); + expect(manager.getConnectionStats().activityEntries).toBe(0); + expect(destroyAgent).toHaveBeenCalled(); + + const recoveryRun = createContext(); + await expect( + callEcho({ + manager, + config: createServerConfig(server.url), + requestScopedConnections: recoveryRun, + runId: 'recovery-run', + value: 'recovered', + serverName: 'broken-request-scoped', + }), + ).resolves.toBeDefined(); + + expect(server.sessionsCreated()).toBe(1); + expect(server.toolCallCount()).toBe(1); + expect(new Set(server.observedRunIds())).toEqual(new Set(['recovery-run'])); + }); +}); diff --git a/packages/api/src/mcp/connection.ts b/packages/api/src/mcp/connection.ts index 35fb985e05..c479ea4bf0 100644 --- a/packages/api/src/mcp/connection.ts +++ b/packages/api/src/mcp/connection.ts @@ -2183,10 +2183,10 @@ export class MCPConnection extends EventEmitter { }; } - private async closeAgents(): Promise { + private async closeAgents(force = false): Promise { const logPrefix = this.getLogPrefix(); const closing = this.agents.map((agent) => - agent.close().catch((err: unknown) => { + (force ? agent.destroy() : agent.close()).catch((err: unknown) => { logger.debug(`${logPrefix} Agent close error (non-fatal):`, err); }), ); @@ -2230,14 +2230,14 @@ export class MCPConnection extends EventEmitter { } } - public async disconnect(resetCycleTracking = true): Promise { + public async disconnect(resetCycleTracking = true, forceAgentClose = false): Promise { try { if (this.transport) { await this.terminateStreamableSession(); await this.client.close(); this.transport = null; } - await this.closeAgents(); + await this.closeAgents(forceAgentClose); if (this.connectionState === 'disconnected') { return; } @@ -2251,6 +2251,13 @@ export class MCPConnection extends EventEmitter { } } + /** Permanently tears down a connection that will never be reused. */ + public async dispose(): Promise { + this.shouldStopReconnecting = true; + this.removeAllListeners(); + await this.disconnect(true, true); + } + async fetchResources(): Promise { try { const { resources } = await this.client.listResources();