feat(web): implement SSE chat streaming with real-time token rendering (#516)

Co-authored-by: Jason Woltje <jason@diversecanvas.com>
Co-committed-by: Jason Woltje <jason@diversecanvas.com>

apps/web/src/hooks/useChat.ts

@@ -4,7 +4,11 @@
  */
 
 import { useState, useCallback, useRef } from "react";
-import { sendChatMessage, type ChatMessage as ApiChatMessage } from "@/lib/api/chat";
+import {
+  sendChatMessage,
+  streamChatMessage,
+  type ChatMessage as ApiChatMessage,
+} from "@/lib/api/chat";
 import { createConversation, updateConversation, getIdea, type Idea } from "@/lib/api/ideas";
 import { safeJsonParse, isMessageArray } from "@/lib/utils/safe-json";
 
@@ -33,10 +37,12 @@ export interface UseChatOptions {
 export interface UseChatReturn {
   messages: Message[];
   isLoading: boolean;
+  isStreaming: boolean;
   error: string | null;
   conversationId: string | null;
   conversationTitle: string | null;
   sendMessage: (content: string) => Promise<void>;
+  abortStream: () => void;
   loadConversation: (ideaId: string) => Promise<void>;
   startNewConversation: (projectId?: string | null) => void;
   setMessages: React.Dispatch<React.SetStateAction<Message[]>>;
@@ -66,6 +72,7 @@ export function useChat(options: UseChatOptions = {}): UseChatReturn {
 
   const [messages, setMessages] = useState<Message[]>([WELCOME_MESSAGE]);
   const [isLoading, setIsLoading] = useState(false);
+  const [isStreaming, setIsStreaming] = useState(false);
   const [error, setError] = useState<string | null>(null);
   const [conversationId, setConversationId] = useState<string | null>(null);
   const [conversationTitle, setConversationTitle] = useState<string | null>(null);
@@ -78,6 +85,16 @@ export function useChat(options: UseChatOptions = {}): UseChatReturn {
   const messagesRef = useRef<Message[]>(messages);
   messagesRef.current = messages;
 
+  // AbortController ref for the active stream
+  const abortControllerRef = useRef<AbortController | null>(null);
+
+  // Track conversation state in refs to avoid stale closures in streaming callbacks
+  const conversationIdRef = useRef<string | null>(conversationId);
+  conversationIdRef.current = conversationId;
+
+  const conversationTitleRef = useRef<string | null>(conversationTitle);
+  conversationTitleRef.current = conversationTitle;
+
   /**
    * Convert our Message format to API ChatMessage format
    */
@@ -119,44 +136,57 @@ export function useChat(options: UseChatOptions = {}): UseChatReturn {
   }, []);
 
   /**
-   * Save conversation to backend
+   * Save conversation to backend.
+   * Uses refs for conversation state to avoid stale closures in streaming callbacks.
    */
   const saveConversation = useCallback(
     async (msgs: Message[], title: string): Promise<string> => {
       const content = serializeMessages(msgs);
+      const currentConvId = conversationIdRef.current;
 
-      if (conversationId) {
-        // Update existing conversation
-        await updateConversation(conversationId, content, title);
-        return conversationId;
+      if (currentConvId) {
+        await updateConversation(currentConvId, content, title);
+        return currentConvId;
       } else {
-        // Create new conversation
         const idea = await createConversation(title, content, projectIdRef.current ?? undefined);
         setConversationId(idea.id);
         setConversationTitle(title);
+        conversationIdRef.current = idea.id;
+        conversationTitleRef.current = title;
         return idea.id;
       }
     },
-    [conversationId, serializeMessages]
+    [serializeMessages]
   );
 
   /**
-   * Send a message to the LLM and save the conversation
+   * Abort an active stream
+   */
+  const abortStream = useCallback((): void => {
+    if (abortControllerRef.current) {
+      abortControllerRef.current.abort();
+      abortControllerRef.current = null;
+    }
+    setIsStreaming(false);
+    setIsLoading(false);
+  }, []);
+
+  /**
+   * Send a message to the LLM using streaming, with fallback to non-streaming
    */
   const sendMessage = useCallback(
     async (content: string): Promise<void> => {
-      if (!content.trim() || isLoading) {
+      if (!content.trim() || isLoading || isStreaming) {
         return;
       }
 
       const userMessage: Message = {
-        id: `user-${Date.now().toString()}`,
+        id: `user-${Date.now().toString()}-${Math.random().toString(36).slice(2, 8)}`,
         role: "user",
         content: content.trim(),
         createdAt: new Date().toISOString(),
       };
 
-      // Add user message immediately using functional update
       setMessages((prev) => {
         const updated = [...prev, userMessage];
         messagesRef.current = updated;
@@ -165,95 +195,186 @@ export function useChat(options: UseChatOptions = {}): UseChatReturn {
       setIsLoading(true);
       setError(null);
 
+      const assistantMessageId = `assistant-${Date.now().toString()}-${Math.random().toString(36).slice(2, 8)}`;
+      const placeholderMessage: Message = {
+        id: assistantMessageId,
+        role: "assistant",
+        content: "",
+        createdAt: new Date().toISOString(),
+        model,
+      };
+
+      const currentMessages = messagesRef.current;
+      const apiMessages = convertToApiMessages(currentMessages);
+
+      const request = {
+        model,
+        messages: apiMessages,
+        ...(temperature !== undefined && { temperature }),
+        ...(maxTokens !== undefined && { maxTokens }),
+        ...(systemPrompt !== undefined && { systemPrompt }),
+      };
+
+      const controller = new AbortController();
+      abortControllerRef.current = controller;
+
+      let streamingSucceeded = false;
+
       try {
-        // Prepare API request - use ref to get current messages (prevents stale closure)
-        const currentMessages = messagesRef.current;
-        const apiMessages = convertToApiMessages(currentMessages);
+        await new Promise<void>((resolve, reject) => {
+          let hasReceivedData = false;
 
-        const request = {
-          model,
-          messages: apiMessages,
-          ...(temperature !== undefined && { temperature }),
-          ...(maxTokens !== undefined && { maxTokens }),
-          ...(systemPrompt !== undefined && { systemPrompt }),
-        };
+          streamChatMessage(
+            request,
+            (chunk: string) => {
+              if (!hasReceivedData) {
+                hasReceivedData = true;
+                setIsLoading(false);
+                setIsStreaming(true);
+                setMessages((prev) => {
+                  const updated = [...prev, { ...placeholderMessage }];
+                  messagesRef.current = updated;
+                  return updated;
+                });
+              }
 
-        // Call LLM API
-        const response = await sendChatMessage(request);
-
-        // Create assistant message
-        const assistantMessage: Message = {
-          id: `assistant-${Date.now().toString()}`,
-          role: "assistant",
-          content: response.message.content,
-          createdAt: new Date().toISOString(),
-          model: response.model,
-          promptTokens: response.promptEvalCount ?? 0,
-          completionTokens: response.evalCount ?? 0,
-          totalTokens: (response.promptEvalCount ?? 0) + (response.evalCount ?? 0),
-        };
-
-        // Add assistant message using functional update
-        let finalMessages: Message[] = [];
-        setMessages((prev) => {
-          finalMessages = [...prev, assistantMessage];
-          messagesRef.current = finalMessages;
-          return finalMessages;
+              setMessages((prev) => {
+                const updated = prev.map((msg) =>
+                  msg.id === assistantMessageId ? { ...msg, content: msg.content + chunk } : msg
+                );
+                messagesRef.current = updated;
+                return updated;
+              });
+            },
+            () => {
+              streamingSucceeded = true;
+              setIsStreaming(false);
+              abortControllerRef.current = null;
+              resolve();
+            },
+            (err: Error) => {
+              reject(err);
+            },
+            controller.signal
+          );
         });
+      } catch (err: unknown) {
+        if (controller.signal.aborted) {
+          setIsStreaming(false);
+          setIsLoading(false);
+          abortControllerRef.current = null;
 
-        // Generate title from first user message if this is a new conversation
-        const isFirstMessage =
-          !conversationId && finalMessages.filter((m) => m.role === "user").length === 1;
-        const title = isFirstMessage
-          ? generateTitle(content)
-          : (conversationTitle ?? "Chat Conversation");
-
-        // Save conversation (separate error handling from LLM errors)
-        try {
-          await saveConversation(finalMessages, title);
-        } catch (saveErr) {
-          const saveErrorMsg =
-            saveErr instanceof Error ? saveErr.message : "Unknown persistence error";
-          setError("Message sent but failed to save. Please try again.");
-          onError?.(saveErr instanceof Error ? saveErr : new Error(saveErrorMsg));
-          console.error("Failed to save conversation", {
-            error: saveErr,
-            errorType: "PERSISTENCE_ERROR",
-            conversationId,
-            detail: saveErrorMsg,
+          // Remove placeholder if no content was received
+          setMessages((prev) => {
+            const assistantMsg = prev.find((m) => m.id === assistantMessageId);
+            if (assistantMsg?.content === "") {
+              const updated = prev.filter((m) => m.id !== assistantMessageId);
+              messagesRef.current = updated;
+              return updated;
+            }
+            messagesRef.current = prev;
+            return prev;
           });
+          return;
         }
-      } catch (err) {
-        const errorMsg = err instanceof Error ? err.message : "Failed to send message";
-        setError("Unable to send message. Please try again.");
-        onError?.(err instanceof Error ? err : new Error(errorMsg));
-        console.error("Failed to send chat message", {
-          error: err,
-          errorType: "LLM_ERROR",
-          conversationId,
-          messageLength: content.length,
-          messagePreview: content.substring(0, 50),
-          model,
-          messageCount: messagesRef.current.length,
-          timestamp: new Date().toISOString(),
+
+        // Streaming failed — fall back to non-streaming
+        console.warn("Streaming failed, falling back to non-streaming", {
+          error: err instanceof Error ? err : new Error(String(err)),
         });
 
-        // Add error message to chat
-        const errorMessage: Message = {
-          id: `error-${String(Date.now())}`,
-          role: "assistant",
-          content: "Something went wrong. Please try again.",
-          createdAt: new Date().toISOString(),
-        };
-        setMessages((prev) => [...prev, errorMessage]);
-      } finally {
-        setIsLoading(false);
+        setMessages((prev) => {
+          const withoutPlaceholder = prev.filter((m) => m.id !== assistantMessageId);
+          messagesRef.current = withoutPlaceholder;
+          return withoutPlaceholder;
+        });
+        setIsStreaming(false);
+
+        try {
+          const response = await sendChatMessage(request);
+
+          const assistantMessage: Message = {
+            id: `assistant-${Date.now().toString()}`,
+            role: "assistant",
+            content: response.message.content,
+            createdAt: new Date().toISOString(),
+            model: response.model,
+            promptTokens: response.promptEvalCount ?? 0,
+            completionTokens: response.evalCount ?? 0,
+            totalTokens: (response.promptEvalCount ?? 0) + (response.evalCount ?? 0),
+          };
+
+          setMessages((prev) => {
+            const updated = [...prev, assistantMessage];
+            messagesRef.current = updated;
+            return updated;
+          });
+
+          streamingSucceeded = true;
+        } catch (fallbackErr: unknown) {
+          const errorMsg =
+            fallbackErr instanceof Error ? fallbackErr.message : "Failed to send message";
+          setError("Unable to send message. Please try again.");
+          onError?.(fallbackErr instanceof Error ? fallbackErr : new Error(errorMsg));
+          console.error("Failed to send chat message", {
+            error: fallbackErr,
+            errorType: "LLM_ERROR",
+            conversationId: conversationIdRef.current,
+            messageLength: content.length,
+            messagePreview: content.substring(0, 50),
+            model,
+            messageCount: messagesRef.current.length,
+            timestamp: new Date().toISOString(),
+          });
+
+          const errorMessage: Message = {
+            id: `error-${String(Date.now())}`,
+            role: "assistant",
+            content: "Something went wrong. Please try again.",
+            createdAt: new Date().toISOString(),
+          };
+          setMessages((prev) => {
+            const updated = [...prev, errorMessage];
+            messagesRef.current = updated;
+            return updated;
+          });
+          setIsLoading(false);
+          return;
+        }
+      }
+
+      setIsLoading(false);
+
+      if (!streamingSucceeded) {
+        return;
+      }
+
+      const finalMessages = messagesRef.current;
+
+      const isFirstMessage =
+        !conversationIdRef.current && finalMessages.filter((m) => m.role === "user").length === 1;
+      const title = isFirstMessage
+        ? generateTitle(content)
+        : (conversationTitleRef.current ?? "Chat Conversation");
+
+      try {
+        await saveConversation(finalMessages, title);
+      } catch (saveErr) {
+        const saveErrorMsg =
+          saveErr instanceof Error ? saveErr.message : "Unknown persistence error";
+        setError("Message sent but failed to save. Please try again.");
+        onError?.(saveErr instanceof Error ? saveErr : new Error(saveErrorMsg));
+        console.error("Failed to save conversation", {
+          error: saveErr,
+          errorType: "PERSISTENCE_ERROR",
+          conversationId: conversationIdRef.current,
+          detail: saveErrorMsg,
+        });
       }
     },
     [
       isLoading,
-      conversationId,
-      conversationTitle,
+      isStreaming,
       model,
       temperature,
       maxTokens,
@@ -280,6 +401,8 @@ export function useChat(options: UseChatOptions = {}): UseChatReturn {
         setMessages(msgs);
         setConversationId(idea.id);
         setConversationTitle(idea.title ?? null);
+        conversationIdRef.current = idea.id;
+        conversationTitleRef.current = idea.title ?? null;
       } catch (err) {
         const errorMsg = err instanceof Error ? err.message : "Failed to load conversation";
         setError("Unable to load conversation. Please try again.");
@@ -305,6 +428,8 @@ export function useChat(options: UseChatOptions = {}): UseChatReturn {
     setConversationId(null);
     setConversationTitle(null);
     setError(null);
+    conversationIdRef.current = null;
+    conversationTitleRef.current = null;
     projectIdRef.current = newProjectId ?? null;
   }, []);
 
@@ -318,10 +443,12 @@ export function useChat(options: UseChatOptions = {}): UseChatReturn {
   return {
     messages,
     isLoading,
+    isStreaming,
     error,
     conversationId,
     conversationTitle,
     sendMessage,
+    abortStream,
     loadConversation,
     startNewConversation,
     setMessages,