Spaces:
Sleeping
Sleeping
File size: 2,494 Bytes
fa81e4d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | import { useEffect, useState } from 'react';
import { useWebSocket } from '@/hooks/useWebSocket';
import type { WebSocketMessage } from '@/types';
interface EpisodeProgressProps {
episodeId: string | null;
onProgressUpdate?: (data: WebSocketMessage) => void;
onCompletion?: (data: WebSocketMessage) => void;
onError?: (error: string) => void;
}
export function useEpisodeProgress({
episodeId,
onProgressUpdate,
onCompletion,
onError,
}: EpisodeProgressProps) {
const [currentStep, setCurrentStep] = useState(0);
const [totalReward, setTotalReward] = useState(0);
const [progress, setProgress] = useState(0);
const [lastAction, setLastAction] = useState<string>('');
const [isCompleted, setIsCompleted] = useState(false);
const [error, setError] = useState<string | null>(null);
const { isConnected, lastMessage } = useWebSocket(
episodeId ? `/ws/episode/${episodeId}` : '',
{
autoConnect: !!episodeId,
onMessage: (message: WebSocketMessage) => {
console.log('WebSocket message:', message);
switch (message.type) {
case 'connected':
console.log('Connected to episode:', message.episode_id);
break;
case 'progress':
if (message.step !== undefined) setCurrentStep(message.step);
if (message.reward !== undefined) setTotalReward(prev => prev + message.reward!);
if (message.progress !== undefined) setProgress(message.progress);
if (message.action_type) setLastAction(message.action_type);
onProgressUpdate?.(message);
break;
case 'completion':
setIsCompleted(true);
if (message.total_reward !== undefined) setTotalReward(message.total_reward);
setProgress(100);
onCompletion?.(message);
break;
case 'error':
const errorMsg = message.error || 'Unknown error occurred';
setError(errorMsg);
onError?.(errorMsg);
break;
}
},
}
);
// Reset state when episode changes
useEffect(() => {
if (episodeId) {
setCurrentStep(0);
setTotalReward(0);
setProgress(0);
setLastAction('');
setIsCompleted(false);
setError(null);
}
}, [episodeId]);
return {
isConnected,
currentStep,
totalReward,
progress,
lastAction,
isCompleted,
error,
lastMessage,
};
}
export default useEpisodeProgress;
|