/** * 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(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(null); // Abort controller for the current fetch const abortControllerRef = useRef(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 => { // 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 }; }