From cbfbdeb78750a4b1f1a78b111cff5abd28934fd3 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 21 Sep 2025 07:55:32 -0400 Subject: [PATCH] refactor: Enhance OAuth polling with gradual backoff and timeout handling; update reconnection tracking --- client/src/hooks/MCP/useMCPServerManager.ts | 71 ++++++++++++++++--- client/src/locales/en/translation.json | 1 + .../src/mcp/oauth/OAuthReconnectionTracker.ts | 26 ++++++- 3 files changed, 88 insertions(+), 10 deletions(-) diff --git a/client/src/hooks/MCP/useMCPServerManager.ts b/client/src/hooks/MCP/useMCPServerManager.ts index c30193919..d721b798c 100644 --- a/client/src/hooks/MCP/useMCPServerManager.ts +++ b/client/src/hooks/MCP/useMCPServerManager.ts @@ -130,7 +130,7 @@ export function useMCPServerManager({ conversationId }: { conversationId?: strin (serverName: string) => { const state = serverStates[serverName]; if (state?.pollInterval) { - clearInterval(state.pollInterval); + clearTimeout(state.pollInterval); } updateServerState(serverName, { isInitializing: false, @@ -145,8 +145,39 @@ export function useMCPServerManager({ conversationId }: { conversationId?: strin const startServerPolling = useCallback( (serverName: string) => { - const pollInterval = setInterval(async () => { + let pollAttempts = 0; + let timeoutId: NodeJS.Timeout | null = null; + + // OAuth typically completes in 5 seconds to 3 minutes + // Poll with gradual backoff from 5 to 8 seconds + const getPollInterval = (attempt: number): number => { + if (attempt < 12) return 5000; // First minute: every 5s (12 polls) + if (attempt < 22) return 6000; // Next minute: every 6s (10 polls) + if (attempt < 32) return 7000; // Next 70s: every 7s (10 polls) + return 8000; // Remaining time: every 8s + }; + + const maxAttempts = 47; // ~5.5 minutes total + + const pollOnce = async () => { try { + pollAttempts++; + + if (pollAttempts > maxAttempts) { + console.warn( + `[MCP Manager] Max polling attempts (${maxAttempts}) reached for ${serverName}`, + ); + showToast({ + message: localize('com_ui_mcp_connection_timeout', { 0: serverName }), + status: 'error', + }); + if (timeoutId) { + clearTimeout(timeoutId); + } + cleanupServerState(serverName); + return; + } + await queryClient.refetchQueries([QueryKeys.mcpConnectionStatus]); const freshConnectionData = queryClient.getQueryData([ @@ -158,7 +189,9 @@ export function useMCPServerManager({ conversationId }: { conversationId?: strin const serverStatus = freshConnectionStatus[serverName]; if (serverStatus?.connectionState === 'connected') { - clearInterval(pollInterval); + if (timeoutId) { + clearTimeout(timeoutId); + } showToast({ message: localize('com_ui_mcp_authenticated_success', { 0: serverName }), @@ -180,12 +213,15 @@ export function useMCPServerManager({ conversationId }: { conversationId?: strin return; } + // Check for OAuth timeout (3 minutes) if (state?.oauthStartTime && Date.now() - state.oauthStartTime > 180000) { showToast({ message: localize('com_ui_mcp_oauth_timeout', { 0: serverName }), status: 'error', }); - clearInterval(pollInterval); + if (timeoutId) { + clearTimeout(timeoutId); + } cleanupServerState(serverName); return; } @@ -195,19 +231,38 @@ export function useMCPServerManager({ conversationId }: { conversationId?: strin message: localize('com_ui_mcp_init_failed'), status: 'error', }); - clearInterval(pollInterval); + if (timeoutId) { + clearTimeout(timeoutId); + } cleanupServerState(serverName); return; } + + // Schedule next poll with smart intervals based on OAuth timing + const nextInterval = getPollInterval(pollAttempts); + + // Log progress periodically + if (pollAttempts % 5 === 0 || pollAttempts <= 2) { + console.debug( + `[MCP Manager] Polling ${serverName} attempt ${pollAttempts}/${maxAttempts}, next in ${nextInterval / 1000}s`, + ); + } + + timeoutId = setTimeout(pollOnce, nextInterval); + updateServerState(serverName, { pollInterval: timeoutId }); } catch (error) { console.error(`[MCP Manager] Error polling server ${serverName}:`, error); - clearInterval(pollInterval); + if (timeoutId) { + clearTimeout(timeoutId); + } cleanupServerState(serverName); return; } - }, 3500); + }; - updateServerState(serverName, { pollInterval }); + // Start the first poll + timeoutId = setTimeout(pollOnce, getPollInterval(0)); + updateServerState(serverName, { pollInterval: timeoutId }); }, [ queryClient, diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index a92c1887c..a03231d41 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -979,6 +979,7 @@ "com_ui_mcp_authenticated_success": "MCP server '{{0}}' authenticated successfully", "com_ui_mcp_configure_server": "Configure {{0}}", "com_ui_mcp_configure_server_description": "Configure custom variables for {{0}}", + "com_ui_mcp_connection_timeout": "Connection timeout for MCP server '{{0}}' - please try again", "com_ui_mcp_enter_var": "Enter value for {{0}}", "com_ui_mcp_init_failed": "Failed to initialize MCP server", "com_ui_mcp_initialize": "Initialize", diff --git a/packages/api/src/mcp/oauth/OAuthReconnectionTracker.ts b/packages/api/src/mcp/oauth/OAuthReconnectionTracker.ts index f18decd1a..b127eec15 100644 --- a/packages/api/src/mcp/oauth/OAuthReconnectionTracker.ts +++ b/packages/api/src/mcp/oauth/OAuthReconnectionTracker.ts @@ -1,14 +1,28 @@ export class OAuthReconnectionTracker { - // Map of userId -> Set of serverNames that have failed reconnection + /** Map of userId -> Set of serverNames that have failed reconnection */ private failed: Map> = new Map(); - // Map of userId -> Set of serverNames that are actively reconnecting + /** Map of userId -> Set of serverNames that are actively reconnecting */ private active: Map> = new Map(); + /** Map of userId:serverName -> timestamp when reconnection started */ + private activeTimestamps: Map = new Map(); + /** Maximum time (ms) a server can be in reconnecting state before auto-cleanup */ + private readonly RECONNECTION_TIMEOUT_MS = 5 * 60 * 1000; // 5 minutes public isFailed(userId: string, serverName: string): boolean { return this.failed.get(userId)?.has(serverName) ?? false; } public isActive(userId: string, serverName: string): boolean { + const key = `${userId}:${serverName}`; + const startTime = this.activeTimestamps.get(key); + + // Check if reconnection has timed out + if (startTime && Date.now() - startTime > this.RECONNECTION_TIMEOUT_MS) { + // Auto-cleanup timed out reconnection + this.removeActive(userId, serverName); + return false; + } + return this.active.get(userId)?.has(serverName) ?? false; } @@ -26,6 +40,10 @@ export class OAuthReconnectionTracker { } this.active.get(userId)?.add(serverName); + + /** Track when reconnection started */ + const key = `${userId}:${serverName}`; + this.activeTimestamps.set(key, Date.now()); } public removeFailed(userId: string, serverName: string): void { @@ -42,5 +60,9 @@ export class OAuthReconnectionTracker { if (userServers?.size === 0) { this.active.delete(userId); } + + /** Clear timestamp tracking */ + const key = `${userId}:${serverName}`; + this.activeTimestamps.delete(key); } }