🔄 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

* 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 2025-10-31 13:00:21 -06:00 committed by GitHub
parent 961f87cfda
commit ce7e6edad8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
45 changed files with 3116 additions and 1150 deletions

View file

@ -1,6 +1,7 @@
import { TokenMethods } from '@librechat/data-schemas';
import { FlowStateManager, MCPConnection, MCPOAuthTokens, MCPOptions } from '../..';
import { MCPManager } from '../MCPManager';
import { mcpServersRegistry } from '../../mcp/registry/MCPServersRegistry';
import { OAuthReconnectionManager } from './OAuthReconnectionManager';
import { OAuthReconnectionTracker } from './OAuthReconnectionTracker';
@ -14,6 +15,12 @@ jest.mock('@librechat/data-schemas', () => ({
}));
jest.mock('../MCPManager');
jest.mock('../../mcp/registry/MCPServersRegistry', () => ({
mcpServersRegistry: {
getServerConfig: jest.fn(),
getOAuthServers: jest.fn(),
},
}));
describe('OAuthReconnectionManager', () => {
let flowManager: jest.Mocked<FlowStateManager<null>>;
@ -51,10 +58,10 @@ describe('OAuthReconnectionManager', () => {
getUserConnection: jest.fn(),
getUserConnections: jest.fn(),
disconnectUserConnection: jest.fn(),
getRawConfig: jest.fn(),
} as unknown as jest.Mocked<MCPManager>;
(MCPManager.getInstance as jest.Mock).mockReturnValue(mockMCPManager);
(mcpServersRegistry.getServerConfig as jest.Mock).mockResolvedValue({});
});
afterEach(() => {
@ -152,7 +159,7 @@ describe('OAuthReconnectionManager', () => {
it('should reconnect eligible servers', async () => {
const userId = 'user-123';
const oauthServers = new Set(['server1', 'server2', 'server3']);
mockMCPManager.getOAuthServers.mockReturnValue(oauthServers);
(mcpServersRegistry.getOAuthServers as jest.Mock).mockResolvedValue(oauthServers);
// server1: has failed reconnection
reconnectionTracker.setFailed(userId, 'server1');
@ -186,7 +193,9 @@ describe('OAuthReconnectionManager', () => {
mockMCPManager.getUserConnection.mockResolvedValue(
mockNewConnection as unknown as MCPConnection,
);
mockMCPManager.getRawConfig.mockReturnValue({ initTimeout: 5000 } as unknown as MCPOptions);
(mcpServersRegistry.getServerConfig as jest.Mock).mockResolvedValue({
initTimeout: 5000,
} as unknown as MCPOptions);
await reconnectionManager.reconnectServers(userId);
@ -215,7 +224,7 @@ describe('OAuthReconnectionManager', () => {
it('should handle failed reconnection attempts', async () => {
const userId = 'user-123';
const oauthServers = new Set(['server1']);
mockMCPManager.getOAuthServers.mockReturnValue(oauthServers);
(mcpServersRegistry.getOAuthServers as jest.Mock).mockResolvedValue(oauthServers);
// server1: has valid token
tokenMethods.findToken.mockResolvedValue({
@ -226,7 +235,9 @@ describe('OAuthReconnectionManager', () => {
// Mock failed connection
mockMCPManager.getUserConnection.mockRejectedValue(new Error('Connection failed'));
mockMCPManager.getRawConfig.mockReturnValue({} as unknown as MCPOptions);
(mcpServersRegistry.getServerConfig as jest.Mock).mockResolvedValue(
{} as unknown as MCPOptions,
);
await reconnectionManager.reconnectServers(userId);
@ -242,7 +253,7 @@ describe('OAuthReconnectionManager', () => {
it('should not reconnect servers with expired tokens', async () => {
const userId = 'user-123';
const oauthServers = new Set(['server1']);
mockMCPManager.getOAuthServers.mockReturnValue(oauthServers);
(mcpServersRegistry.getOAuthServers as jest.Mock).mockResolvedValue(oauthServers);
// server1: has expired token
tokenMethods.findToken.mockResolvedValue({
@ -261,7 +272,7 @@ describe('OAuthReconnectionManager', () => {
it('should handle connection that returns but is not connected', async () => {
const userId = 'user-123';
const oauthServers = new Set(['server1']);
mockMCPManager.getOAuthServers.mockReturnValue(oauthServers);
(mcpServersRegistry.getOAuthServers as jest.Mock).mockResolvedValue(oauthServers);
tokenMethods.findToken.mockResolvedValue({
userId,
@ -277,7 +288,9 @@ describe('OAuthReconnectionManager', () => {
mockMCPManager.getUserConnection.mockResolvedValue(
mockConnection as unknown as MCPConnection,
);
mockMCPManager.getRawConfig.mockReturnValue({} as unknown as MCPOptions);
(mcpServersRegistry.getServerConfig as jest.Mock).mockResolvedValue(
{} as unknown as MCPOptions,
);
await reconnectionManager.reconnectServers(userId);
@ -359,7 +372,7 @@ describe('OAuthReconnectionManager', () => {
it('should not attempt to reconnect servers that have timed out during reconnection', async () => {
const userId = 'user-123';
const oauthServers = new Set(['server1', 'server2']);
mockMCPManager.getOAuthServers.mockReturnValue(oauthServers);
(mcpServersRegistry.getOAuthServers as jest.Mock).mockResolvedValue(oauthServers);
const now = Date.now();
jest.setSystemTime(now);
@ -414,7 +427,7 @@ describe('OAuthReconnectionManager', () => {
const userId = 'user-123';
const serverName = 'server1';
const oauthServers = new Set([serverName]);
mockMCPManager.getOAuthServers.mockReturnValue(oauthServers);
(mcpServersRegistry.getOAuthServers as jest.Mock).mockResolvedValue(oauthServers);
const now = Date.now();
jest.setSystemTime(now);
@ -428,7 +441,9 @@ describe('OAuthReconnectionManager', () => {
// First reconnect attempt - will fail
mockMCPManager.getUserConnection.mockRejectedValueOnce(new Error('Connection failed'));
mockMCPManager.getRawConfig.mockReturnValue({} as unknown as MCPOptions);
(mcpServersRegistry.getServerConfig as jest.Mock).mockResolvedValue(
{} as unknown as MCPOptions,
);
await reconnectionManager.reconnectServers(userId);
await jest.runAllTimersAsync();

View file

@ -5,6 +5,7 @@ import type { MCPOAuthTokens } from './types';
import { OAuthReconnectionTracker } from './OAuthReconnectionTracker';
import { FlowStateManager } from '~/flow/manager';
import { MCPManager } from '~/mcp/MCPManager';
import { mcpServersRegistry } from '~/mcp/registry/MCPServersRegistry';
const DEFAULT_CONNECTION_TIMEOUT_MS = 10_000; // ms
@ -72,7 +73,7 @@ export class OAuthReconnectionManager {
// 1. derive the servers to reconnect
const serversToReconnect = [];
for (const serverName of this.mcpManager.getOAuthServers()) {
for (const serverName of await mcpServersRegistry.getOAuthServers()) {
const canReconnect = await this.canReconnect(userId, serverName);
if (canReconnect) {
serversToReconnect.push(serverName);
@ -104,7 +105,7 @@ export class OAuthReconnectionManager {
logger.info(`${logPrefix} Attempting reconnection`);
const config = this.mcpManager.getRawConfig(serverName);
const config = await mcpServersRegistry.getServerConfig(serverName, userId);
const cleanupOnFailedReconnect = () => {
this.reconnectionsTracker.setFailed(userId, serverName);