import { useState, useCallback, useRef, useEffect } from 'react' import { apiClient } from '../api/client' import type { SegmentationResult, JobStatus } from '../types' // Polling interval in milliseconds const POLLING_INTERVAL = 2000 /** * Hook for running segmentation with async job polling. * * Instead of waiting for the full inference to complete (which can timeout * on HuggingFace Spaces), this hook: * 1. Creates a job that returns immediately with a job ID * 2. Polls for job status/progress every 2 seconds * 3. Returns results when the job completes * * This avoids the ~60s gateway timeout on HF Spaces while providing * real-time progress updates to the user. */ export function useSegmentation() { // Result state const [result, setResult] = useState(null) const [error, setError] = useState(null) // Job tracking state const [jobId, setJobId] = useState(null) const [jobStatus, setJobStatus] = useState(null) const [progress, setProgress] = useState(0) const [progressMessage, setProgressMessage] = useState('') const [elapsedSeconds, setElapsedSeconds] = useState( undefined ) // Loading state - true from job creation until completion/failure const [isLoading, setIsLoading] = useState(false) // Refs for managing async operations const currentJobRef = useRef(null) const pollingIntervalRef = useRef(null) const abortControllerRef = useRef(null) /** * Stop polling for job status */ const stopPolling = useCallback(() => { if (pollingIntervalRef.current) { clearInterval(pollingIntervalRef.current) pollingIntervalRef.current = null } }, []) /** * Poll for job status and update state */ const pollJobStatus = useCallback( async (id: string, signal: AbortSignal) => { // Don't poll if this isn't the current job if (id !== currentJobRef.current) { stopPolling() return } try { const response = await apiClient.getJobStatus(id, signal) // Ignore results if job changed if (id !== currentJobRef.current) return // Update progress state setJobStatus(response.status) setProgress(response.progress) setProgressMessage(response.progressMessage) setElapsedSeconds(response.elapsedSeconds) // Handle completion if (response.status === 'completed' && response.result) { stopPolling() setIsLoading(false) setResult({ dwiUrl: response.result.dwiUrl, predictionUrl: response.result.predictionUrl, metrics: { caseId: response.result.caseId, diceScore: response.result.diceScore, volumeMl: response.result.volumeMl, elapsedSeconds: response.result.elapsedSeconds, }, }) } // Handle failure if (response.status === 'failed') { stopPolling() setIsLoading(false) setError(response.error || 'Job failed') setResult(null) } } catch (err) { // Ignore abort errors if (err instanceof Error && err.name === 'AbortError') return // Don't stop polling on transient network errors - retry next interval console.warn('Polling error (will retry):', err) } }, [stopPolling] ) /** * Start segmentation job and begin polling */ const runSegmentation = useCallback( async (caseId: string, fastMode = true) => { // Cancel any existing job/polling stopPolling() abortControllerRef.current?.abort() const abortController = new AbortController() abortControllerRef.current = abortController // Reset state setError(null) setResult(null) setProgress(0) setProgressMessage('Creating job...') setJobStatus('pending') setElapsedSeconds(undefined) setIsLoading(true) try { // Create the job const response = await apiClient.createSegmentJob( caseId, fastMode, abortController.signal ) // Store job reference const newJobId = response.jobId setJobId(newJobId) currentJobRef.current = newJobId setJobStatus(response.status) setProgressMessage(response.message) // Start polling pollingIntervalRef.current = window.setInterval(() => { pollJobStatus(newJobId, abortController.signal) }, POLLING_INTERVAL) // Do an initial poll immediately await pollJobStatus(newJobId, abortController.signal) } catch (err) { // Ignore abort errors if (err instanceof Error && err.name === 'AbortError') return const message = err instanceof Error ? err.message : 'Failed to start job' setError(message) setIsLoading(false) setJobStatus('failed') } }, [pollJobStatus, stopPolling] ) /** * Cancel the current job (stops polling, clears loading state) */ const cancelJob = useCallback(() => { stopPolling() abortControllerRef.current?.abort() currentJobRef.current = null setIsLoading(false) setJobStatus(null) setProgress(0) setProgressMessage('') }, [stopPolling]) // Cleanup on unmount useEffect(() => { return () => { stopPolling() abortControllerRef.current?.abort() } }, [stopPolling]) return { // Result data result, error, // Job status jobId, jobStatus, progress, progressMessage, elapsedSeconds, // Loading state isLoading, // Actions runSegmentation, cancelJob, } }