Production Streaming UI with Gemini API + TanStack Query — Cancellation, Retries, and Cache Coherence

// app/api/chat/stream/route.ts
import { GoogleGenerativeAI } from "@google/generative-ai";
 
export const runtime = "nodejs";
export const dynamic = "force-dynamic"; // disable caching
 
const genAI = new GoogleGenerativeAI(process.env.GEMINI_API_KEY!);
 
export async function POST(req: Request) {
  const { messages, conversationId } = await req.json();
 
  // Detect client disconnects
  const abortSignal = req.signal;
 
  const model = genAI.getGenerativeModel({
    model: "gemini-2.5-flash",
    generationConfig: { temperature: 0.7, maxOutputTokens: 2048 },
  });
 
  const encoder = new TextEncoder();
  const stream = new ReadableStream({
    async start(controller) {
      try {
        const result = await model.generateContentStream({
          contents: messages.map((m: { role: string; content: string }) => ({
            role: m.role === "assistant" ? "model" : "user",
            parts: [{ text: m.content }],
          })),
        });
 
        for await (const chunk of result.stream) {
          // Bail immediately if the client disconnected
          if (abortSignal.aborted) {
            controller.close();
            return;
          }
          const text = chunk.text();
          if (!text) continue;
          // SSE format: "data: <JSON>\n\n"
          const payload = JSON.stringify({ type: "chunk", text });
          controller.enqueue(encoder.encode(`data: ${payload}\n\n`));
        }
        // Completion signal
        controller.enqueue(
          encoder.encode(`data: ${JSON.stringify({ type: "done", conversationId })}\n\n`)
        );
        controller.close();
      } catch (err: unknown) {
        const message = err instanceof Error ? err.message : "stream error";
        controller.enqueue(
          encoder.encode(`data: ${JSON.stringify({ type: "error", message })}\n\n`)
        );
        controller.close();
      }
    },
  });
 
  return new Response(stream, {
    headers: {
      "Content-Type": "text/event-stream",
      "Cache-Control": "no-cache, no-transform",
      Connection: "keep-alive",
      "X-Accel-Buffering": "no", // disables intermediate proxy buffering
    },
  });
}

// hooks/useStreamingChat.ts
import { useMutation, useQueryClient } from "@tanstack/react-query";
import { useState, useRef, useCallback, useEffect } from "react";
 
type Message = { role: "user" | "assistant"; content: string };
 
type StreamRequest = { conversationId: string; messages: Message[] };
 
async function* readSSE(response: Response, signal: AbortSignal) {
  if (!response.body) throw new Error("No response body");
  const reader = response.body.getReader();
  const decoder = new TextDecoder();
  let buffer = "";
  try {
    while (true) {
      if (signal.aborted) return;
      const { done, value } = await reader.read();
      if (done) return;
      buffer += decoder.decode(value, { stream: true });
      // Split on double-newline to find SSE event boundaries
      const events = buffer.split("\n\n");
      buffer = events.pop() ?? "";
      for (const ev of events) {
        const line = ev.split("\n").find((l) => l.startsWith("data: "));
        if (!line) continue;
        try {
          yield JSON.parse(line.slice(6));
        } catch {
          // Skip malformed events (common on flaky mobile networks)
        }
      }
    }
  } finally {
    reader.releaseLock();
  }
}
 
export function useStreamingChat(conversationId: string) {
  const qc = useQueryClient();
  const [streamingText, setStreamingText] = useState("");
  const abortRef = useRef<AbortController | null>(null);
 
  const mutation = useMutation({
    mutationKey: ["chat-stream", conversationId],
    mutationFn: async (req: StreamRequest) => {
      // Cancel any in-flight stream (handles rapid clicks)
      abortRef.current?.abort();
      const ac = new AbortController();
      abortRef.current = ac;
      setStreamingText("");
 
      const res = await fetch("/api/chat/stream", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify(req),
        signal: ac.signal,
      });
      if (!res.ok) throw new Error(`HTTP ${res.status}`);
 
      let full = "";
      for await (const event of readSSE(res, ac.signal)) {
        if (event.type === "chunk") {
          full += event.text;
          // Show partial text via React state
          setStreamingText(full);
        } else if (event.type === "error") {
          throw new Error(event.message);
        } else if (event.type === "done") {
          return { conversationId: event.conversationId, text: full };
        }
      }
      // Loop ended without seeing "done" (network drop, etc.)
      throw new Error("stream ended unexpectedly");
    },
    onSuccess: (data) => {
      // Sync finalized result into the query cache
      qc.setQueryData<Message[]>(["conversation", conversationId], (old = []) => [
        ...old,
        { role: "assistant", content: data.text },
      ]);
      setStreamingText("");
    },
    onError: () => {
      // Whether to keep partial text is a UX call.
      // I prefer to keep what was rendered and append a notice.
      if (streamingText) {
        qc.setQueryData<Message[]>(["conversation", conversationId], (old = []) => [
          ...old,
          { role: "assistant", content: streamingText + "\n\n[Response was interrupted]" },
        ]);
      }
      setStreamingText("");
    },
  });
 
  // Mandatory cleanup: cancel on unmount to stop runaway streams
  useEffect(() => {
    return () => {
      abortRef.current?.abort();
    };
  }, []);
 
  const cancel = useCallback(() => {
    abortRef.current?.abort();
  }, []);
 
  return { ...mutation, streamingText, cancel };
}

