import { useRef, useCallback } from 'react'; interface PollingState { activeWorkflowId: string | null; isPolling: boolean; isPollInProgress: boolean; isPaused: boolean; currentInterval: number; failureCount: number; rateLimitFailureCount: number; timeoutId: NodeJS.Timeout | null; } const BASE_INTERVAL = 5000; // 5 seconds const MAX_INTERVAL = 10000; // 10 seconds const BACKOFF_MULTIPLIER = 1.5; const RATE_LIMIT_BACKOFF_MULTIPLIER = 2.0; const MAX_RATE_LIMIT_FAILURES = 5; export type PollCallback = (workflowId: string) => Promise; export function useWorkflowPolling() { const stateRef = useRef({ activeWorkflowId: null, isPolling: false, isPollInProgress: false, isPaused: false, currentInterval: BASE_INTERVAL, failureCount: 0, rateLimitFailureCount: 0, timeoutId: null }); const pollCallbackRef = useRef(null); const calculateInterval = useCallback((isRateLimit: boolean = false): number => { const state = stateRef.current; const multiplier = isRateLimit ? RATE_LIMIT_BACKOFF_MULTIPLIER : BACKOFF_MULTIPLIER; const newInterval = Math.min( BASE_INTERVAL * Math.pow(multiplier, state.failureCount), MAX_INTERVAL ); return Math.floor(newInterval); }, []); const scheduleNextPoll = useCallback((interval: number) => { const state = stateRef.current; // Clear any existing timeout if (state.timeoutId) { clearTimeout(state.timeoutId); state.timeoutId = null; } // Don't schedule if not polling or paused if (!state.isPolling || state.isPaused || !state.activeWorkflowId) { return; } // Schedule next poll state.timeoutId = setTimeout(() => { state.timeoutId = null; doPolling(); }, interval); }, []); const doPolling = useCallback(async () => { const state = stateRef.current; // Prevent concurrent polls if (state.isPollInProgress) { return; } // Validate workflow is still active if (!state.activeWorkflowId || !state.isPolling || state.isPaused) { return; } const workflowId = state.activeWorkflowId; state.isPollInProgress = true; try { if (pollCallbackRef.current) { await pollCallbackRef.current(workflowId); } // Success - reset failure counts and interval state.failureCount = 0; state.rateLimitFailureCount = 0; state.currentInterval = BASE_INTERVAL; // Schedule next poll scheduleNextPoll(state.currentInterval); } catch (error: any) { // Handle errors const isRateLimit = error?.status === 429 || error?.response?.status === 429; if (isRateLimit) { state.rateLimitFailureCount++; // Stop polling after too many rate limit errors if (state.rateLimitFailureCount >= MAX_RATE_LIMIT_FAILURES) { console.error('Too many rate limit errors, stopping polling'); stopPolling(); return; } } else { state.rateLimitFailureCount = 0; // Reset rate limit count on non-rate-limit errors } state.failureCount++; const nextInterval = calculateInterval(isRateLimit); state.currentInterval = nextInterval; console.warn(`Polling error (attempt ${state.failureCount}):`, error); // Schedule next poll with backoff scheduleNextPoll(nextInterval); } finally { state.isPollInProgress = false; } }, [scheduleNextPoll, calculateInterval]); const startPolling = useCallback((workflowId: string, callback: PollCallback) => { const state = stateRef.current; // Stop any existing polling if (state.isPolling) { stopPolling(); } // Validate workflow ID if (!workflowId || typeof workflowId !== 'string') { console.error('Invalid workflow ID for polling:', workflowId); return; } // Set up polling state state.activeWorkflowId = workflowId; state.isPolling = true; state.isPaused = false; state.failureCount = 0; state.rateLimitFailureCount = 0; state.currentInterval = BASE_INTERVAL; pollCallbackRef.current = callback; // Execute immediate first poll (no delay) doPolling(); }, [doPolling]); const stopPolling = useCallback(() => { const state = stateRef.current; // Clear timeout if (state.timeoutId) { clearTimeout(state.timeoutId); state.timeoutId = null; } // Reset state state.isPolling = false; state.isPollInProgress = false; state.activeWorkflowId = null; state.failureCount = 0; state.rateLimitFailureCount = 0; state.currentInterval = BASE_INTERVAL; state.isPaused = false; pollCallbackRef.current = null; }, []); const pausePolling = useCallback(() => { const state = stateRef.current; state.isPaused = true; }, []); const resumePolling = useCallback(() => { const state = stateRef.current; if (state.isPolling && state.isPaused) { state.isPaused = false; // Resume polling immediately if (!state.isPollInProgress) { scheduleNextPoll(0); } } }, [scheduleNextPoll]); const isPolling = useCallback((): boolean => { return stateRef.current.isPolling && !stateRef.current.isPaused; }, []); const getActiveWorkflowId = useCallback((): string | null => { return stateRef.current.activeWorkflowId; }, []); return { startPolling, stopPolling, pausePolling, resumePolling, isPolling, getActiveWorkflowId }; }