🔒 feat: Idempotency Check for OAuth Flow Completion (#10468)

* 🔒 feat: Implement idempotency check for OAuth flow completion

- Added a check to prevent duplicate token exchanges if the OAuth flow has already been completed.
- Updated the OAuth callback route to redirect appropriately when a completed flow is detected.
- Refactored token storage logic to use original flow state credentials instead of updated ones.
- Enhanced tests to cover the new idempotency behavior and ensure correct handling of OAuth flow states.

* chore: add back scope for logging

* refactor: Add isFlowStale method to FlowStateManager for stale flow detection

- Implemented a new method to check if a flow is stale based on its age and status.
- Updated MCPConnectionFactory to utilize the isFlowStale method for cleaning up stale OAuth flows.
- Enhanced logging to provide more informative messages regarding flow status and age during cleanup.

* test: Add unit tests for isFlowStale method in FlowStateManager

- Implemented comprehensive tests for the isFlowStale method to verify its behavior across various flow statuses (PENDING, COMPLETED, FAILED) and age thresholds.
- Ensured correct handling of edge cases, including flows with missing timestamps and custom stale thresholds.
- Enhanced test coverage to validate the logic for determining flow staleness based on createdAt, completedAt, and failedAt timestamps.
This commit is contained in:
Danny Avila 2025-11-12 08:44:45 -05:00 committed by GitHub
parent a49c509ebc
commit dd35f42073
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 380 additions and 72 deletions

View file

@ -7,9 +7,9 @@ import type { FlowMetadata } from '~/flow/types';
import type * as t from './types';
import { MCPTokenStorage, MCPOAuthHandler } from '~/mcp/oauth';
import { sanitizeUrlForLogging } from './utils';
import { withTimeout } from '~/utils/promise';
import { MCPConnection } from './connection';
import { processMCPEnv } from '~/utils';
import { withTimeout } from '~/utils/promise';
/**
* Factory for creating MCP connections with optional OAuth authentication.
@ -343,28 +343,42 @@ export class MCPConnectionFactory {
`${this.logPrefix} OAuth flow completed, tokens received for ${this.serverName}`,
);
// Re-fetch flow state after completion to get updated credentials
const updatedFlowState = await MCPOAuthHandler.getFlowState(
flowId,
this.flowManager as FlowStateManager<MCPOAuthTokens>,
);
/** Client information from the updated flow metadata */
/** Client information from the existing flow metadata */
const existingMetadata = existingFlow.metadata as unknown as MCPOAuthFlowMetadata;
const clientInfo = updatedFlowState?.clientInfo || existingMetadata?.clientInfo;
const clientInfo = existingMetadata?.clientInfo;
return { tokens, clientInfo };
}
// Clean up old completed flows: createFlow() may return cached results otherwise
// 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') {
try {
await this.flowManager.deleteFlow(flowId, 'mcp_oauth');
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 {
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} Cleared stale ${existingFlow.status} OAuth flow for ${flowId}`,
`${this.logPrefix} Skipping cleanup of recent ${status} flow (age: ${Math.round(age / 1000)}s, threshold: ${STALE_FLOW_THRESHOLD / 1000}s)`,
);
} catch (error) {
logger.warn(`${this.logPrefix} Failed to clear stale OAuth flow`, error);
// 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`,
);
}
}
}
@ -402,15 +416,9 @@ export class MCPConnectionFactory {
}
logger.info(`${this.logPrefix} OAuth flow completed, tokens received for ${this.serverName}`);
// Re-fetch flow state after completion to get updated credentials
const updatedFlowState = await MCPOAuthHandler.getFlowState(
newFlowId,
this.flowManager as FlowStateManager<MCPOAuthTokens>,
);
/** Client information from the updated flow state */
const clientInfo = updatedFlowState?.clientInfo || flowMetadata?.clientInfo;
const metadata = updatedFlowState?.metadata || flowMetadata?.metadata;
/** Client information from the flow metadata */
const clientInfo = flowMetadata?.clientInfo;
const metadata = flowMetadata?.metadata;
return {
tokens,