rehearshalhub/web/src/hooks/useWebSocket.ts

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

type WsEvent = { event: string; data: unknown };
type EventHandler = (data: unknown) => void;

export function useVersionWebSocket(
  versionId: string | null,
  handlers: Record<string, EventHandler>
) {
  const handlersRef = useRef(handlers);
  useLayoutEffect(() => {
    handlersRef.current = handlers;
  });

  useEffect(() => {
    if (!versionId) return;
    const protocol = window.location.protocol === "https:" ? "wss" : "ws";
    const ws = new WebSocket(`${protocol}://${window.location.host}/ws/versions/${versionId}`);

    ws.onmessage = (evt) => {
      try {
        const msg: WsEvent = JSON.parse(evt.data);
        handlersRef.current[msg.event]?.(msg.data);
      } catch {
        // ignore malformed messages
      }
    };

    const pingInterval = setInterval(() => {
      if (ws.readyState === WebSocket.OPEN) {
        ws.send(JSON.stringify({ event: "ping" }));
      }
    }, 30_000);

    return () => {
      clearInterval(pingInterval);
      ws.close();
    };
  }, [versionId]);
}