109 lines
3.2 KiB
TypeScript
109 lines
3.2 KiB
TypeScript
import { useEffect, useRef } from 'react'
|
|
import { useTaskStore } from '../stores/taskStore'
|
|
import { useUIStore } from '../stores/uiStore'
|
|
|
|
export const useWebSocket = (taskId: string | null) => {
|
|
const wsRef = useRef<WebSocket | null>(null)
|
|
const reconnectTimeoutRef = useRef<NodeJS.Timeout | null>(null)
|
|
const reconnectAttempts = useRef(0)
|
|
|
|
const { setStatus, setProgress, setMessage, setError, setTracks, addLog } = useTaskStore()
|
|
const { setWsConnected } = useUIStore()
|
|
|
|
useEffect(() => {
|
|
if (!taskId) {
|
|
if (wsRef.current) {
|
|
wsRef.current.close()
|
|
wsRef.current = null
|
|
}
|
|
setWsConnected(false)
|
|
return
|
|
}
|
|
|
|
const connect = () => {
|
|
const wsUrl = `/ws/${taskId}`
|
|
const ws = new WebSocket(wsUrl)
|
|
|
|
ws.onopen = () => {
|
|
console.log(`WebSocket connected for task ${taskId}`)
|
|
setWsConnected(true)
|
|
reconnectAttempts.current = 0
|
|
if (reconnectTimeoutRef.current) {
|
|
clearTimeout(reconnectTimeoutRef.current)
|
|
reconnectTimeoutRef.current = null
|
|
}
|
|
}
|
|
|
|
ws.onmessage = (event) => {
|
|
try {
|
|
const data = JSON.parse(event.data)
|
|
|
|
if (data.type === 'status') {
|
|
const status = data.data
|
|
setStatus(status.status)
|
|
setProgress(status.progress)
|
|
setMessage(status.message)
|
|
if (status.error) {
|
|
setError(status.error)
|
|
}
|
|
if (status.tracks) {
|
|
setTracks(status.tracks)
|
|
}
|
|
} else if (data.type === 'progress') {
|
|
const progressData = data.data
|
|
setStatus(progressData.status)
|
|
setProgress(progressData.progress)
|
|
setMessage(progressData.message)
|
|
if (progressData.status === 'done') {
|
|
addLog('✅ Split complete!')
|
|
} else if (progressData.status === 'error') {
|
|
setError(progressData.message)
|
|
addLog(`❌ Error: ${progressData.message}`)
|
|
} else {
|
|
addLog(`🔄 ${progressData.message} (${progressData.progress}%)`)
|
|
}
|
|
}
|
|
} catch (error) {
|
|
console.error('Failed to parse WebSocket message:', error)
|
|
}
|
|
}
|
|
|
|
ws.onclose = () => {
|
|
console.log(`WebSocket disconnected for task ${taskId}`)
|
|
setWsConnected(false)
|
|
|
|
// Try to reconnect if the task is still processing
|
|
// We'll check status via polling if needed
|
|
if (reconnectAttempts.current < 5) {
|
|
reconnectTimeoutRef.current = setTimeout(() => {
|
|
reconnectAttempts.current += 1
|
|
connect()
|
|
}, 3000)
|
|
}
|
|
}
|
|
|
|
ws.onerror = (error) => {
|
|
console.error('WebSocket error:', error)
|
|
// The onclose will handle reconnection
|
|
}
|
|
|
|
wsRef.current = ws
|
|
}
|
|
|
|
connect()
|
|
|
|
return () => {
|
|
if (wsRef.current) {
|
|
wsRef.current.close()
|
|
wsRef.current = null
|
|
}
|
|
if (reconnectTimeoutRef.current) {
|
|
clearTimeout(reconnectTimeoutRef.current)
|
|
reconnectTimeoutRef.current = null
|
|
}
|
|
setWsConnected(false)
|
|
}
|
|
}, [taskId, setStatus, setProgress, setMessage, setError, setTracks, addLog, setWsConnected])
|
|
|
|
return wsRef.current
|
|
} |