🔄 refactor: Optimize MCP Tool Initialization

🔄 refactor: Optimize MCP Tool Initialization

fix: update tool caching to use separated mcp logic

refactor: Replace `req.user` with `userId` in MCP handling functions

refactor: Replace `req` parameter with `userId` in file search tool functions

fix: Update user connection parameter to use object format in reinitMCPServer

refactor: Simplify MCP tool creation logic and improve handling of tool configurations to avoid capturing too much in closures

refactor: ensure MCP available tools are fetched from cache only when needed
This commit is contained in:
Danny Avila 2025-09-21 16:52:43 -04:00
parent 386900fb4f
commit 5b1a31ef4d
No known key found for this signature in database
GPG key ID: BF31EEB2C5CA0956
15 changed files with 111 additions and 134 deletions

View file

@ -33,7 +33,7 @@ const { createFileSearchTool, primeFiles: primeSearchFiles } = require('./fileSe
const { getUserPluginAuthValue } = require('~/server/services/PluginService');
const { createMCPTool, createMCPTools } = require('~/server/services/MCP');
const { loadAuthValues } = require('~/server/services/Tools/credentials');
const { getCachedTools } = require('~/server/services/Config');
const { getMCPServerTools } = require('~/server/services/Config');
const { getRoleByName } = require('~/models/Role');
/**
@ -250,7 +250,6 @@ const loadTools = async ({
/** @type {Record<string, string>} */
const toolContextMap = {};
const cachedTools = (await getCachedTools({ userId: user, includeGlobal: true })) ?? {};
const requestedMCPTools = {};
for (const tool of tools) {
@ -307,7 +306,7 @@ const loadTools = async ({
}
return createFileSearchTool({
req: options.req,
userId: user,
files,
entity_id: agent?.id,
fileCitations,
@ -340,7 +339,7 @@ Current Date & Time: ${replaceSpecialVars({ text: '{{iso_datetime}}' })}
});
};
continue;
} else if (tool && cachedTools && mcpToolPattern.test(tool)) {
} else if (tool && mcpToolPattern.test(tool)) {
const [toolName, serverName] = tool.split(Constants.mcp_delimiter);
if (toolName === Constants.mcp_server) {
/** Placeholder used for UI purposes */
@ -353,33 +352,21 @@ Current Date & Time: ${replaceSpecialVars({ text: '{{iso_datetime}}' })}
continue;
}
if (toolName === Constants.mcp_all) {
const currentMCPGenerator = async (index) =>
createMCPTools({
req: options.req,
res: options.res,
index,
requestedMCPTools[serverName] = [
{
type: 'all',
serverName,
userMCPAuthMap,
model: agent?.model ?? model,
provider: agent?.provider ?? endpoint,
signal,
});
requestedMCPTools[serverName] = [currentMCPGenerator];
},
];
continue;
}
const currentMCPGenerator = async (index) =>
createMCPTool({
index,
req: options.req,
res: options.res,
toolKey: tool,
userMCPAuthMap,
model: agent?.model ?? model,
provider: agent?.provider ?? endpoint,
signal,
});
requestedMCPTools[serverName] = requestedMCPTools[serverName] || [];
requestedMCPTools[serverName].push(currentMCPGenerator);
requestedMCPTools[serverName].push({
type: 'single',
toolKey: tool,
serverName,
});
continue;
}
@ -422,20 +409,51 @@ Current Date & Time: ${replaceSpecialVars({ text: '{{iso_datetime}}' })}
const mcpToolPromises = [];
/** MCP server tools are initialized sequentially by server */
let index = -1;
for (const [serverName, generators] of Object.entries(requestedMCPTools)) {
for (const [serverName, toolConfigs] of Object.entries(requestedMCPTools)) {
index++;
for (const generator of generators) {
/** @type {LCAvailableTools} */
let availableTools;
for (const config of toolConfigs) {
try {
if (generator && generators.length === 1) {
const mcpParams = {
res: options.res,
userId: user,
index,
serverName: config.serverName,
userMCPAuthMap,
model: agent?.model ?? model,
provider: agent?.provider ?? endpoint,
signal,
};
if (config.type === 'all' && toolConfigs.length === 1) {
/** Handle async loading for single 'all' tool config */
mcpToolPromises.push(
generator(index).catch((error) => {
createMCPTools(mcpParams).catch((error) => {
logger.error(`Error loading ${serverName} tools:`, error);
return null;
}),
);
continue;
}
const mcpTool = await generator(index);
if (!availableTools) {
try {
availableTools = await getMCPServerTools(serverName);
} catch (error) {
logger.error(`Error fetching available tools for MCP server ${serverName}:`, error);
}
}
/** Handle synchronous loading */
const mcpTool =
config.type === 'all'
? await createMCPTools(mcpParams)
: await createMCPTool({
...mcpParams,
availableTools,
toolKey: config.toolKey,
});
if (Array.isArray(mcpTool)) {
loadedTools.push(...mcpTool);
} else if (mcpTool) {