From 139d61c437e62c97bc531c9af5cc402e9487eb9e Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Thu, 11 Jun 2026 01:17:14 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=9A=90=20fix:=20Reuse=20Request-Scoped=20?= =?UTF-8?q?MCP=20Connections=20per=20Run=20(#13673)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(mcp): reuse request-scoped connections per run * test(mcp): update connection factory defaults --- api/app/clients/tools/util/handleTools.js | 4 ++ api/server/controllers/agents/openai.js | 2 + api/server/controllers/agents/responses.js | 3 + .../services/Endpoints/agents/initialize.js | 2 + .../services/Endpoints/agents/skillDeps.js | 2 + api/server/services/MCP.js | 14 ++++ api/server/services/MCPRequestContext.js | 69 +++++++++++++++++++ api/server/services/ToolService.js | 14 +++- api/server/services/Tools/mcp.js | 5 +- .../services/__tests__/ToolService.spec.js | 2 +- packages/api/src/agents/initialize.ts | 8 ++- packages/api/src/mcp/MCPConnectionFactory.ts | 5 ++ packages/api/src/mcp/MCPManager.ts | 6 +- packages/api/src/mcp/UserConnectionManager.ts | 67 ++++++++++++++++++ .../__tests__/MCPConnectionFactory.test.ts | 6 ++ .../api/src/mcp/__tests__/MCPManager.test.ts | 49 +++++++++++++ packages/api/src/mcp/__tests__/utils.test.ts | 19 ++--- packages/api/src/mcp/connection.ts | 12 +++- packages/api/src/mcp/types/index.ts | 8 +++ packages/api/src/mcp/utils.ts | 19 ++--- 20 files changed, 291 insertions(+), 25 deletions(-) create mode 100644 api/server/services/MCPRequestContext.js diff --git a/api/app/clients/tools/util/handleTools.js b/api/app/clients/tools/util/handleTools.js index b29e5bcad1..3f0dc8ab9b 100644 --- a/api/app/clients/tools/util/handleTools.js +++ b/api/app/clients/tools/util/handleTools.js @@ -42,6 +42,7 @@ const { createMCPPermissionContext, resolveConfigServers, } = require('~/server/services/MCP'); +const { getMCPRequestContext } = require('~/server/services/MCPRequestContext'); const { createFileSearchTool, primeFiles: primeSearchFiles } = require('./fileSearch'); const { primeFiles: primeCodeFiles } = require('~/server/services/Files/Code/process'); const { getUserPluginAuthValue } = require('~/server/services/PluginService'); @@ -452,6 +453,8 @@ const loadTools = async ({ let index = -1; const failedMCPServers = new Set(); const safeUser = createSafeUser(options.req?.user); + const requestScopedConnections = + options.requestScopedConnections ?? getMCPRequestContext(options.req, options.res); for (const [serverName, toolConfigs] of Object.entries(requestedMCPTools)) { index++; @@ -470,6 +473,7 @@ const loadTools = async ({ userMCPAuthMap, configServers, requestBody: options.req?.body, + requestScopedConnections, res: options.res, streamId: options.req?._resumableStreamId || null, model: agent?.model ?? model, diff --git a/api/server/controllers/agents/openai.js b/api/server/controllers/agents/openai.js index ff4d5f6dc9..5d50ae4d9d 100644 --- a/api/server/controllers/agents/openai.js +++ b/api/server/controllers/agents/openai.js @@ -348,6 +348,7 @@ const OpenAIChatCompletionController = async (req, res) => { * @type {Map>, * tool_resources?: object, * actionsEnabled?: boolean, @@ -493,6 +494,7 @@ const OpenAIChatCompletionController = async (req, res) => { signal: abortController.signal, toolRegistry: ctx.toolRegistry, mcpAvailableTools: ctx.mcpAvailableTools, + requestScopedConnections: ctx.requestScopedConnections, userMCPAuthMap: ctx.userMCPAuthMap, tool_resources: ctx.tool_resources, actionsEnabled: ctx.actionsEnabled, diff --git a/api/server/controllers/agents/responses.js b/api/server/controllers/agents/responses.js index 4247c99702..220209fff8 100644 --- a/api/server/controllers/agents/responses.js +++ b/api/server/controllers/agents/responses.js @@ -467,6 +467,7 @@ const createResponse = async (req, res) => { * @type {Map>, * tool_resources?: object, * actionsEnabled?: boolean, @@ -685,6 +686,7 @@ const createResponse = async (req, res) => { signal: abortController.signal, toolRegistry: ctx.toolRegistry, mcpAvailableTools: ctx.mcpAvailableTools, + requestScopedConnections: ctx.requestScopedConnections, userMCPAuthMap: ctx.userMCPAuthMap, tool_resources: ctx.tool_resources, actionsEnabled: ctx.actionsEnabled, @@ -860,6 +862,7 @@ const createResponse = async (req, res) => { signal: abortController.signal, toolRegistry: ctx.toolRegistry, mcpAvailableTools: ctx.mcpAvailableTools, + requestScopedConnections: ctx.requestScopedConnections, userMCPAuthMap: ctx.userMCPAuthMap, tool_resources: ctx.tool_resources, actionsEnabled: ctx.actionsEnabled, diff --git a/api/server/services/Endpoints/agents/initialize.js b/api/server/services/Endpoints/agents/initialize.js index 478c6c4e12..eb0655d886 100644 --- a/api/server/services/Endpoints/agents/initialize.js +++ b/api/server/services/Endpoints/agents/initialize.js @@ -184,6 +184,7 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => { * agent?: object, * tool_resources?: object, * toolRegistry?: import('@librechat/agents').LCToolRegistry, + * requestScopedConnections?: import('@librechat/api').RequestScopedMCPConnectionStore, * openAIApiKey?: string * }>} */ @@ -204,6 +205,7 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => { agent: ctx.agent, toolRegistry: ctx.toolRegistry, mcpAvailableTools: ctx.mcpAvailableTools, + requestScopedConnections: ctx.requestScopedConnections, userMCPAuthMap: ctx.userMCPAuthMap, tool_resources: ctx.tool_resources, actionsEnabled: ctx.actionsEnabled, diff --git a/api/server/services/Endpoints/agents/skillDeps.js b/api/server/services/Endpoints/agents/skillDeps.js index f47541bf05..83d287fc08 100644 --- a/api/server/services/Endpoints/agents/skillDeps.js +++ b/api/server/services/Endpoints/agents/skillDeps.js @@ -266,6 +266,7 @@ function buildSkillPrimedIdsByName(manualSkillPrimes, alwaysApplySkillPrimes) { * @param {object} params.agent * @param {object} params.config * @param {Record} [params.config.mcpAvailableTools] + * @param {import('@librechat/api').RequestScopedMCPConnectionStore} [params.config.requestScopedConnections] * @returns {object} */ function buildAgentToolContext({ agent, config }) { @@ -273,6 +274,7 @@ function buildAgentToolContext({ agent, config }) { agent, toolRegistry: config.toolRegistry, mcpAvailableTools: config.mcpAvailableTools, + requestScopedConnections: config.requestScopedConnections, userMCPAuthMap: config.userMCPAuthMap, tool_resources: config.tool_resources, actionsEnabled: config.actionsEnabled, diff --git a/api/server/services/MCP.js b/api/server/services/MCP.js index 4c6daa3682..46971a416e 100644 --- a/api/server/services/MCP.js +++ b/api/server/services/MCP.js @@ -387,6 +387,7 @@ function createOAuthCallback({ runStepEmitter, runStepDeltaEmitter }) { * @param {number} [params.index] * @param {string | null} [params.streamId] - The stream ID for resumable mode. * @param {Record>} [params.userMCPAuthMap] + * @param {import('@librechat/api').RequestScopedMCPConnectionStore} [params.requestScopedConnections] * @param {import('@librechat/api').ParsedServerConfig} [params.serverConfig] - Used to bypass reconnect throttling for request-scoped servers. * @returns { Promise unknown}>> } An object with `_call` method to execute the tool input. */ @@ -400,6 +401,7 @@ async function reconnectServer({ configServers, userMCPAuthMap, requestBody, + requestScopedConnections, streamId = null, }) { logger.debug( @@ -477,6 +479,7 @@ async function reconnectServer({ flowManager, userMCPAuthMap, requestBody, + requestScopedConnections, forceNew: true, returnOnOAuth: false, connectionTimeout: Time.THIRTY_SECONDS, @@ -507,6 +510,7 @@ async function reconnectServer({ * @param {string | null} [params.streamId] - The stream ID for resumable mode. * @param {import('@librechat/api').ParsedServerConfig} [params.config] * @param {import('@librechat/api').RequestBody} [params.requestBody] + * @param {import('@librechat/api').RequestScopedMCPConnectionStore} [params.requestScopedConnections] * @param {Record>} [params.userMCPAuthMap] * @returns { Promise unknown}>> } An object with `_call` method to execute the tool input. */ @@ -522,6 +526,7 @@ async function createMCPTools({ configServers, userMCPAuthMap, requestBody, + requestScopedConnections, streamId = null, }) { const serverConfig = @@ -560,6 +565,7 @@ async function createMCPTools({ configServers, userMCPAuthMap, requestBody, + requestScopedConnections, streamId, }); if (result === null) { @@ -584,6 +590,7 @@ async function createMCPTools({ availableTools: result.availableTools, toolKey: `${tool.name}${Constants.mcp_delimiter}${serverName}`, requestBody, + requestScopedConnections, config: serverConfig, }); if (toolInstance) { @@ -608,6 +615,7 @@ async function createMCPTools({ * @param {Providers | EModelEndpoint} params.provider - The provider for the tool. * @param {LCAvailableTools} [params.availableTools] * @param {import('@librechat/api').RequestBody} [params.requestBody] + * @param {import('@librechat/api').RequestScopedMCPConnectionStore} [params.requestScopedConnections] * @param {Record>} [params.userMCPAuthMap] * @param {import('@librechat/api').ParsedServerConfig} [params.config] * @param {(availableTools: LCAvailableTools) => void} [params.onAvailableTools] @@ -624,6 +632,7 @@ async function createMCPTool({ userMCPAuthMap, availableTools, requestBody, + requestScopedConnections, config, configServers, onAvailableTools, @@ -683,6 +692,7 @@ async function createMCPTool({ configServers, userMCPAuthMap, requestBody, + requestScopedConnections, streamId, }); if (result?.availableTools) { @@ -708,6 +718,7 @@ async function createMCPTool({ mcpPermissionContext, user, requestBody, + requestScopedConnections, provider, toolName, serverName, @@ -722,6 +733,7 @@ function createToolInstance({ mcpPermissionContext, user: capturedUser = null, requestBody: capturedRequestBody, + requestScopedConnections: capturedRequestScopedConnections, toolName, serverName, serverConfig: capturedServerConfig, @@ -816,6 +828,8 @@ function createToolInstance({ }, user: effectiveUser, requestBody: config?.configurable?.requestBody ?? capturedRequestBody, + requestScopedConnections: + config?.configurable?.requestScopedConnections ?? capturedRequestScopedConnections, customUserVars, flowManager, tokenMethods: { diff --git a/api/server/services/MCPRequestContext.js b/api/server/services/MCPRequestContext.js new file mode 100644 index 0000000000..e4e7936d2d --- /dev/null +++ b/api/server/services/MCPRequestContext.js @@ -0,0 +1,69 @@ +const { logger } = require('@librechat/data-schemas'); + +const MCP_REQUEST_CONTEXT = Symbol.for('librechat.mcpRequestContext'); + +function createMCPRequestContext() { + return { + connections: new Map(), + pending: new Map(), + cleanupStarted: false, + }; +} + +async function cleanupMCPRequestContext(context) { + if (!context || context.cleanupStarted) { + return; + } + context.cleanupStarted = true; + + const connections = new Set(context.connections.values()); + const pending = Array.from(context.pending.values()); + if (pending.length > 0) { + const settled = await Promise.allSettled(pending); + for (const result of settled) { + if (result.status === 'fulfilled' && result.value) { + connections.add(result.value); + } + } + } + + await Promise.allSettled( + Array.from(connections).map(async (connection) => { + try { + await connection.disconnect(); + } catch (error) { + logger.warn('[MCP Request Context] Failed to disconnect request-scoped connection', error); + } + }), + ); + + context.connections.clear(); + context.pending.clear(); +} + +function getMCPRequestContext(req, res) { + if (!req) { + return undefined; + } + + if (!req[MCP_REQUEST_CONTEXT]) { + const context = createMCPRequestContext(); + req[MCP_REQUEST_CONTEXT] = context; + + const cleanup = () => { + cleanupMCPRequestContext(context).catch((error) => { + logger.warn('[MCP Request Context] Cleanup failed', error); + }); + }; + res?.once?.('finish', cleanup); + res?.once?.('close', cleanup); + } + + return req[MCP_REQUEST_CONTEXT]; +} + +module.exports = { + cleanupMCPRequestContext, + createMCPRequestContext, + getMCPRequestContext, +}; diff --git a/api/server/services/ToolService.js b/api/server/services/ToolService.js index 8cf71c0c02..52d9f713ad 100644 --- a/api/server/services/ToolService.js +++ b/api/server/services/ToolService.js @@ -70,6 +70,7 @@ const { manifestToolMap, toolkits } = require('~/app/clients/tools/manifest'); const { createOnSearchResults } = require('~/server/services/Tools/search'); const { reinitMCPServer } = require('~/server/services/Tools/mcp'); const { createMCPPermissionContext, resolveConfigServers } = require('~/server/services/MCP'); +const { getMCPRequestContext } = require('~/server/services/MCPRequestContext'); const { recordUsage } = require('~/server/services/Threads'); const { loadTools } = require('~/app/clients/tools/util'); const { redactMessage } = require('~/config/parsers'); @@ -608,6 +609,7 @@ async function loadToolDefinitionsWrapper({ req, res, agent, streamId = null, to const oauthStepIndexes = new Map(); /** @type {Record} */ const mcpAvailableTools = {}; + const requestScopedConnections = getMCPRequestContext(req, res); const rememberMCPAvailableTools = (serverName, availableTools) => { if (!availableTools || Object.keys(availableTools).length === 0) { return; @@ -784,6 +786,7 @@ async function loadToolDefinitionsWrapper({ req, res, agent, streamId = null, to configServers, userMCPAuthMap, requestBody: req.body, + requestScopedConnections, }); rememberMCPAvailableTools(serverName, result?.availableTools); @@ -1041,6 +1044,7 @@ async function loadToolDefinitionsWrapper({ req, res, agent, streamId = null, to return { toolRegistry, mcpAvailableTools, + requestScopedConnections, userMCPAuthMap, toolContextMap, dynamicToolContextMap, @@ -1167,6 +1171,7 @@ async function loadAgentTools({ uploadImageBuffer, returnMetadata: true, mcpPermissionContext, + requestScopedConnections: getMCPRequestContext(req, res), [Tools.web_search]: webSearchCallbacks, }, webSearch: appConfig.webSearch, @@ -1242,6 +1247,7 @@ async function loadAgentTools({ if (!hasActionTools) { return { toolRegistry, + requestScopedConnections: getMCPRequestContext(req, res), userMCPAuthMap, toolContextMap, dynamicToolContextMap, @@ -1260,6 +1266,7 @@ async function loadAgentTools({ } return { toolRegistry, + requestScopedConnections: getMCPRequestContext(req, res), userMCPAuthMap, toolContextMap, dynamicToolContextMap, @@ -1388,6 +1395,7 @@ async function loadAgentTools({ return { toolRegistry, + requestScopedConnections: getMCPRequestContext(req, res), toolContextMap, dynamicToolContextMap, userMCPAuthMap, @@ -1414,6 +1422,7 @@ async function loadAgentTools({ * @param {string[]} params.toolNames - Names of tools to load * @param {Map} [params.toolRegistry] - Tool registry * @param {Record} [params.mcpAvailableTools] - Run-scoped MCP tool definitions + * @param {import('@librechat/api').RequestScopedMCPConnectionStore} [params.requestScopedConnections] - Run-scoped MCP connections * @param {Record>} [params.userMCPAuthMap] - User MCP auth map * @param {Object} [params.tool_resources] - Tool resources * @param {string|null} [params.streamId] - Stream ID for web search callbacks @@ -1428,6 +1437,7 @@ async function loadToolsForExecution({ toolNames, toolRegistry, mcpAvailableTools, + requestScopedConnections, userMCPAuthMap, tool_resources, streamId = null, @@ -1435,7 +1445,8 @@ async function loadToolsForExecution({ }) { const appConfig = req.config; const allLoadedTools = []; - const configurable = { userMCPAuthMap }; + const mcpRequestScopedConnections = requestScopedConnections ?? getMCPRequestContext(req, res); + const configurable = { userMCPAuthMap, requestScopedConnections: mcpRequestScopedConnections }; const isToolSearch = toolNames.includes(AgentConstants.TOOL_SEARCH); const ptcToolNames = [ @@ -1556,6 +1567,7 @@ async function loadToolsForExecution({ uploadImageBuffer, returnMetadata: true, mcpAvailableTools, + requestScopedConnections: mcpRequestScopedConnections, [Tools.web_search]: webSearchCallbacks, }, webSearch: appConfig?.webSearch, diff --git a/api/server/services/Tools/mcp.js b/api/server/services/Tools/mcp.js index d0b37a0b87..5a2dcf8fcc 100644 --- a/api/server/services/Tools/mcp.js +++ b/api/server/services/Tools/mcp.js @@ -24,6 +24,7 @@ const { getLogStores } = require('~/cache'); * @param {(authURL: string, options?: { expiresAt?: number }) => Promise} [params.oauthStart] * @param {() => Promise} [params.oauthEnd] * @param {import('@librechat/api').RequestBody} [params.requestBody] + * @param {import('@librechat/api').RequestScopedMCPConnectionStore} [params.requestScopedConnections] * @param {Record>} [params.userMCPAuthMap] */ async function reinitMCPServer({ @@ -39,6 +40,7 @@ async function reinitMCPServer({ flowManager: _flowManager, serverConfig: providedConfig, requestBody, + requestScopedConnections, oauthEnd, }) { /** @type {MCPConnection | null} */ @@ -144,6 +146,7 @@ async function reinitMCPServer({ oauthEnd, customUserVars, requestBody, + requestScopedConnections, connectionTimeout, serverConfig, graphTokenResolver: getGraphApiToken, @@ -265,7 +268,7 @@ async function reinitMCPServer({ error, ); } finally { - if (connection && ephemeralServer) { + if (connection && ephemeralServer && !requestScopedConnections) { try { await connection.disconnect(); } catch (error) { diff --git a/api/server/services/__tests__/ToolService.spec.js b/api/server/services/__tests__/ToolService.spec.js index 80b7e8c317..cd8d05478c 100644 --- a/api/server/services/__tests__/ToolService.spec.js +++ b/api/server/services/__tests__/ToolService.spec.js @@ -778,7 +778,7 @@ describe('ToolService - Action Capability Gating', () => { mockGetEndpointsConfig.mockResolvedValue(createEndpointsConfig(capabilities)); mockGetServerConfig.mockResolvedValue({ type: 'streamable-http', - url: 'https://mcp.example.com/{{LIBRECHAT_OPENID_ACCESS_TOKEN}}/mcp', + url: 'https://mcp.example.com/{{LIBRECHAT_BODY_MESSAGEID}}/mcp', source: 'yaml', }); mockGetMCPServerTools.mockResolvedValue(null); diff --git a/packages/api/src/agents/initialize.ts b/packages/api/src/agents/initialize.ts index 82ca0e4ad3..9dc39176ec 100644 --- a/packages/api/src/agents/initialize.ts +++ b/packages/api/src/agents/initialize.ts @@ -24,9 +24,9 @@ import type { GenericTool, LCToolRegistry, ToolMap, LCTool } from '@librechat/ag import type { Response as ServerResponse } from 'express'; import type { IMongoFile } from '@librechat/data-schemas'; import type { InitializeResultBase, ServerRequest, EndpointDbMethods } from '~/types'; +import type { LCAvailableTools, RequestScopedMCPConnectionStore } from '../mcp/types'; import type { ResolvedManualSkill, ResolvedAlwaysApplySkill } from './skills'; import type { TFilterFilesByAgentAccess } from './resources'; -import type { LCAvailableTools } from '../mcp/types'; import { injectSkillCatalog, resolveManualSkills, @@ -244,6 +244,8 @@ export type InitializedAgent = Agent & { toolRegistry?: LCToolRegistry; /** Run-scoped MCP tool definitions for request-scoped servers. */ mcpAvailableTools?: Record; + /** Run-scoped MCP connections for request-scoped servers. */ + requestScopedConnections?: RequestScopedMCPConnectionStore; /** Serializable tool definitions for event-driven execution */ toolDefinitions?: LCTool[]; /** Precomputed flag indicating if any tools have defer_loading enabled (for efficient runtime checks) */ @@ -348,6 +350,7 @@ export interface InitializeAgentParams { userMCPAuthMap?: Record>; toolRegistry?: LCToolRegistry; mcpAvailableTools?: Record; + requestScopedConnections?: RequestScopedMCPConnectionStore; /** Serializable tool definitions for event-driven mode */ toolDefinitions?: LCTool[]; hasDeferredTools?: boolean; @@ -880,6 +883,7 @@ export async function initializeAgent( userMCPAuthMap, toolDefinitions: loadedToolDefinitions, mcpAvailableTools, + requestScopedConnections, hasDeferredTools, actionsEnabled, tools: structuredTools, @@ -891,6 +895,7 @@ export async function initializeAgent( userMCPAuthMap: undefined, toolRegistry: undefined, mcpAvailableTools: undefined, + requestScopedConnections: undefined, toolDefinitions: [], hasDeferredTools: false, actionsEnabled: undefined, @@ -1196,6 +1201,7 @@ export async function initializeAgent( resendFiles, toolRegistry, mcpAvailableTools, + requestScopedConnections, tool_resources, userMCPAuthMap, toolDefinitions, diff --git a/packages/api/src/mcp/MCPConnectionFactory.ts b/packages/api/src/mcp/MCPConnectionFactory.ts index be1639048e..a7808feaab 100644 --- a/packages/api/src/mcp/MCPConnectionFactory.ts +++ b/packages/api/src/mcp/MCPConnectionFactory.ts @@ -54,6 +54,7 @@ export class MCPConnectionFactory { protected readonly useSSRFProtection: boolean; protected readonly allowedDomains?: string[] | null; protected readonly allowedAddresses?: string[] | null; + protected readonly ephemeralConnection: boolean; // OAuth-related properties (only set when useOAuth is true) protected readonly userId?: string; @@ -175,6 +176,7 @@ export class MCPConnectionFactory { oauthTokens, useSSRFProtection: this.useSSRFProtection, allowedAddresses: this.allowedAddresses, + ephemeralConnection: this.ephemeralConnection, }); const oauthHandler = () => { @@ -248,6 +250,7 @@ export class MCPConnectionFactory { oauthTokens: null, useSSRFProtection: this.useSSRFProtection, allowedAddresses: this.allowedAddresses, + ephemeralConnection: this.ephemeralConnection, }); unauthConnection.on('oauthRequired', () => { @@ -299,6 +302,7 @@ export class MCPConnectionFactory { this.useSSRFProtection = basic.useSSRFProtection === true; this.allowedDomains = basic.allowedDomains; this.allowedAddresses = basic.allowedAddresses; + this.ephemeralConnection = basic.ephemeralConnection === true; this.connectionTimeout = options?.connectionTimeout; this.tenantContext = tenantStorage?.getStore?.(); this.tenantId = this.tenantContext?.tenantId ?? getTenantId(); @@ -396,6 +400,7 @@ export class MCPConnectionFactory { oauthTokens, useSSRFProtection: this.useSSRFProtection, allowedAddresses: this.allowedAddresses, + ephemeralConnection: this.ephemeralConnection, }); let cleanupOAuthHandlers: (() => void) | null = null; diff --git a/packages/api/src/mcp/MCPManager.ts b/packages/api/src/mcp/MCPManager.ts index e8663c598e..6b26bc709a 100644 --- a/packages/api/src/mcp/MCPManager.ts +++ b/packages/api/src/mcp/MCPManager.ts @@ -350,6 +350,7 @@ Please follow these instructions when using tools from the respective MCP server options, tokenMethods, requestBody, + requestScopedConnections, flowManager, oauthStart, oauthEnd, @@ -367,6 +368,7 @@ Please follow these instructions when using tools from the respective MCP server toolArguments?: Record; options?: RequestOptions; requestBody?: RequestBody; + requestScopedConnections?: t.RequestScopedMCPConnectionStore; tokenMethods?: TokenMethods; customUserVars?: Record; flowManager: FlowStateManager; @@ -399,6 +401,7 @@ Please follow these instructions when using tools from the respective MCP server signal: options?.signal, customUserVars, requestBody, + requestScopedConnections, serverConfig: providedConfig, }); @@ -419,7 +422,8 @@ Please follow these instructions when using tools from the respective MCP server ); } const isDbSourced = isUserSourced(rawConfig); - disconnectAfterCall = !!userId && requiresEphemeralUserConnection(rawConfig); + disconnectAfterCall = + !!userId && requiresEphemeralUserConnection(rawConfig) && !requestScopedConnections; /** Pre-process Graph token placeholders (async) before the synchronous processMCPEnv pass */ const graphProcessedConfig = isDbSourced diff --git a/packages/api/src/mcp/UserConnectionManager.ts b/packages/api/src/mcp/UserConnectionManager.ts index 6c02852586..e12720dc01 100644 --- a/packages/api/src/mcp/UserConnectionManager.ts +++ b/packages/api/src/mcp/UserConnectionManager.ts @@ -85,6 +85,72 @@ export abstract class UserConnectionManager { ); } const ephemeralConnection = config ? requiresEphemeralUserConnection(config) : false; + const requestScopedConnections = ephemeralConnection + ? opts.requestScopedConnections + : undefined; + if (requestScopedConnections) { + const requestConnectionKey = `${userId}:${serverName}`; + const existing = requestScopedConnections.connections.get(requestConnectionKey) as + | MCPConnection + | undefined; + if (existing) { + if (!config || (config.updatedAt && existing.isStale(config.updatedAt))) { + await existing.disconnect().catch((error) => { + logger.warn( + `[MCP][User: ${userId}][${serverName}] Failed to disconnect stale request-scoped connection`, + error, + ); + }); + 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); + } + } + + const pending = requestScopedConnections.pending.get(requestConnectionKey) as + | Promise + | undefined; + if (pending) { + logger.debug( + `[MCP][User: ${userId}][${serverName}] Joining in-flight request-scoped connection attempt`, + ); + return pending; + } + + const pendingOAuth = this.createPendingOAuthState(opts.oauthStart); + const connectionPromise = this.createUserConnectionInternal( + { + ...opts, + forceNew: true, + ephemeralConnection: true, + serverConfig: config, + oauthStart: this.createPendingOAuthStart(serverName, userId, pendingOAuth), + }, + userId, + forceNew === true, + ).then((connection) => { + requestScopedConnections.connections.set(requestConnectionKey, connection); + return connection; + }); + + requestScopedConnections.pending.set( + requestConnectionKey, + connectionPromise as Promise, + ); + + try { + return await connectionPromise; + } finally { + if (requestScopedConnections.pending.get(requestConnectionKey) === connectionPromise) { + requestScopedConnections.pending.delete(requestConnectionKey); + } + } + } + const forceNewConnection = forceNew || ephemeralConnection; const clearCooldown = forceNew === true; @@ -406,6 +472,7 @@ export abstract class UserConnectionManager { useSSRFProtection: registry.shouldEnableSSRFProtection(), allowedDomains, allowedAddresses, + ephemeralConnection, }; const useOAuth = requiresOAuthMachinery(runtimeConfig); diff --git a/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts b/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts index fdf9c472ac..73048c7ab3 100644 --- a/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts +++ b/packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts @@ -167,6 +167,8 @@ describe('MCPConnectionFactory', () => { userId: undefined, oauthTokens: null, useSSRFProtection: false, + allowedAddresses: undefined, + ephemeralConnection: false, }); expect(mockConnectionInstance.connect).toHaveBeenCalled(); }); @@ -286,6 +288,8 @@ describe('MCPConnectionFactory', () => { userId: 'user123', oauthTokens: mockTokens, useSSRFProtection: false, + allowedAddresses: undefined, + ephemeralConnection: false, }); }); @@ -380,6 +384,8 @@ describe('MCPConnectionFactory', () => { userId: 'user123', oauthTokens: null, useSSRFProtection: false, + allowedAddresses: undefined, + ephemeralConnection: false, }); expect(mockLogger.debug).toHaveBeenCalledWith( expect.stringContaining('No existing tokens found or error loading tokens'), diff --git a/packages/api/src/mcp/__tests__/MCPManager.test.ts b/packages/api/src/mcp/__tests__/MCPManager.test.ts index 04b6ae2ac0..0a0b519f24 100644 --- a/packages/api/src/mcp/__tests__/MCPManager.test.ts +++ b/packages/api/src/mcp/__tests__/MCPManager.test.ts @@ -1815,6 +1815,55 @@ describe('MCPManager', () => { expect(MCPConnectionFactory.create).toHaveBeenCalledTimes(2); }); + it('should reuse BODY-scoped connections within a request-scoped connection store', async () => { + const bodyUrlConfig: t.ParsedServerConfig = { + type: 'streamable-http', + url: 'https://api.example.com/messages/{{LIBRECHAT_BODY_MESSAGEID}}/mcp', + source: 'yaml', + requiresOAuth: false, + }; + const requestScopedConnection = { + isConnected: jest.fn().mockResolvedValue(true), + } as unknown as MCPConnection; + const requestScopedConnections: t.RequestScopedMCPConnectionStore = { + connections: new Map(), + pending: new Map(), + }; + + mockAppConnections({ + has: jest.fn().mockResolvedValue(false), + }); + (mockRegistryInstance.getServerConfig as jest.Mock).mockResolvedValue(bodyUrlConfig); + mockProcessMCPEnv.mockImplementation(({ options, body }) => ({ + ...options, + ...('url' in options && { + url: options.url?.replace('{{LIBRECHAT_BODY_MESSAGEID}}', body?.messageId ?? ''), + }), + })); + (MCPConnectionFactory.create as jest.Mock).mockResolvedValue(requestScopedConnection); + + const manager = await MCPManager.createInstance(newMCPServersConfig()); + const first = await manager.getUserConnection({ + serverName, + user: mockUser, + requestBody: { messageId: 'message-1' }, + requestScopedConnections, + }); + const second = await manager.getUserConnection({ + serverName, + user: mockUser, + requestBody: { messageId: 'message-1' }, + requestScopedConnections, + }); + + expect(first).toBe(requestScopedConnection); + expect(second).toBe(requestScopedConnection); + expect(MCPConnectionFactory.create).toHaveBeenCalledTimes(1); + expect(requestScopedConnections.connections.get(`${mockUser.id}:${serverName}`)).toBe( + requestScopedConnection, + ); + }); + it('should not clear server cooldowns for ephemeral runtime connections', async () => { const bodyUrlConfig: t.ParsedServerConfig = { type: 'streamable-http', diff --git a/packages/api/src/mcp/__tests__/utils.test.ts b/packages/api/src/mcp/__tests__/utils.test.ts index 199f259434..bd6baa2fa0 100644 --- a/packages/api/src/mcp/__tests__/utils.test.ts +++ b/packages/api/src/mcp/__tests__/utils.test.ts @@ -557,26 +557,28 @@ describe('getMissingRuntimeBodyPlaceholderFields', () => { }); describe('requiresEphemeralUserConnection', () => { - it('returns true when request-varying placeholders affect oauth_headers', () => { + it('returns true when BODY placeholders affect oauth_headers', () => { expect( requiresEphemeralUserConnection({ source: 'yaml', url: 'https://example.com/mcp', oauth_headers: { - Authorization: 'Bearer {{LIBRECHAT_OPENID_ACCESS_TOKEN}}', + 'X-Message': '{{LIBRECHAT_BODY_MESSAGEID}}', }, }), ).toBe(true); }); - it('returns true when request-varying placeholders affect connection fields', () => { + it('returns true when BODY placeholders affect connection fields', () => { expect( requiresEphemeralUserConnection({ source: 'yaml', url: 'https://example.com/messages/{{LIBRECHAT_BODY_MESSAGEID}}/mcp', }), ).toBe(true); + }); + it('does not treat Graph placeholders as request-scoped by themselves', () => { expect( requiresEphemeralUserConnection({ source: 'config', @@ -584,16 +586,16 @@ describe('requiresEphemeralUserConnection', () => { GRAPH_TOKEN: '{{LIBRECHAT_GRAPH_ACCESS_TOKEN}}', }, }), - ).toBe(true); + ).toBe(false); }); - it('returns true when OpenID token placeholders affect connection fields', () => { + it('does not treat OpenID token placeholders as request-scoped by themselves', () => { expect( requiresEphemeralUserConnection({ source: 'yaml', args: ['--id-token={{LIBRECHAT_OPENID_ID_TOKEN}}'], }), - ).toBe(true); + ).toBe(false); expect( requiresEphemeralUserConnection({ @@ -602,16 +604,15 @@ describe('requiresEphemeralUserConnection', () => { Authorization: 'Bearer {{LIBRECHAT_OPENID_ACCESS_TOKEN}}', }, }), - ).toBe(true); + ).toBe(false); }); - it('returns true when request-varying placeholders affect remote transport headers', () => { + it('returns true when BODY placeholders affect remote transport headers', () => { expect( requiresEphemeralUserConnection({ source: 'yaml', headers: { 'X-Message': '{{LIBRECHAT_BODY_MESSAGEID}}', - 'X-Graph': '{{LIBRECHAT_GRAPH_ACCESS_TOKEN}}', }, }), ).toBe(true); diff --git a/packages/api/src/mcp/connection.ts b/packages/api/src/mcp/connection.ts index 217a44d89a..2366181684 100644 --- a/packages/api/src/mcp/connection.ts +++ b/packages/api/src/mcp/connection.ts @@ -1092,6 +1092,7 @@ interface MCPConnectionParams { oauthTokens?: MCPOAuthTokens | null; useSSRFProtection?: boolean; allowedAddresses?: string[] | null; + ephemeralConnection?: boolean; } export class MCPConnection extends EventEmitter { @@ -1116,6 +1117,7 @@ export class MCPConnection extends EventEmitter { private oauthRecovery = false; private readonly useSSRFProtection: boolean; private readonly allowedAddresses?: string[] | null; + private readonly ephemeralConnection: boolean; private readonly proxyConfig?: MCPProxyConfig; iconPath?: string; timeout?: number; @@ -1232,6 +1234,7 @@ export class MCPConnection extends EventEmitter { this.userId = params.userId; this.useSSRFProtection = params.useSSRFProtection === true; this.allowedAddresses = params.allowedAddresses ?? null; + this.ephemeralConnection = params.ephemeralConnection === true; this.proxyConfig = getMCPProxyConfig(params.serverConfig); this.iconPath = params.serverConfig.iconPath; this.timeout = params.serverConfig.timeout; @@ -2026,8 +2029,13 @@ export class MCPConnection extends EventEmitter { async connect(): Promise { try { - // preserve cycle tracking across reconnects so the circuit breaker can detect rapid cycling - await this.disconnect(false); + /** + * Persistent connections preserve cycle tracking across reconnects so the + * circuit breaker can detect storms. Request-scoped connections are + * intentionally short-lived per tool call, so their clean lifecycle should + * not consume the reconnect-storm cycle budget. + */ + await this.disconnect(this.ephemeralConnection); await this.connectClient(); if (!(await this.isConnected())) { throw new Error('Connection not established'); diff --git a/packages/api/src/mcp/types/index.ts b/packages/api/src/mcp/types/index.ts index 253d7c43c2..852dc0b837 100644 --- a/packages/api/src/mcp/types/index.ts +++ b/packages/api/src/mcp/types/index.ts @@ -194,6 +194,8 @@ export interface BasicConnectionOptions { dbSourced?: boolean; /** When true, serverConfig has already gone through processMCPEnv for this request */ skipEnvProcessing?: boolean; + /** When true, the connection is intentionally short-lived for a single request/tool call */ + ephemeralConnection?: boolean; } /** User context for placeholder resolution in MCP connections (non-OAuth and OAuth alike) */ @@ -201,10 +203,16 @@ export interface UserConnectionContext { user?: IUser; customUserVars?: Record; requestBody?: RequestBody; + requestScopedConnections?: RequestScopedMCPConnectionStore; graphTokenResolver?: GraphTokenResolver; connectionTimeout?: number; } +export interface RequestScopedMCPConnectionStore { + connections: Map; + pending: Map>; +} + export interface OAuthStartOptions { expiresAt?: number; } diff --git a/packages/api/src/mcp/utils.ts b/packages/api/src/mcp/utils.ts index 9d58c10358..12829a93e1 100644 --- a/packages/api/src/mcp/utils.ts +++ b/packages/api/src/mcp/utils.ts @@ -5,7 +5,6 @@ import type { RequestBody } from '~/types'; export const mcpToolPattern: RegExp = new RegExp(`^.+${Constants.mcp_delimiter}.+$`); const RUNTIME_CONTEXT_PLACEHOLDER_PATTERN = /\{\{LIBRECHAT_(?:USER|OPENID|GRAPH|BODY)_[^}]+\}\}/; -const EPHEMERAL_CONNECTION_PLACEHOLDER_PATTERN = /\{\{LIBRECHAT_(?:OPENID|GRAPH|BODY)_[^}]+\}\}/; const RUNTIME_BODY_PLACEHOLDER_PATTERN = /\{\{LIBRECHAT_BODY_[^}]+\}\}/; const RUNTIME_BODY_PLACEHOLDER_CAPTURE_PATTERN = /\{\{LIBRECHAT_BODY_([^}]+)\}\}/g; @@ -76,10 +75,6 @@ function hasRuntimeContextPlaceholder(value: PlaceholderValue): boolean { return hasPlaceholder(value, RUNTIME_CONTEXT_PLACEHOLDER_PATTERN); } -function hasEphemeralConnectionPlaceholder(value: PlaceholderValue): boolean { - return hasPlaceholder(value, EPHEMERAL_CONNECTION_PLACEHOLDER_PATTERN); -} - function hasPlaceholder(value: PlaceholderValue, pattern: RegExp): boolean { if (typeof value === 'string') { return pattern.test(value); @@ -176,19 +171,25 @@ export function getMissingRuntimeBodyPlaceholderFields( } /** - * `GRAPH` and `BODY` placeholders can change per request. If they affect the - * connection-defining parts of a config, the normal userId:serverName cache - * would reuse a connection built with stale request context. + * `BODY` placeholders vary by chat request, so the normal userId:serverName + * cache would reuse a connection built with stale request context. * * Ephemeral connections are created and torn down per tool call — configs using * these placeholders pay a full connect + initialize on every invocation. + * + * User/OpenID/Graph placeholders still require user-scoped connections, but they + * are not request-scoped by themselves. HTTP transports refresh resolved headers + * before each tool call, so token/user headers can remain on the cached user + * connection without forcing a reconnect for every invocation. */ export function requiresEphemeralUserConnection(config: UserScopedConnectionConfig): boolean { if (isUserSourced(config)) { return false; } - return placeholderBearingFields(config).some(hasEphemeralConnectionPlaceholder); + return placeholderBearingFields(config).some((value) => + hasPlaceholder(value, RUNTIME_BODY_PLACEHOLDER_PATTERN), + ); } /**