mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-01-17 07:55:32 +01:00
♻️ refactor: MCPManager for Scalability, Fix App-Level Detection, Add Lazy Connections (#8930)
* 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
GitHub
parent
9dbf153489
commit
8780a78165
32 changed files with 2571 additions and 1468 deletions
347
packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts
Normal file
347
packages/api/src/mcp/__tests__/MCPConnectionFactory.test.ts
Normal file
|
|
@ -0,0 +1,347 @@
|
|||
import { logger } from '@librechat/data-schemas';
|
||||
import type { TUser } from 'librechat-data-provider';
|
||||
import type { FlowStateManager } from '~/flow/manager';
|
||||
import type { MCPOAuthTokens } from '~/mcp/oauth';
|
||||
import { MCPConnectionFactory } from '../MCPConnectionFactory';
|
||||
import { MCPOAuthHandler } from '~/mcp/oauth';
|
||||
import { MCPConnection } from '../connection';
|
||||
import { processMCPEnv } from '~/utils';
|
||||
import type * as t from '../types';
|
||||
|
||||
jest.mock('../connection');
|
||||
jest.mock('~/mcp/oauth');
|
||||
jest.mock('~/utils');
|
||||
jest.mock('@librechat/data-schemas', () => ({
|
||||
logger: {
|
||||
info: jest.fn(),
|
||||
warn: jest.fn(),
|
||||
error: jest.fn(),
|
||||
debug: jest.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
const mockLogger = logger as jest.Mocked<typeof logger>;
|
||||
const mockProcessMCPEnv = processMCPEnv as jest.MockedFunction<typeof processMCPEnv>;
|
||||
const mockMCPConnection = MCPConnection as jest.MockedClass<typeof MCPConnection>;
|
||||
const mockMCPOAuthHandler = MCPOAuthHandler as jest.Mocked<typeof MCPOAuthHandler>;
|
||||
|
||||
describe('MCPConnectionFactory', () => {
|
||||
let mockUser: TUser;
|
||||
let mockServerConfig: t.MCPOptions;
|
||||
let mockFlowManager: jest.Mocked<FlowStateManager<MCPOAuthTokens | null>>;
|
||||
let mockConnectionInstance: jest.Mocked<MCPConnection>;
|
||||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks();
|
||||
mockUser = {
|
||||
id: 'user123',
|
||||
email: 'test@example.com',
|
||||
} as TUser;
|
||||
|
||||
mockServerConfig = {
|
||||
command: 'node',
|
||||
args: ['server.js'],
|
||||
initTimeout: 5000,
|
||||
} as t.MCPOptions;
|
||||
|
||||
mockFlowManager = {
|
||||
createFlow: jest.fn(),
|
||||
createFlowWithHandler: jest.fn(),
|
||||
getFlowState: jest.fn(),
|
||||
} as unknown as jest.Mocked<FlowStateManager<MCPOAuthTokens | null>>;
|
||||
|
||||
mockConnectionInstance = {
|
||||
connect: jest.fn(),
|
||||
isConnected: jest.fn(),
|
||||
setOAuthTokens: jest.fn(),
|
||||
on: jest.fn().mockReturnValue(mockConnectionInstance),
|
||||
emit: jest.fn(),
|
||||
} as unknown as jest.Mocked<MCPConnection>;
|
||||
|
||||
mockMCPConnection.mockImplementation(() => mockConnectionInstance);
|
||||
mockProcessMCPEnv.mockReturnValue(mockServerConfig);
|
||||
});
|
||||
|
||||
describe('static create method', () => {
|
||||
it('should create a basic connection without OAuth', async () => {
|
||||
const basicOptions = {
|
||||
serverName: 'test-server',
|
||||
serverConfig: mockServerConfig,
|
||||
};
|
||||
|
||||
mockConnectionInstance.isConnected.mockResolvedValue(true);
|
||||
|
||||
const connection = await MCPConnectionFactory.create(basicOptions);
|
||||
|
||||
expect(connection).toBe(mockConnectionInstance);
|
||||
expect(mockProcessMCPEnv).toHaveBeenCalledWith(mockServerConfig, undefined, undefined);
|
||||
expect(mockMCPConnection).toHaveBeenCalledWith({
|
||||
serverName: 'test-server',
|
||||
serverConfig: mockServerConfig,
|
||||
userId: undefined,
|
||||
oauthTokens: null,
|
||||
});
|
||||
expect(mockConnectionInstance.connect).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should create a connection with OAuth', async () => {
|
||||
const basicOptions = {
|
||||
serverName: 'test-server',
|
||||
serverConfig: mockServerConfig,
|
||||
};
|
||||
|
||||
const oauthOptions = {
|
||||
useOAuth: true as const,
|
||||
user: mockUser,
|
||||
flowManager: mockFlowManager,
|
||||
tokenMethods: {
|
||||
findToken: jest.fn(),
|
||||
createToken: jest.fn(),
|
||||
updateToken: jest.fn(),
|
||||
deleteTokens: jest.fn(),
|
||||
},
|
||||
};
|
||||
|
||||
const mockTokens: MCPOAuthTokens = {
|
||||
access_token: 'access123',
|
||||
refresh_token: 'refresh123',
|
||||
token_type: 'Bearer',
|
||||
obtained_at: Date.now(),
|
||||
};
|
||||
|
||||
mockFlowManager.createFlowWithHandler.mockResolvedValue(mockTokens);
|
||||
mockConnectionInstance.isConnected.mockResolvedValue(true);
|
||||
|
||||
const connection = await MCPConnectionFactory.create(basicOptions, oauthOptions);
|
||||
|
||||
expect(connection).toBe(mockConnectionInstance);
|
||||
expect(mockProcessMCPEnv).toHaveBeenCalledWith(mockServerConfig, mockUser, undefined);
|
||||
expect(mockMCPConnection).toHaveBeenCalledWith({
|
||||
serverName: 'test-server',
|
||||
serverConfig: mockServerConfig,
|
||||
userId: 'user123',
|
||||
oauthTokens: mockTokens,
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('OAuth token handling', () => {
|
||||
it('should return null when no findToken method is provided', async () => {
|
||||
const basicOptions = {
|
||||
serverName: 'test-server',
|
||||
serverConfig: mockServerConfig,
|
||||
};
|
||||
|
||||
const oauthOptions = {
|
||||
useOAuth: true as const,
|
||||
user: mockUser,
|
||||
flowManager: mockFlowManager,
|
||||
tokenMethods: {
|
||||
findToken: undefined as unknown as () => Promise<any>,
|
||||
createToken: jest.fn(),
|
||||
updateToken: jest.fn(),
|
||||
deleteTokens: jest.fn(),
|
||||
},
|
||||
};
|
||||
|
||||
mockConnectionInstance.isConnected.mockResolvedValue(true);
|
||||
|
||||
await MCPConnectionFactory.create(basicOptions, oauthOptions);
|
||||
|
||||
expect(mockFlowManager.createFlowWithHandler).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should handle token retrieval errors gracefully', async () => {
|
||||
const basicOptions = {
|
||||
serverName: 'test-server',
|
||||
serverConfig: mockServerConfig,
|
||||
};
|
||||
|
||||
const oauthOptions = {
|
||||
useOAuth: true as const,
|
||||
user: mockUser,
|
||||
flowManager: mockFlowManager,
|
||||
tokenMethods: {
|
||||
findToken: jest.fn(),
|
||||
createToken: jest.fn(),
|
||||
updateToken: jest.fn(),
|
||||
deleteTokens: jest.fn(),
|
||||
},
|
||||
};
|
||||
|
||||
mockFlowManager.createFlowWithHandler.mockRejectedValue(new Error('Token fetch failed'));
|
||||
mockConnectionInstance.isConnected.mockResolvedValue(true);
|
||||
|
||||
const connection = await MCPConnectionFactory.create(basicOptions, oauthOptions);
|
||||
|
||||
expect(connection).toBe(mockConnectionInstance);
|
||||
expect(mockMCPConnection).toHaveBeenCalledWith({
|
||||
serverName: 'test-server',
|
||||
serverConfig: mockServerConfig,
|
||||
userId: 'user123',
|
||||
oauthTokens: null,
|
||||
});
|
||||
expect(mockLogger.debug).toHaveBeenCalledWith(
|
||||
expect.stringContaining('No existing tokens found or error loading tokens'),
|
||||
expect.any(Error),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('OAuth event handling', () => {
|
||||
it('should handle oauthRequired event for returnOnOAuth scenario', async () => {
|
||||
const basicOptions = {
|
||||
serverName: 'test-server',
|
||||
serverConfig: {
|
||||
...mockServerConfig,
|
||||
url: 'https://api.example.com',
|
||||
type: 'sse' as const,
|
||||
} as t.SSEOptions,
|
||||
};
|
||||
|
||||
const oauthOptions = {
|
||||
useOAuth: true as const,
|
||||
user: mockUser,
|
||||
flowManager: mockFlowManager,
|
||||
returnOnOAuth: true,
|
||||
oauthStart: jest.fn(),
|
||||
tokenMethods: {
|
||||
findToken: jest.fn(),
|
||||
createToken: jest.fn(),
|
||||
updateToken: jest.fn(),
|
||||
deleteTokens: jest.fn(),
|
||||
},
|
||||
};
|
||||
|
||||
const mockFlowData = {
|
||||
authorizationUrl: 'https://auth.example.com',
|
||||
flowId: 'flow123',
|
||||
flowMetadata: {
|
||||
serverName: 'test-server',
|
||||
userId: 'user123',
|
||||
serverUrl: 'https://api.example.com',
|
||||
state: 'random-state',
|
||||
clientInfo: { client_id: 'client123' },
|
||||
},
|
||||
};
|
||||
|
||||
mockMCPOAuthHandler.initiateOAuthFlow.mockResolvedValue(mockFlowData);
|
||||
mockFlowManager.createFlow.mockRejectedValue(new Error('Timeout expected'));
|
||||
mockConnectionInstance.isConnected.mockResolvedValue(false);
|
||||
|
||||
let oauthRequiredHandler: (data: Record<string, unknown>) => Promise<void>;
|
||||
mockConnectionInstance.on.mockImplementation((event, handler) => {
|
||||
if (event === 'oauthRequired') {
|
||||
oauthRequiredHandler = handler as (data: Record<string, unknown>) => Promise<void>;
|
||||
}
|
||||
return mockConnectionInstance;
|
||||
});
|
||||
|
||||
try {
|
||||
await MCPConnectionFactory.create(basicOptions, oauthOptions);
|
||||
} catch {
|
||||
// Expected to fail due to connection not established
|
||||
}
|
||||
|
||||
expect(oauthRequiredHandler!).toBeDefined();
|
||||
|
||||
await oauthRequiredHandler!({ serverUrl: 'https://api.example.com' });
|
||||
|
||||
expect(mockMCPOAuthHandler.initiateOAuthFlow).toHaveBeenCalledWith(
|
||||
'test-server',
|
||||
'https://api.example.com',
|
||||
'user123',
|
||||
undefined,
|
||||
);
|
||||
expect(oauthOptions.oauthStart).toHaveBeenCalledWith('https://auth.example.com');
|
||||
expect(mockConnectionInstance.emit).toHaveBeenCalledWith(
|
||||
'oauthFailed',
|
||||
expect.objectContaining({
|
||||
message: 'OAuth flow initiated - return early',
|
||||
}),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('connection retry logic', () => {
|
||||
it('should establish connection successfully', async () => {
|
||||
const basicOptions = {
|
||||
serverName: 'test-server',
|
||||
serverConfig: mockServerConfig, // Use default 5000ms timeout
|
||||
};
|
||||
|
||||
mockConnectionInstance.connect.mockResolvedValue(undefined);
|
||||
mockConnectionInstance.isConnected.mockResolvedValue(true);
|
||||
|
||||
const connection = await MCPConnectionFactory.create(basicOptions);
|
||||
|
||||
expect(connection).toBe(mockConnectionInstance);
|
||||
expect(mockConnectionInstance.connect).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('should handle OAuth errors during connection attempts', async () => {
|
||||
const basicOptions = {
|
||||
serverName: 'test-server',
|
||||
serverConfig: mockServerConfig,
|
||||
};
|
||||
|
||||
const oauthOptions = {
|
||||
useOAuth: true as const,
|
||||
user: mockUser,
|
||||
flowManager: mockFlowManager,
|
||||
oauthStart: jest.fn(),
|
||||
tokenMethods: {
|
||||
findToken: jest.fn(),
|
||||
createToken: jest.fn(),
|
||||
updateToken: jest.fn(),
|
||||
deleteTokens: jest.fn(),
|
||||
},
|
||||
};
|
||||
|
||||
const oauthError = new Error('Non-200 status code (401)');
|
||||
(oauthError as unknown as Record<string, unknown>).isOAuthError = true;
|
||||
|
||||
mockConnectionInstance.connect.mockRejectedValue(oauthError);
|
||||
mockConnectionInstance.isConnected.mockResolvedValue(false);
|
||||
|
||||
await expect(MCPConnectionFactory.create(basicOptions, oauthOptions)).rejects.toThrow(
|
||||
'Non-200 status code (401)',
|
||||
);
|
||||
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
expect.stringContaining('OAuth required, stopping connection attempts'),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('isOAuthError method', () => {
|
||||
it('should identify OAuth errors by message content', async () => {
|
||||
const basicOptions = {
|
||||
serverName: 'test-server',
|
||||
serverConfig: mockServerConfig,
|
||||
};
|
||||
|
||||
const oauthOptions = {
|
||||
useOAuth: true as const,
|
||||
user: mockUser,
|
||||
flowManager: mockFlowManager,
|
||||
tokenMethods: {
|
||||
findToken: jest.fn(),
|
||||
createToken: jest.fn(),
|
||||
updateToken: jest.fn(),
|
||||
deleteTokens: jest.fn(),
|
||||
},
|
||||
};
|
||||
|
||||
const error401 = new Error('401 Unauthorized');
|
||||
|
||||
mockConnectionInstance.connect.mockRejectedValue(error401);
|
||||
mockConnectionInstance.isConnected.mockResolvedValue(false);
|
||||
|
||||
await expect(MCPConnectionFactory.create(basicOptions, oauthOptions)).rejects.toThrow('401');
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
expect.stringContaining('OAuth required, stopping connection attempts'),
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
Loading…
Add table
Add a link
Reference in a new issue