🔌 fix: Shared MCP Server Connection Management (#9822)

- Fixed a bug in reinitMCPServer where a user connection was created for an app-level server whenever this server is reinitialized
- Made MCPManager.getUserConnection to return an error if the connection is app-level
- Add MCPManager.getConnection to return either an app connection or a user connection based on the serverName
- Made MCPManager.appConnections public to avoid unnecessary wrapper methods.
This commit is contained in:
Theo N. Truong 2025-09-26 06:24:36 -06:00 committed by GitHub
parent 4f3683fd9a
commit 3219734b9e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 56 additions and 42 deletions

View file

@ -20,8 +20,6 @@ import { processMCPEnv } from '~/utils/env';
*/
export class MCPManager extends UserConnectionManager {
private static instance: MCPManager | null;
// Connections shared by all users.
private appConnections: ConnectionsRepository | null = null;
/** Creates and initializes the singleton MCPManager instance */
public static async createInstance(configs: t.MCPServers): Promise<MCPManager> {
@ -43,9 +41,25 @@ export class MCPManager extends UserConnectionManager {
this.appConnections = new ConnectionsRepository(this.serversRegistry.appServerConfigs!);
}
/** Returns all app-level connections */
public async getAllConnections(): Promise<Map<string, MCPConnection> | null> {
return this.appConnections!.getAll();
/** Retrieves an app-level or user-specific connection based on provided arguments */
public async getConnection(
args: {
serverName: string;
user?: TUser;
forceNew?: boolean;
flowManager?: FlowStateManager<MCPOAuthTokens | null>;
} & Omit<t.OAuthConnectionOptions, 'useOAuth' | 'user' | 'flowManager'>,
): Promise<MCPConnection> {
if (this.appConnections!.has(args.serverName)) {
return this.appConnections!.get(args.serverName);
} else if (args.user?.id) {
return this.getUserConnection(args as Parameters<typeof this.getUserConnection>[0]);
} else {
throw new McpError(
ErrorCode.InvalidRequest,
`No connection found for server ${args.serverName}`,
);
}
}
/** Get servers that require OAuth */
@ -180,30 +194,19 @@ Please follow these instructions when using tools from the respective MCP server
const logPrefix = userId ? `[MCP][User: ${userId}][${serverName}]` : `[MCP][${serverName}]`;
try {
if (!this.appConnections?.has(serverName) && userId && user) {
this.updateUserLastActivity(userId);
/** Get or create user-specific connection */
connection = await this.getUserConnection({
user,
serverName,
flowManager,
tokenMethods,
oauthStart,
oauthEnd,
signal: options?.signal,
customUserVars,
requestBody,
});
} else {
/** App-level connection */
connection = await this.appConnections!.get(serverName);
if (!connection) {
throw new McpError(
ErrorCode.InvalidRequest,
`${logPrefix} No app-level connection found. Cannot execute tool ${toolName}.`,
);
}
}
if (userId && user) this.updateUserLastActivity(userId);
connection = await this.getConnection({
serverName,
user,
flowManager,
tokenMethods,
oauthStart,
oauthEnd,
signal: options?.signal,
customUserVars,
requestBody,
});
if (!(await connection.isConnected())) {
/** May happen if getUserConnection failed silently or app connection dropped */