mirror of
https://github.com/danny-avila/LibreChat.git
synced 2025-12-20 02:10:15 +01:00
🔐 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:
parent
e3a645e8fb
commit
51f2d43fed
2 changed files with 39 additions and 6 deletions
|
|
@ -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 */
|
||||||
|
|
|
||||||
|
|
@ -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>;
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue