From da4aa37493024aef255091ec5a83f07aafcfd0f6 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Wed, 20 Aug 2025 12:19:29 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=9B=A0=EF=B8=8F=20refactor:=20Consolidate?= =?UTF-8?q?=20MCP=20Tool=20Caching=20(#9172)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 🛠️ refactor: Consolidate MCP Tool Caching * 🐍 fix: Correctly mock and utilize updateMCPUserTools in MCP route tests --- api/cache/getLogStores.js | 1 - api/server/controllers/PluginController.js | 151 ++++++++---------- .../controllers/PluginController.spec.js | 32 ++-- api/server/routes/__tests__/mcp.spec.js | 22 ++- api/server/routes/mcp.js | 60 ++----- api/server/services/Config/index.js | 2 + api/server/services/Config/mcpToolsCache.js | 142 ++++++++++++++++ api/server/services/initializeMCPs.js | 23 +-- packages/api/src/mcp/MCPManager.ts | 47 +++--- packages/data-provider/src/config.ts | 4 - 10 files changed, 293 insertions(+), 191 deletions(-) create mode 100644 api/server/services/Config/mcpToolsCache.js diff --git a/api/cache/getLogStores.js b/api/cache/getLogStores.js index 675e16a76..dfb201744 100644 --- a/api/cache/getLogStores.js +++ b/api/cache/getLogStores.js @@ -31,7 +31,6 @@ const namespaces = { [CacheKeys.SAML_SESSION]: sessionCache(CacheKeys.SAML_SESSION), [CacheKeys.ROLES]: standardCache(CacheKeys.ROLES), - [CacheKeys.MCP_TOOLS]: standardCache(CacheKeys.MCP_TOOLS), [CacheKeys.CONFIG_STORE]: standardCache(CacheKeys.CONFIG_STORE), [CacheKeys.STATIC_CONFIG]: standardCache(CacheKeys.STATIC_CONFIG), [CacheKeys.PENDING_REQ]: standardCache(CacheKeys.PENDING_REQ), diff --git a/api/server/controllers/PluginController.js b/api/server/controllers/PluginController.js index 06b4585a0..0329834ad 100644 --- a/api/server/controllers/PluginController.js +++ b/api/server/controllers/PluginController.js @@ -6,9 +6,15 @@ const { filterUniquePlugins, convertMCPToolsToPlugins, } = require('@librechat/api'); -const { getCustomConfig, getCachedTools } = require('~/server/services/Config'); +const { + getCachedTools, + setCachedTools, + mergeUserTools, + getCustomConfig, +} = require('~/server/services/Config'); +const { loadAndFormatTools } = require('~/server/services/ToolService'); const { availableTools, toolkits } = require('~/app/clients/tools'); -const { getMCPManager, getFlowStateManager } = require('~/config'); +const { getMCPManager } = require('~/config'); const { getLogStores } = require('~/cache'); const getAvailablePluginsController = async (req, res) => { @@ -22,6 +28,7 @@ const getAvailablePluginsController = async (req, res) => { /** @type {{ filteredTools: string[], includedTools: string[] }} */ const { filteredTools = [], includedTools = [] } = req.app.locals; + /** @type {import('@librechat/api').LCManifestTool[]} */ const pluginManifest = availableTools; const uniquePlugins = filterUniquePlugins(pluginManifest); @@ -47,45 +54,6 @@ const getAvailablePluginsController = async (req, res) => { } }; -function createServerToolsCallback() { - /** - * @param {string} serverName - * @param {TPlugin[] | null} serverTools - */ - return async function (serverName, serverTools) { - try { - const mcpToolsCache = getLogStores(CacheKeys.MCP_TOOLS); - if (!serverName || !mcpToolsCache) { - return; - } - await mcpToolsCache.set(serverName, serverTools); - logger.debug(`MCP tools for ${serverName} added to cache.`); - } catch (error) { - logger.error('Error retrieving MCP tools from cache:', error); - } - }; -} - -function createGetServerTools() { - /** - * Retrieves cached server tools - * @param {string} serverName - * @returns {Promise} - */ - return async function (serverName) { - try { - const mcpToolsCache = getLogStores(CacheKeys.MCP_TOOLS); - if (!mcpToolsCache) { - return null; - } - return await mcpToolsCache.get(serverName); - } catch (error) { - logger.error('Error retrieving MCP tools from cache:', error); - return null; - } - }; -} - /** * Retrieves and returns a list of available tools, either from a cache or by reading a plugin manifest file. * @@ -101,11 +69,18 @@ function createGetServerTools() { const getAvailableTools = async (req, res) => { try { const userId = req.user?.id; + if (!userId) { + logger.warn('[getAvailableTools] User ID not found in request'); + return res.status(401).json({ message: 'Unauthorized' }); + } const customConfig = await getCustomConfig(); const cache = getLogStores(CacheKeys.CONFIG_STORE); const cachedToolsArray = await cache.get(CacheKeys.TOOLS); const cachedUserTools = await getCachedTools({ userId }); - const userPlugins = convertMCPToolsToPlugins({ functionTools: cachedUserTools, customConfig }); + const userPlugins = + cachedUserTools != null + ? convertMCPToolsToPlugins({ functionTools: cachedUserTools, customConfig }) + : undefined; if (cachedToolsArray != null && userPlugins != null) { const dedupedTools = filterUniquePlugins([...userPlugins, ...cachedToolsArray]); @@ -113,31 +88,47 @@ const getAvailableTools = async (req, res) => { return; } + /** @type {Record | null} Get tool definitions to filter which tools are actually available */ + let toolDefinitions = await getCachedTools({ includeGlobal: true }); + let prelimCachedTools; + + // TODO: this is a temp fix until app config is refactored + if (!toolDefinitions) { + toolDefinitions = loadAndFormatTools({ + adminFilter: req.app.locals?.filteredTools, + adminIncluded: req.app.locals?.includedTools, + directory: req.app.locals?.paths.structuredTools, + }); + prelimCachedTools = toolDefinitions; + } + + /** @type {import('@librechat/api').LCManifestTool[]} */ let pluginManifest = availableTools; if (customConfig?.mcpServers != null) { try { const mcpManager = getMCPManager(); - const flowsCache = getLogStores(CacheKeys.FLOWS); - const flowManager = flowsCache ? getFlowStateManager(flowsCache) : null; - const serverToolsCallback = createServerToolsCallback(); - const getServerTools = createGetServerTools(); - const mcpTools = await mcpManager.loadManifestTools({ - flowManager, - serverToolsCallback, - getServerTools, - }); - pluginManifest = [...mcpTools, ...pluginManifest]; + const mcpTools = await mcpManager.loadAllManifestTools(userId); + const mcpToolsRecord = mcpTools.reduce((acc, tool) => { + pluginManifest.push(tool); + acc[tool.pluginKey] = tool; + if (!toolDefinitions[tool.pluginKey]) { + toolDefinitions[tool.pluginKey] = tool; + } + return acc; + }, prelimCachedTools ?? {}); + await mergeUserTools({ userId, cachedUserTools, userTools: mcpToolsRecord }); } catch (error) { logger.error( '[getAvailableTools] Error loading MCP Tools, servers may still be initializing:', error, ); } + } else if (prelimCachedTools != null) { + await setCachedTools(prelimCachedTools, { isGlobal: true }); } - /** @type {TPlugin[]} */ + /** @type {TPlugin[]} Deduplicate and authenticate plugins */ const uniquePlugins = filterUniquePlugins(pluginManifest); - const authenticatedPlugins = uniquePlugins.map((plugin) => { if (checkPluginAuth(plugin)) { return { ...plugin, authenticated: true }; @@ -146,8 +137,7 @@ const getAvailableTools = async (req, res) => { } }); - const toolDefinitions = (await getCachedTools({ includeGlobal: true })) || {}; - + /** Filter plugins based on availability and add MCP-specific auth config */ const toolsOutput = []; for (const plugin of authenticatedPlugins) { const isToolDefined = toolDefinitions[plugin.pluginKey] !== undefined; @@ -163,41 +153,36 @@ const getAvailableTools = async (req, res) => { const toolToAdd = { ...plugin }; - if (!plugin.pluginKey.includes(Constants.mcp_delimiter)) { - toolsOutput.push(toolToAdd); - continue; - } + if (plugin.pluginKey.includes(Constants.mcp_delimiter)) { + const parts = plugin.pluginKey.split(Constants.mcp_delimiter); + const serverName = parts[parts.length - 1]; + const serverConfig = customConfig?.mcpServers?.[serverName]; - const parts = plugin.pluginKey.split(Constants.mcp_delimiter); - const serverName = parts[parts.length - 1]; - const serverConfig = customConfig?.mcpServers?.[serverName]; - - if (!serverConfig?.customUserVars) { - toolsOutput.push(toolToAdd); - continue; - } - - const customVarKeys = Object.keys(serverConfig.customUserVars); - - if (customVarKeys.length === 0) { - toolToAdd.authConfig = []; - toolToAdd.authenticated = true; - } else { - toolToAdd.authConfig = Object.entries(serverConfig.customUserVars).map(([key, value]) => ({ - authField: key, - label: value.title || key, - description: value.description || '', - })); - toolToAdd.authenticated = false; + if (serverConfig?.customUserVars) { + const customVarKeys = Object.keys(serverConfig.customUserVars); + if (customVarKeys.length === 0) { + toolToAdd.authConfig = []; + toolToAdd.authenticated = true; + } else { + toolToAdd.authConfig = Object.entries(serverConfig.customUserVars).map( + ([key, value]) => ({ + authField: key, + label: value.title || key, + description: value.description || '', + }), + ); + toolToAdd.authenticated = false; + } + } } toolsOutput.push(toolToAdd); } + const finalTools = filterUniquePlugins(toolsOutput); await cache.set(CacheKeys.TOOLS, finalTools); - const dedupedTools = filterUniquePlugins([...userPlugins, ...finalTools]); - + const dedupedTools = filterUniquePlugins([...(userPlugins ?? []), ...finalTools]); res.status(200).json(dedupedTools); } catch (error) { logger.error('[getAvailableTools]', error); diff --git a/api/server/controllers/PluginController.spec.js b/api/server/controllers/PluginController.spec.js index 89dfd72e4..07dd02b1d 100644 --- a/api/server/controllers/PluginController.spec.js +++ b/api/server/controllers/PluginController.spec.js @@ -13,15 +13,18 @@ jest.mock('@librechat/data-schemas', () => ({ jest.mock('~/server/services/Config', () => ({ getCustomConfig: jest.fn(), getCachedTools: jest.fn(), + setCachedTools: jest.fn(), + mergeUserTools: jest.fn(), })); jest.mock('~/server/services/ToolService', () => ({ getToolkitKey: jest.fn(), + loadAndFormatTools: jest.fn(), })); jest.mock('~/config', () => ({ getMCPManager: jest.fn(() => ({ - loadManifestTools: jest.fn().mockResolvedValue([]), + loadAllManifestTools: jest.fn().mockResolvedValue([]), })), getFlowStateManager: jest.fn(), })); @@ -50,6 +53,7 @@ const { convertMCPToolsToPlugins, getToolkitKey, } = require('@librechat/api'); +const { loadAndFormatTools } = require('~/server/services/ToolService'); describe('PluginController', () => { let mockReq, mockRes, mockCache; @@ -310,7 +314,7 @@ describe('PluginController', () => { // Mock the MCP manager to return tools const mockMCPManager = { - loadManifestTools: jest.fn().mockResolvedValue(mcpManagerTools), + loadAllManifestTools: jest.fn().mockResolvedValue(mcpManagerTools), }; require('~/config').getMCPManager.mockReturnValue(mockMCPManager); @@ -379,10 +383,8 @@ describe('PluginController', () => { await getAvailableTools(mockReq, mockRes); - expect(convertMCPToolsToPlugins).toHaveBeenCalledWith({ - functionTools: null, - customConfig: null, - }); + // When cachedUserTools is null, convertMCPToolsToPlugins is not called + expect(convertMCPToolsToPlugins).not.toHaveBeenCalled(); }); it('should handle when getCachedTools returns undefined', async () => { @@ -399,10 +401,8 @@ describe('PluginController', () => { await getAvailableTools(mockReq, mockRes); - expect(convertMCPToolsToPlugins).toHaveBeenCalledWith({ - functionTools: undefined, - customConfig: null, - }); + // When cachedUserTools is undefined, convertMCPToolsToPlugins is not called + expect(convertMCPToolsToPlugins).not.toHaveBeenCalled(); }); it('should handle cachedToolsArray and userPlugins both being defined', async () => { @@ -500,6 +500,15 @@ describe('PluginController', () => { toolkit: true, }; + // Ensure req.app.locals is properly mocked + mockReq.app = { + locals: { + filteredTools: [], + includedTools: [], + paths: { structuredTools: '/mock/path' }, + }, + }; + mockCache.get.mockResolvedValue(null); getCachedTools.mockResolvedValue({}); convertMCPToolsToPlugins.mockReturnValue([]); @@ -508,6 +517,9 @@ describe('PluginController', () => { getToolkitKey.mockReturnValue(undefined); getCustomConfig.mockResolvedValue(null); + // Mock loadAndFormatTools to return an empty object when toolDefinitions is null + loadAndFormatTools.mockReturnValue({}); + // Mock getCachedTools second call to return null getCachedTools.mockResolvedValueOnce({}).mockResolvedValueOnce(null); diff --git a/api/server/routes/__tests__/mcp.spec.js b/api/server/routes/__tests__/mcp.spec.js index 243711d76..25d17c51d 100644 --- a/api/server/routes/__tests__/mcp.spec.js +++ b/api/server/routes/__tests__/mcp.spec.js @@ -1,9 +1,10 @@ -const { MongoMemoryServer } = require('mongodb-memory-server'); +const express = require('express'); const request = require('supertest'); const mongoose = require('mongoose'); -const express = require('express'); +const { MongoMemoryServer } = require('mongodb-memory-server'); jest.mock('@librechat/api', () => ({ + ...jest.requireActual('@librechat/api'), MCPOAuthHandler: { initiateOAuthFlow: jest.fn(), getFlowState: jest.fn(), @@ -44,6 +45,10 @@ jest.mock('~/server/services/Config', () => ({ loadCustomConfig: jest.fn(), })); +jest.mock('~/server/services/Config/mcpToolsCache', () => ({ + updateMCPUserTools: jest.fn(), +})); + jest.mock('~/server/services/MCP', () => ({ getMCPSetupData: jest.fn(), getServerConnectionStatus: jest.fn(), @@ -759,8 +764,10 @@ describe('MCP Routes', () => { require('~/cache').getLogStores.mockReturnValue({}); const { getCachedTools, setCachedTools } = require('~/server/services/Config'); + const { updateMCPUserTools } = require('~/server/services/Config/mcpToolsCache'); getCachedTools.mockResolvedValue({}); setCachedTools.mockResolvedValue(); + updateMCPUserTools.mockResolvedValue(); const response = await request(app).post('/api/mcp/test-server/reinitialize'); @@ -776,7 +783,14 @@ describe('MCP Routes', () => { 'test-user-id', 'test-server', ); - expect(setCachedTools).toHaveBeenCalled(); + expect(updateMCPUserTools).toHaveBeenCalledWith({ + userId: 'test-user-id', + serverName: 'test-server', + tools: [ + { name: 'tool1', description: 'Test tool 1', inputSchema: { type: 'object' } }, + { name: 'tool2', description: 'Test tool 2', inputSchema: { type: 'object' } }, + ], + }); }); it('should handle server with custom user variables', async () => { @@ -803,8 +817,10 @@ describe('MCP Routes', () => { ); const { getCachedTools, setCachedTools } = require('~/server/services/Config'); + const { updateMCPUserTools } = require('~/server/services/Config/mcpToolsCache'); getCachedTools.mockResolvedValue({}); setCachedTools.mockResolvedValue(); + updateMCPUserTools.mockResolvedValue(); const response = await request(app).post('/api/mcp/test-server/reinitialize'); diff --git a/api/server/routes/mcp.js b/api/server/routes/mcp.js index 5931ca02f..0a068cd8f 100644 --- a/api/server/routes/mcp.js +++ b/api/server/routes/mcp.js @@ -3,7 +3,7 @@ const { MCPOAuthHandler } = require('@librechat/api'); const { Router } = require('express'); const { getMCPSetupData, getServerConnectionStatus } = require('~/server/services/MCP'); const { findToken, updateToken, createToken, deleteTokens } = require('~/models'); -const { setCachedTools, getCachedTools } = require('~/server/services/Config'); +const { updateMCPUserTools } = require('~/server/services/Config/mcpToolsCache'); const { getUserPluginAuthValue } = require('~/server/services/PluginService'); const { CacheKeys, Constants } = require('librechat-data-provider'); const { getMCPManager, getFlowStateManager } = require('~/config'); @@ -142,33 +142,12 @@ router.get('/:serverName/oauth/callback', async (req, res) => { `[MCP OAuth] Successfully reconnected ${serverName} for user ${flowState.userId}`, ); - const userTools = (await getCachedTools({ userId: flowState.userId })) || {}; - - const mcpDelimiter = Constants.mcp_delimiter; - for (const key of Object.keys(userTools)) { - if (key.endsWith(`${mcpDelimiter}${serverName}`)) { - delete userTools[key]; - } - } - const tools = await userConnection.fetchTools(); - for (const tool of tools) { - const name = `${tool.name}${Constants.mcp_delimiter}${serverName}`; - userTools[name] = { - type: 'function', - ['function']: { - name, - description: tool.description, - parameters: tool.inputSchema, - }, - }; - } - - await setCachedTools(userTools, { userId: flowState.userId }); - - logger.debug( - `[MCP OAuth] Cached ${tools.length} tools for ${serverName} user ${flowState.userId}`, - ); + await updateMCPUserTools({ + userId: flowState.userId, + serverName, + tools, + }); } else { logger.debug(`[MCP OAuth] System-level OAuth completed for ${serverName}`); } @@ -396,29 +375,12 @@ router.post('/:serverName/reinitialize', requireJwtAuth, async (req, res) => { } if (userConnection && !oauthRequired) { - const userTools = (await getCachedTools({ userId: user.id })) || {}; - - const mcpDelimiter = Constants.mcp_delimiter; - for (const key of Object.keys(userTools)) { - if (key.endsWith(`${mcpDelimiter}${serverName}`)) { - delete userTools[key]; - } - } - const tools = await userConnection.fetchTools(); - for (const tool of tools) { - const name = `${tool.name}${Constants.mcp_delimiter}${serverName}`; - userTools[name] = { - type: 'function', - ['function']: { - name, - description: tool.description, - parameters: tool.inputSchema, - }, - }; - } - - await setCachedTools(userTools, { userId: user.id }); + await updateMCPUserTools({ + userId: user.id, + serverName, + tools, + }); } logger.debug( diff --git a/api/server/services/Config/index.js b/api/server/services/Config/index.js index ad25e5799..aafb2dcb4 100644 --- a/api/server/services/Config/index.js +++ b/api/server/services/Config/index.js @@ -1,6 +1,7 @@ const { config } = require('./EndpointService'); const getCachedTools = require('./getCachedTools'); const getCustomConfig = require('./getCustomConfig'); +const mcpToolsCache = require('./mcpToolsCache'); const loadCustomConfig = require('./loadCustomConfig'); const loadConfigModels = require('./loadConfigModels'); const loadDefaultModels = require('./loadDefaultModels'); @@ -17,5 +18,6 @@ module.exports = { loadAsyncEndpoints, ...getCachedTools, ...getCustomConfig, + ...mcpToolsCache, ...getEndpointsConfig, }; diff --git a/api/server/services/Config/mcpToolsCache.js b/api/server/services/Config/mcpToolsCache.js new file mode 100644 index 000000000..bd21807cd --- /dev/null +++ b/api/server/services/Config/mcpToolsCache.js @@ -0,0 +1,142 @@ +const { logger } = require('@librechat/data-schemas'); +const { CacheKeys, Constants } = require('librechat-data-provider'); +const { getCachedTools, setCachedTools } = require('./getCachedTools'); +const { getLogStores } = require('~/cache'); + +/** + * Updates MCP tools in the cache for a specific server and user + * @param {Object} params - Parameters for updating MCP tools + * @param {string} params.userId - User ID + * @param {string} params.serverName - MCP server name + * @param {Array} params.tools - Array of tool objects from MCP server + * @returns {Promise} + */ +async function updateMCPUserTools({ userId, serverName, tools }) { + try { + const userTools = await getCachedTools({ userId }); + + const mcpDelimiter = Constants.mcp_delimiter; + for (const key of Object.keys(userTools)) { + if (key.endsWith(`${mcpDelimiter}${serverName}`)) { + delete userTools[key]; + } + } + + for (const tool of tools) { + const name = `${tool.name}${Constants.mcp_delimiter}${serverName}`; + userTools[name] = { + type: 'function', + ['function']: { + name, + description: tool.description, + parameters: tool.inputSchema, + }, + }; + } + + await setCachedTools(userTools, { userId }); + + const cache = getLogStores(CacheKeys.CONFIG_STORE); + await cache.delete(CacheKeys.TOOLS); + logger.debug(`[MCP Cache] Updated ${tools.length} tools for ${serverName} user ${userId}`); + } catch (error) { + logger.error(`[MCP Cache] Failed to update tools for ${serverName}:`, error); + throw error; + } +} + +/** + * Merges app-level tools with global tools + * @param {import('@librechat/api').LCAvailableTools} appTools + * @returns {Promise} + */ +async function mergeAppTools(appTools) { + try { + const count = Object.keys(appTools).length; + if (!count) { + return; + } + const cachedTools = await getCachedTools({ includeGlobal: true }); + const mergedTools = { ...cachedTools, ...appTools }; + await setCachedTools(mergedTools, { isGlobal: true }); + const cache = getLogStores(CacheKeys.CONFIG_STORE); + await cache.delete(CacheKeys.TOOLS); + logger.debug(`Merged ${count} app-level tools`); + } catch (error) { + logger.error('Failed to merge app-level tools:', error); + throw error; + } +} + +/** + * Merges user-level tools with global tools + * @param {object} params + * @param {string} params.userId + * @param {Record} params.cachedUserTools + * @param {import('@librechat/api').LCAvailableTools} params.userTools + * @returns {Promise} + */ +async function mergeUserTools({ userId, cachedUserTools, userTools }) { + try { + if (!userId) { + return; + } + const count = Object.keys(userTools).length; + if (!count) { + return; + } + const cachedTools = cachedUserTools ?? (await getCachedTools({ userId })); + const mergedTools = { ...cachedTools, ...userTools }; + await setCachedTools(mergedTools, { userId }); + const cache = getLogStores(CacheKeys.CONFIG_STORE); + await cache.delete(CacheKeys.TOOLS); + logger.debug(`Merged ${count} user-level tools`); + } catch (error) { + logger.error('Failed to merge user-level tools:', error); + throw error; + } +} + +/** + * Clears all MCP tools for a specific server + * @param {Object} params - Parameters for clearing MCP tools + * @param {string} [params.userId] - User ID (if clearing user-specific tools) + * @param {string} params.serverName - MCP server name + * @returns {Promise} + */ +async function clearMCPServerTools({ userId, serverName }) { + try { + const tools = await getCachedTools({ userId, includeGlobal: !userId }); + + // Remove all tools for this server + const mcpDelimiter = Constants.mcp_delimiter; + let removedCount = 0; + for (const key of Object.keys(tools)) { + if (key.endsWith(`${mcpDelimiter}${serverName}`)) { + delete tools[key]; + removedCount++; + } + } + + if (removedCount > 0) { + await setCachedTools(tools, userId ? { userId } : { isGlobal: true }); + + const cache = getLogStores(CacheKeys.CONFIG_STORE); + await cache.delete(CacheKeys.TOOLS); + + logger.debug( + `[MCP Cache] Removed ${removedCount} tools for ${serverName}${userId ? ` user ${userId}` : ' (global)'}`, + ); + } + } catch (error) { + logger.error(`[MCP Cache] Failed to clear tools for ${serverName}:`, error); + throw error; + } +} + +module.exports = { + mergeAppTools, + mergeUserTools, + updateMCPUserTools, + clearMCPServerTools, +}; diff --git a/api/server/services/initializeMCPs.js b/api/server/services/initializeMCPs.js index 0ca402766..696373384 100644 --- a/api/server/services/initializeMCPs.js +++ b/api/server/services/initializeMCPs.js @@ -1,8 +1,6 @@ const { logger } = require('@librechat/data-schemas'); -const { getCachedTools, setCachedTools } = require('./Config'); -const { CacheKeys } = require('librechat-data-provider'); const { createMCPManager } = require('~/config'); -const { getLogStores } = require('~/cache'); +const { mergeAppTools } = require('./Config'); /** * Initialize MCP servers @@ -18,21 +16,12 @@ async function initializeMCPs(app) { try { delete app.locals.mcpConfig; - const cachedTools = await getCachedTools(); + const mcpTools = mcpManager.getAppToolFunctions() || {}; + await mergeAppTools(mcpTools); - if (!cachedTools) { - logger.warn('No available tools found in cache during MCP initialization'); - return; - } - - const mcpTools = mcpManager.getAppToolFunctions() ?? {}; - await setCachedTools({ ...cachedTools, ...mcpTools }, { isGlobal: true }); - - const cache = getLogStores(CacheKeys.CONFIG_STORE); - await cache.delete(CacheKeys.TOOLS); - logger.debug('Cleared tools array cache after MCP initialization'); - - logger.info('MCP servers initialized successfully'); + logger.info( + `MCP servers initialized successfully. Added ${Object.keys(mcpTools).length} MCP tools.`, + ); } catch (error) { logger.error('Failed to initialize MCP servers:', error); } diff --git a/packages/api/src/mcp/MCPManager.ts b/packages/api/src/mcp/MCPManager.ts index 41b4f1a0f..b561e71fd 100644 --- a/packages/api/src/mcp/MCPManager.ts +++ b/packages/api/src/mcp/MCPManager.ts @@ -101,34 +101,36 @@ ${formattedInstructions} Please follow these instructions when using tools from the respective MCP servers.`; } - /** Loads tools from all app-level connections into the manifest. */ - public async loadManifestTools({ - serverToolsCallback, - getServerTools, - }: { - flowManager: FlowStateManager; - serverToolsCallback?: (serverName: string, tools: t.LCManifestTool[]) => Promise; - getServerTools?: (serverName: string) => Promise; - }): Promise { - const mcpTools: t.LCManifestTool[] = []; + private async loadAppManifestTools(): Promise { const connections = await this.appConnections!.getAll(); + return await this.loadManifestTools(connections); + } + + private async loadUserManifestTools(userId: string): Promise { + const connections = this.getUserConnections(userId); + return await this.loadManifestTools(connections); + } + + public async loadAllManifestTools(userId: string): Promise { + const appTools = await this.loadAppManifestTools(); + const userTools = await this.loadUserManifestTools(userId); + return [...appTools, ...userTools]; + } + + /** Loads tools from all app-level connections into the manifest. */ + private async loadManifestTools( + connections?: Map | null, + ): Promise { + const mcpTools: t.LCManifestTool[] = []; + if (!connections || connections.size === 0) { + return mcpTools; + } for (const [serverName, connection] of connections.entries()) { try { if (!(await connection.isConnected())) { logger.warn( `[MCP][${serverName}] Connection not available for ${serverName} manifest tools.`, ); - if (typeof getServerTools !== 'function') { - logger.warn( - `[MCP][${serverName}] No \`getServerTools\` function provided, skipping tool loading.`, - ); - continue; - } - const serverTools = await getServerTools(serverName); - if (serverTools && serverTools.length > 0) { - logger.info(`[MCP][${serverName}] Loaded tools from cache for manifest`); - mcpTools.push(...serverTools); - } continue; } @@ -157,9 +159,6 @@ Please follow these instructions when using tools from the respective MCP server mcpTools.push(manifestTool); serverTools.push(manifestTool); } - if (typeof serverToolsCallback === 'function') { - await serverToolsCallback(serverName, serverTools); - } } catch (error) { logger.error(`[MCP][${serverName}] Error fetching tools for manifest:`, error); } diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts index bb333d736..0be8ce3ff 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -1242,10 +1242,6 @@ export enum CacheKeys { * Key for in-progress flow states. */ FLOWS = 'FLOWS', - /** - * Key for individual MCP Tool Manifests. - */ - MCP_TOOLS = 'MCP_TOOLS', /** * Key for pending chat requests (concurrency check) */