🚦 feat: Auto-reinitialize MCP Servers on Request (#9226)

This commit is contained in:
Danny Avila 2025-08-23 03:27:05 -04:00
parent ac608ded46
commit c827fdd10e
No known key found for this signature in database
GPG key ID: BF31EEB2C5CA0956
28 changed files with 871 additions and 312 deletions

View file

@ -28,6 +28,7 @@ export class MCPConnectionFactory {
protected readonly oauthStart?: (authURL: string) => Promise<void>;
protected readonly oauthEnd?: () => Promise<void>;
protected readonly returnOnOAuth?: boolean;
protected readonly connectionTimeout?: number;
/** Creates a new MCP connection with optional OAuth support */
static async create(
@ -47,6 +48,7 @@ export class MCPConnectionFactory {
});
this.serverName = basic.serverName;
this.useOAuth = !!oauth?.useOAuth;
this.connectionTimeout = oauth?.connectionTimeout;
this.logPrefix = oauth?.user
? `[MCP][${basic.serverName}][${oauth.user.id}]`
: `[MCP][${basic.serverName}]`;
@ -82,8 +84,9 @@ export class MCPConnectionFactory {
if (!this.tokenMethods?.findToken) return null;
try {
const flowId = MCPOAuthHandler.generateFlowId(this.userId!, this.serverName);
const tokens = await this.flowManager!.createFlowWithHandler(
`tokens:${this.userId}:${this.serverName}`,
flowId,
'mcp_get_tokens',
async () => {
return await MCPTokenStorage.getTokens({
@ -203,7 +206,7 @@ export class MCPConnectionFactory {
/** Attempts to establish connection with timeout handling */
protected async attemptToConnect(connection: MCPConnection): Promise<void> {
const connectTimeout = this.serverConfig.initTimeout ?? 30000;
const connectTimeout = this.connectionTimeout ?? this.serverConfig.initTimeout ?? 30000;
const connectionTimeout = new Promise<void>((_, reject) =>
setTimeout(
() => reject(new Error(`Connection timeout after ${connectTimeout}ms`)),
@ -347,6 +350,7 @@ export class MCPConnectionFactory {
newFlowId,
'mcp_oauth',
flowMetadata as FlowMetadata,
this.signal,
);
if (typeof this.oauthEnd === 'function') {
await this.oauthEnd();