File size: 2,113 Bytes
32bba92 | 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 | import { useState, useEffect, useRef, useCallback } from 'react'
export function useWebSocket(onMessage) {
const [status, setStatus] = useState('disconnected')
const wsRef = useRef(null)
const reconnectTimeoutRef = useRef(null)
const connect = useCallback(() => {
const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'
// Always use root WebSocket path (works with proxy)
const wsUrl = `${protocol}//${window.location.host}/ws/stream`
console.log('[WS] Connecting to:', wsUrl)
const ws = new WebSocket(wsUrl)
ws.binaryType = 'arraybuffer'
ws.onopen = () => {
console.log('[WS] Connected')
setStatus('connected')
}
ws.onclose = () => {
console.log('[WS] Disconnected')
setStatus('disconnected')
reconnectTimeoutRef.current = setTimeout(connect, 3000)
}
ws.onerror = (e) => {
console.error('[WS] Error:', e)
setStatus('error')
}
ws.onmessage = async (event) => {
if (event.data instanceof ArrayBuffer) {
onMessage?.({ type: 'binary', data: event.data })
} else {
try {
const data = JSON.parse(event.data)
onMessage?.(data)
} catch (e) {
console.error('[WS] Parse error:', e)
}
}
}
wsRef.current = ws
// Keep alive
const pingInterval = setInterval(() => {
if (ws.readyState === WebSocket.OPEN) {
ws.send(JSON.stringify({ type: 'ping' }))
}
}, 10000)
return () => {
clearInterval(pingInterval)
ws.close()
}
}, [onMessage])
useEffect(() => {
connect()
return () => {
if (reconnectTimeoutRef.current) {
clearTimeout(reconnectTimeoutRef.current)
}
if (wsRef.current) {
wsRef.current.close()
}
}
}, [connect])
const send = useCallback((data) => {
if (wsRef.current?.readyState === WebSocket.OPEN) {
if (typeof data === 'string') {
wsRef.current.send(data)
} else {
wsRef.current.send(JSON.stringify(data))
}
}
}, [])
return { status, send }
}
|