Submatter/apps/web/lib/use-chat-subscription.ts

"use client";

import { useEffect, useRef } from "react";
import { useQueryClient } from "@tanstack/react-query";
import { getCollabUrl } from "@/lib/yjs-provider";

/** Invalidation-only transport. Durable chat history must be fetched through the authenticated API. */
export function useChatSubscription(
  workspace: string,
  channelIds: readonly string[],
  onHint?: () => void,
): void {
  const hint = useRef(onHint);
  hint.current = onHint;
  const queryClient = useQueryClient();
  const channels = [...new Set(channelIds)].sort().join(",");
  useEffect(() => {
    if (!workspace || !channels) return;
    let socket: WebSocket | undefined;
    let timer: ReturnType<typeof setTimeout> | undefined;
    let stopped = false;
    let attempts = 0;
    let revoked = false;
    const invalidate = () => {
      if (stopped) return;
      hint.current?.();
      void queryClient.invalidateQueries({ predicate: (query) => {
        const [path, options] = query.queryKey;
        return Array.isArray(path) && path[0] === "chat" &&
          !!options && typeof options === "object" && "input" in options &&
          !!options.input && typeof options.input === "object" &&
          "workspace" in options.input && options.input.workspace === workspace;
      } });
    };
    async function connect() {
      revoked = false;
      try {
        const response = await fetch("/api/chat/ticket", {
          method: "POST",
          credentials: "same-origin",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify({ workspace }),
          cache: "no-store",
        });
        if (!response.ok) {
          if (
            response.status === 401 ||
            response.status === 403 ||
            response.status === 404
          )
            return;
          throw new Error("Chat ticket unavailable");
        }
        const data: unknown = await response.json();
        if (
          !data ||
          typeof data !== "object" ||
          !("ticket" in data) ||
          typeof data.ticket !== "string"
        )
          throw new Error("Invalid chat ticket");
        if (stopped) return;
        const url = new URL(
          process.env.NEXT_PUBLIC_CHAT_SERVER_URL || getCollabUrl(),
        );
        if (!process.env.NEXT_PUBLIC_CHAT_SERVER_URL) url.port = "1235";
        url.pathname = "/chat";
        url.search = new URLSearchParams({ ticket: data.ticket }).toString();
        socket = new WebSocket(url);
        socket.onopen = () => {
          attempts = 0;
          invalidate(); // catch up after missed Redis notifications
          for (const channelId of channels.split(","))
            socket?.send(JSON.stringify({ type: "subscribe", channelId }));
        };
        socket.onmessage = (event) => {
          try {
            const message: unknown = JSON.parse(String(event.data));
            if (
              message &&
              typeof message === "object" &&
              "type" in message &&
              (message.type === "invalidate" ||
                message.type === "revoked" ||
                message.type === "subscribed")
            ) {
              if (message.type === "revoked") revoked = true;
              invalidate();
            }
          } catch {
            /* ignore malformed hints */
          }
        };
        socket.onclose = (event) => {
          if (event.code === 1008 && revoked) {
            invalidate();
          }
          if (!stopped)
            timer = setTimeout(
              () => {
                void connect();
              },
              Math.min(30_000, 1000 * 2 ** Math.min(attempts++, 5)),
            );
        };
      } catch {
        if (!stopped)
          timer = setTimeout(
            () => {
              void connect();
            },
            Math.min(30_000, 1000 * 2 ** Math.min(attempts++, 5)),
          );
      }
    }
    void connect();
    return () => {
      stopped = true;
      clearTimeout(timer);
      socket?.close();
    };
  }, [workspace, channels, queryClient]);
}