frontend/src/hooks/useStream.ts

/**
 * useStream — SSE streaming hook for POST /api/query/stream.
 *
 * Accumulates SSE `token` events into a `useRef` (not `useState`) to avoid
 * per-token re-renders. A `requestAnimationFrame` loop flushes the ref into
 * rendered state at ~60fps.
 *
 * Retries up to 3 times with exponential back-off (1s, 2s, 4s) on
 * network/non-2xx failure before setting an error state.
 *
 * On an `error` SSE event from the server the connection is closed
 * immediately (no retry — the error is from the pipeline, not the network).
 *
 * Requirements: 4.2, 4.3, 4.5, 4.6
 */

import { useRef, useState, useCallback, useEffect } from "react";
import { flushSync } from "react-dom";
import { streamQuery } from "../lib/api";
import type { Message, QueryResult } from "../types";

export interface StreamPayload {
  query: string;
  history: Message[];
  profession: string | null;
  language: string | null;
}

export interface UseStreamResult {
  /** Start a new streaming request. */
  startStream: (
    payload: StreamPayload,
    onDone: (metadata: QueryResult, finalContent: string) => void,
    onError: (message: string, finalContent: string) => void
  ) => void;
  /** True while an SSE connection is open. */
  isStreaming: boolean;
  /** Accumulated token text (ref-backed; flushed to state). */
  streamingContent: string;
  /** Error message if the stream failed after all retries. */
  error: string | null;
  /** Reset streaming state (call before starting a new session). */
  reset: () => void;
  /** Abort the in-flight request immediately. */
  abort: () => void;
}

const RETRY_DELAYS_MS = [1000, 2000, 4000];
const FLUSH_INTERVAL_MS = 50;

export function useStream(token: string): UseStreamResult {
  const [isStreaming, setIsStreaming] = useState(false);
  const [streamingContent, setStreamingContent] = useState("");
  const [error, setError] = useState<string | null>(null);

  // Accumulate tokens in a ref to avoid per-token re-renders
  const tokenBufferRef = useRef("");
  // Interval handle for the continuous flush loop (number in browser)
  const intervalRef = useRef<number | null>(null);
  // Abort controller for the current fetch
  const abortControllerRef = useRef<AbortController | null>(null);

  /** Flush the token buffer into React state synchronously. */
  const flush = useCallback(() => {
    flushSync(() => {
      setStreamingContent(tokenBufferRef.current);
    });
  }, []);

  /** Start a continuous flush loop that updates state every 50ms. */
  const startFlushLoop = useCallback(() => {
    if (intervalRef.current !== null) return; // already running
    intervalRef.current = setInterval(flush, FLUSH_INTERVAL_MS);
  }, [flush]);

  /** Stop the continuous flush loop. */
  const stopFlushLoop = useCallback(() => {
    if (intervalRef.current !== null) {
      clearInterval(intervalRef.current);
      intervalRef.current = null;
    }
  }, []);

  const reset = useCallback(() => {
    stopFlushLoop();
    tokenBufferRef.current = "";
    setStreamingContent("");
    setError(null);
    setIsStreaming(false);
  }, [stopFlushLoop]);

  // Cancel any in-flight request on unmount
  useEffect(() => {
    return () => {
      abortControllerRef.current?.abort();
      stopFlushLoop();
    };
  }, [stopFlushLoop]);

  const startStream = useCallback(
    (
      payload: StreamPayload,
      onDone: (metadata: QueryResult, finalContent: string) => void,
      onError: (message: string, finalContent: string) => void
    ) => {
      // Reset state for the new stream
      tokenBufferRef.current = "";
      setStreamingContent("");
      setError(null);
      setIsStreaming(true);

      // Start the continuous flush loop so tokens appear progressively
      startFlushLoop();

      const attempt = async (retriesLeft: number): Promise<void> => {
        // Cancel any previous request
        abortControllerRef.current?.abort();
        const controller = new AbortController();
        abortControllerRef.current = controller;

        try {
          await streamQuery(
            payload,
            token,
            controller.signal,
            // onToken: accumulate into ref (no setState — flush loop handles it)
            (text) => {
              tokenBufferRef.current += text;
            },
            // onDone
            (metadata) => {
              stopFlushLoop();
              const finalContent = tokenBufferRef.current;
              // Final flush
              setStreamingContent(finalContent);
              setIsStreaming(false);
              onDone(metadata, finalContent);
            },
            // onError (SSE error event — no retry)
            (message) => {
              stopFlushLoop();
              const finalContent = tokenBufferRef.current;
              setStreamingContent(finalContent);
              setIsStreaming(false);
              setError(message);
              onError(message, finalContent);
            }
          );
        } catch (err) {
          if (controller.signal.aborted) {
            // Intentional abort — do not retry or set error
            return;
          }

          // Do not retry on 401 — the token is invalid/expired. App.tsx will
          // call handle401() when it receives "Unauthorized" in onError.
          const isAuthError =
            err instanceof Error && err.message === "Unauthorized";

          if (retriesLeft > 0 && !isAuthError) {
            const delayIndex = RETRY_DELAYS_MS.length - retriesLeft;
            const delay = RETRY_DELAYS_MS[delayIndex] ?? 1000;
            await new Promise((resolve) => setTimeout(resolve, delay));
            return attempt(retriesLeft - 1);
          }

          // All retries exhausted
          stopFlushLoop();
          const finalContent = tokenBufferRef.current;
          setStreamingContent(finalContent);
          setIsStreaming(false);
          const message =
            err instanceof Error
              ? err.message
              : "Connection failed after 3 attempts. Please try again.";
          setError(message);
          onError(message, finalContent);
        }
      };

      void attempt(RETRY_DELAYS_MS.length);
    },
    [token, startFlushLoop, stopFlushLoop]
  );

  const abort = useCallback(() => {
    abortControllerRef.current?.abort();
    stopFlushLoop();
    setIsStreaming(false);
  }, [stopFlushLoop]);

  return { startStream, isStreaming, streamingContent, error, reset, abort };
}