FusionAGI/frontend/src/hooks/useWebSocket.ts

import { useState, useCallback, useRef, useEffect } from 'react'
import type { WSEvent } from '../types'

type WSStatus = 'disconnected' | 'connecting' | 'connected' | 'error'

export function useWebSocket(sessionId: string | null) {
  const [status, setStatus] = useState<WSStatus>('disconnected')
  const [events, setEvents] = useState<WSEvent[]>([])
  const wsRef = useRef<WebSocket | null>(null)

  const connect = useCallback((sid: string) => {
    if (wsRef.current) wsRef.current.close()
    setStatus('connecting')
    const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'
    const ws = new WebSocket(`${protocol}//${window.location.host}/v1/sessions/${sid}/stream`)
    wsRef.current = ws

    ws.onopen = () => setStatus('connected')
    ws.onclose = () => setStatus('disconnected')
    ws.onerror = () => setStatus('error')
    ws.onmessage = (e) => {
      try {
        const event: WSEvent = JSON.parse(e.data)
        setEvents((prev) => [...prev, event])
      } catch { /* ignore malformed */ }
    }
  }, [])

  const send = useCallback((data: Record<string, unknown>) => {
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify(data))
    }
  }, [])

  const disconnect = useCallback(() => {
    wsRef.current?.close()
    wsRef.current = null
    setStatus('disconnected')
  }, [])

  const clearEvents = useCallback(() => setEvents([]), [])

  useEffect(() => () => { wsRef.current?.close() }, [])

  return { status, events, connect, send, disconnect, clearEvents }
}