Agent/frontend/react-template/src/hooks/useWebSocket.ts

import { useEffect, useRef, useState } from "react";

interface UseWebSocketOptions {
  onMessage?: (data: unknown) => void;
  onError?: (error: Event) => void;
  onClose?: () => void;
}

export const useWebSocket = (traceId: string | null, options: UseWebSocketOptions = {}) => {
  const wsRef = useRef<WebSocket | null>(null);
  const [connected, setConnected] = useState(false);
  const { onMessage, onError, onClose } = options;

  useEffect(() => {
    if (!traceId) return;

    const url = `ws://localhost:8000/api/traces/${traceId}/watch?since_event_id=0`;
    const ws = new WebSocket(url);

    ws.onopen = () => {
      setConnected(true);
    };

    ws.onmessage = (event) => {
      try {
        const data = JSON.parse(event.data);
        onMessage?.(data);
      } catch {
        onMessage?.(event.data);
      }
    };

    ws.onerror = (error) => {
      onError?.(error);
    };

    ws.onclose = () => {
      setConnected(false);
      onClose?.();
    };

    wsRef.current = ws;

    return () => {
      ws.close();
    };
  }, [traceId, onMessage, onError, onClose]);

  return { connected, ws: wsRef.current };
};