🔐 refactor: Improve MCP OAuth Event Handler Cleanup (#9584)

* 🔐 refactor: Improve MCP OAuth event handling and cleanup

* ci: MCPConnection mock with additional event handling methods
This commit is contained in:
Danny Avila 2025-09-11 18:54:43 -04:00 committed by GitHub
parent e3a645e8fb
commit 51f2d43fed
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 39 additions and 6 deletions

View file

@ -74,9 +74,23 @@ export class MCPConnectionFactory {
oauthTokens, oauthTokens,
}); });
if (this.useOAuth) this.handleOAuthEvents(connection); let cleanupOAuthHandlers: (() => void) | null = null;
if (this.useOAuth) {
cleanupOAuthHandlers = this.handleOAuthEvents(connection);
}
try {
await this.attemptToConnect(connection); await this.attemptToConnect(connection);
if (cleanupOAuthHandlers) {
cleanupOAuthHandlers();
}
return connection; return connection;
} catch (error) {
if (cleanupOAuthHandlers) {
cleanupOAuthHandlers();
}
throw error;
}
} }
/** Retrieves existing OAuth tokens from storage or returns null */ /** Retrieves existing OAuth tokens from storage or returns null */
@ -133,8 +147,8 @@ export class MCPConnectionFactory {
} }
/** Sets up OAuth event handlers for the connection */ /** Sets up OAuth event handlers for the connection */
protected handleOAuthEvents(connection: MCPConnection): void { protected handleOAuthEvents(connection: MCPConnection): () => void {
connection.on('oauthRequired', async (data) => { const oauthHandler = async (data: { serverUrl?: string }) => {
logger.info(`${this.logPrefix} oauthRequired event received`); logger.info(`${this.logPrefix} oauthRequired event received`);
// If we just want to initiate OAuth and return, handle it differently // If we just want to initiate OAuth and return, handle it differently
@ -202,7 +216,23 @@ export class MCPConnectionFactory {
logger.warn(`${this.logPrefix} OAuth failed, emitting oauthFailed event`); logger.warn(`${this.logPrefix} OAuth failed, emitting oauthFailed event`);
connection.emit('oauthFailed', new Error('OAuth authentication failed')); connection.emit('oauthFailed', new Error('OAuth authentication failed'));
} }
}); };
connection.on('oauthRequired', oauthHandler);
/** Handler reference for cleanup when connection state changes to disconnected */
const cleanupHandler = (state: string) => {
if (state === 'disconnected') {
connection.removeListener('oauthRequired', oauthHandler);
connection.removeListener('connectionChange', cleanupHandler);
}
};
connection.on('connectionChange', cleanupHandler);
return () => {
connection.removeListener('oauthRequired', oauthHandler);
connection.removeListener('connectionChange', cleanupHandler);
};
} }
/** Attempts to establish connection with timeout handling */ /** Attempts to establish connection with timeout handling */

View file

@ -56,6 +56,9 @@ describe('MCPConnectionFactory', () => {
isConnected: jest.fn(), isConnected: jest.fn(),
setOAuthTokens: jest.fn(), setOAuthTokens: jest.fn(),
on: jest.fn().mockReturnValue(mockConnectionInstance), on: jest.fn().mockReturnValue(mockConnectionInstance),
once: jest.fn().mockReturnValue(mockConnectionInstance),
off: jest.fn().mockReturnValue(mockConnectionInstance),
removeListener: jest.fn().mockReturnValue(mockConnectionInstance),
emit: jest.fn(), emit: jest.fn(),
} as unknown as jest.Mocked<MCPConnection>; } as unknown as jest.Mocked<MCPConnection>;