Spaces:
Build error
Build error
File size: 2,032 Bytes
6423ff2 a3c3ef2 6423ff2 a3c3ef2 6423ff2 5bf2d26 6423ff2 5bf2d26 6423ff2 5bf2d26 6423ff2 5bf2d26 6423ff2 | 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 | import { useEffect, useRef, useState } from 'react';
type WebSocketStatus = 'connecting' | 'connected' | 'disconnected' | 'error';
interface WebSocketMessage {
status: 'pending' | 'processing' | 'completed' | 'failed';
stage?: string;
progress?: number;
message?: string;
audio_url?: string;
error?: string;
}
export function useGenerationWebSocket(generationId: string, isActive: boolean) {
const [status, setStatus] = useState<WebSocketStatus>('disconnected');
const [lastMessage, setLastMessage] = useState<WebSocketMessage | null>(null);
const wsRef = useRef<WebSocket | null>(null);
useEffect(() => {
// Only connect if the generation is active (pending/processing)
if (!isActive) {
return;
}
// Prevent duplicate connections
if (wsRef.current && (wsRef.current.readyState === WebSocket.OPEN || wsRef.current.readyState === WebSocket.CONNECTING)) {
return;
}
const connect = () => {
try {
setStatus('connecting');
const apiBase = process.env.NEXT_PUBLIC_API_URL || 'http://localhost:8000';
// Convert http(s) to ws(s)
const wsBase = apiBase.replace(/^http/, 'ws');
const wsUrl = `${wsBase}/api/v1/ws/generations/${generationId}`;
const ws = new WebSocket(wsUrl);
ws.onopen = () => {
setStatus('connected');
};
ws.onmessage = (event) => {
try {
const data = JSON.parse(event.data);
setLastMessage(data);
} catch {
// Invalid JSON - ignore
}
};
ws.onclose = () => {
setStatus('disconnected');
};
ws.onerror = () => {
setStatus('error');
};
wsRef.current = ws;
} catch {
setStatus('error');
}
};
connect();
return () => {
const ws = wsRef.current;
wsRef.current = null;
if (ws) {
ws.close();
}
};
}, [generationId, isActive]);
return { status, lastMessage };
}
|