mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-01-06 10:38:50 +01:00
🔄 refactor: MCP Registry System with Distributed Caching (#10191)
Some checks failed
Docker Dev Branch Images Build / build (Dockerfile, lc-dev, node) (push) Has been cancelled
Docker Dev Branch Images Build / build (Dockerfile.multi, lc-dev-api, api-build) (push) Has been cancelled
Docker Dev Images Build / build (Dockerfile, librechat-dev, node) (push) Has been cancelled
Docker Dev Images Build / build (Dockerfile.multi, librechat-dev-api, api-build) (push) Has been cancelled
Sync Locize Translations & Create Translation PR / Sync Translation Keys with Locize (push) Has been cancelled
Sync Locize Translations & Create Translation PR / Create Translation PR on Version Published (push) Has been cancelled
Some checks failed
Docker Dev Branch Images Build / build (Dockerfile, lc-dev, node) (push) Has been cancelled
Docker Dev Branch Images Build / build (Dockerfile.multi, lc-dev-api, api-build) (push) Has been cancelled
Docker Dev Images Build / build (Dockerfile, librechat-dev, node) (push) Has been cancelled
Docker Dev Images Build / build (Dockerfile.multi, librechat-dev-api, api-build) (push) Has been cancelled
Sync Locize Translations & Create Translation PR / Sync Translation Keys with Locize (push) Has been cancelled
Sync Locize Translations & Create Translation PR / Create Translation PR on Version Published (push) Has been cancelled
* refactor: Restructure MCP registry system with caching - Split MCPServersRegistry into modular components: - MCPServerInspector: handles server inspection and health checks - MCPServersInitializer: manages server initialization logic - MCPServersRegistry: simplified registry coordination - Add distributed caching layer: - ServerConfigsCacheRedis: Redis-backed configuration cache - ServerConfigsCacheInMemory: in-memory fallback cache - RegistryStatusCache: distributed leader election state - Add promise utilities (withTimeout) replacing Promise.race patterns - Add comprehensive cache integration tests for all cache implementations - Remove unused MCPManager.getAllToolFunctions method * fix: Update OAuth flow to include user-specific headers * chore: Update Jest configuration to ignore additional test files - Added patterns to ignore files ending with .helper.ts and .helper.d.ts in testPathIgnorePatterns for cleaner test runs. * fix: oauth headers in callback * chore: Update Jest testPathIgnorePatterns to exclude helper files - Modified testPathIgnorePatterns in package.json to ignore files ending with .helper.ts and .helper.d.ts for cleaner test execution. * ci: update test mocks --------- Co-authored-by: Danny Avila <danny@librechat.ai>
This commit is contained in:
Theo N. Truong
GitHub
parent
961f87cfda
commit
ce7e6edad8
45 changed files with 3116 additions and 1150 deletions
|
|
@ -1,7 +1,9 @@
|
|||
import { logger } from '@librechat/data-schemas';
|
||||
import type * as t from '~/mcp/types';
|
||||
import { MCPManager } from '~/mcp/MCPManager';
|
||||
import { MCPServersRegistry } from '~/mcp/MCPServersRegistry';
|
||||
import { mcpServersRegistry } from '~/mcp/registry/MCPServersRegistry';
|
||||
import { MCPServersInitializer } from '~/mcp/registry/MCPServersInitializer';
|
||||
import { MCPServerInspector } from '~/mcp/registry/MCPServerInspector';
|
||||
import { ConnectionsRepository } from '~/mcp/ConnectionsRepository';
|
||||
import { MCPConnection } from '../connection';
|
||||
|
||||
|
|
@ -15,7 +17,24 @@ jest.mock('@librechat/data-schemas', () => ({
|
|||
},
|
||||
}));
|
||||
|
||||
jest.mock('~/mcp/MCPServersRegistry');
|
||||
jest.mock('~/mcp/registry/MCPServersRegistry', () => ({
|
||||
mcpServersRegistry: {
|
||||
sharedAppServers: {
|
||||
getAll: jest.fn(),
|
||||
},
|
||||
getServerConfig: jest.fn(),
|
||||
getAllServerConfigs: jest.fn(),
|
||||
getOAuthServers: jest.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
jest.mock('~/mcp/registry/MCPServersInitializer', () => ({
|
||||
MCPServersInitializer: {
|
||||
initialize: jest.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
jest.mock('~/mcp/registry/MCPServerInspector');
|
||||
jest.mock('~/mcp/ConnectionsRepository');
|
||||
|
||||
const mockLogger = logger as jest.Mocked<typeof logger>;
|
||||
|
|
@ -28,20 +47,12 @@ describe('MCPManager', () => {
|
|||
// Reset MCPManager singleton state
|
||||
(MCPManager as unknown as { instance: null }).instance = null;
|
||||
jest.clearAllMocks();
|
||||
});
|
||||
|
||||
function mockRegistry(
|
||||
registryConfig: Partial<MCPServersRegistry>,
|
||||
): jest.MockedClass<typeof MCPServersRegistry> {
|
||||
const mock = {
|
||||
initialize: jest.fn().mockResolvedValue(undefined),
|
||||
getToolFunctions: jest.fn().mockResolvedValue(null),
|
||||
...registryConfig,
|
||||
};
|
||||
return (MCPServersRegistry as jest.MockedClass<typeof MCPServersRegistry>).mockImplementation(
|
||||
() => mock as unknown as MCPServersRegistry,
|
||||
);
|
||||
}
|
||||
// Set up default mock implementations
|
||||
(MCPServersInitializer.initialize as jest.Mock).mockResolvedValue(undefined);
|
||||
(mcpServersRegistry.sharedAppServers.getAll as jest.Mock).mockResolvedValue({});
|
||||
(mcpServersRegistry.getAllServerConfigs as jest.Mock).mockResolvedValue({});
|
||||
});
|
||||
|
||||
function mockAppConnections(
|
||||
appConnectionsConfig: Partial<ConnectionsRepository>,
|
||||
|
|
@ -66,12 +77,229 @@ describe('MCPManager', () => {
|
|||
};
|
||||
}
|
||||
|
||||
describe('getAppToolFunctions', () => {
|
||||
it('should return empty object when no servers have tool functions', async () => {
|
||||
(mcpServersRegistry.getAllServerConfigs as jest.Mock).mockResolvedValue({
|
||||
server1: { type: 'stdio', command: 'test', args: [] },
|
||||
server2: { type: 'stdio', command: 'test2', args: [] },
|
||||
});
|
||||
|
||||
const manager = await MCPManager.createInstance(newMCPServersConfig());
|
||||
const result = await manager.getAppToolFunctions();
|
||||
|
||||
expect(result).toEqual({});
|
||||
});
|
||||
|
||||
it('should collect tool functions from multiple servers', async () => {
|
||||
const toolFunctions1 = {
|
||||
tool1_mcp_server1: {
|
||||
type: 'function' as const,
|
||||
function: {
|
||||
name: 'tool1_mcp_server1',
|
||||
description: 'Tool 1',
|
||||
parameters: { type: 'object' as const },
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
const toolFunctions2 = {
|
||||
tool2_mcp_server2: {
|
||||
type: 'function' as const,
|
||||
function: {
|
||||
name: 'tool2_mcp_server2',
|
||||
description: 'Tool 2',
|
||||
parameters: { type: 'object' as const },
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
(mcpServersRegistry.getAllServerConfigs as jest.Mock).mockResolvedValue({
|
||||
server1: {
|
||||
type: 'stdio',
|
||||
command: 'test',
|
||||
args: [],
|
||||
toolFunctions: toolFunctions1,
|
||||
},
|
||||
server2: {
|
||||
type: 'stdio',
|
||||
command: 'test2',
|
||||
args: [],
|
||||
toolFunctions: toolFunctions2,
|
||||
},
|
||||
});
|
||||
|
||||
const manager = await MCPManager.createInstance(newMCPServersConfig());
|
||||
const result = await manager.getAppToolFunctions();
|
||||
|
||||
expect(result).toEqual({
|
||||
...toolFunctions1,
|
||||
...toolFunctions2,
|
||||
});
|
||||
});
|
||||
|
||||
it('should handle servers with null or undefined toolFunctions', async () => {
|
||||
const toolFunctions1 = {
|
||||
tool1_mcp_server1: {
|
||||
type: 'function' as const,
|
||||
function: {
|
||||
name: 'tool1_mcp_server1',
|
||||
description: 'Tool 1',
|
||||
parameters: { type: 'object' as const },
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
(mcpServersRegistry.getAllServerConfigs as jest.Mock).mockResolvedValue({
|
||||
server1: {
|
||||
type: 'stdio',
|
||||
command: 'test',
|
||||
args: [],
|
||||
toolFunctions: toolFunctions1,
|
||||
},
|
||||
server2: {
|
||||
type: 'stdio',
|
||||
command: 'test2',
|
||||
args: [],
|
||||
toolFunctions: null,
|
||||
},
|
||||
server3: {
|
||||
type: 'stdio',
|
||||
command: 'test3',
|
||||
args: [],
|
||||
},
|
||||
});
|
||||
|
||||
const manager = await MCPManager.createInstance(newMCPServersConfig());
|
||||
const result = await manager.getAppToolFunctions();
|
||||
|
||||
expect(result).toEqual(toolFunctions1);
|
||||
});
|
||||
});
|
||||
|
||||
describe('formatInstructionsForContext', () => {
|
||||
it('should return empty string when no servers have instructions', async () => {
|
||||
(mcpServersRegistry.getAllServerConfigs as jest.Mock).mockResolvedValue({
|
||||
server1: { type: 'stdio', command: 'test', args: [] },
|
||||
server2: { type: 'stdio', command: 'test2', args: [] },
|
||||
});
|
||||
|
||||
const manager = await MCPManager.createInstance(newMCPServersConfig());
|
||||
const result = await manager.formatInstructionsForContext();
|
||||
|
||||
expect(result).toBe('');
|
||||
});
|
||||
|
||||
it('should format instructions from multiple servers', async () => {
|
||||
(mcpServersRegistry.getAllServerConfigs as jest.Mock).mockResolvedValue({
|
||||
github: {
|
||||
type: 'sse',
|
||||
url: 'https://api.github.com',
|
||||
serverInstructions: 'Use GitHub API with care',
|
||||
},
|
||||
files: {
|
||||
type: 'stdio',
|
||||
command: 'node',
|
||||
args: ['files.js'],
|
||||
serverInstructions: 'Only read/write files in allowed directories',
|
||||
},
|
||||
});
|
||||
|
||||
const manager = await MCPManager.createInstance(newMCPServersConfig());
|
||||
const result = await manager.formatInstructionsForContext();
|
||||
|
||||
expect(result).toContain('# MCP Server Instructions');
|
||||
expect(result).toContain('## github MCP Server Instructions');
|
||||
expect(result).toContain('Use GitHub API with care');
|
||||
expect(result).toContain('## files MCP Server Instructions');
|
||||
expect(result).toContain('Only read/write files in allowed directories');
|
||||
});
|
||||
|
||||
it('should filter instructions by server names when provided', async () => {
|
||||
(mcpServersRegistry.getAllServerConfigs as jest.Mock).mockResolvedValue({
|
||||
github: {
|
||||
type: 'sse',
|
||||
url: 'https://api.github.com',
|
||||
serverInstructions: 'Use GitHub API with care',
|
||||
},
|
||||
files: {
|
||||
type: 'stdio',
|
||||
command: 'node',
|
||||
args: ['files.js'],
|
||||
serverInstructions: 'Only read/write files in allowed directories',
|
||||
},
|
||||
database: {
|
||||
type: 'stdio',
|
||||
command: 'node',
|
||||
args: ['db.js'],
|
||||
serverInstructions: 'Be careful with database operations',
|
||||
},
|
||||
});
|
||||
|
||||
const manager = await MCPManager.createInstance(newMCPServersConfig());
|
||||
const result = await manager.formatInstructionsForContext(['github', 'database']);
|
||||
|
||||
expect(result).toContain('## github MCP Server Instructions');
|
||||
expect(result).toContain('Use GitHub API with care');
|
||||
expect(result).toContain('## database MCP Server Instructions');
|
||||
expect(result).toContain('Be careful with database operations');
|
||||
expect(result).not.toContain('files');
|
||||
expect(result).not.toContain('Only read/write files in allowed directories');
|
||||
});
|
||||
|
||||
it('should handle servers with null or undefined instructions', async () => {
|
||||
(mcpServersRegistry.getAllServerConfigs as jest.Mock).mockResolvedValue({
|
||||
github: {
|
||||
type: 'sse',
|
||||
url: 'https://api.github.com',
|
||||
serverInstructions: 'Use GitHub API with care',
|
||||
},
|
||||
files: {
|
||||
type: 'stdio',
|
||||
command: 'node',
|
||||
args: ['files.js'],
|
||||
serverInstructions: null,
|
||||
},
|
||||
database: {
|
||||
type: 'stdio',
|
||||
command: 'node',
|
||||
args: ['db.js'],
|
||||
},
|
||||
});
|
||||
|
||||
const manager = await MCPManager.createInstance(newMCPServersConfig());
|
||||
const result = await manager.formatInstructionsForContext();
|
||||
|
||||
expect(result).toContain('## github MCP Server Instructions');
|
||||
expect(result).toContain('Use GitHub API with care');
|
||||
expect(result).not.toContain('files');
|
||||
expect(result).not.toContain('database');
|
||||
});
|
||||
|
||||
it('should return empty string when filtered servers have no instructions', async () => {
|
||||
(mcpServersRegistry.getAllServerConfigs as jest.Mock).mockResolvedValue({
|
||||
github: {
|
||||
type: 'sse',
|
||||
url: 'https://api.github.com',
|
||||
serverInstructions: 'Use GitHub API with care',
|
||||
},
|
||||
files: {
|
||||
type: 'stdio',
|
||||
command: 'node',
|
||||
args: ['files.js'],
|
||||
},
|
||||
});
|
||||
|
||||
const manager = await MCPManager.createInstance(newMCPServersConfig());
|
||||
const result = await manager.formatInstructionsForContext(['files']);
|
||||
|
||||
expect(result).toBe('');
|
||||
});
|
||||
});
|
||||
|
||||
describe('getServerToolFunctions', () => {
|
||||
it('should catch and handle errors gracefully', async () => {
|
||||
mockRegistry({
|
||||
getToolFunctions: jest.fn(() => {
|
||||
throw new Error('Connection failed');
|
||||
}),
|
||||
(MCPServerInspector.getToolFunctions as jest.Mock) = jest.fn(() => {
|
||||
throw new Error('Connection failed');
|
||||
});
|
||||
|
||||
mockAppConnections({
|
||||
|
|
@ -90,9 +318,7 @@ describe('MCPManager', () => {
|
|||
});
|
||||
|
||||
it('should catch synchronous errors from getUserConnections', async () => {
|
||||
mockRegistry({
|
||||
getToolFunctions: jest.fn().mockResolvedValue({}),
|
||||
});
|
||||
(MCPServerInspector.getToolFunctions as jest.Mock) = jest.fn().mockResolvedValue({});
|
||||
|
||||
mockAppConnections({
|
||||
has: jest.fn().mockReturnValue(false),
|
||||
|
|
@ -126,9 +352,9 @@ describe('MCPManager', () => {
|
|||
},
|
||||
};
|
||||
|
||||
mockRegistry({
|
||||
getToolFunctions: jest.fn().mockResolvedValue(expectedTools),
|
||||
});
|
||||
(MCPServerInspector.getToolFunctions as jest.Mock) = jest
|
||||
.fn()
|
||||
.mockResolvedValue(expectedTools);
|
||||
|
||||
mockAppConnections({
|
||||
has: jest.fn().mockReturnValue(true),
|
||||
|
|
@ -145,10 +371,8 @@ describe('MCPManager', () => {
|
|||
it('should include specific server name in error messages', async () => {
|
||||
const specificServerName = 'github_mcp_server';
|
||||
|
||||
mockRegistry({
|
||||
getToolFunctions: jest.fn(() => {
|
||||
throw new Error('Server specific error');
|
||||
}),
|
||||
(MCPServerInspector.getToolFunctions as jest.Mock) = jest.fn(() => {
|
||||
throw new Error('Server specific error');
|
||||
});
|
||||
|
||||
mockAppConnections({
|
||||
|
|
|
|||
|
|
@ -1,595 +0,0 @@
|
|||
import { join } from 'path';
|
||||
import { readFileSync } from 'fs';
|
||||
import { load as yamlLoad } from 'js-yaml';
|
||||
import { logger } from '@librechat/data-schemas';
|
||||
import type { OAuthDetectionResult } from '~/mcp/oauth/detectOAuth';
|
||||
import type * as t from '~/mcp/types';
|
||||
import { ConnectionsRepository } from '~/mcp/ConnectionsRepository';
|
||||
import { MCPServersRegistry } from '~/mcp/MCPServersRegistry';
|
||||
import { detectOAuthRequirement } from '~/mcp/oauth';
|
||||
import { MCPConnection } from '~/mcp/connection';
|
||||
|
||||
// Mock external dependencies
|
||||
jest.mock('../oauth/detectOAuth');
|
||||
jest.mock('../ConnectionsRepository');
|
||||
jest.mock('../connection');
|
||||
jest.mock('@librechat/data-schemas', () => ({
|
||||
logger: {
|
||||
info: jest.fn(),
|
||||
warn: jest.fn(),
|
||||
error: jest.fn(),
|
||||
debug: jest.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
// Mock processMCPEnv to verify it's called and adds a processed marker
|
||||
jest.mock('~/utils', () => ({
|
||||
...jest.requireActual('~/utils'),
|
||||
processMCPEnv: jest.fn(({ options }) => ({
|
||||
...options,
|
||||
_processed: true, // Simple marker to verify processing occurred
|
||||
})),
|
||||
}));
|
||||
|
||||
const mockDetectOAuthRequirement = detectOAuthRequirement as jest.MockedFunction<
|
||||
typeof detectOAuthRequirement
|
||||
>;
|
||||
const mockLogger = logger as jest.Mocked<typeof logger>;
|
||||
|
||||
describe('MCPServersRegistry - Initialize Function', () => {
|
||||
let rawConfigs: t.MCPServers;
|
||||
let expectedParsedConfigs: Record<string, t.ParsedServerConfig>;
|
||||
let mockConnectionsRepo: jest.Mocked<ConnectionsRepository>;
|
||||
let mockConnections: Map<string, jest.Mocked<MCPConnection>>;
|
||||
|
||||
beforeEach(() => {
|
||||
// Load fixtures
|
||||
const rawConfigsPath = join(__dirname, 'fixtures', 'MCPServersRegistry.rawConfigs.yml');
|
||||
const parsedConfigsPath = join(__dirname, 'fixtures', 'MCPServersRegistry.parsedConfigs.yml');
|
||||
|
||||
rawConfigs = yamlLoad(readFileSync(rawConfigsPath, 'utf8')) as t.MCPServers;
|
||||
expectedParsedConfigs = yamlLoad(readFileSync(parsedConfigsPath, 'utf8')) as Record<
|
||||
string,
|
||||
t.ParsedServerConfig
|
||||
>;
|
||||
|
||||
// Setup mock connections
|
||||
mockConnections = new Map();
|
||||
const serverNames = Object.keys(rawConfigs);
|
||||
|
||||
serverNames.forEach((serverName) => {
|
||||
const mockClient = {
|
||||
listTools: jest.fn(),
|
||||
getInstructions: jest.fn(),
|
||||
getServerCapabilities: jest.fn(),
|
||||
};
|
||||
const mockConnection = {
|
||||
client: mockClient,
|
||||
} as unknown as jest.Mocked<MCPConnection>;
|
||||
|
||||
// Setup mock responses based on expected configs
|
||||
const expectedConfig = expectedParsedConfigs[serverName];
|
||||
|
||||
// Mock listTools response
|
||||
if (expectedConfig.tools) {
|
||||
const toolNames = expectedConfig.tools.split(', ');
|
||||
const tools = toolNames.map((name: string) => ({
|
||||
name,
|
||||
description: `Description for ${name}`,
|
||||
inputSchema: {
|
||||
type: 'object' as const,
|
||||
properties: {
|
||||
input: { type: 'string' },
|
||||
},
|
||||
},
|
||||
}));
|
||||
(mockClient.listTools as jest.Mock).mockResolvedValue({ tools });
|
||||
} else {
|
||||
(mockClient.listTools as jest.Mock).mockResolvedValue({ tools: [] });
|
||||
}
|
||||
|
||||
// Mock getInstructions response
|
||||
if (expectedConfig.serverInstructions) {
|
||||
(mockClient.getInstructions as jest.Mock).mockReturnValue(
|
||||
expectedConfig.serverInstructions as string,
|
||||
);
|
||||
} else {
|
||||
(mockClient.getInstructions as jest.Mock).mockReturnValue(undefined);
|
||||
}
|
||||
|
||||
// Mock getServerCapabilities response
|
||||
if (expectedConfig.capabilities) {
|
||||
const capabilities = JSON.parse(expectedConfig.capabilities) as Record<string, unknown>;
|
||||
(mockClient.getServerCapabilities as jest.Mock).mockReturnValue(capabilities);
|
||||
} else {
|
||||
(mockClient.getServerCapabilities as jest.Mock).mockReturnValue(undefined);
|
||||
}
|
||||
|
||||
mockConnections.set(serverName, mockConnection);
|
||||
});
|
||||
|
||||
// Setup ConnectionsRepository mock
|
||||
mockConnectionsRepo = {
|
||||
get: jest.fn(),
|
||||
getLoaded: jest.fn(),
|
||||
disconnectAll: jest.fn(),
|
||||
disconnect: jest.fn().mockResolvedValue(undefined),
|
||||
} as unknown as jest.Mocked<ConnectionsRepository>;
|
||||
|
||||
mockConnectionsRepo.get.mockImplementation((serverName: string) => {
|
||||
const connection = mockConnections.get(serverName);
|
||||
if (!connection) {
|
||||
throw new Error(`Connection not found for server: ${serverName}`);
|
||||
}
|
||||
return Promise.resolve(connection);
|
||||
});
|
||||
|
||||
mockConnectionsRepo.getLoaded.mockResolvedValue(mockConnections);
|
||||
|
||||
(ConnectionsRepository as jest.Mock).mockImplementation(() => mockConnectionsRepo);
|
||||
|
||||
// Setup OAuth detection mock with deterministic results
|
||||
mockDetectOAuthRequirement.mockImplementation((url: string) => {
|
||||
const oauthResults: Record<string, OAuthDetectionResult> = {
|
||||
'https://api.github.com/mcp': {
|
||||
requiresOAuth: true,
|
||||
method: 'protected-resource-metadata',
|
||||
metadata: {
|
||||
authorization_url: 'https://github.com/login/oauth/authorize',
|
||||
token_url: 'https://github.com/login/oauth/access_token',
|
||||
},
|
||||
},
|
||||
'https://api.disabled.com/mcp': {
|
||||
requiresOAuth: false,
|
||||
method: 'no-metadata-found',
|
||||
metadata: null,
|
||||
},
|
||||
'https://api.public.com/mcp': {
|
||||
requiresOAuth: false,
|
||||
method: 'no-metadata-found',
|
||||
metadata: null,
|
||||
},
|
||||
};
|
||||
|
||||
return Promise.resolve(
|
||||
oauthResults[url] || { requiresOAuth: false, method: 'no-metadata-found', metadata: null },
|
||||
);
|
||||
});
|
||||
|
||||
// Clear all mocks
|
||||
jest.clearAllMocks();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
delete process.env.MCP_INIT_TIMEOUT_MS;
|
||||
jest.clearAllMocks();
|
||||
});
|
||||
|
||||
describe('initialize() method', () => {
|
||||
it('should only run initialization once', async () => {
|
||||
const registry = new MCPServersRegistry(rawConfigs);
|
||||
|
||||
await registry.initialize();
|
||||
await registry.initialize(); // Second call should not re-run
|
||||
|
||||
// Verify that connections are only requested for servers that need them
|
||||
// (servers with serverInstructions=true or all servers for capabilities)
|
||||
expect(mockConnectionsRepo.get).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should set all public properties correctly after initialization', async () => {
|
||||
const registry = new MCPServersRegistry(rawConfigs);
|
||||
|
||||
// Verify initial state
|
||||
expect(registry.oauthServers.size).toBe(0);
|
||||
expect(registry.serverInstructions).toEqual({});
|
||||
expect(registry.toolFunctions).toEqual({});
|
||||
expect(registry.appServerConfigs).toEqual({});
|
||||
|
||||
await registry.initialize();
|
||||
|
||||
// Test oauthServers Set
|
||||
expect(registry.oauthServers).toEqual(
|
||||
new Set(['oauth_server', 'oauth_predefined', 'oauth_startup_enabled']),
|
||||
);
|
||||
|
||||
// Test serverInstructions - OAuth servers keep their original boolean value, non-OAuth fetch actual strings
|
||||
expect(registry.serverInstructions).toEqual({
|
||||
stdio_server: 'Follow these instructions for stdio server',
|
||||
oauth_server: true,
|
||||
non_oauth_server: 'Public API instructions',
|
||||
});
|
||||
|
||||
// Test appServerConfigs (startup enabled, non-OAuth servers only)
|
||||
expect(registry.appServerConfigs).toEqual({
|
||||
stdio_server: rawConfigs.stdio_server,
|
||||
websocket_server: rawConfigs.websocket_server,
|
||||
non_oauth_server: rawConfigs.non_oauth_server,
|
||||
});
|
||||
|
||||
// Test toolFunctions (only non-OAuth servers get their tools fetched during initialization)
|
||||
const expectedToolFunctions = {
|
||||
file_read_mcp_stdio_server: {
|
||||
type: 'function',
|
||||
function: {
|
||||
name: 'file_read_mcp_stdio_server',
|
||||
description: 'Description for file_read',
|
||||
parameters: { type: 'object', properties: { input: { type: 'string' } } },
|
||||
},
|
||||
},
|
||||
file_write_mcp_stdio_server: {
|
||||
type: 'function',
|
||||
function: {
|
||||
name: 'file_write_mcp_stdio_server',
|
||||
description: 'Description for file_write',
|
||||
parameters: { type: 'object', properties: { input: { type: 'string' } } },
|
||||
},
|
||||
},
|
||||
};
|
||||
expect(registry.toolFunctions).toEqual(expectedToolFunctions);
|
||||
});
|
||||
|
||||
it('should handle errors gracefully and continue initialization of other servers', async () => {
|
||||
const registry = new MCPServersRegistry(rawConfigs);
|
||||
|
||||
// Make one specific server throw an error during OAuth detection
|
||||
mockDetectOAuthRequirement.mockImplementation((url: string) => {
|
||||
if (url === 'https://api.github.com/mcp') {
|
||||
return Promise.reject(new Error('OAuth detection failed'));
|
||||
}
|
||||
// Return normal responses for other servers
|
||||
const oauthResults: Record<string, OAuthDetectionResult> = {
|
||||
'https://api.disabled.com/mcp': {
|
||||
requiresOAuth: false,
|
||||
method: 'no-metadata-found',
|
||||
metadata: null,
|
||||
},
|
||||
'https://api.public.com/mcp': {
|
||||
requiresOAuth: false,
|
||||
method: 'no-metadata-found',
|
||||
metadata: null,
|
||||
},
|
||||
};
|
||||
return Promise.resolve(
|
||||
oauthResults[url] ?? {
|
||||
requiresOAuth: false,
|
||||
method: 'no-metadata-found',
|
||||
metadata: null,
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
await registry.initialize();
|
||||
|
||||
// Should still initialize successfully for other servers
|
||||
expect(registry.oauthServers).toBeInstanceOf(Set);
|
||||
expect(registry.toolFunctions).toBeDefined();
|
||||
|
||||
// The failed server should not be in oauthServers (since it failed OAuth detection)
|
||||
expect(registry.oauthServers.has('oauth_server')).toBe(false);
|
||||
|
||||
// But other servers should still be processed successfully
|
||||
expect(registry.appServerConfigs).toHaveProperty('stdio_server');
|
||||
expect(registry.appServerConfigs).toHaveProperty('non_oauth_server');
|
||||
|
||||
// Error should be logged as a warning at the higher level
|
||||
expect(mockLogger.warn).toHaveBeenCalledWith(
|
||||
expect.stringContaining('[MCP][oauth_server] Failed to initialize server:'),
|
||||
expect.any(Error),
|
||||
);
|
||||
});
|
||||
|
||||
it('should disconnect individual connections after each server initialization', async () => {
|
||||
const registry = new MCPServersRegistry(rawConfigs);
|
||||
|
||||
await registry.initialize();
|
||||
|
||||
// Verify disconnect was called for each server during initialization
|
||||
// All servers attempt to connect during initialization for metadata gathering
|
||||
const serverNames = Object.keys(rawConfigs);
|
||||
expect(mockConnectionsRepo.disconnect).toHaveBeenCalledTimes(serverNames.length);
|
||||
});
|
||||
|
||||
it('should log configuration updates for each startup-enabled server', async () => {
|
||||
const registry = new MCPServersRegistry(rawConfigs);
|
||||
|
||||
await registry.initialize();
|
||||
|
||||
const serverNames = Object.keys(rawConfigs);
|
||||
serverNames.forEach((serverName) => {
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
expect.stringContaining(`[MCP][${serverName}] URL:`),
|
||||
);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
expect.stringContaining(`[MCP][${serverName}] OAuth Required:`),
|
||||
);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
expect.stringContaining(`[MCP][${serverName}] Capabilities:`),
|
||||
);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
expect.stringContaining(`[MCP][${serverName}] Tools:`),
|
||||
);
|
||||
expect(mockLogger.info).toHaveBeenCalledWith(
|
||||
expect.stringContaining(`[MCP][${serverName}] Server Instructions:`),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
it('should have parsedConfigs matching the expected fixture after initialization', async () => {
|
||||
const registry = new MCPServersRegistry(rawConfigs);
|
||||
|
||||
await registry.initialize();
|
||||
|
||||
// Compare the actual parsedConfigs against the expected fixture
|
||||
expect(registry.parsedConfigs).toEqual(expectedParsedConfigs);
|
||||
});
|
||||
|
||||
it('should handle serverInstructions as string "true" correctly and fetch from server', async () => {
|
||||
// Create test config with serverInstructions as string "true"
|
||||
const testConfig: t.MCPServers = {
|
||||
test_server_string_true: {
|
||||
type: 'stdio',
|
||||
args: [],
|
||||
command: 'test-command',
|
||||
serverInstructions: 'true', // Simulating string "true" from YAML parsing
|
||||
},
|
||||
test_server_custom_string: {
|
||||
type: 'stdio',
|
||||
args: [],
|
||||
command: 'test-command',
|
||||
serverInstructions: 'Custom instructions here',
|
||||
},
|
||||
test_server_bool_true: {
|
||||
type: 'stdio',
|
||||
args: [],
|
||||
command: 'test-command',
|
||||
serverInstructions: true,
|
||||
},
|
||||
};
|
||||
|
||||
const registry = new MCPServersRegistry(testConfig);
|
||||
|
||||
// Setup mock connection for servers that should fetch
|
||||
const mockClient = {
|
||||
listTools: jest.fn().mockResolvedValue({ tools: [] }),
|
||||
getInstructions: jest.fn().mockReturnValue('Fetched instructions from server'),
|
||||
getServerCapabilities: jest.fn().mockReturnValue({ tools: {} }),
|
||||
};
|
||||
const mockConnection = {
|
||||
client: mockClient,
|
||||
} as unknown as jest.Mocked<MCPConnection>;
|
||||
|
||||
mockConnectionsRepo.get.mockResolvedValue(mockConnection);
|
||||
mockConnectionsRepo.getLoaded.mockResolvedValue(
|
||||
new Map([
|
||||
['test_server_string_true', mockConnection],
|
||||
['test_server_bool_true', mockConnection],
|
||||
]),
|
||||
);
|
||||
mockDetectOAuthRequirement.mockResolvedValue({
|
||||
requiresOAuth: false,
|
||||
method: 'no-metadata-found',
|
||||
metadata: null,
|
||||
});
|
||||
|
||||
await registry.initialize();
|
||||
|
||||
// Verify that string "true" was treated as fetch-from-server
|
||||
expect(registry.parsedConfigs['test_server_string_true'].serverInstructions).toBe(
|
||||
'Fetched instructions from server',
|
||||
);
|
||||
|
||||
// Verify that custom string was kept as-is
|
||||
expect(registry.parsedConfigs['test_server_custom_string'].serverInstructions).toBe(
|
||||
'Custom instructions here',
|
||||
);
|
||||
|
||||
// Verify that boolean true also fetched from server
|
||||
expect(registry.parsedConfigs['test_server_bool_true'].serverInstructions).toBe(
|
||||
'Fetched instructions from server',
|
||||
);
|
||||
|
||||
// Verify getInstructions was called for both "true" cases
|
||||
expect(mockClient.getInstructions).toHaveBeenCalledTimes(2);
|
||||
});
|
||||
|
||||
it('should use Promise.allSettled for individual server initialization', async () => {
|
||||
const registry = new MCPServersRegistry(rawConfigs);
|
||||
|
||||
// Spy on Promise.allSettled to verify it's being used
|
||||
const allSettledSpy = jest.spyOn(Promise, 'allSettled');
|
||||
|
||||
await registry.initialize();
|
||||
|
||||
// Verify Promise.allSettled was called with an array of server initialization promises
|
||||
expect(allSettledSpy).toHaveBeenCalledWith(expect.arrayContaining([expect.any(Promise)]));
|
||||
|
||||
// Verify it was called with the correct number of server promises
|
||||
const serverNames = Object.keys(rawConfigs);
|
||||
expect(allSettledSpy).toHaveBeenCalledWith(
|
||||
expect.arrayContaining(new Array(serverNames.length).fill(expect.any(Promise))),
|
||||
);
|
||||
|
||||
allSettledSpy.mockRestore();
|
||||
});
|
||||
|
||||
it('should isolate server failures and not affect other servers', async () => {
|
||||
const registry = new MCPServersRegistry(rawConfigs);
|
||||
|
||||
// Make multiple servers fail in different ways
|
||||
mockConnectionsRepo.get.mockImplementation((serverName: string) => {
|
||||
if (serverName === 'stdio_server') {
|
||||
// First server fails
|
||||
throw new Error('Connection failed for stdio_server');
|
||||
}
|
||||
if (serverName === 'websocket_server') {
|
||||
// Second server fails
|
||||
throw new Error('Connection failed for websocket_server');
|
||||
}
|
||||
// Other servers succeed
|
||||
const connection = mockConnections.get(serverName);
|
||||
if (!connection) {
|
||||
throw new Error(`Connection not found for server: ${serverName}`);
|
||||
}
|
||||
return Promise.resolve(connection);
|
||||
});
|
||||
|
||||
await registry.initialize();
|
||||
|
||||
// Despite failures, initialization should complete
|
||||
expect(registry.oauthServers).toBeInstanceOf(Set);
|
||||
expect(registry.toolFunctions).toBeDefined();
|
||||
|
||||
// Successful servers should still be processed
|
||||
expect(registry.appServerConfigs).toHaveProperty('non_oauth_server');
|
||||
|
||||
// Failed servers should not crash the whole initialization
|
||||
expect(mockLogger.warn).toHaveBeenCalledWith(
|
||||
expect.stringContaining('[MCP][stdio_server] Failed to fetch server capabilities:'),
|
||||
expect.any(Error),
|
||||
);
|
||||
expect(mockLogger.warn).toHaveBeenCalledWith(
|
||||
expect.stringContaining('[MCP][websocket_server] Failed to fetch server capabilities:'),
|
||||
expect.any(Error),
|
||||
);
|
||||
});
|
||||
|
||||
it('should properly clean up connections even when some servers fail', async () => {
|
||||
const registry = new MCPServersRegistry(rawConfigs);
|
||||
|
||||
// Track disconnect failures but suppress unhandled rejections
|
||||
const disconnectErrors: Error[] = [];
|
||||
mockConnectionsRepo.disconnect.mockImplementation((serverName: string) => {
|
||||
if (serverName === 'stdio_server') {
|
||||
const error = new Error('Disconnect failed');
|
||||
disconnectErrors.push(error);
|
||||
return Promise.reject(error).catch(() => {}); // Suppress unhandled rejection
|
||||
}
|
||||
return Promise.resolve();
|
||||
});
|
||||
|
||||
await registry.initialize();
|
||||
|
||||
// Should still attempt to disconnect all servers during initialization
|
||||
const serverNames = Object.keys(rawConfigs);
|
||||
expect(mockConnectionsRepo.disconnect).toHaveBeenCalledTimes(serverNames.length);
|
||||
expect(disconnectErrors).toHaveLength(1);
|
||||
});
|
||||
|
||||
it('should timeout individual server initialization after configured timeout', async () => {
|
||||
const timeout = 2000;
|
||||
// Create registry with a short timeout for testing
|
||||
process.env.MCP_INIT_TIMEOUT_MS = `${timeout}`;
|
||||
|
||||
const registry = new MCPServersRegistry(rawConfigs);
|
||||
|
||||
// Make one server hang indefinitely during OAuth detection
|
||||
mockDetectOAuthRequirement.mockImplementation((url: string) => {
|
||||
if (url === 'https://api.github.com/mcp') {
|
||||
// Slow init
|
||||
return new Promise((res) => setTimeout(res, timeout * 2));
|
||||
}
|
||||
// Return normal responses for other servers
|
||||
return Promise.resolve({
|
||||
requiresOAuth: false,
|
||||
method: 'no-metadata-found',
|
||||
metadata: null,
|
||||
});
|
||||
});
|
||||
|
||||
const start = Date.now();
|
||||
await registry.initialize();
|
||||
const duration = Date.now() - start;
|
||||
|
||||
// Should complete within reasonable time despite one server hanging
|
||||
// Allow some buffer for test execution overhead
|
||||
expect(duration).toBeLessThan(timeout * 1.5);
|
||||
|
||||
// The timeout should prevent the hanging server from blocking initialization
|
||||
// Other servers should still be processed successfully
|
||||
expect(registry.appServerConfigs).toHaveProperty('stdio_server');
|
||||
expect(registry.appServerConfigs).toHaveProperty('non_oauth_server');
|
||||
}, 10_000); // 10 second Jest timeout
|
||||
|
||||
it('should skip tool function fetching if connection was not established', async () => {
|
||||
const testConfig: t.MCPServers = {
|
||||
server_with_connection: {
|
||||
type: 'stdio',
|
||||
args: [],
|
||||
command: 'test-command',
|
||||
},
|
||||
server_without_connection: {
|
||||
type: 'stdio',
|
||||
args: [],
|
||||
command: 'failing-command',
|
||||
},
|
||||
};
|
||||
|
||||
const registry = new MCPServersRegistry(testConfig);
|
||||
|
||||
const mockClient = {
|
||||
listTools: jest.fn().mockResolvedValue({
|
||||
tools: [
|
||||
{
|
||||
name: 'test_tool',
|
||||
description: 'Test tool',
|
||||
inputSchema: { type: 'object', properties: {} },
|
||||
},
|
||||
],
|
||||
}),
|
||||
getInstructions: jest.fn().mockReturnValue(undefined),
|
||||
getServerCapabilities: jest.fn().mockReturnValue({ tools: {} }),
|
||||
};
|
||||
const mockConnection = {
|
||||
client: mockClient,
|
||||
} as unknown as jest.Mocked<MCPConnection>;
|
||||
|
||||
mockConnectionsRepo.get.mockImplementation((serverName: string) => {
|
||||
if (serverName === 'server_with_connection') {
|
||||
return Promise.resolve(mockConnection);
|
||||
}
|
||||
throw new Error('Connection failed');
|
||||
});
|
||||
|
||||
// Mock getLoaded to return connections map - the real implementation returns all loaded connections at once
|
||||
mockConnectionsRepo.getLoaded.mockResolvedValue(
|
||||
new Map([['server_with_connection', mockConnection]]),
|
||||
);
|
||||
|
||||
mockDetectOAuthRequirement.mockResolvedValue({
|
||||
requiresOAuth: false,
|
||||
method: 'no-metadata-found',
|
||||
metadata: null,
|
||||
});
|
||||
|
||||
await registry.initialize();
|
||||
|
||||
expect(registry.toolFunctions).toHaveProperty('test_tool_mcp_server_with_connection');
|
||||
expect(Object.keys(registry.toolFunctions)).toHaveLength(1);
|
||||
});
|
||||
|
||||
it('should handle getLoaded returning empty map gracefully', async () => {
|
||||
const testConfig: t.MCPServers = {
|
||||
test_server: {
|
||||
type: 'stdio',
|
||||
args: [],
|
||||
command: 'test-command',
|
||||
},
|
||||
};
|
||||
|
||||
const registry = new MCPServersRegistry(testConfig);
|
||||
|
||||
mockConnectionsRepo.get.mockRejectedValue(new Error('All connections failed'));
|
||||
mockConnectionsRepo.getLoaded.mockResolvedValue(new Map());
|
||||
mockDetectOAuthRequirement.mockResolvedValue({
|
||||
requiresOAuth: false,
|
||||
method: 'no-metadata-found',
|
||||
metadata: null,
|
||||
});
|
||||
|
||||
await registry.initialize();
|
||||
|
||||
expect(registry.toolFunctions).toEqual({});
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -1,67 +0,0 @@
|
|||
# Expected parsed MCP server configurations after running initialize()
|
||||
# These represent the expected state of parsedConfigs after all fetch functions complete
|
||||
|
||||
oauth_server:
|
||||
_processed: true
|
||||
type: "streamable-http"
|
||||
url: "https://api.github.com/mcp"
|
||||
headers:
|
||||
Authorization: "Bearer {{GITHUB_TOKEN}}"
|
||||
serverInstructions: true
|
||||
requiresOAuth: true
|
||||
oauthMetadata:
|
||||
authorization_url: "https://github.com/login/oauth/authorize"
|
||||
token_url: "https://github.com/login/oauth/access_token"
|
||||
|
||||
oauth_predefined:
|
||||
_processed: true
|
||||
type: "sse"
|
||||
url: "https://api.example.com/sse"
|
||||
requiresOAuth: true
|
||||
oauthMetadata:
|
||||
authorization_url: "https://example.com/oauth/authorize"
|
||||
token_url: "https://example.com/oauth/token"
|
||||
|
||||
stdio_server:
|
||||
_processed: true
|
||||
command: "node"
|
||||
args: ["server.js"]
|
||||
env:
|
||||
API_KEY: "${TEST_API_KEY}"
|
||||
startup: true
|
||||
serverInstructions: "Follow these instructions for stdio server"
|
||||
requiresOAuth: false
|
||||
capabilities: '{"tools":{"listChanged":true},"resources":{"listChanged":true},"prompts":{}}'
|
||||
tools: "file_read, file_write"
|
||||
|
||||
websocket_server:
|
||||
_processed: true
|
||||
type: "websocket"
|
||||
url: "ws://localhost:3001/mcp"
|
||||
startup: true
|
||||
requiresOAuth: false
|
||||
oauthMetadata: null
|
||||
capabilities: '{"tools":{},"resources":{},"prompts":{}}'
|
||||
tools: ""
|
||||
|
||||
disabled_server:
|
||||
_processed: true
|
||||
requiresOAuth: false
|
||||
type: "streamable-http"
|
||||
url: "https://api.disabled.com/mcp"
|
||||
startup: false
|
||||
|
||||
non_oauth_server:
|
||||
_processed: true
|
||||
type: "streamable-http"
|
||||
url: "https://api.public.com/mcp"
|
||||
requiresOAuth: false
|
||||
serverInstructions: "Public API instructions"
|
||||
capabilities: '{"tools":{},"resources":{},"prompts":{}}'
|
||||
tools: ""
|
||||
|
||||
oauth_startup_enabled:
|
||||
_processed: true
|
||||
type: "sse"
|
||||
url: "https://api.oauth-startup.com/sse"
|
||||
requiresOAuth: true
|
||||
|
|
@ -1,53 +0,0 @@
|
|||
# Raw MCP server configurations used as input to MCPServersRegistry constructor
|
||||
# These configs test different code paths in the initialization process
|
||||
|
||||
# Test OAuth detection with URL - should trigger fetchOAuthRequirement
|
||||
oauth_server:
|
||||
type: "streamable-http"
|
||||
url: "https://api.github.com/mcp"
|
||||
headers:
|
||||
Authorization: "Bearer {{GITHUB_TOKEN}}"
|
||||
serverInstructions: true
|
||||
|
||||
# Test OAuth already specified - should skip OAuth detection
|
||||
oauth_predefined:
|
||||
type: "sse"
|
||||
url: "https://api.example.com/sse"
|
||||
requiresOAuth: true
|
||||
oauthMetadata:
|
||||
authorization_url: "https://example.com/oauth/authorize"
|
||||
token_url: "https://example.com/oauth/token"
|
||||
|
||||
# Test stdio server without URL - should set requiresOAuth to false
|
||||
stdio_server:
|
||||
command: "node"
|
||||
args: ["server.js"]
|
||||
env:
|
||||
API_KEY: "${TEST_API_KEY}"
|
||||
startup: true
|
||||
serverInstructions: "Follow these instructions for stdio server"
|
||||
|
||||
# Test websocket server with capabilities but no tools
|
||||
websocket_server:
|
||||
type: "websocket"
|
||||
url: "ws://localhost:3001/mcp"
|
||||
startup: true
|
||||
|
||||
# Test server with startup disabled - should not be included in appServerConfigs
|
||||
disabled_server:
|
||||
type: "streamable-http"
|
||||
url: "https://api.disabled.com/mcp"
|
||||
startup: false
|
||||
|
||||
# Test non-OAuth server - should be included in appServerConfigs
|
||||
non_oauth_server:
|
||||
type: "streamable-http"
|
||||
url: "https://api.public.com/mcp"
|
||||
requiresOAuth: false
|
||||
serverInstructions: true
|
||||
|
||||
# Test server with OAuth but startup enabled - should not be in appServerConfigs
|
||||
oauth_startup_enabled:
|
||||
type: "sse"
|
||||
url: "https://api.oauth-startup.com/sse"
|
||||
requiresOAuth: true
|
||||
Loading…
Add table
Add a link
Reference in a new issue