mirror of
https://github.com/danny-avila/LibreChat.git
synced 2025-12-17 08:50:15 +01:00
🛠️ refactor: Consolidate MCP Tool Caching (#9172)
* 🛠️ refactor: Consolidate MCP Tool Caching * 🐍 fix: Correctly mock and utilize updateMCPUserTools in MCP route tests
This commit is contained in:
parent
5a14ee9c6a
commit
da4aa37493
10 changed files with 293 additions and 191 deletions
1
api/cache/getLogStores.js
vendored
1
api/cache/getLogStores.js
vendored
|
|
@ -31,7 +31,6 @@ const namespaces = {
|
||||||
[CacheKeys.SAML_SESSION]: sessionCache(CacheKeys.SAML_SESSION),
|
[CacheKeys.SAML_SESSION]: sessionCache(CacheKeys.SAML_SESSION),
|
||||||
|
|
||||||
[CacheKeys.ROLES]: standardCache(CacheKeys.ROLES),
|
[CacheKeys.ROLES]: standardCache(CacheKeys.ROLES),
|
||||||
[CacheKeys.MCP_TOOLS]: standardCache(CacheKeys.MCP_TOOLS),
|
|
||||||
[CacheKeys.CONFIG_STORE]: standardCache(CacheKeys.CONFIG_STORE),
|
[CacheKeys.CONFIG_STORE]: standardCache(CacheKeys.CONFIG_STORE),
|
||||||
[CacheKeys.STATIC_CONFIG]: standardCache(CacheKeys.STATIC_CONFIG),
|
[CacheKeys.STATIC_CONFIG]: standardCache(CacheKeys.STATIC_CONFIG),
|
||||||
[CacheKeys.PENDING_REQ]: standardCache(CacheKeys.PENDING_REQ),
|
[CacheKeys.PENDING_REQ]: standardCache(CacheKeys.PENDING_REQ),
|
||||||
|
|
|
||||||
|
|
@ -6,9 +6,15 @@ const {
|
||||||
filterUniquePlugins,
|
filterUniquePlugins,
|
||||||
convertMCPToolsToPlugins,
|
convertMCPToolsToPlugins,
|
||||||
} = require('@librechat/api');
|
} = 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 { availableTools, toolkits } = require('~/app/clients/tools');
|
||||||
const { getMCPManager, getFlowStateManager } = require('~/config');
|
const { getMCPManager } = require('~/config');
|
||||||
const { getLogStores } = require('~/cache');
|
const { getLogStores } = require('~/cache');
|
||||||
|
|
||||||
const getAvailablePluginsController = async (req, res) => {
|
const getAvailablePluginsController = async (req, res) => {
|
||||||
|
|
@ -22,6 +28,7 @@ const getAvailablePluginsController = async (req, res) => {
|
||||||
|
|
||||||
/** @type {{ filteredTools: string[], includedTools: string[] }} */
|
/** @type {{ filteredTools: string[], includedTools: string[] }} */
|
||||||
const { filteredTools = [], includedTools = [] } = req.app.locals;
|
const { filteredTools = [], includedTools = [] } = req.app.locals;
|
||||||
|
/** @type {import('@librechat/api').LCManifestTool[]} */
|
||||||
const pluginManifest = availableTools;
|
const pluginManifest = availableTools;
|
||||||
|
|
||||||
const uniquePlugins = filterUniquePlugins(pluginManifest);
|
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<TPlugin[] | null>}
|
|
||||||
*/
|
|
||||||
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.
|
* 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) => {
|
const getAvailableTools = async (req, res) => {
|
||||||
try {
|
try {
|
||||||
const userId = req.user?.id;
|
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 customConfig = await getCustomConfig();
|
||||||
const cache = getLogStores(CacheKeys.CONFIG_STORE);
|
const cache = getLogStores(CacheKeys.CONFIG_STORE);
|
||||||
const cachedToolsArray = await cache.get(CacheKeys.TOOLS);
|
const cachedToolsArray = await cache.get(CacheKeys.TOOLS);
|
||||||
const cachedUserTools = await getCachedTools({ userId });
|
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) {
|
if (cachedToolsArray != null && userPlugins != null) {
|
||||||
const dedupedTools = filterUniquePlugins([...userPlugins, ...cachedToolsArray]);
|
const dedupedTools = filterUniquePlugins([...userPlugins, ...cachedToolsArray]);
|
||||||
|
|
@ -113,31 +88,47 @@ const getAvailableTools = async (req, res) => {
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/** @type {Record<string, FunctionTool> | 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;
|
let pluginManifest = availableTools;
|
||||||
if (customConfig?.mcpServers != null) {
|
if (customConfig?.mcpServers != null) {
|
||||||
try {
|
try {
|
||||||
const mcpManager = getMCPManager();
|
const mcpManager = getMCPManager();
|
||||||
const flowsCache = getLogStores(CacheKeys.FLOWS);
|
const mcpTools = await mcpManager.loadAllManifestTools(userId);
|
||||||
const flowManager = flowsCache ? getFlowStateManager(flowsCache) : null;
|
const mcpToolsRecord = mcpTools.reduce((acc, tool) => {
|
||||||
const serverToolsCallback = createServerToolsCallback();
|
pluginManifest.push(tool);
|
||||||
const getServerTools = createGetServerTools();
|
acc[tool.pluginKey] = tool;
|
||||||
const mcpTools = await mcpManager.loadManifestTools({
|
if (!toolDefinitions[tool.pluginKey]) {
|
||||||
flowManager,
|
toolDefinitions[tool.pluginKey] = tool;
|
||||||
serverToolsCallback,
|
}
|
||||||
getServerTools,
|
return acc;
|
||||||
});
|
}, prelimCachedTools ?? {});
|
||||||
pluginManifest = [...mcpTools, ...pluginManifest];
|
await mergeUserTools({ userId, cachedUserTools, userTools: mcpToolsRecord });
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
logger.error(
|
logger.error(
|
||||||
'[getAvailableTools] Error loading MCP Tools, servers may still be initializing:',
|
'[getAvailableTools] Error loading MCP Tools, servers may still be initializing:',
|
||||||
error,
|
error,
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
} else if (prelimCachedTools != null) {
|
||||||
|
await setCachedTools(prelimCachedTools, { isGlobal: true });
|
||||||
}
|
}
|
||||||
|
|
||||||
/** @type {TPlugin[]} */
|
/** @type {TPlugin[]} Deduplicate and authenticate plugins */
|
||||||
const uniquePlugins = filterUniquePlugins(pluginManifest);
|
const uniquePlugins = filterUniquePlugins(pluginManifest);
|
||||||
|
|
||||||
const authenticatedPlugins = uniquePlugins.map((plugin) => {
|
const authenticatedPlugins = uniquePlugins.map((plugin) => {
|
||||||
if (checkPluginAuth(plugin)) {
|
if (checkPluginAuth(plugin)) {
|
||||||
return { ...plugin, authenticated: true };
|
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 = [];
|
const toolsOutput = [];
|
||||||
for (const plugin of authenticatedPlugins) {
|
for (const plugin of authenticatedPlugins) {
|
||||||
const isToolDefined = toolDefinitions[plugin.pluginKey] !== undefined;
|
const isToolDefined = toolDefinitions[plugin.pluginKey] !== undefined;
|
||||||
|
|
@ -163,41 +153,36 @@ const getAvailableTools = async (req, res) => {
|
||||||
|
|
||||||
const toolToAdd = { ...plugin };
|
const toolToAdd = { ...plugin };
|
||||||
|
|
||||||
if (!plugin.pluginKey.includes(Constants.mcp_delimiter)) {
|
if (plugin.pluginKey.includes(Constants.mcp_delimiter)) {
|
||||||
toolsOutput.push(toolToAdd);
|
const parts = plugin.pluginKey.split(Constants.mcp_delimiter);
|
||||||
continue;
|
const serverName = parts[parts.length - 1];
|
||||||
}
|
const serverConfig = customConfig?.mcpServers?.[serverName];
|
||||||
|
|
||||||
const parts = plugin.pluginKey.split(Constants.mcp_delimiter);
|
if (serverConfig?.customUserVars) {
|
||||||
const serverName = parts[parts.length - 1];
|
const customVarKeys = Object.keys(serverConfig.customUserVars);
|
||||||
const serverConfig = customConfig?.mcpServers?.[serverName];
|
if (customVarKeys.length === 0) {
|
||||||
|
toolToAdd.authConfig = [];
|
||||||
if (!serverConfig?.customUserVars) {
|
toolToAdd.authenticated = true;
|
||||||
toolsOutput.push(toolToAdd);
|
} else {
|
||||||
continue;
|
toolToAdd.authConfig = Object.entries(serverConfig.customUserVars).map(
|
||||||
}
|
([key, value]) => ({
|
||||||
|
authField: key,
|
||||||
const customVarKeys = Object.keys(serverConfig.customUserVars);
|
label: value.title || key,
|
||||||
|
description: value.description || '',
|
||||||
if (customVarKeys.length === 0) {
|
}),
|
||||||
toolToAdd.authConfig = [];
|
);
|
||||||
toolToAdd.authenticated = true;
|
toolToAdd.authenticated = false;
|
||||||
} 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);
|
toolsOutput.push(toolToAdd);
|
||||||
}
|
}
|
||||||
|
|
||||||
const finalTools = filterUniquePlugins(toolsOutput);
|
const finalTools = filterUniquePlugins(toolsOutput);
|
||||||
await cache.set(CacheKeys.TOOLS, finalTools);
|
await cache.set(CacheKeys.TOOLS, finalTools);
|
||||||
|
|
||||||
const dedupedTools = filterUniquePlugins([...userPlugins, ...finalTools]);
|
const dedupedTools = filterUniquePlugins([...(userPlugins ?? []), ...finalTools]);
|
||||||
|
|
||||||
res.status(200).json(dedupedTools);
|
res.status(200).json(dedupedTools);
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
logger.error('[getAvailableTools]', error);
|
logger.error('[getAvailableTools]', error);
|
||||||
|
|
|
||||||
|
|
@ -13,15 +13,18 @@ jest.mock('@librechat/data-schemas', () => ({
|
||||||
jest.mock('~/server/services/Config', () => ({
|
jest.mock('~/server/services/Config', () => ({
|
||||||
getCustomConfig: jest.fn(),
|
getCustomConfig: jest.fn(),
|
||||||
getCachedTools: jest.fn(),
|
getCachedTools: jest.fn(),
|
||||||
|
setCachedTools: jest.fn(),
|
||||||
|
mergeUserTools: jest.fn(),
|
||||||
}));
|
}));
|
||||||
|
|
||||||
jest.mock('~/server/services/ToolService', () => ({
|
jest.mock('~/server/services/ToolService', () => ({
|
||||||
getToolkitKey: jest.fn(),
|
getToolkitKey: jest.fn(),
|
||||||
|
loadAndFormatTools: jest.fn(),
|
||||||
}));
|
}));
|
||||||
|
|
||||||
jest.mock('~/config', () => ({
|
jest.mock('~/config', () => ({
|
||||||
getMCPManager: jest.fn(() => ({
|
getMCPManager: jest.fn(() => ({
|
||||||
loadManifestTools: jest.fn().mockResolvedValue([]),
|
loadAllManifestTools: jest.fn().mockResolvedValue([]),
|
||||||
})),
|
})),
|
||||||
getFlowStateManager: jest.fn(),
|
getFlowStateManager: jest.fn(),
|
||||||
}));
|
}));
|
||||||
|
|
@ -50,6 +53,7 @@ const {
|
||||||
convertMCPToolsToPlugins,
|
convertMCPToolsToPlugins,
|
||||||
getToolkitKey,
|
getToolkitKey,
|
||||||
} = require('@librechat/api');
|
} = require('@librechat/api');
|
||||||
|
const { loadAndFormatTools } = require('~/server/services/ToolService');
|
||||||
|
|
||||||
describe('PluginController', () => {
|
describe('PluginController', () => {
|
||||||
let mockReq, mockRes, mockCache;
|
let mockReq, mockRes, mockCache;
|
||||||
|
|
@ -310,7 +314,7 @@ describe('PluginController', () => {
|
||||||
|
|
||||||
// Mock the MCP manager to return tools
|
// Mock the MCP manager to return tools
|
||||||
const mockMCPManager = {
|
const mockMCPManager = {
|
||||||
loadManifestTools: jest.fn().mockResolvedValue(mcpManagerTools),
|
loadAllManifestTools: jest.fn().mockResolvedValue(mcpManagerTools),
|
||||||
};
|
};
|
||||||
require('~/config').getMCPManager.mockReturnValue(mockMCPManager);
|
require('~/config').getMCPManager.mockReturnValue(mockMCPManager);
|
||||||
|
|
||||||
|
|
@ -379,10 +383,8 @@ describe('PluginController', () => {
|
||||||
|
|
||||||
await getAvailableTools(mockReq, mockRes);
|
await getAvailableTools(mockReq, mockRes);
|
||||||
|
|
||||||
expect(convertMCPToolsToPlugins).toHaveBeenCalledWith({
|
// When cachedUserTools is null, convertMCPToolsToPlugins is not called
|
||||||
functionTools: null,
|
expect(convertMCPToolsToPlugins).not.toHaveBeenCalled();
|
||||||
customConfig: null,
|
|
||||||
});
|
|
||||||
});
|
});
|
||||||
|
|
||||||
it('should handle when getCachedTools returns undefined', async () => {
|
it('should handle when getCachedTools returns undefined', async () => {
|
||||||
|
|
@ -399,10 +401,8 @@ describe('PluginController', () => {
|
||||||
|
|
||||||
await getAvailableTools(mockReq, mockRes);
|
await getAvailableTools(mockReq, mockRes);
|
||||||
|
|
||||||
expect(convertMCPToolsToPlugins).toHaveBeenCalledWith({
|
// When cachedUserTools is undefined, convertMCPToolsToPlugins is not called
|
||||||
functionTools: undefined,
|
expect(convertMCPToolsToPlugins).not.toHaveBeenCalled();
|
||||||
customConfig: null,
|
|
||||||
});
|
|
||||||
});
|
});
|
||||||
|
|
||||||
it('should handle cachedToolsArray and userPlugins both being defined', async () => {
|
it('should handle cachedToolsArray and userPlugins both being defined', async () => {
|
||||||
|
|
@ -500,6 +500,15 @@ describe('PluginController', () => {
|
||||||
toolkit: true,
|
toolkit: true,
|
||||||
};
|
};
|
||||||
|
|
||||||
|
// Ensure req.app.locals is properly mocked
|
||||||
|
mockReq.app = {
|
||||||
|
locals: {
|
||||||
|
filteredTools: [],
|
||||||
|
includedTools: [],
|
||||||
|
paths: { structuredTools: '/mock/path' },
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
mockCache.get.mockResolvedValue(null);
|
mockCache.get.mockResolvedValue(null);
|
||||||
getCachedTools.mockResolvedValue({});
|
getCachedTools.mockResolvedValue({});
|
||||||
convertMCPToolsToPlugins.mockReturnValue([]);
|
convertMCPToolsToPlugins.mockReturnValue([]);
|
||||||
|
|
@ -508,6 +517,9 @@ describe('PluginController', () => {
|
||||||
getToolkitKey.mockReturnValue(undefined);
|
getToolkitKey.mockReturnValue(undefined);
|
||||||
getCustomConfig.mockResolvedValue(null);
|
getCustomConfig.mockResolvedValue(null);
|
||||||
|
|
||||||
|
// Mock loadAndFormatTools to return an empty object when toolDefinitions is null
|
||||||
|
loadAndFormatTools.mockReturnValue({});
|
||||||
|
|
||||||
// Mock getCachedTools second call to return null
|
// Mock getCachedTools second call to return null
|
||||||
getCachedTools.mockResolvedValueOnce({}).mockResolvedValueOnce(null);
|
getCachedTools.mockResolvedValueOnce({}).mockResolvedValueOnce(null);
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,9 +1,10 @@
|
||||||
const { MongoMemoryServer } = require('mongodb-memory-server');
|
const express = require('express');
|
||||||
const request = require('supertest');
|
const request = require('supertest');
|
||||||
const mongoose = require('mongoose');
|
const mongoose = require('mongoose');
|
||||||
const express = require('express');
|
const { MongoMemoryServer } = require('mongodb-memory-server');
|
||||||
|
|
||||||
jest.mock('@librechat/api', () => ({
|
jest.mock('@librechat/api', () => ({
|
||||||
|
...jest.requireActual('@librechat/api'),
|
||||||
MCPOAuthHandler: {
|
MCPOAuthHandler: {
|
||||||
initiateOAuthFlow: jest.fn(),
|
initiateOAuthFlow: jest.fn(),
|
||||||
getFlowState: jest.fn(),
|
getFlowState: jest.fn(),
|
||||||
|
|
@ -44,6 +45,10 @@ jest.mock('~/server/services/Config', () => ({
|
||||||
loadCustomConfig: jest.fn(),
|
loadCustomConfig: jest.fn(),
|
||||||
}));
|
}));
|
||||||
|
|
||||||
|
jest.mock('~/server/services/Config/mcpToolsCache', () => ({
|
||||||
|
updateMCPUserTools: jest.fn(),
|
||||||
|
}));
|
||||||
|
|
||||||
jest.mock('~/server/services/MCP', () => ({
|
jest.mock('~/server/services/MCP', () => ({
|
||||||
getMCPSetupData: jest.fn(),
|
getMCPSetupData: jest.fn(),
|
||||||
getServerConnectionStatus: jest.fn(),
|
getServerConnectionStatus: jest.fn(),
|
||||||
|
|
@ -759,8 +764,10 @@ describe('MCP Routes', () => {
|
||||||
require('~/cache').getLogStores.mockReturnValue({});
|
require('~/cache').getLogStores.mockReturnValue({});
|
||||||
|
|
||||||
const { getCachedTools, setCachedTools } = require('~/server/services/Config');
|
const { getCachedTools, setCachedTools } = require('~/server/services/Config');
|
||||||
|
const { updateMCPUserTools } = require('~/server/services/Config/mcpToolsCache');
|
||||||
getCachedTools.mockResolvedValue({});
|
getCachedTools.mockResolvedValue({});
|
||||||
setCachedTools.mockResolvedValue();
|
setCachedTools.mockResolvedValue();
|
||||||
|
updateMCPUserTools.mockResolvedValue();
|
||||||
|
|
||||||
const response = await request(app).post('/api/mcp/test-server/reinitialize');
|
const response = await request(app).post('/api/mcp/test-server/reinitialize');
|
||||||
|
|
||||||
|
|
@ -776,7 +783,14 @@ describe('MCP Routes', () => {
|
||||||
'test-user-id',
|
'test-user-id',
|
||||||
'test-server',
|
'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 () => {
|
it('should handle server with custom user variables', async () => {
|
||||||
|
|
@ -803,8 +817,10 @@ describe('MCP Routes', () => {
|
||||||
);
|
);
|
||||||
|
|
||||||
const { getCachedTools, setCachedTools } = require('~/server/services/Config');
|
const { getCachedTools, setCachedTools } = require('~/server/services/Config');
|
||||||
|
const { updateMCPUserTools } = require('~/server/services/Config/mcpToolsCache');
|
||||||
getCachedTools.mockResolvedValue({});
|
getCachedTools.mockResolvedValue({});
|
||||||
setCachedTools.mockResolvedValue();
|
setCachedTools.mockResolvedValue();
|
||||||
|
updateMCPUserTools.mockResolvedValue();
|
||||||
|
|
||||||
const response = await request(app).post('/api/mcp/test-server/reinitialize');
|
const response = await request(app).post('/api/mcp/test-server/reinitialize');
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -3,7 +3,7 @@ const { MCPOAuthHandler } = require('@librechat/api');
|
||||||
const { Router } = require('express');
|
const { Router } = require('express');
|
||||||
const { getMCPSetupData, getServerConnectionStatus } = require('~/server/services/MCP');
|
const { getMCPSetupData, getServerConnectionStatus } = require('~/server/services/MCP');
|
||||||
const { findToken, updateToken, createToken, deleteTokens } = require('~/models');
|
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 { getUserPluginAuthValue } = require('~/server/services/PluginService');
|
||||||
const { CacheKeys, Constants } = require('librechat-data-provider');
|
const { CacheKeys, Constants } = require('librechat-data-provider');
|
||||||
const { getMCPManager, getFlowStateManager } = require('~/config');
|
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}`,
|
`[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();
|
const tools = await userConnection.fetchTools();
|
||||||
for (const tool of tools) {
|
await updateMCPUserTools({
|
||||||
const name = `${tool.name}${Constants.mcp_delimiter}${serverName}`;
|
userId: flowState.userId,
|
||||||
userTools[name] = {
|
serverName,
|
||||||
type: 'function',
|
tools,
|
||||||
['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}`,
|
|
||||||
);
|
|
||||||
} else {
|
} else {
|
||||||
logger.debug(`[MCP OAuth] System-level OAuth completed for ${serverName}`);
|
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) {
|
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();
|
const tools = await userConnection.fetchTools();
|
||||||
for (const tool of tools) {
|
await updateMCPUserTools({
|
||||||
const name = `${tool.name}${Constants.mcp_delimiter}${serverName}`;
|
userId: user.id,
|
||||||
userTools[name] = {
|
serverName,
|
||||||
type: 'function',
|
tools,
|
||||||
['function']: {
|
});
|
||||||
name,
|
|
||||||
description: tool.description,
|
|
||||||
parameters: tool.inputSchema,
|
|
||||||
},
|
|
||||||
};
|
|
||||||
}
|
|
||||||
|
|
||||||
await setCachedTools(userTools, { userId: user.id });
|
|
||||||
}
|
}
|
||||||
|
|
||||||
logger.debug(
|
logger.debug(
|
||||||
|
|
|
||||||
|
|
@ -1,6 +1,7 @@
|
||||||
const { config } = require('./EndpointService');
|
const { config } = require('./EndpointService');
|
||||||
const getCachedTools = require('./getCachedTools');
|
const getCachedTools = require('./getCachedTools');
|
||||||
const getCustomConfig = require('./getCustomConfig');
|
const getCustomConfig = require('./getCustomConfig');
|
||||||
|
const mcpToolsCache = require('./mcpToolsCache');
|
||||||
const loadCustomConfig = require('./loadCustomConfig');
|
const loadCustomConfig = require('./loadCustomConfig');
|
||||||
const loadConfigModels = require('./loadConfigModels');
|
const loadConfigModels = require('./loadConfigModels');
|
||||||
const loadDefaultModels = require('./loadDefaultModels');
|
const loadDefaultModels = require('./loadDefaultModels');
|
||||||
|
|
@ -17,5 +18,6 @@ module.exports = {
|
||||||
loadAsyncEndpoints,
|
loadAsyncEndpoints,
|
||||||
...getCachedTools,
|
...getCachedTools,
|
||||||
...getCustomConfig,
|
...getCustomConfig,
|
||||||
|
...mcpToolsCache,
|
||||||
...getEndpointsConfig,
|
...getEndpointsConfig,
|
||||||
};
|
};
|
||||||
|
|
|
||||||
142
api/server/services/Config/mcpToolsCache.js
Normal file
142
api/server/services/Config/mcpToolsCache.js
Normal file
|
|
@ -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<void>}
|
||||||
|
*/
|
||||||
|
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<void>}
|
||||||
|
*/
|
||||||
|
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<string, FunctionTool>} params.cachedUserTools
|
||||||
|
* @param {import('@librechat/api').LCAvailableTools} params.userTools
|
||||||
|
* @returns {Promise<void>}
|
||||||
|
*/
|
||||||
|
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<void>}
|
||||||
|
*/
|
||||||
|
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,
|
||||||
|
};
|
||||||
|
|
@ -1,8 +1,6 @@
|
||||||
const { logger } = require('@librechat/data-schemas');
|
const { logger } = require('@librechat/data-schemas');
|
||||||
const { getCachedTools, setCachedTools } = require('./Config');
|
|
||||||
const { CacheKeys } = require('librechat-data-provider');
|
|
||||||
const { createMCPManager } = require('~/config');
|
const { createMCPManager } = require('~/config');
|
||||||
const { getLogStores } = require('~/cache');
|
const { mergeAppTools } = require('./Config');
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Initialize MCP servers
|
* Initialize MCP servers
|
||||||
|
|
@ -18,21 +16,12 @@ async function initializeMCPs(app) {
|
||||||
|
|
||||||
try {
|
try {
|
||||||
delete app.locals.mcpConfig;
|
delete app.locals.mcpConfig;
|
||||||
const cachedTools = await getCachedTools();
|
const mcpTools = mcpManager.getAppToolFunctions() || {};
|
||||||
|
await mergeAppTools(mcpTools);
|
||||||
|
|
||||||
if (!cachedTools) {
|
logger.info(
|
||||||
logger.warn('No available tools found in cache during MCP initialization');
|
`MCP servers initialized successfully. Added ${Object.keys(mcpTools).length} MCP tools.`,
|
||||||
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');
|
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
logger.error('Failed to initialize MCP servers:', error);
|
logger.error('Failed to initialize MCP servers:', error);
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -101,34 +101,36 @@ ${formattedInstructions}
|
||||||
Please follow these instructions when using tools from the respective MCP servers.`;
|
Please follow these instructions when using tools from the respective MCP servers.`;
|
||||||
}
|
}
|
||||||
|
|
||||||
/** Loads tools from all app-level connections into the manifest. */
|
private async loadAppManifestTools(): Promise<t.LCManifestTool[]> {
|
||||||
public async loadManifestTools({
|
|
||||||
serverToolsCallback,
|
|
||||||
getServerTools,
|
|
||||||
}: {
|
|
||||||
flowManager: FlowStateManager<MCPOAuthTokens | null>;
|
|
||||||
serverToolsCallback?: (serverName: string, tools: t.LCManifestTool[]) => Promise<void>;
|
|
||||||
getServerTools?: (serverName: string) => Promise<t.LCManifestTool[] | undefined>;
|
|
||||||
}): Promise<t.LCToolManifest> {
|
|
||||||
const mcpTools: t.LCManifestTool[] = [];
|
|
||||||
const connections = await this.appConnections!.getAll();
|
const connections = await this.appConnections!.getAll();
|
||||||
|
return await this.loadManifestTools(connections);
|
||||||
|
}
|
||||||
|
|
||||||
|
private async loadUserManifestTools(userId: string): Promise<t.LCManifestTool[]> {
|
||||||
|
const connections = this.getUserConnections(userId);
|
||||||
|
return await this.loadManifestTools(connections);
|
||||||
|
}
|
||||||
|
|
||||||
|
public async loadAllManifestTools(userId: string): Promise<t.LCManifestTool[]> {
|
||||||
|
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<string, MCPConnection> | null,
|
||||||
|
): Promise<t.LCToolManifest> {
|
||||||
|
const mcpTools: t.LCManifestTool[] = [];
|
||||||
|
if (!connections || connections.size === 0) {
|
||||||
|
return mcpTools;
|
||||||
|
}
|
||||||
for (const [serverName, connection] of connections.entries()) {
|
for (const [serverName, connection] of connections.entries()) {
|
||||||
try {
|
try {
|
||||||
if (!(await connection.isConnected())) {
|
if (!(await connection.isConnected())) {
|
||||||
logger.warn(
|
logger.warn(
|
||||||
`[MCP][${serverName}] Connection not available for ${serverName} manifest tools.`,
|
`[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;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -157,9 +159,6 @@ Please follow these instructions when using tools from the respective MCP server
|
||||||
mcpTools.push(manifestTool);
|
mcpTools.push(manifestTool);
|
||||||
serverTools.push(manifestTool);
|
serverTools.push(manifestTool);
|
||||||
}
|
}
|
||||||
if (typeof serverToolsCallback === 'function') {
|
|
||||||
await serverToolsCallback(serverName, serverTools);
|
|
||||||
}
|
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
logger.error(`[MCP][${serverName}] Error fetching tools for manifest:`, error);
|
logger.error(`[MCP][${serverName}] Error fetching tools for manifest:`, error);
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -1242,10 +1242,6 @@ export enum CacheKeys {
|
||||||
* Key for in-progress flow states.
|
* Key for in-progress flow states.
|
||||||
*/
|
*/
|
||||||
FLOWS = 'FLOWS',
|
FLOWS = 'FLOWS',
|
||||||
/**
|
|
||||||
* Key for individual MCP Tool Manifests.
|
|
||||||
*/
|
|
||||||
MCP_TOOLS = 'MCP_TOOLS',
|
|
||||||
/**
|
/**
|
||||||
* Key for pending chat requests (concurrency check)
|
* Key for pending chat requests (concurrency check)
|
||||||
*/
|
*/
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue