🔒 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.

packages/api/src/mcp/MCPConnectionFactory.ts

@@ -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,

packages/api/src/mcp/oauth/handler.ts

@@ -1,5 +1,7 @@
 import { randomBytes } from 'crypto';
 import { logger } from '@librechat/data-schemas';
+import { FetchLike } from '@modelcontextprotocol/sdk/shared/transport';
+import { OAuthMetadataSchema } from '@modelcontextprotocol/sdk/shared/auth.js';
 import {
   registerClient,
   startAuthorization,
@@ -7,7 +9,6 @@ import {
   discoverAuthorizationServerMetadata,
   discoverOAuthProtectedResourceMetadata,
 } from '@modelcontextprotocol/sdk/client/auth.js';
-import { OAuthMetadataSchema } from '@modelcontextprotocol/sdk/shared/auth.js';
 import type { MCPOptions } from 'librechat-data-provider';
 import type { FlowStateManager } from '~/flow/manager';
 import type {
@@ -18,7 +19,6 @@ import type {
   OAuthMetadata,
 } from './types';
 import { sanitizeUrlForLogging } from '~/mcp/utils';
-import { FetchLike } from '@modelcontextprotocol/sdk/shared/transport';
 
 /** Type for the OAuth metadata from the SDK */
 type SDKOAuthMetadata = Parameters<typeof registerClient>[1]['metadata'];
@@ -439,9 +439,10 @@ export class MCPOAuthHandler {
         fetchFn: this.createOAuthFetch(oauthHeaders),
       });
 
-      logger.debug('[MCPOAuth] Raw tokens from exchange:', {
-        access_token: tokens.access_token ? '[REDACTED]' : undefined,
-        refresh_token: tokens.refresh_token ? '[REDACTED]' : undefined,
+      logger.debug('[MCPOAuth] Token exchange successful', {
+        flowId,
+        has_access_token: !!tokens.access_token,
+        has_refresh_token: !!tokens.refresh_token,
         expires_in: tokens.expires_in,
         token_type: tokens.token_type,
         scope: tokens.scope,