claude-collab/ui/src/hooks/useWebSocket.ts

Claude Collab - The AI collaboration framework that prevents echo chambers

import { useEffect, useState, useRef, useCallback } from 'react'
import type { ConnectionStatus, DashboardData } from '../types/claude-collab'

interface UseWebSocketOptions {
  url?: string
  onMessage?: (data: DashboardData) => void
  reconnectInterval?: number
  maxReconnectAttempts?: number
}

export function useWebSocket({
  url = 'ws://localhost:8765',
  onMessage,
  reconnectInterval = 1000,
  maxReconnectAttempts = 10
}: UseWebSocketOptions = {}) {
  const [isConnected, setIsConnected] = useState(false)
  const [connectionStatus, setConnectionStatus] = useState<ConnectionStatus>('connecting')
  const wsRef = useRef<WebSocket | null>(null)
  const reconnectCount = useRef(0)
  const reconnectTimeoutRef = useRef<NodeJS.Timeout>()

  const connect = useCallback(() => {
    try {
      const ws = new WebSocket(url)
      
      ws.onopen = () => {
        console.log('WebSocket connected')
        setIsConnected(true)
        setConnectionStatus('connected')
        reconnectCount.current = 0
        
        // TODO: Implement dashboard authentication flow
        // For now, this will be rejected by server if not authenticated
        ws.send(JSON.stringify({ type: 'subscribe-dashboard' }))
      }

      ws.onmessage = (event) => {
        try {
          const data = JSON.parse(event.data)
          onMessage?.(data)
        } catch (error) {
          console.error('Failed to parse WebSocket message:', error)
        }
      }

      ws.onerror = (error) => {
        console.error('WebSocket error:', error)
        setConnectionStatus('error')
      }

      ws.onclose = () => {
        console.log('WebSocket disconnected')
        setIsConnected(false)
        setConnectionStatus('disconnected')
        wsRef.current = null

        // Attempt reconnection
        if (reconnectCount.current < maxReconnectAttempts) {
          setConnectionStatus('reconnecting')
          reconnectCount.current++
          reconnectTimeoutRef.current = setTimeout(() => {
            connect()
          }, reconnectInterval * Math.pow(1.5, reconnectCount.current))
        }
      }

      wsRef.current = ws
    } catch (error) {
      console.error('Failed to create WebSocket connection:', error)
      setConnectionStatus('error')
    }
  }, [url, onMessage, reconnectInterval, maxReconnectAttempts])

  useEffect(() => {
    connect()

    return () => {
      if (reconnectTimeoutRef.current) {
        clearTimeout(reconnectTimeoutRef.current)
      }
      if (wsRef.current) {
        wsRef.current.close()
      }
    }
  }, [connect])

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

  return {
    isConnected,
    connectionStatus,
    sendMessage
  }
}