mirror of
https://github.com/danny-avila/LibreChat.git
synced 2025-12-18 17:30:16 +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
236
packages/api/src/mcp/UserConnectionManager.ts
Normal file
236
packages/api/src/mcp/UserConnectionManager.ts
Normal file
|
|
@ -0,0 +1,236 @@
|
|||
import { ErrorCode, McpError } from '@modelcontextprotocol/sdk/types.js';
|
||||
import { logger } from '@librechat/data-schemas';
|
||||
import type { TokenMethods } 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 '~/mcp/MCPConnectionFactory';
|
||||
import { MCPServersRegistry } from '~/mcp/MCPServersRegistry';
|
||||
import { MCPConnection } from './connection';
|
||||
import type * as t from './types';
|
||||
|
||||
/**
|
||||
* Abstract base class for managing user-specific MCP connections with lifecycle management.
|
||||
* Only meant to be extended by MCPManager.
|
||||
* Much of the logic was move here from the old MCPManager to make it more manageable.
|
||||
* User connections will soon be ephemeral and not cached anymore:
|
||||
* https://github.com/danny-avila/LibreChat/discussions/8790
|
||||
*/
|
||||
export abstract class UserConnectionManager {
|
||||
protected readonly serversRegistry: MCPServersRegistry;
|
||||
protected userConnections: Map<string, Map<string, MCPConnection>> = new Map();
|
||||
/** Last activity timestamp for users (not per server) */
|
||||
protected userLastActivity: Map<string, number> = new Map();
|
||||
protected readonly USER_CONNECTION_IDLE_TIMEOUT = 15 * 60 * 1000; // 15 minutes (TODO: make configurable)
|
||||
|
||||
constructor(serverConfigs: t.MCPServers) {
|
||||
this.serversRegistry = new MCPServersRegistry(serverConfigs);
|
||||
}
|
||||
|
||||
/** fetches am MCP Server config from the registry */
|
||||
public getRawConfig(serverName: string): t.MCPOptions | undefined {
|
||||
return this.serversRegistry.rawConfigs[serverName];
|
||||
}
|
||||
|
||||
/** Updates the last activity timestamp for a user */
|
||||
protected updateUserLastActivity(userId: string): void {
|
||||
const now = Date.now();
|
||||
this.userLastActivity.set(userId, now);
|
||||
logger.debug(
|
||||
`[MCP][User: ${userId}] Updated last activity timestamp: ${new Date(now).toISOString()}`,
|
||||
);
|
||||
}
|
||||
|
||||
/** Gets or creates a connection for a specific user */
|
||||
public async getUserConnection({
|
||||
user,
|
||||
serverName,
|
||||
flowManager,
|
||||
customUserVars,
|
||||
tokenMethods,
|
||||
oauthStart,
|
||||
oauthEnd,
|
||||
signal,
|
||||
returnOnOAuth = false,
|
||||
}: {
|
||||
user: TUser;
|
||||
serverName: string;
|
||||
flowManager: FlowStateManager<MCPOAuthTokens | null>;
|
||||
customUserVars?: Record<string, string>;
|
||||
tokenMethods?: TokenMethods;
|
||||
oauthStart?: (authURL: string) => Promise<void>;
|
||||
oauthEnd?: () => Promise<void>;
|
||||
signal?: AbortSignal;
|
||||
returnOnOAuth?: boolean;
|
||||
}): Promise<MCPConnection> {
|
||||
const userId = user.id;
|
||||
if (!userId) {
|
||||
throw new McpError(ErrorCode.InvalidRequest, `[MCP] User object missing id property`);
|
||||
}
|
||||
|
||||
const userServerMap = this.userConnections.get(userId);
|
||||
let connection = userServerMap?.get(serverName);
|
||||
const now = Date.now();
|
||||
|
||||
// Check if user is idle
|
||||
const lastActivity = this.userLastActivity.get(userId);
|
||||
if (lastActivity && now - lastActivity > this.USER_CONNECTION_IDLE_TIMEOUT) {
|
||||
logger.info(`[MCP][User: ${userId}] User idle for too long. Disconnecting all connections.`);
|
||||
// Disconnect all user connections
|
||||
try {
|
||||
await this.disconnectUserConnections(userId);
|
||||
} catch (err) {
|
||||
logger.error(`[MCP][User: ${userId}] Error disconnecting idle connections:`, err);
|
||||
}
|
||||
connection = undefined; // Force creation of a new connection
|
||||
} else if (connection) {
|
||||
if (await connection.isConnected()) {
|
||||
logger.debug(`[MCP][User: ${userId}][${serverName}] Reusing active connection`);
|
||||
this.updateUserLastActivity(userId);
|
||||
return connection;
|
||||
} else {
|
||||
// Connection exists but is not connected, attempt to remove potentially stale entry
|
||||
logger.warn(
|
||||
`[MCP][User: ${userId}][${serverName}] Found existing but disconnected connection object. Cleaning up.`,
|
||||
);
|
||||
this.removeUserConnection(userId, serverName); // Clean up maps
|
||||
connection = undefined;
|
||||
}
|
||||
}
|
||||
|
||||
// If no valid connection exists, create a new one
|
||||
if (!connection) {
|
||||
logger.info(`[MCP][User: ${userId}][${serverName}] Establishing new connection`);
|
||||
}
|
||||
|
||||
const config = this.serversRegistry.parsedConfigs[serverName];
|
||||
if (!config) {
|
||||
throw new McpError(
|
||||
ErrorCode.InvalidRequest,
|
||||
`[MCP][User: ${userId}] Configuration for server "${serverName}" not found.`,
|
||||
);
|
||||
}
|
||||
|
||||
try {
|
||||
connection = await MCPConnectionFactory.create(
|
||||
{
|
||||
serverName: serverName,
|
||||
serverConfig: config,
|
||||
},
|
||||
{
|
||||
useOAuth: true,
|
||||
user: user,
|
||||
customUserVars: customUserVars,
|
||||
flowManager: flowManager,
|
||||
tokenMethods: tokenMethods,
|
||||
signal: signal,
|
||||
oauthStart: oauthStart,
|
||||
oauthEnd: oauthEnd,
|
||||
returnOnOAuth: returnOnOAuth,
|
||||
},
|
||||
);
|
||||
|
||||
if (!(await connection?.isConnected())) {
|
||||
throw new Error('Failed to establish connection after initialization attempt.');
|
||||
}
|
||||
|
||||
if (!this.userConnections.has(userId)) {
|
||||
this.userConnections.set(userId, new Map());
|
||||
}
|
||||
this.userConnections.get(userId)?.set(serverName, connection);
|
||||
|
||||
logger.info(`[MCP][User: ${userId}][${serverName}] Connection successfully established`);
|
||||
// Update timestamp on creation
|
||||
this.updateUserLastActivity(userId);
|
||||
return connection;
|
||||
} catch (error) {
|
||||
logger.error(`[MCP][User: ${userId}][${serverName}] Failed to establish connection`, error);
|
||||
// Ensure partial connection state is cleaned up if initialization fails
|
||||
await connection?.disconnect().catch((disconnectError) => {
|
||||
logger.error(
|
||||
`[MCP][User: ${userId}][${serverName}] Error during cleanup after failed connection`,
|
||||
disconnectError,
|
||||
);
|
||||
});
|
||||
// Ensure cleanup even if connection attempt fails
|
||||
this.removeUserConnection(userId, serverName);
|
||||
throw error; // Re-throw the error to the caller
|
||||
}
|
||||
}
|
||||
|
||||
/** Returns all connections for a specific user */
|
||||
public getUserConnections(userId: string) {
|
||||
return this.userConnections.get(userId);
|
||||
}
|
||||
|
||||
/** Removes a specific user connection entry */
|
||||
protected removeUserConnection(userId: string, serverName: string): void {
|
||||
const userMap = this.userConnections.get(userId);
|
||||
if (userMap) {
|
||||
userMap.delete(serverName);
|
||||
if (userMap.size === 0) {
|
||||
this.userConnections.delete(userId);
|
||||
// Only remove user activity timestamp if all connections are gone
|
||||
this.userLastActivity.delete(userId);
|
||||
}
|
||||
}
|
||||
|
||||
logger.debug(`[MCP][User: ${userId}][${serverName}] Removed connection entry.`);
|
||||
}
|
||||
|
||||
/** Disconnects and removes a specific user connection */
|
||||
public async disconnectUserConnection(userId: string, serverName: string): Promise<void> {
|
||||
const userMap = this.userConnections.get(userId);
|
||||
const connection = userMap?.get(serverName);
|
||||
if (connection) {
|
||||
logger.info(`[MCP][User: ${userId}][${serverName}] Disconnecting...`);
|
||||
await connection.disconnect();
|
||||
this.removeUserConnection(userId, serverName);
|
||||
}
|
||||
}
|
||||
|
||||
/** Disconnects and removes all connections for a specific user */
|
||||
public async disconnectUserConnections(userId: string): Promise<void> {
|
||||
const userMap = this.userConnections.get(userId);
|
||||
const disconnectPromises: Promise<void>[] = [];
|
||||
if (userMap) {
|
||||
logger.info(`[MCP][User: ${userId}] Disconnecting all servers...`);
|
||||
const userServers = Array.from(userMap.keys());
|
||||
for (const serverName of userServers) {
|
||||
disconnectPromises.push(
|
||||
this.disconnectUserConnection(userId, serverName).catch((error) => {
|
||||
logger.error(
|
||||
`[MCP][User: ${userId}][${serverName}] Error during disconnection:`,
|
||||
error,
|
||||
);
|
||||
}),
|
||||
);
|
||||
}
|
||||
await Promise.allSettled(disconnectPromises);
|
||||
// Ensure user activity timestamp is removed
|
||||
this.userLastActivity.delete(userId);
|
||||
logger.info(`[MCP][User: ${userId}] All connections processed for disconnection.`);
|
||||
}
|
||||
}
|
||||
|
||||
/** Check for and disconnect idle connections */
|
||||
protected checkIdleConnections(currentUserId?: string): void {
|
||||
const now = Date.now();
|
||||
|
||||
// Iterate through all users to check for idle ones
|
||||
for (const [userId, lastActivity] of this.userLastActivity.entries()) {
|
||||
if (currentUserId && currentUserId === userId) {
|
||||
continue;
|
||||
}
|
||||
if (now - lastActivity > this.USER_CONNECTION_IDLE_TIMEOUT) {
|
||||
logger.info(
|
||||
`[MCP][User: ${userId}] User idle for too long. Disconnecting all connections...`,
|
||||
);
|
||||
// Disconnect all user connections asynchronously (fire and forget)
|
||||
this.disconnectUserConnections(userId).catch((err) =>
|
||||
logger.error(`[MCP][User: ${userId}] Error disconnecting idle connections:`, err),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue