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

api/routers/source_chat.py

@@ -1,10 +1,10 @@
 import asyncio
 import json
+import time
 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
 from loguru import logger
 from pydantic import BaseModel, Field
@@ -16,6 +16,7 @@
     get_source_or_404,
     get_verified_source_session,
 )
+from api.source_chat_service import SessionTurnLease, persist_pending_human_turn
 from open_notebook.database.repository import ensure_record_id, repo_query
 from open_notebook.domain.notebook import ChatSession
 from open_notebook.exceptions import (
@@ -28,6 +29,15 @@
 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
+# Poll client disconnect more often than keepalive so dropped connections stop
+# generation promptly instead of waiting for the next SSE comment interval.
+DISCONNECT_POLL_INTERVAL_SECONDS = 1.0
+
+
 # Request/Response models
 class CreateSourceChatSessionRequest(BaseModel):
     source_id: str = Field(..., description="Source ID to create chat session for")
@@ -76,6 +86,14 @@ class SourceChatSessionWithMessagesResponse(SourceChatSessionResponse):
 
 class SendMessageRequest(BaseModel):
     message: str = Field(..., description="User message content")
+    message_id: Optional[str] = Field(
+        None,
+        description=(
+            "Client-generated message identity, used to deduplicate a retry of "
+            "the same turn (same id) while still keeping distinct identical "
+            "messages (different ids)."
+        ),
+    )
     model_override: Optional[str] = Field(
         None, description="Optional model override for this message"
     )
@@ -330,70 +348,117 @@ 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,
+    message_id: 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}
+    )
+    # The turn lock is queued for by polling rather than blocking so the
+    # connection can emit keepalives and notice disconnects while another turn
+    # for this session is still generating — while queued the stream is silent,
+    # and an idle connection is dropped by the same idle-timeout proxies the
+    # generation loop's keepalive guards against.
+    lease = SessionTurnLease(session_id)
     try:
-        # Get current state
-        # Use sync get_state() in a thread since SqliteSaver doesn't support async
-        current_state = await asyncio.to_thread(
-            source_chat_graph.get_state,
-            config=RunnableConfig(configurable={"thread_id": session_id}),
+        last_keepalive = time.monotonic()
+        while not await lease.wait_for_lock(DISCONNECT_POLL_INTERVAL_SECONDS):
+            if await request.is_disconnected():
+                # Client went away while queued — stop waiting instead of
+                # holding a queue slot.
+                return
+            now = time.monotonic()
+            if now - last_keepalive >= KEEPALIVE_INTERVAL_SECONDS:
+                # SSE comment — ignored by clients, keeps the connection alive.
+                yield ": ping\n\n"
+                last_keepalive = now
+
+        # Serializes the turn for this session and persists the pending user
+        # message before generation starts. The lock is held for the whole
+        # stream, so a concurrent request queues (with keepalives) instead of
+        # starting a second generation on the same thread.
+        await persist_pending_human_turn(
+            graph=source_chat_graph,
+            config=config,
+            message=message,
+            message_id=message_id,
         )
 
-        # 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)
-
         # 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
-        # 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}
-                ),
+        # 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.
+        invoke_task = asyncio.create_task(
+            source_chat_graph.ainvoke(
+                input={"source_id": source_id, "model_override": model_override},  # type: ignore[call-overload]
+                config=config,
             )
         )
-
-        # Stream the complete AI response
-        if "messages" in result:
-            for msg in result["messages"]:
-                if hasattr(msg, "type") and msg.type == "ai":
-                    ai_event = {
-                        "type": "ai_message",
-                        "content": msg.content if hasattr(msg, "content") else str(msg),
-                        "timestamp": None,
-                    }
-                    yield f"data: {json.dumps(ai_event)}\n\n"
-
-        # Stream context indicators
-        if "context_indicators" in result:
-            context_event = {
-                "type": "context_indicators",
-                "data": result["context_indicators"],
-            }
-            yield f"data: {json.dumps(context_event)}\n\n"
-
-        # Send completion signal
-        completion_event = {"type": "complete"}
-        yield f"data: {json.dumps(completion_event)}\n\n"
+        try:
+            while True:
+                done, _ = await asyncio.wait(
+                    {invoke_task}, timeout=DISCONNECT_POLL_INTERVAL_SECONDS
+                )
+                # A due keepalive is emitted even when the task just finished:
+                # the final response otherwise has to cover the whole gap since
+                # the last keepalive, which can exceed the proxy idle timeout
+                # when generation ends just past the interval.
+                now = time.monotonic()
+                if now - last_keepalive >= KEEPALIVE_INTERVAL_SECONDS:
+                    # SSE comment — ignored by clients, keeps the connection
+                    # alive.
+                    yield ": ping\n\n"
+                    last_keepalive = now
+                if done:
+                    # Re-raises on graph error, caught by the except below.
+                    result = invoke_task.result()
+                    break
+                if await request.is_disconnected():
+                    # Client went away — stop generating instead of burning
+                    # tokens.
+                    return
+
+            # Stream the complete AI response
+            if "messages" in result:
+                for msg in result["messages"]:
+                    if hasattr(msg, "type") and msg.type == "ai":
+                        ai_event = {
+                            "type": "ai_message",
+                            "content": msg.content
+                            if hasattr(msg, "content")
+                            else str(msg),
+                            "timestamp": None,
+                        }
+                        yield f"data: {json.dumps(ai_event)}\n\n"
+
+            # Stream context indicators
+            if "context_indicators" in result:
+                context_event = {
+                    "type": "context_indicators",
+                    "data": result["context_indicators"],
+                }
+                yield f"data: {json.dumps(context_event)}\n\n"
+
+            # Send completion signal
+            completion_event = {"type": "complete"}
+            yield f"data: {json.dumps(completion_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. Runs before the turn lock is released.
+            if not invoke_task.done():
+                invoke_task.cancel()
+                await asyncio.gather(invoke_task, return_exceptions=True)
 
     except Exception as e:
         from open_notebook.utils.error_classifier import classify_error
@@ -402,10 +467,13 @@ 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:
+        await lease.release()
 
 
 @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,10 +499,12 @@ 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,
                 model_override=model_override,
+                message_id=request.message_id,
             ),
             media_type="text/event-stream",
             headers={