1
import { useEffect, useRef, useState } from "react";
2
import { createWebSocketTicket, type AuthDescriptor } from "@/client/lib/api";
3
 
4
type SocketState = "idle" | "connecting" | "open" | "closed" | "error";
5
 
6
interface Decoder<T> {
7
  assertDecode(value: unknown): T;
8
}
9
 
10
interface UseWebSocketOptions<TMessage> {
11
  address: string;
12
  auth: AuthDescriptor | null;
13
  enabled: boolean;
14
  messageCodec: Decoder<TMessage>;
15
  onMessage?: (message: TMessage) => void;
16
}
17
 
18
export function useWebSocket<TMessage>({
19
  address,
20
  auth,
21
  enabled,
22
  messageCodec,
23
  onMessage,
24
}: UseWebSocketOptions<TMessage>) {
25
  const callbackRef = useRef(onMessage);
26
  const [status, setStatus] = useState<SocketState>(enabled && auth ? "connecting" : "idle");
27
 
28
  // Reduce auth to a stable key so callers passing freshly-allocated
29
  // descriptor objects on each render do not retrigger the connect
30
  // effect. The descriptor itself is reconstructed inside the effect.
31
  const authMode = auth?.mode ?? "none";
32
  const userToken = auth?.mode === "user" ? auth.token : "";
33
 
34
  useEffect(() => {
35
    callbackRef.current = onMessage;
36
  }, [onMessage]);
37
 
38
  useEffect(() => {
39
    if (!enabled || authMode === "none") {
40
      setStatus("idle");
41
      return;
42
    }
43
    const effectiveAuth: AuthDescriptor = authMode === "user" ? { mode: "user", token: userToken } : { mode: "admin" };
44
 
45
    let closedByHook = false;
46
    let reconnectDelay = 1_000;
47
    let reconnectTimer: number | undefined;
48
    let heartbeatTimer: number | undefined;
49
    let socket: WebSocket | null = null;
50
 
51
    const scheduleReconnect = () => {
52
      if (closedByHook) {
53
        return;
54
      }
55
 
56
      setStatus("closed");
57
      reconnectTimer = window.setTimeout(() => {
58
        reconnectDelay = Math.min(reconnectDelay * 1.8, 10_000);
59
        void connect();
60
      }, reconnectDelay);
61
    };
62
 
63
    const connect = async () => {
64
      setStatus("connecting");
65
 
66
      try {
67
        const { ticket } = await createWebSocketTicket(address, effectiveAuth);
68
        if (closedByHook) {
69
          return;
70
        }
71
 
72
        const scheme = window.location.protocol === "https:" ? "wss" : "ws";
73
        socket = new WebSocket(
74
          `${scheme}://${window.location.host}/ws?address=${encodeURIComponent(address)}&ticket=${encodeURIComponent(ticket)}`,
75
        );
76
 
77
        socket.addEventListener("open", () => {
78
          reconnectDelay = 1_000;
79
          setStatus("open");
80
          heartbeatTimer = window.setInterval(() => {
81
            if (socket?.readyState === WebSocket.OPEN) {
82
              socket.send("ping");
83
            }
84
          }, 25_000);
85
        });
86
 
87
        socket.addEventListener("message", (event) => {
88
          try {
89
            const payload = messageCodec.assertDecode(JSON.parse(event.data as string));
90
            callbackRef.current?.(payload);
91
          } catch {
92
            // Ignore malformed payloads.
93
          }
94
        });
95
 
96
        socket.addEventListener("close", () => {
97
          if (heartbeatTimer) {
98
            window.clearInterval(heartbeatTimer);
99
          }
100
 
101
          if (closedByHook) {
102
            setStatus("closed");
103
            return;
104
          }
105
 
106
          scheduleReconnect();
107
        });
108
 
109
        socket.addEventListener("error", () => {
110
          setStatus("error");
111
        });
112
      } catch {
113
        if (closedByHook) {
114
          return;
115
        }
116
 
117
        setStatus("error");
118
        scheduleReconnect();
119
      }
120
    };
121
 
122
    void connect();
123
 
124
    return () => {
125
      closedByHook = true;
126
 
127
      if (reconnectTimer) {
128
        window.clearTimeout(reconnectTimer);
129
      }
130
 
131
      if (heartbeatTimer) {
132
        window.clearInterval(heartbeatTimer);
133
      }
134
 
135
      socket?.close();
136
    };
137
  }, [address, authMode, userToken, enabled, messageCodec]);
138
 
139
  return status;
140
}