refactor: Enhance OAuth polling with gradual backoff and timeout handling; update reconnection tracking

This commit is contained in:
Danny Avila 2025-09-21 07:55:32 -04:00
parent 95fb9436fe
commit cbfbdeb787
No known key found for this signature in database
GPG key ID: BF31EEB2C5CA0956
3 changed files with 88 additions and 10 deletions

View file

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

View file

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

View file

@ -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<string, Set<string>> = 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<string, Set<string>> = new Map();
/** Map of userId:serverName -> timestamp when reconnection started */
private activeTimestamps: Map<string, number> = 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);
}
}