// CancelButton.tsx
import { useStreamingChat } from "@/hooks/useStreamingChat";
 
export function CancelButton({ conversationId }: { conversationId: string }) {
  const { cancel, isPending } = useStreamingChat(conversationId);
  if (!isPending) return null;
  return (
    <button
      type="button"
      onClick={cancel}
      className="rounded bg-red-500 px-3 py-1 text-white"
    >
      Stop
    </button>
  );
}

// ❌ Bad: only checks at the start
async function bad(stream: AsyncIterable<unknown>, abortSignal: AbortSignal) {
  if (abortSignal.aborted) return;
  for await (const chunk of stream) {
    // No way to cancel once we're here
    void chunk;
  }
}
 
// ✅ Good: checks every chunk
async function good(stream: AsyncIterable<unknown>, abortSignal: AbortSignal) {
  for await (const chunk of stream) {
    if (abortSignal.aborted) return;
    void chunk;
  }
}

// useMutation options excerpt
{
  retry: (failureCount, error) => {
    // If a stream had already started, don't auto-retry
    if (error instanceof Error && error.message === "stream ended unexpectedly") {
      return false; // avoid double-billing
    }
    // Pre-connection failures: retry up to 3 times
    return failureCount < 3;
  },
  retryDelay: (attempt) => Math.min(1000 * 2 ** attempt, 30000),
}

// hooks/useConversation.ts
import { useQuery } from "@tanstack/react-query";
 
type Message = { role: "user" | "assistant"; content: string };
 
export function useConversation(conversationId: string) {
  return useQuery<Message[]>({
    queryKey: ["conversation", conversationId],
    queryFn: async () => {
      const res = await fetch(`/api/conversations/${conversationId}`);
      if (!res.ok) throw new Error("Failed to load conversation");
      return res.json();
    },
    // Suppress duplicate fetches during streaming
    staleTime: 60_000,
    // Auto-refetch on focus needs care (see below)
    refetchOnWindowFocus: false,
  });
}

// app/providers.tsx
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
 
const queryClient = new QueryClient({
  defaultOptions: {
    queries: {
      refetchOnWindowFocus: false,
      staleTime: 30_000,
    },
    mutations: {
      // Default is 0; written explicitly to make intent clear
      retry: 0,
    },
  },
});

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

// __tests__/streamingChat.test.tsx
import { setupServer } from "msw/node";
import { http, HttpResponse } from "msw";
import { renderHook, act, waitFor } from "@testing-library/react";
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
import { useStreamingChat } from "@/hooks/useStreamingChat";
 
const server = setupServer(
  http.post("/api/chat/stream", () => {
    const stream = new ReadableStream({
      async start(controller) {
        const enc = new TextEncoder();
        const events = [
          { type: "chunk", text: "Hello" },
          { type: "chunk", text: ", I'm Gemini." },
          { type: "done", conversationId: "conv-1" },
        ];
        for (const e of events) {
          controller.enqueue(enc.encode(`data: ${JSON.stringify(e)}
 
`));
          await new Promise((r) => setTimeout(r, 10));
        }
        controller.close();
      },
    });
    return new HttpResponse(stream, {
      headers: { "Content-Type": "text/event-stream" },
    });
  })
);
 
beforeAll(() => server.listen());
afterAll(() => server.close());
 
const wrapper = ({ children }: { children: React.ReactNode }) => {
  const qc = new QueryClient({ defaultOptions: { mutations: { retry: 0 } } });
  return <QueryClientProvider client={qc}>{children}</QueryClientProvider>;
};
 
test("writes finalized result into the cache after streaming completes", async () => {
  const { result } = renderHook(() => useStreamingChat("conv-1"), { wrapper });
  act(() => {
    result.current.mutate({
      conversationId: "conv-1",
      messages: [{ role: "user", content: "test" }],
    });
  });
  await waitFor(() => expect(result.current.isSuccess).toBe(true));
  expect(result.current.data?.text).toBe("Hello, I'm Gemini.");
});