Andreas/LibreChat
1
0
Fork 0
mirror of https://github.com/danny-avila/LibreChat.git synced 2025-12-17 00:40:14 +01:00
Code Issues Projects Releases Packages Wiki Activity Actions 6

♻️ refactor: MCPManager for Scalability, Fix App-Level Detection, Add Lazy Connections (#8930)

Browse source 
* feat: MCP Connection management overhaul - Making MCPManager manageable

Refactor the monolithic MCPManager into focused, single-responsibility classes:

• MCPServersRegistry: Server configuration discovery and metadata management
• UserConnectionManager: Manages user-level connections
• ConnectionsRepository: Low-level connection pool with lazy loading
• MCPConnectionFactory: Handles MCP connection creation with OAuth support

New Features:
• Lazy loading of app-level connections for horizontal scaling
• Automatic reconnection for app-level connections
• Enhanced OAuth detection with explicit requiresOAuth flag
• Centralized MCP configuration management

Bug Fixes:
• App-level connection detection in MCPManager.callTool
• MCP Connection Reinitialization route behavior

Optimizations:
• MCPConnection.isConnected() caching to reduce overhead
• Concurrent server metadata retrieval instead of sequential

This refactoring addresses scalability bottlenecks and improves reliability
while maintaining backward compatibility with existing configurations.

* feat: Enabled import order in eslint.

* # Moved tests to __tests__ folder
# added tests for MCPServersRegistry.ts

* # Add unit tests for ConnectionsRepository functionality

* # Add unit tests for MCPConnectionFactory functionality

* # Reorganize MCP connection tests and improve error handling

* # reordering imports

* # Update testPathIgnorePatterns in jest.config.mjs to exclude development TypeScript files

* # removed mcp/manager.ts
This commit is contained in:
Theo N. Truong 2025-08-13 09:45:06 -06:00 • committed by GitHub
parent 9dbf153489
commit 8780a78165
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
32 changed files with 2571 additions and 1468 deletions
Download patch file Download diff file Expand all files Collapse all files

19
api/config/index.js
View file

@ -1,27 +1,13 @@
const { MCPManager, FlowStateManager } = require('@librechat/api');
const { EventSource } = require('eventsource');
const { Time } = require('librechat-data-provider');
const { MCPManager, FlowStateManager } = require('@librechat/api');
const logger = require('./winston');
global.EventSource = EventSource;
/** @type {MCPManager} */
let mcpManager = null;
let flowManager = null;
/**
* @param {string} [userId] - Optional user ID, to avoid disconnecting the current user.
* @returns {MCPManager}
*/
function getMCPManager(userId) {
if (!mcpManager) {
mcpManager = MCPManager.getInstance();
} else {
mcpManager.checkIdleConnections(userId);
}
return mcpManager;
}
/**
* @param {Keyv} flowsCache
* @returns {FlowStateManager}
@ -37,6 +23,7 @@ function getFlowStateManager(flowsCache) {
module.exports = {
logger,
getMCPManager,
createMCPManager: MCPManager.createInstance,
getMCPManager: MCPManager.getInstance,
getFlowStateManager,
};

295
api/server/routes/__tests__/mcp.spec.js
View file

@ -1,7 +1,7 @@
const express = require('express');
const { MongoMemoryServer } = require('mongodb-memory-server');
const request = require('supertest');
const mongoose = require('mongoose');
const { MongoMemoryServer } = require('mongodb-memory-server');
const express = require('express');
jest.mock('@librechat/api', () => ({
MCPOAuthHandler: {
@ -494,12 +494,9 @@ describe('MCP Routes', () => {
});
it('should return 500 when token retrieval throws an unexpected error', async () => {
const mockFlowManager = {
getFlowState: jest.fn().mockRejectedValue(new Error('Database connection failed')),
};
getLogStores.mockReturnValue({});
require('~/config').getFlowStateManager.mockReturnValue(mockFlowManager);
getLogStores.mockImplementation(() => {
throw new Error('Database connection failed');
});
const response = await request(app).get('/api/mcp/oauth/tokens/test-user-id:error-flow');
@ -563,8 +560,8 @@ describe('MCP Routes', () => {
});
describe('POST /oauth/cancel/:serverName', () => {
const { getLogStores } = require('~/cache');
const { MCPOAuthHandler } = require('@librechat/api');
const { getLogStores } = require('~/cache');
it('should cancel OAuth flow successfully', async () => {
const mockFlowManager = {
@ -644,15 +641,15 @@ describe('MCP Routes', () => {
});
describe('POST /:serverName/reinitialize', () => {
const { loadCustomConfig } = require('~/server/services/Config');
const { getUserPluginAuthValue } = require('~/server/services/PluginService');
it('should return 404 when server is not found in configuration', async () => {
loadCustomConfig.mockResolvedValue({
mcpServers: {
'other-server': {},
},
});
const mockMcpManager = {
getRawConfig: jest.fn().mockReturnValue(null),
disconnectUserConnection: jest.fn().mockResolvedValue(),
};
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
require('~/config').getFlowStateManager.mockReturnValue({});
require('~/cache').getLogStores.mockReturnValue({});
const response = await request(app).post('/api/mcp/non-existent-server/reinitialize');
@ -663,16 +660,11 @@ describe('MCP Routes', () => {
});
it('should handle OAuth requirement during reinitialize', async () => {
loadCustomConfig.mockResolvedValue({
mcpServers: {
'oauth-server': {
customUserVars: {},
},
},
});
const mockMcpManager = {
disconnectServer: jest.fn().mockResolvedValue(),
getRawConfig: jest.fn().mockReturnValue({
customUserVars: {},
}),
disconnectUserConnection: jest.fn().mockResolvedValue(),
mcpConfigs: {},
getUserConnection: jest.fn().mockImplementation(async ({ oauthStart }) => {
if (oauthStart) {
@ -690,7 +682,7 @@ describe('MCP Routes', () => {
expect(response.status).toBe(200);
expect(response.body).toEqual({
success: 'https://oauth.example.com/auth',
success: true,
message: "MCP server 'oauth-server' ready for OAuth authentication",
serverName: 'oauth-server',
oauthRequired: true,
@ -699,14 +691,9 @@ describe('MCP Routes', () => {
});
it('should return 500 when reinitialize fails with non-OAuth error', async () => {
loadCustomConfig.mockResolvedValue({
mcpServers: {
'error-server': {},
},
});
const mockMcpManager = {
disconnectServer: jest.fn().mockResolvedValue(),
getRawConfig: jest.fn().mockReturnValue({}),
disconnectUserConnection: jest.fn().mockResolvedValue(),
mcpConfigs: {},
getUserConnection: jest.fn().mockRejectedValue(new Error('Connection failed')),
};
@ -724,7 +711,13 @@ describe('MCP Routes', () => {
});
it('should return 500 when unexpected error occurs', async () => {
loadCustomConfig.mockRejectedValue(new Error('Config loading failed'));
const mockMcpManager = {
getRawConfig: jest.fn().mockImplementation(() => {
throw new Error('Config loading failed');
}),
};
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
const response = await request(app).post('/api/mcp/test-server/reinitialize');
@ -747,29 +740,17 @@ describe('MCP Routes', () => {
expect(response.body).toEqual({ error: 'User not authenticated' });
});
it('should handle errors when fetching custom user variables', async () => {
loadCustomConfig.mockResolvedValue({
mcpServers: {
'test-server': {
customUserVars: {
API_KEY: 'test-key-var',
SECRET_TOKEN: 'test-secret-var',
},
},
},
});
getUserPluginAuthValue
.mockResolvedValueOnce('test-api-key-value')
.mockRejectedValueOnce(new Error('Database error'));
it('should successfully reinitialize server and cache tools', async () => {
const mockUserConnection = {
fetchTools: jest.fn().mockResolvedValue([]),
fetchTools: jest.fn().mockResolvedValue([
{ name: 'tool1', description: 'Test tool 1', inputSchema: { type: 'object' } },
{ name: 'tool2', description: 'Test tool 2', inputSchema: { type: 'object' } },
]),
};
const mockMcpManager = {
disconnectServer: jest.fn().mockResolvedValue(),
mcpConfigs: {},
getRawConfig: jest.fn().mockReturnValue({ endpoint: 'http://test-server.com' }),
disconnectUserConnection: jest.fn().mockResolvedValue(),
getUserConnection: jest.fn().mockResolvedValue(mockUserConnection),
};
@ -784,38 +765,54 @@ describe('MCP Routes', () => {
const response = await request(app).post('/api/mcp/test-server/reinitialize');
expect(response.status).toBe(200);
expect(response.body.success).toBe(true);
expect(response.body).toEqual({
success: true,
message: "MCP server 'test-server' reinitialized successfully",
serverName: 'test-server',
oauthRequired: false,
oauthUrl: null,
});
expect(mockMcpManager.disconnectUserConnection).toHaveBeenCalledWith(
'test-user-id',
'test-server',
);
expect(setCachedTools).toHaveBeenCalled();
});
it('should return failure message when reinitialize completely fails', async () => {
loadCustomConfig.mockResolvedValue({
mcpServers: {
'test-server': {},
},
});
it('should handle server with custom user variables', async () => {
const mockUserConnection = {
fetchTools: jest.fn().mockResolvedValue([]),
};
const mockMcpManager = {
disconnectServer: jest.fn().mockResolvedValue(),
mcpConfigs: {},
getUserConnection: jest.fn().mockResolvedValue(null),
getRawConfig: jest.fn().mockReturnValue({
endpoint: 'http://test-server.com',
customUserVars: {
API_KEY: 'some-env-var',
},
}),
disconnectUserConnection: jest.fn().mockResolvedValue(),
getUserConnection: jest.fn().mockResolvedValue(mockUserConnection),
};
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
require('~/config').getFlowStateManager.mockReturnValue({});
require('~/cache').getLogStores.mockReturnValue({});
require('~/server/services/PluginService').getUserPluginAuthValue.mockResolvedValue(
'api-key-value',
);
const { getCachedTools, setCachedTools } = require('~/server/services/Config');
const { Constants } = require('librechat-data-provider');
getCachedTools.mockResolvedValue({
[`existing-tool${Constants.mcp_delimiter}test-server`]: { type: 'function' },
});
getCachedTools.mockResolvedValue({});
setCachedTools.mockResolvedValue();
const response = await request(app).post('/api/mcp/test-server/reinitialize');
expect(response.status).toBe(200);
expect(response.body.success).toBe(false);
expect(response.body.message).toBe("Failed to reinitialize MCP server 'test-server'");
expect(response.body.success).toBe(true);
expect(
require('~/server/services/PluginService').getUserPluginAuthValue,
).toHaveBeenCalledWith('test-user-id', 'API_KEY', false);
});
});
@ -984,21 +981,19 @@ describe('MCP Routes', () => {
});
describe('GET /:serverName/auth-values', () => {
const { loadCustomConfig } = require('~/server/services/Config');
const { getUserPluginAuthValue } = require('~/server/services/PluginService');
it('should return auth value flags for server', async () => {
loadCustomConfig.mockResolvedValue({
mcpServers: {
'test-server': {
customUserVars: {
API_KEY: 'some-env-var',
SECRET_TOKEN: 'another-env-var',
},
const mockMcpManager = {
getRawConfig: jest.fn().mockReturnValue({
customUserVars: {
API_KEY: 'some-env-var',
SECRET_TOKEN: 'another-env-var',
},
},
});
}),
};
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
getUserPluginAuthValue.mockResolvedValueOnce('some-api-key-value').mockResolvedValueOnce('');
const response = await request(app).get('/api/mcp/test-server/auth-values');
@ -1017,11 +1012,11 @@ describe('MCP Routes', () => {
});
it('should return 404 when server is not found in configuration', async () => {
loadCustomConfig.mockResolvedValue({
mcpServers: {
'other-server': {},
},
});
const mockMcpManager = {
getRawConfig: jest.fn().mockReturnValue(null),
};
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
const response = await request(app).get('/api/mcp/non-existent-server/auth-values');
@ -1032,16 +1027,15 @@ describe('MCP Routes', () => {
});
it('should handle errors when checking auth values', async () => {
loadCustomConfig.mockResolvedValue({
mcpServers: {
'test-server': {
customUserVars: {
API_KEY: 'some-env-var',
},
const mockMcpManager = {
getRawConfig: jest.fn().mockReturnValue({
customUserVars: {
API_KEY: 'some-env-var',
},
},
});
}),
};
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
getUserPluginAuthValue.mockRejectedValue(new Error('Database error'));
const response = await request(app).get('/api/mcp/test-server/auth-values');
@ -1057,7 +1051,13 @@ describe('MCP Routes', () => {
});
it('should return 500 when auth values check throws unexpected error', async () => {
loadCustomConfig.mockRejectedValue(new Error('Config loading failed'));
const mockMcpManager = {
getRawConfig: jest.fn().mockImplementation(() => {
throw new Error('Config loading failed');
}),
};
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
const response = await request(app).get('/api/mcp/test-server/auth-values');
@ -1066,14 +1066,13 @@ describe('MCP Routes', () => {
});
it('should handle customUserVars that is not an object', async () => {
const { loadCustomConfig } = require('~/server/services/Config');
loadCustomConfig.mockResolvedValue({
mcpServers: {
'test-server': {
customUserVars: 'not-an-object',
},
},
});
const mockMcpManager = {
getRawConfig: jest.fn().mockReturnValue({
customUserVars: 'not-an-object',
}),
};
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
const response = await request(app).get('/api/mcp/test-server/auth-values');
@ -1097,98 +1096,6 @@ describe('MCP Routes', () => {
});
});
describe('POST /:serverName/reinitialize - Tool Deletion Coverage', () => {
it('should handle null cached tools during reinitialize (triggers || {} fallback)', async () => {
const { loadCustomConfig, getCachedTools } = require('~/server/services/Config');
const mockUserConnection = {
fetchTools: jest.fn().mockResolvedValue([{ name: 'new-tool', description: 'A new tool' }]),
};
const mockMcpManager = {
getUserConnection: jest.fn().mockResolvedValue(mockUserConnection),
disconnectServer: jest.fn(),
initializeServer: jest.fn(),
mcpConfigs: {},
};
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
loadCustomConfig.mockResolvedValue({
mcpServers: {
'test-server': { env: { API_KEY: 'test-key' } },
},
});
getCachedTools.mockResolvedValue(null);
const response = await request(app).post('/api/mcp/test-server/reinitialize').expect(200);
expect(response.body).toEqual({
message: "MCP server 'test-server' reinitialized successfully",
success: true,
oauthRequired: false,
oauthUrl: null,
serverName: 'test-server',
});
});
it('should delete existing cached tools during successful reinitialize', async () => {
const {
loadCustomConfig,
getCachedTools,
setCachedTools,
} = require('~/server/services/Config');
const mockUserConnection = {
fetchTools: jest.fn().mockResolvedValue([{ name: 'new-tool', description: 'A new tool' }]),
};
const mockMcpManager = {
getUserConnection: jest.fn().mockResolvedValue(mockUserConnection),
disconnectServer: jest.fn(),
initializeServer: jest.fn(),
mcpConfigs: {},
};
require('~/config').getMCPManager.mockReturnValue(mockMcpManager);
loadCustomConfig.mockResolvedValue({
mcpServers: {
'test-server': { env: { API_KEY: 'test-key' } },
},
});
const existingTools = {
'old-tool_mcp_test-server': { type: 'function' },
'other-tool_mcp_other-server': { type: 'function' },
};
getCachedTools.mockResolvedValue(existingTools);
const response = await request(app).post('/api/mcp/test-server/reinitialize').expect(200);
expect(response.body).toEqual({
message: "MCP server 'test-server' reinitialized successfully",
success: true,
oauthRequired: false,
oauthUrl: null,
serverName: 'test-server',
});
expect(setCachedTools).toHaveBeenCalledWith(
expect.objectContaining({
'new-tool_mcp_test-server': expect.any(Object),
'other-tool_mcp_other-server': { type: 'function' },
}),
{ userId: 'test-user-id' },
);
expect(setCachedTools).toHaveBeenCalledWith(
expect.not.objectContaining({
'old-tool_mcp_test-server': expect.anything(),
}),
{ userId: 'test-user-id' },
);
});
});
describe('GET /:serverName/oauth/callback - Edge Cases', () => {
it('should handle OAuth callback without toolFlowId (falsy toolFlowId)', async () => {
const { MCPOAuthHandler } = require('@librechat/api');

32
api/server/routes/mcp.js
View file

@ -1,11 +1,11 @@
const { Router } = require('express');
const { logger } = require('@librechat/data-schemas');
const { MCPOAuthHandler } = require('@librechat/api');
const { CacheKeys, Constants } = require('librechat-data-provider');
const { findToken, updateToken, createToken, deleteTokens } = require('~/models');
const { setCachedTools, getCachedTools, loadCustomConfig } = require('~/server/services/Config');
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 { getUserPluginAuthValue } = require('~/server/services/PluginService');
const { CacheKeys, Constants } = require('librechat-data-provider');
const { getMCPManager, getFlowStateManager } = require('~/config');
const { requireJwtAuth } = require('~/server/middleware');
const { getLogStores } = require('~/cache');
@ -315,9 +315,9 @@ router.post('/:serverName/reinitialize', requireJwtAuth, async (req, res) => {
logger.info(`[MCP Reinitialize] Reinitializing server: ${serverName}`);
const printConfig = false;
const config = await loadCustomConfig(printConfig);
if (!config || !config.mcpServers || !config.mcpServers[serverName]) {
const mcpManager = getMCPManager();
const serverConfig = mcpManager.getRawConfig(serverName);
if (!serverConfig) {
return res.status(404).json({
error: `MCP server '${serverName}' not found in configuration`,
});
@ -325,13 +325,12 @@ router.post('/:serverName/reinitialize', requireJwtAuth, async (req, res) => {
const flowsCache = getLogStores(CacheKeys.FLOWS);
const flowManager = getFlowStateManager(flowsCache);
const mcpManager = getMCPManager();
await mcpManager.disconnectServer(serverName);
logger.info(`[MCP Reinitialize] Disconnected existing server: ${serverName}`);
await mcpManager.disconnectUserConnection(user.id, serverName);
logger.info(
`[MCP Reinitialize] Disconnected existing user connection for server: ${serverName}`,
);
const serverConfig = config.mcpServers[serverName];
mcpManager.mcpConfigs[serverName] = serverConfig;
let customUserVars = {};
if (serverConfig.customUserVars && typeof serverConfig.customUserVars === 'object') {
for (const varName of Object.keys(serverConfig.customUserVars)) {
@ -437,7 +436,7 @@ router.post('/:serverName/reinitialize', requireJwtAuth, async (req, res) => {
};
res.json({
success: (userConnection && !oauthRequired) || (oauthRequired && oauthUrl),
success: Boolean((userConnection && !oauthRequired) || (oauthRequired && oauthUrl)),
message: getResponseMessage(),
serverName,
oauthRequired,
@ -551,15 +550,14 @@ router.get('/:serverName/auth-values', requireJwtAuth, async (req, res) => {
return res.status(401).json({ error: 'User not authenticated' });
}
const printConfig = false;
const config = await loadCustomConfig(printConfig);
if (!config || !config.mcpServers || !config.mcpServers[serverName]) {
const mcpManager = getMCPManager();
const serverConfig = mcpManager.getRawConfig(serverName);
if (!serverConfig) {
return res.status(404).json({
error: `MCP server '${serverName}' not found in configuration`,
});
}
const serverConfig = config.mcpServers[serverName];
const pluginKey = `${Constants.mcp_prefix}${serverName}`;
const authValueFlags = {};

29
api/server/services/initializeMCPs.js
View file

@ -1,8 +1,7 @@
const { logger } = require('@librechat/data-schemas');
const { CacheKeys } = require('librechat-data-provider');
const { findToken, updateToken, createToken, deleteTokens } = require('~/models');
const { getMCPManager, getFlowStateManager } = require('~/config');
const { getCachedTools, setCachedTools } = require('./Config');
const { CacheKeys } = require('librechat-data-provider');
const { createMCPManager } = require('~/config');
const { getLogStores } = require('~/cache');
/**
@ -31,33 +30,19 @@ async function initializeMCPs(app) {
}
logger.info('Initializing MCP servers...');
const mcpManager = getMCPManager();
const flowsCache = getLogStores(CacheKeys.FLOWS);
const flowManager = flowsCache ? getFlowStateManager(flowsCache) : null;
const mcpManager = await createMCPManager(mcpServers);
try {
await mcpManager.initializeMCPs({
mcpServers: filteredServers,
flowManager,
tokenMethods: {
findToken,
updateToken,
createToken,
deleteTokens,
},
});
delete app.locals.mcpConfig;
const availableTools = await getCachedTools();
const cachedTools = await getCachedTools();
if (!availableTools) {
if (!cachedTools) {
logger.warn('No available tools found in cache during MCP initialization');
return;
}
const toolsCopy = { ...availableTools };
await mcpManager.mapAvailableTools(toolsCopy, flowManager);
await setCachedTools(toolsCopy, { isGlobal: true });
const mcpTools = mcpManager.getAppToolFunctions();
await setCachedTools({ ...cachedTools, ...mcpTools }, { isGlobal: true });
const cache = getLogStores(CacheKeys.CONFIG_STORE);
await cache.delete(CacheKeys.TOOLS);