Fix LLM Streaming in React: Prevent Re-renders, Jank, and Token Flooding

import { useEffect, useMemo, useRef, useState, startTransition } from "react";

export function useLLMStream() {
  const [text, setText] = useState("");
  const [isStreaming, setIsStreaming] = useState(false);

  const abortRef = useRef(null);
  const bufferRef = useRef("");
  const flushTimerRef = useRef(null);

  const flush = useMemo(() => {
    let last = 0;
    return (force = false) => {
      const now = performance.now();
      if (!force && now - last < 50) return; // cadence: 50ms
      last = now;

      const chunk = bufferRef.current;
      if (!chunk) return;
      bufferRef.current = "";

      startTransition(() => {
        setText((t) => t + chunk);
      });
    };
  }, []);

  useEffect(() => {
    return () => {
      if (abortRef.current) abortRef.current.abort();
      if (flushTimerRef.current) clearInterval(flushTimerRef.current);
    };
  }, []);

  async function start(prompt) {
    // cancel any in-flight stream
    if (abortRef.current) abortRef.current.abort();

    const ac = new AbortController();
    abortRef.current = ac;

    setText("");
    bufferRef.current = "";
    setIsStreaming(true);

    // flush buffer on a steady cadence
    if (flushTimerRef.current) clearInterval(flushTimerRef.current);
    flushTimerRef.current = setInterval(() => flush(false), 50);

    try {
      const res = await fetch("/api/llm/stream", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ prompt }),
        signal: ac.signal
      });

      if (!res.ok) throw new Error(`HTTP ${res.status}`);
      if (!res.body) throw new Error("Missing response body stream");

      const reader = res.body.getReader();
      const decoder = new TextDecoder();

      while (true) {
        const { value, done } = await reader.read();
        if (done) break;

        // Decode bytes to string (stream-safe)
        const str = decoder.decode(value, { stream: true });
        bufferRef.current += str;

        // Safety: if buffer is huge, flush immediately
        if (bufferRef.current.length > 32_000) flush(true);
      }

      // final flush
      flush(true);
    } finally {
      setIsStreaming(false);
      if (flushTimerRef.current) {
        clearInterval(flushTimerRef.current);
        flushTimerRef.current = null;
      }
    }
  }

  function stop() {
    abortRef.current?.abort();
  }

  return { text, isStreaming, start, stop };
}