mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-08-03 22:32:42 +00:00
🚐 fix: Reuse Request-Scoped MCP Connections per Run (#13673)
Some checks are pending
Docker Dev Branch Images Build / build (Dockerfile, lc-dev, node) (push) Waiting to run
Docker Dev Branch Images Build / build (Dockerfile.multi, lc-dev-api, api-build) (push) Waiting to run
GitNexus Index / index (push) Waiting to run
GitNexus Index / post-index (push) Blocked by required conditions
Docker Dev Images Build / build (Dockerfile, librechat-dev, node) (push) Waiting to run
Docker Dev Images Build / build (Dockerfile.multi, librechat-dev-api, api-build) (push) Waiting to run
Sync Locize Translations & Create Translation PR / Sync Translation Keys with Locize (push) Waiting to run
Sync Locize Translations & Create Translation PR / Create Translation PR on Version Published (push) Blocked by required conditions
Sync Helm Chart Tags / Ignore non-main push (push) Waiting to run
Sync Helm Chart Tags / Sync chart tags (push) Waiting to run
Some checks are pending
Docker Dev Branch Images Build / build (Dockerfile, lc-dev, node) (push) Waiting to run
Docker Dev Branch Images Build / build (Dockerfile.multi, lc-dev-api, api-build) (push) Waiting to run
GitNexus Index / index (push) Waiting to run
GitNexus Index / post-index (push) Blocked by required conditions
Docker Dev Images Build / build (Dockerfile, librechat-dev, node) (push) Waiting to run
Docker Dev Images Build / build (Dockerfile.multi, librechat-dev-api, api-build) (push) Waiting to run
Sync Locize Translations & Create Translation PR / Sync Translation Keys with Locize (push) Waiting to run
Sync Locize Translations & Create Translation PR / Create Translation PR on Version Published (push) Blocked by required conditions
Sync Helm Chart Tags / Ignore non-main push (push) Waiting to run
Sync Helm Chart Tags / Sync chart tags (push) Waiting to run
* fix(mcp): reuse request-scoped connections per run * test(mcp): update connection factory defaults
This commit is contained in:
parent
65bca95023
commit
139d61c437
20 changed files with 291 additions and 25 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -348,6 +348,7 @@ const OpenAIChatCompletionController = async (req, res) => {
|
|||
* @type {Map<string, {
|
||||
* agent: object,
|
||||
* toolRegistry?: import('@librechat/agents').LCToolRegistry,
|
||||
* requestScopedConnections?: import('@librechat/api').RequestScopedMCPConnectionStore,
|
||||
* userMCPAuthMap?: Record<string, Record<string, string>>,
|
||||
* 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,
|
||||
|
|
|
|||
|
|
@ -467,6 +467,7 @@ const createResponse = async (req, res) => {
|
|||
* @type {Map<string, {
|
||||
* agent: object,
|
||||
* toolRegistry?: import('@librechat/agents').LCToolRegistry,
|
||||
* requestScopedConnections?: import('@librechat/api').RequestScopedMCPConnectionStore,
|
||||
* userMCPAuthMap?: Record<string, Record<string, string>>,
|
||||
* 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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -266,6 +266,7 @@ function buildSkillPrimedIdsByName(manualSkillPrimes, alwaysApplySkillPrimes) {
|
|||
* @param {object} params.agent
|
||||
* @param {object} params.config
|
||||
* @param {Record<string, import('@librechat/api').LCAvailableTools>} [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,
|
||||
|
|
|
|||
|
|
@ -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<string, Record<string, string>>} [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<Array<typeof tool | { _call: (toolInput: Object | string) => 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<string, Record<string, string>>} [params.userMCPAuthMap]
|
||||
* @returns { Promise<Array<typeof tool | { _call: (toolInput: Object | string) => 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<string, Record<string, string>>} [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: {
|
||||
|
|
|
|||
69
api/server/services/MCPRequestContext.js
Normal file
69
api/server/services/MCPRequestContext.js
Normal file
|
|
@ -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,
|
||||
};
|
||||
|
|
@ -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<string, import('@librechat/api').LCAvailableTools>} */
|
||||
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<string, import('@librechat/api').LCAvailableTools>} [params.mcpAvailableTools] - Run-scoped MCP tool definitions
|
||||
* @param {import('@librechat/api').RequestScopedMCPConnectionStore} [params.requestScopedConnections] - Run-scoped MCP connections
|
||||
* @param {Record<string, Record<string, string>>} [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,
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ const { getLogStores } = require('~/cache');
|
|||
* @param {(authURL: string, options?: { expiresAt?: number }) => Promise<void>} [params.oauthStart]
|
||||
* @param {() => Promise<void>} [params.oauthEnd]
|
||||
* @param {import('@librechat/api').RequestBody} [params.requestBody]
|
||||
* @param {import('@librechat/api').RequestScopedMCPConnectionStore} [params.requestScopedConnections]
|
||||
* @param {Record<string, Record<string, string>>} [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) {
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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<string, LCAvailableTools>;
|
||||
/** 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<string, Record<string, string>>;
|
||||
toolRegistry?: LCToolRegistry;
|
||||
mcpAvailableTools?: Record<string, LCAvailableTools>;
|
||||
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,
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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<string, unknown>;
|
||||
options?: RequestOptions;
|
||||
requestBody?: RequestBody;
|
||||
requestScopedConnections?: t.RequestScopedMCPConnectionStore;
|
||||
tokenMethods?: TokenMethods;
|
||||
customUserVars?: Record<string, string>;
|
||||
flowManager: FlowStateManager<MCPOAuthTokens | null>;
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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<MCPConnection>
|
||||
| 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<unknown>,
|
||||
);
|
||||
|
||||
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);
|
||||
|
|
|
|||
|
|
@ -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'),
|
||||
|
|
|
|||
|
|
@ -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',
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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<void> {
|
||||
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');
|
||||
|
|
|
|||
|
|
@ -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<string, string>;
|
||||
requestBody?: RequestBody;
|
||||
requestScopedConnections?: RequestScopedMCPConnectionStore;
|
||||
graphTokenResolver?: GraphTokenResolver;
|
||||
connectionTimeout?: number;
|
||||
}
|
||||
|
||||
export interface RequestScopedMCPConnectionStore {
|
||||
connections: Map<string, unknown>;
|
||||
pending: Map<string, Promise<unknown>>;
|
||||
}
|
||||
|
||||
export interface OAuthStartOptions {
|
||||
expiresAt?: number;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue