Andreas/LibreChat
1
0
Fork 0
mirror of https://github.com/danny-avila/LibreChat.git synced 2025-12-29 14:48:51 +01:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

fix(mcp): clean up OAuth flows on abort and simplify flow handling

Browse source 
- Add abort handler in reconnectServer to clean up mcp_oauth and mcp_get_tokens flows
- Update createAbortHandler to clean up both flow types on tool call abort
- Pass abort signal to createFlow in returnOnOAuth path
- Simplify handleOAuthRequired to always cancel existing flows and start fresh
- This ensures user always gets a new OAuth URL instead of waiting for stale flows
This commit is contained in:
Danny Avila 2025-12-18 19:26:00 -05:00
parent cc931bcf51
commit 39adeac86e
No known key found for this signature in database
GPG key ID: BF31EEB2C5CA0956
2 changed files with 75 additions and 80 deletions
Download patch file Download diff file Expand all files Collapse all files

67
packages/api/src/mcp/MCPConnectionFactory.ts
View file

@ -1,7 +1,7 @@
import { logger } from '@librechat/data-schemas';
import type { OAuthClientInformation } from '@modelcontextprotocol/sdk/shared/auth.js';
import type { TokenMethods } from '@librechat/data-schemas';
import type { MCPOAuthTokens, MCPOAuthFlowMetadata, OAuthMetadata } from '~/mcp/oauth';
import type { MCPOAuthTokens, OAuthMetadata } from '~/mcp/oauth';
import type { FlowStateManager } from '~/flow/manager';
import type { FlowMetadata } from '~/flow/types';
import type * as t from './types';
@ -173,9 +173,10 @@ export class MCPConnectionFactory {
// Create the flow state so the OAuth callback can find it
// We spawn this in the background without waiting for it
this.flowManager!.createFlow(flowId, 'mcp_oauth', flowMetadata).catch(() => {
// Pass signal so the flow can be aborted if the request is cancelled
this.flowManager!.createFlow(flowId, 'mcp_oauth', flowMetadata, this.signal).catch(() => {
// The OAuth callback will resolve this flow, so we expect it to timeout here
// which is fine - we just need the flow state to exist
// or it will be aborted if the request is cancelled - both are fine
});
if (this.oauthStart) {
@ -354,56 +355,26 @@ export class MCPConnectionFactory {
/** Check if there's already an ongoing OAuth flow for this flowId */
const existingFlow = await this.flowManager.getFlowState(flowId, 'mcp_oauth');
if (existingFlow && existingFlow.status === 'PENDING') {
// If any flow exists (PENDING, COMPLETED, FAILED), cancel it and start fresh
// This ensures the user always gets a new OAuth URL instead of waiting for stale flows
if (existingFlow) {
logger.debug(
`${this.logPrefix} OAuth flow already exists for ${flowId}, waiting for completion`,
`${this.logPrefix} Found existing OAuth flow (status: ${existingFlow.status}), cancelling to start fresh`,
);
/** Tokens from existing flow to complete */
const tokens = await this.flowManager.createFlow(flowId, 'mcp_oauth');
if (typeof this.oauthEnd === 'function') {
await this.oauthEnd();
}
logger.info(
`${this.logPrefix} OAuth flow completed, tokens received for ${this.serverName}`,
);
/** Client information from the existing flow metadata */
const existingMetadata = existingFlow.metadata as unknown as MCPOAuthFlowMetadata;
const clientInfo = existingMetadata?.clientInfo;
return { tokens, clientInfo };
}
// Clean up old completed/failed flows, but only if they're actually stale
// This prevents race conditions where we delete a flow that's still being processed
if (existingFlow && existingFlow.status !== 'PENDING') {
const STALE_FLOW_THRESHOLD = 2 * 60 * 1000; // 2 minutes
const { isStale, age, status } = await this.flowManager.isFlowStale(
flowId,
'mcp_oauth',
STALE_FLOW_THRESHOLD,
);
if (isStale) {
try {
try {
if (existingFlow.status === 'PENDING') {
await this.flowManager.failFlow(
flowId,
'mcp_oauth',
new Error('Cancelled for new OAuth request'),
);
} else {
await this.flowManager.deleteFlow(flowId, 'mcp_oauth');
logger.debug(
`${this.logPrefix} Cleared stale ${status} OAuth flow (age: ${Math.round(age / 1000)}s)`,
);
} catch (error) {
logger.warn(`${this.logPrefix} Failed to clear stale OAuth flow`, error);
}
} else {
logger.debug(
`${this.logPrefix} Skipping cleanup of recent ${status} flow (age: ${Math.round(age / 1000)}s, threshold: ${STALE_FLOW_THRESHOLD / 1000}s)`,
);
// If flow is recent but not pending, something might be wrong
if (status === 'FAILED') {
logger.warn(
`${this.logPrefix} Recent OAuth flow failed, will retry after ${Math.round((STALE_FLOW_THRESHOLD - age) / 1000)}s`,
);
}
} catch (error) {
logger.warn(`${this.logPrefix} Failed to cancel existing OAuth flow`, error);
}
// Continue to start a new flow below
}
logger.debug(`${this.logPrefix} Initiating new OAuth flow for ${this.serverName}...`);