Spaces:
Sleeping
Sleeping
| import { useState, useEffect, useRef, useCallback } from 'react' | |
| export function useProcessingProgress(sessionId) { | |
| const [progress, setProgress] = useState({}) | |
| const wsRef = useRef(null) | |
| useEffect(() => { | |
| if (!sessionId) return | |
| // Connect to WebSocket | |
| const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:' | |
| const wsUrl = `${protocol}//${window.location.host}/api/ws/${sessionId}` | |
| const ws = new WebSocket(wsUrl) | |
| wsRef.current = ws | |
| ws.onopen = () => { | |
| console.log('WebSocket connected for session:', sessionId) | |
| } | |
| ws.onmessage = (event) => { | |
| try { | |
| const data = JSON.parse(event.data) | |
| if (data.type === 'progress') { | |
| setProgress(prev => ({ | |
| ...prev, | |
| [data.stem]: data.status | |
| })) | |
| } else if (data.type === 'complete') { | |
| console.log('Processing complete in', data.total_time, 'seconds') | |
| } | |
| } catch (err) { | |
| console.error('WebSocket message error:', err) | |
| } | |
| } | |
| ws.onerror = (error) => { | |
| console.error('WebSocket error:', error) | |
| } | |
| ws.onclose = () => { | |
| console.log('WebSocket disconnected') | |
| } | |
| // Cleanup | |
| return () => { | |
| if (ws.readyState === WebSocket.OPEN) { | |
| ws.close() | |
| } | |
| } | |
| }, [sessionId]) | |
| const resetProgress = useCallback(() => { | |
| setProgress({}) | |
| }, []) | |
| return { progress, resetProgress } | |
| } | |