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;
  sinceEventId?: number;
}

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

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

    const httpBase =
      (typeof window !== "undefined"
        ? (window as unknown as { CONFIG?: { API_BASE_URL?: string } }).CONFIG?.API_BASE_URL
        : undefined) ||
      (typeof import.meta !== "undefined" && import.meta.env && import.meta.env.VITE_API_BASE_URL
        ? import.meta.env.VITE_API_BASE_URL
        : `${window.location.protocol}//${window.location.hostname}:8000`);

    const wsBase = httpBase.replace(/^http(s?):\/\//, "ws$1://").replace(/\/+$/, "");
    const url = `${wsBase}/api/traces/${traceId}/watch?since_event_id=${sinceEventId}`;
    const ws = new WebSocket(url);
    let pingTimer: number | null = null;

    ws.onopen = () => {
      setConnected(true);
      pingTimer = window.setInterval(() => {
        if (ws.readyState === WebSocket.OPEN) {
          ws.send("ping");
        }
      }, 15000);
    };

    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);
      if (pingTimer) {
        window.clearInterval(pingTimer);
        pingTimer = null;
      }
      onClose?.();
    };

    wsRef.current = ws;

    return () => {
      if (pingTimer) {
        window.clearInterval(pingTimer);
        pingTimer = null;
      }
      ws.close();
    };
  }, [traceId, onMessage, onError, onClose, sinceEventId]);

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