Spaces:
Sleeping
Sleeping
| import { useState, useEffect, useRef, useCallback } from 'react' | |
| export function useWebSocket(sessionId) { | |
| const [messages, setMessages] = useState([]) | |
| const [isConnected, setIsConnected] = useState(false) | |
| const wsRef = useRef(null) | |
| const reconnectTimeoutRef = useRef(null) | |
| const connect = useCallback(() => { | |
| if (!sessionId) return | |
| const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:' | |
| const wsUrl = `${protocol}//${window.location.host}/api/ws/${sessionId}` | |
| try { | |
| wsRef.current = new WebSocket(wsUrl) | |
| wsRef.current.onopen = () => { | |
| setIsConnected(true) | |
| console.log('WebSocket connected') | |
| } | |
| wsRef.current.onmessage = (event) => { | |
| try { | |
| const data = JSON.parse(event.data) | |
| setMessages(prev => [...prev, data]) | |
| } catch (e) { | |
| console.error('Failed to parse WebSocket message:', e) | |
| } | |
| } | |
| wsRef.current.onclose = () => { | |
| setIsConnected(false) | |
| console.log('WebSocket disconnected') | |
| // Attempt to reconnect after 3 seconds | |
| reconnectTimeoutRef.current = setTimeout(() => { | |
| connect() | |
| }, 3000) | |
| } | |
| wsRef.current.onerror = (error) => { | |
| console.error('WebSocket error:', error) | |
| } | |
| } catch (e) { | |
| console.error('Failed to create WebSocket:', e) | |
| } | |
| }, [sessionId]) | |
| const disconnect = useCallback(() => { | |
| if (reconnectTimeoutRef.current) { | |
| clearTimeout(reconnectTimeoutRef.current) | |
| } | |
| if (wsRef.current) { | |
| wsRef.current.close() | |
| } | |
| }, []) | |
| useEffect(() => { | |
| if (sessionId) { | |
| connect() | |
| } | |
| return () => { | |
| disconnect() | |
| } | |
| }, [sessionId, connect, disconnect]) | |
| const clearMessages = useCallback(() => { | |
| setMessages([]) | |
| }, []) | |
| return { | |
| messages, | |
| isConnected, | |
| clearMessages | |
| } | |
| } | |