🔒 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

@ -151,9 +151,25 @@ export class FlowStateManager<T = unknown> {
const flowState = (await this.keyv.get(flowKey)) as FlowState<T> | undefined;
if (!flowState) {
logger.warn('[FlowStateManager] Cannot complete flow - flow state not found', {
flowId,
type,
});
return false;
}
/** Prevent duplicate completion */
if (flowState.status === 'COMPLETED') {
logger.debug(
'[FlowStateManager] Flow already completed, skipping to prevent duplicate completion',
{
flowId,
type,
},
);
return true;
}
const updatedState: FlowState<T> = {
...flowState,
status: 'COMPLETED',
@ -162,9 +178,55 @@ export class FlowStateManager<T = unknown> {
};
await this.keyv.set(flowKey, updatedState, this.ttl);
logger.debug('[FlowStateManager] Flow completed successfully', {
flowId,
type,
});
return true;
}
/**
* Checks if a flow is stale based on its age and status
* @param flowId - The flow identifier
* @param type - The flow type
* @param staleThresholdMs - Age in milliseconds after which a non-pending flow is considered stale (default: 2 minutes)
* @returns Object with isStale boolean and age in milliseconds
*/
async isFlowStale(
flowId: string,
type: string,
staleThresholdMs: number = 2 * 60 * 1000,
): Promise<{ isStale: boolean; age: number; status?: string }> {
const flowKey = this.getFlowKey(flowId, type);
const flowState = (await this.keyv.get(flowKey)) as FlowState<T> | undefined;
if (!flowState) {
return { isStale: false, age: 0 };
}
if (flowState.status === 'PENDING') {
return { isStale: false, age: 0, status: flowState.status };
}
const completedAt = flowState.completedAt || flowState.failedAt;
const createdAt = flowState.createdAt;
let flowAge = 0;
if (completedAt) {
flowAge = Date.now() - completedAt;
} else if (createdAt) {
flowAge = Date.now() - createdAt;
}
return {
isStale: flowAge > staleThresholdMs,
age: flowAge,
status: flowState.status,
};
}
/**
* Marks a flow as failed
*/