Download section didn't appear. Resolved

web/frontend/src/hooks/useWebSocket.ts

@@ -1,11 +1,13 @@
 import { useEffect, useRef } from 'react'
 import { useTaskStore } from '../stores/taskStore'
 import { useUIStore } from '../stores/uiStore'
+import { getStatus } from '../api/client'
 
 export const useWebSocket = (taskId: string | null) => {
   const wsRef = useRef<WebSocket | null>(null)
-  const reconnectTimeoutRef = useRef<NodeJS.Timeout | null>(null)
+  const reconnectTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null)
   const reconnectAttempts = useRef(0)
+  const pollingRef = useRef<ReturnType<typeof setTimeout> | null>(null)
 
   const {
     setStatus,
@@ -19,6 +21,55 @@ export const useWebSocket = (taskId: string | null) => {
   } = useTaskStore()
   const { setWsConnected } = useUIStore()
 
+  const pollStatus = async () => {
+    if (!taskId) return
+  
+    try {
+      const response = await getStatus(taskId)
+      console.log('[Polling] Status response:', response)
+  
+      // Update all state fields together
+      setStatus(response.status)
+      setProgress(response.progress)
+      setMessage(response.message)
+  
+      // Explicitly set tracks if present
+      if (response.tracks && response.tracks.length > 0) {
+        console.log('[Polling] Setting tracks:', response.tracks)
+        setTracks(response.tracks)
+      }
+  
+      if (response.status === 'done') {
+        console.log('[Polling] Split complete, tracks set:', response.tracks)
+        setIsProcessing(false)
+        // Ensure tracks are set one more time (safety)
+        if (response.tracks && response.tracks.length > 0) {
+          setTracks(response.tracks)
+        }
+        return
+      }
+  
+      if (response.status === 'error') {
+        setError(response.error || 'Split failed')
+        setIsProcessing(false)
+        return
+      }
+  
+      // Still processing – poll again
+      if (pollingRef.current) {
+        clearTimeout(pollingRef.current)
+      }
+      pollingRef.current = setTimeout(pollStatus, 2000)
+    } catch (error) {
+      console.error('[Polling] Error:', error)
+      if (pollingRef.current) {
+        clearTimeout(pollingRef.current)
+      }
+      pollingRef.current = setTimeout(pollStatus, 3000)
+    }
+  }
+
+
   useEffect(() => {
     if (!taskId) {
       if (wsRef.current) {
@@ -26,6 +77,11 @@ export const useWebSocket = (taskId: string | null) => {
         wsRef.current = null
       }
       setWsConnected(false)
+      // Clear polling
+      if (pollingRef.current) {
+        clearTimeout(pollingRef.current)
+        pollingRef.current = null
+      }
       return
     }
 
@@ -41,11 +97,17 @@ export const useWebSocket = (taskId: string | null) => {
           clearTimeout(reconnectTimeoutRef.current)
           reconnectTimeoutRef.current = null
         }
+        // Start polling when connection is established
+        // This ensures we get the final status even if WebSocket fails
+        setTimeout(() => {
+          pollStatus()
+        }, 1000)
       }
 
       ws.onmessage = (event) => {
         try {
           const data = JSON.parse(event.data)
+          console.log('[WebSocket] Message:', data)
 
           if (data.type === 'status') {
             const statusData = data.data
@@ -58,7 +120,6 @@ export const useWebSocket = (taskId: string | null) => {
             if (statusData.tracks) {
               setTracks(statusData.tracks)
             }
-            // If status is done or error, stop processing
             if (statusData.status === 'done' || statusData.status === 'error') {
               setIsProcessing(false)
             }
@@ -67,13 +128,12 @@ export const useWebSocket = (taskId: string | null) => {
             setStatus(progressData.status)
             setProgress(progressData.progress)
             setMessage(progressData.message)
+            if (progressData.tracks) {
+              setTracks(progressData.tracks)
+            }
             if (progressData.status === 'done') {
               addLog('✅ Split complete!')
               setIsProcessing(false)
-              // If tracks are included, update them
-              if (progressData.tracks) {
-                setTracks(progressData.tracks)
-              }
             } else if (progressData.status === 'error') {
               setError(progressData.message)
               addLog(`❌ Error: ${progressData.message}`)
@@ -83,7 +143,7 @@ export const useWebSocket = (taskId: string | null) => {
             }
           }
         } catch (error) {
-          console.error('Failed to parse WebSocket message:', error)
+          console.error('[WebSocket] Failed to parse message:', error)
         }
       }
 
@@ -91,18 +151,16 @@ export const useWebSocket = (taskId: string | null) => {
         console.log(`WebSocket disconnected for task ${taskId}`)
         setWsConnected(false)
 
-        // If task is still in processing state, reconnect
-        // We'll check via polling if needed
-        if (reconnectAttempts.current < 5) {
-          reconnectTimeoutRef.current = setTimeout(() => {
-            reconnectAttempts.current += 1
-            connect()
-          }, 3000)
+        // If task is not done and we have a taskId, start polling
+        // We check the status store to see if it's already done
+        if (taskId && status !== 'done' && status !== 'error') {
+          console.log('[WebSocket] Disconnected while processing, starting polling...')
+          setTimeout(pollStatus, 1000)
         }
       }
 
       ws.onerror = (error) => {
-        console.error('WebSocket error:', error)
+        console.error('[WebSocket] Error:', error)
         // onclose will handle reconnection
       }
 
@@ -120,9 +178,13 @@ export const useWebSocket = (taskId: string | null) => {
         clearTimeout(reconnectTimeoutRef.current)
         reconnectTimeoutRef.current = null
       }
+      if (pollingRef.current) {
+        clearTimeout(pollingRef.current)
+        pollingRef.current = null
+      }
       setWsConnected(false)
     }
-  }, [taskId, setStatus, setProgress, setMessage, setError, setTracks, addLog, setWsConnected, setIsProcessing])
+  }, [taskId, setStatus, setProgress, setMessage, setError, setTracks, addLog, setWsConnected, setIsProcessing, status])
 
   return wsRef.current
 }