fix(source-chat): emit SSE keepalives during generation and add stop-…

api/routers/source_chat.py

@@ -2,7 +2,7 @@
 import json
 from typing import AsyncGenerator, List, Optional
 
-from fastapi import APIRouter, HTTPException, Path
+from fastapi import APIRouter, HTTPException, Path, Request
 from fastapi.responses import StreamingResponse
 from langchain_core.messages import HumanMessage
 from langchain_core.runnables import RunnableConfig
@@ -28,6 +28,12 @@
 router = APIRouter()
 
 
+# Seconds between SSE keepalive comments while the LLM generates. Keeps the
+# connection from going idle so proxies (incl. the Next.js rewrite in front of
+# FastAPI) don't drop it mid-generation.
+KEEPALIVE_INTERVAL_SECONDS = 15.0
+
+
 # Request/Response models
 class CreateSourceChatSessionRequest(BaseModel):
     source_id: str = Field(..., description="Source ID to create chat session for")
@@ -330,47 +336,69 @@ async def delete_source_chat_session(
 
 
 async def stream_source_chat_response(
-    session_id: str, source_id: str, message: str, model_override: Optional[str] = None
+    request: Request,
+    session_id: str,
+    source_id: str,
+    message: str,
+    model_override: Optional[str] = None,
 ) -> AsyncGenerator[str, None]:
     """Stream the source chat response as Server-Sent Events."""
+    config = RunnableConfig(
+        configurable={"thread_id": session_id, "model_id": model_override}
+    )
+    invoke_task: Optional[asyncio.Task] = None
     try:
-        # Get current state
-        # Use sync get_state() in a thread since SqliteSaver doesn't support async
+        # Persist the user message to the checkpoint up front so it survives a
+        # mid-generation disconnect (the frontend refetches the checkpoint on
+        # cancel/complete and would otherwise drop the user's message). Skip the
+        # append when this exact message is already the trailing (unanswered)
+        # turn — a retry after a failed generation would otherwise duplicate it.
+        # A completed exchange always ends with an AI message, so a trailing
+        # human turn is necessarily a pending one.
         current_state = await asyncio.to_thread(
-            source_chat_graph.get_state,
-            config=RunnableConfig(configurable={"thread_id": session_id}),
+            source_chat_graph.get_state, config=config
         )
-
-        # Prepare state for execution
-        state_values = current_state.values if current_state else {}
-        state_values["messages"] = state_values.get("messages", [])
-        state_values["source_id"] = source_id
-        state_values["model_override"] = model_override
-
-        # Add user message to state
-        user_message = HumanMessage(content=message)
-        state_values["messages"].append(user_message)
+        already_pending = False
+        if current_state and current_state.values and "messages" in current_state.values:
+            existing_messages = current_state.values["messages"]
+            last_message = existing_messages[-1] if existing_messages else None
+            already_pending = (
+                isinstance(last_message, HumanMessage)
+                and last_message.content == message
+            )
+        if not already_pending:
+            await source_chat_graph.aupdate_state(
+                config, {"messages": [HumanMessage(content=message)]}
+            )
 
         # Send user message event
         user_event = {"type": "user_message", "content": message, "timestamp": None}
         yield f"data: {json.dumps(user_event)}\n\n"
 
-        # Run the synchronous LangGraph invoke in a thread so it doesn't block the
-        # event loop. While blocked, even the already-yielded SSE events can't
-        # flush and every other request stalls until the LLM finishes. Mirrors the
-        # get_state() calls above.
-        # The lambda pins down which `invoke` overload is used; asyncio.to_thread
-        # can't resolve overloaded callables on its own. The ignore is a langgraph
+        # Run the async graph with ainvoke so generation is cancellable. Only the
+        # per-message config is passed as input; the messages (incl. the user
+        # message above) are read from the checkpoint. The ignore is a langgraph
         # typing limitation: it accepts a partial state dict at runtime, but the
         # signature requires the full state type.
-        result = await asyncio.to_thread(
-            lambda: source_chat_graph.invoke(
-                input=state_values,  # type: ignore[arg-type]
-                config=RunnableConfig(
-                    configurable={"thread_id": session_id, "model_id": model_override}
-                ),
+        invoke_task = asyncio.create_task(
+            source_chat_graph.ainvoke(
+                input={"source_id": source_id, "model_override": model_override},  # type: ignore[call-overload]
+                config=config,
             )
         )
+        while True:
+            done, _ = await asyncio.wait(
+                {invoke_task}, timeout=KEEPALIVE_INTERVAL_SECONDS
+            )
+            if done:
+                # Re-raises on graph error, caught by the outer try/except below.
+                result = invoke_task.result()
+                break
+            if await request.is_disconnected():
+                # Client went away — stop generating instead of burning tokens.
+                return
+            # SSE comment — ignored by clients, keeps the connection alive.
+            yield ": ping\n\n"
 
         # Stream the complete AI response
         if "messages" in result:
@@ -402,10 +430,17 @@ async def stream_source_chat_response(
         logger.error(f"Error in source chat streaming: {str(e)}")
         error_event = {"type": "error", "message": error_message}
         yield f"data: {json.dumps(error_event)}\n\n"
+    finally:
+        # Stop generation if the generator is torn down mid-flight (client
+        # disconnect or server cancellation) so the model doesn't keep running.
+        if invoke_task is not None and not invoke_task.done():
+            invoke_task.cancel()
+            await asyncio.gather(invoke_task, return_exceptions=True)
 
 
 @router.post("/sources/{source_id}/chat/sessions/{session_id}/messages")
 async def send_message_to_source_chat(
+    http_request: Request,
     request: SendMessageRequest,
     source_id: str = Path(..., description="Source ID"),
     session_id: str = Path(..., description="Session ID"),
@@ -431,6 +466,7 @@ async def send_message_to_source_chat(
         # Return streaming response
         return StreamingResponse(
             stream_source_chat_response(
+                http_request,
                 session_id=full_session_id,
                 source_id=full_source_id,
                 message=request.message,

frontend/src/app/(dashboard)/sources/[id]/page.tsx

@@ -57,6 +57,7 @@ export default function SourceDetailPage() {
             isStreaming={chat.isStreaming}
             contextIndicators={chat.contextIndicators}
             onSendMessage={(message, model) => chat.sendMessage(message, model)}
+            onCancel={chat.cancelStreaming}
             modelOverride={chat.currentSession?.model_override}
             onModelChange={(model) => {
               if (chat.currentSessionId) {

frontend/src/components/sources/ChatPanel.tsx

@@ -7,7 +7,7 @@ import { ScrollArea } from '@/components/ui/scroll-area'
 import { Card, CardContent, CardHeader, CardTitle } from '@/components/ui/card'
 import { Badge } from '@/components/ui/badge'
 import { Dialog, DialogContent, DialogTitle } from '@/components/ui/dialog'
-import { Bot, User, Send, Loader2, FileText, Lightbulb, StickyNote, Clock } from 'lucide-react'
+import { Bot, User, Send, Square, Loader2, FileText, Lightbulb, StickyNote, Clock } from 'lucide-react'
 import { MarkdownRenderer } from '@/components/ui/markdown-renderer'
 import {
   SourceChatMessage,
@@ -36,6 +36,7 @@ interface ChatPanelProps {
   isStreaming: boolean
   contextIndicators: SourceChatContextIndicator | null
   onSendMessage: (message: string, modelOverride?: string) => void
+  onCancel?: () => void
   modelOverride?: string
   onModelChange?: (model?: string) => void
   // Session management props
@@ -60,6 +61,7 @@ export function ChatPanel({
   isStreaming,
   contextIndicators,
   onSendMessage,
+  onCancel,
   modelOverride,
   onModelChange,
   sessions = [],
@@ -218,6 +220,7 @@ export function ChatPanel({
         {/* Input Area */}
         <ChatComposer
           onSendMessage={onSendMessage}
+          onCancel={onCancel}
           isStreaming={isStreaming}
           modelOverride={modelOverride}
           onModelChange={onModelChange}
@@ -233,13 +236,15 @@ export function ChatPanel({
 // re-render this small component instead of the whole message history.
 interface ChatComposerProps {
   onSendMessage: (message: string, modelOverride?: string) => void
+  onCancel?: () => void
   isStreaming: boolean
   modelOverride?: string
   onModelChange?: (model?: string) => void
 }
 
 function ChatComposer({
   onSendMessage,
+  onCancel,
   isStreaming,
   modelOverride,
   onModelChange
@@ -297,18 +302,40 @@ function ChatComposer({
           className="flex-1 min-h-[40px] max-h-[100px] resize-none py-2 px-3 min-w-0"
           rows={1}
         />
-        <Button
-          onClick={handleSend}
-          disabled={!input.trim() || isStreaming}
-          size="icon"
-          className="h-[40px] w-[40px] flex-shrink-0"
-        >
-          {isStreaming ? (
-            <Loader2 className="h-4 w-4 animate-spin" />
+        {isStreaming ? (
+          onCancel ? (
+            <Button
+              onClick={onCancel}
+              size="icon"
+              variant="secondary"
+              aria-label={t('chat.stop')}
+              title={t('chat.stop')}
+              className="h-[40px] w-[40px] flex-shrink-0"
+            >
+              <Square className="h-4 w-4" />
+            </Button>
           ) : (
+            // No cancel callback (e.g. notebook chat, which is not streamed) —
+            // show a disabled spinner rather than a dead Stop button.
+            <Button
+              disabled
+              size="icon"
+              aria-label={t('common.saving')}
+              className="h-[40px] w-[40px] flex-shrink-0"
+            >
+              <Loader2 className="h-4 w-4 animate-spin" />
+            </Button>
+          )
+        ) : (
+          <Button
+            onClick={handleSend}
+            disabled={!input.trim()}
+            size="icon"
+            className="h-[40px] w-[40px] flex-shrink-0"
+          >
             <Send className="h-4 w-4" />
-          )}
-        </Button>
+          </Button>
+        )}
       </div>
     </div>
   )

frontend/src/lib/api/source-chat.ts

@@ -47,7 +47,7 @@ export const sourceChatApi = {
   },
 
   // Messaging with streaming
-  sendMessage: (sourceId: string, sessionId: string, data: SendMessageRequest) => {
+  sendMessage: (sourceId: string, sessionId: string, data: SendMessageRequest, signal?: AbortSignal) => {
     // Get auth token using the same logic as apiClient interceptor
     const token = getAuthToken()
 
@@ -62,7 +62,8 @@ export const sourceChatApi = {
         'Content-Type': 'application/json',
         ...(token && { 'Authorization': `Bearer ${token}` })
       },
-      body: JSON.stringify(data)
+      body: JSON.stringify(data),
+      signal
     }).then(response => {
       if (!response.ok) {
         throw new Error(`HTTP error! status: ${response.status}`)

frontend/src/lib/hooks/use-source-chat.ts

@@ -23,6 +23,13 @@ export function useSourceChat(sourceId: string) {
   const [contextIndicators, setContextIndicators] = useState<SourceChatContextIndicator | null>(null)
   const abortControllerRef = useRef<AbortController | null>(null)
 
+  // Abort any in-flight stream when the component unmounts.
+  useEffect(() => {
+    return () => {
+      abortControllerRef.current?.abort()
+    }
+  }, [])
+
   // Fetch sessions
   const { data: sessions = [], isLoading: loadingSessions, refetch: refetchSessions } = useQuery<SourceChatSession[]>({
     queryKey: ['sourceChatSessions', sourceId],
@@ -103,6 +110,12 @@ export function useSourceChat(sourceId: string) {
 
   // Send message with streaming
   const sendMessage = useCallback(async (message: string, modelOverride?: string) => {
+    // Abort any previous in-flight request
+    abortControllerRef.current?.abort()
+    const controller = new AbortController()
+    abortControllerRef.current = controller
+    const signal = controller.signal
+
     let sessionId = currentSessionId
 
     // Auto-create session if none exists
@@ -135,7 +148,7 @@ export function useSourceChat(sourceId: string) {
       const response = await sourceChatApi.sendMessage(sourceId, sessionId, {
         message,
         model_override: modelOverride
-      })
+      }, signal)
 
       if (!response) {
         throw new Error('No response body')
@@ -199,24 +212,35 @@ export function useSourceChat(sourceId: string) {
         }
       }
     } catch (err: unknown) {
+      // Cancelled by the user — the finally block still refetches persisted messages.
+      if (err instanceof DOMException && err.name === 'AbortError') {
+        return
+      }
       const error = err as { response?: { data?: { detail?: string } }, message?: string };
       console.error('Error sending message:', error)
       toast.error(getApiErrorMessage(error.response?.data?.detail || error.message, (key) => t(key), 'apiErrors.failedToSendMessage'))
       // Remove optimistic messages on error
       setMessages(prev => prev.filter(msg => !msg.id.startsWith('temp-')))
     } finally {
-      setIsStreaming(false)
-      // Refetch session to get persisted messages
-      refetchCurrentSession()
+      // A superseded send (replaced by a newer one) must not clear the newer
+      // stream's loading state or refetch over its messages. A user-initiated
+      // cancel sets the ref to null (and clears isStreaming itself), so it still
+      // falls through to refetch the persisted user message.
+      const superseded =
+        abortControllerRef.current !== null && abortControllerRef.current !== controller
+      if (!superseded) {
+        setIsStreaming(false)
+        // Refetch session to get persisted messages
+        refetchCurrentSession()
+      }
     }
   }, [sourceId, currentSessionId, refetchCurrentSession, queryClient, t])
 
   // Cancel streaming
   const cancelStreaming = useCallback(() => {
-    if (abortControllerRef.current) {
-      abortControllerRef.current.abort()
-      setIsStreaming(false)
-    }
+    abortControllerRef.current?.abort()
+    abortControllerRef.current = null
+    setIsStreaming(false)
   }, [])
 
   // Switch session

frontend/src/lib/locales/bn-IN/index.ts

@@ -461,6 +461,7 @@ export const bnIN = {
     sessionCreated: "চ্যাট সেশন তৈরি হয়েছে",
     sessionUpdated: "সেশন আপডেট হয়েছে",
     sessionDeleted: "সেশন মুছে ফেলা হয়েছে",
+    stop: "থামান",
   },
   searchPage: {
     askAndSearch: "জিজ্ঞাসা ও অনুসন্ধান",

frontend/src/lib/locales/ca-ES/index.ts

@@ -461,6 +461,7 @@ export const caES = {
     sessionCreated: "S'ha creat la sessió de xat",
     sessionUpdated: "S'ha actualitzat la sessió",
     sessionDeleted: "S'ha suprimit la sessió",
+    stop: "Atura",
   },
   searchPage: {
     askAndSearch: "Pregunta i cerca",

frontend/src/lib/locales/de-DE/index.ts

@@ -464,6 +464,7 @@ export const deDE = {
     sessionCreated: "Chat-Sitzung erstellt",
     sessionUpdated: "Sitzung aktualisiert",
     sessionDeleted: "Sitzung gelöscht",
+    stop: "Stoppen",
   },
   searchPage: {
     askAndSearch: "Fragen und Suchen",

frontend/src/lib/locales/en-US/index.ts

@@ -459,6 +459,7 @@ export const enUS = {
     sessionCreated: "Chat session created",
     sessionUpdated: "Session updated",
     sessionDeleted: "Session deleted",
+    stop: "Stop",
   },
   searchPage: {
     askAndSearch: "Ask and Search",

frontend/src/lib/locales/es-ES/index.ts

@@ -461,6 +461,7 @@ export const esES = {
     sessionCreated: "Sesión de chat creada",
     sessionUpdated: "Sesión actualizada",
     sessionDeleted: "Sesión eliminada",
+    stop: "Detener",
   },
   searchPage: {
     askAndSearch: "Preguntar y buscar",

frontend/src/lib/locales/fr-FR/index.ts

@@ -461,6 +461,7 @@ export const frFR = {
     sessionCreated: "Session de chat créée",
     sessionUpdated: "Session mise à jour",
     sessionDeleted: "Session supprimée",
+    stop: "Arrêter",
   },
   searchPage: {
     askAndSearch: "Poser une question et Rechercher",