fix(source-chat): make streaming cancellable on client disconnect and persist user message

Convert the source chat graph node and streaming endpoint to async so
generation can be cancelled when the client disconnects. Persist the
user message to the checkpoint up front via aupdate_state so it survives
a mid-generation disconnect. Add a HybridSqliteSaver to delegate
LangGraph async checkpointer calls to the existing sync SQLite connection.
Update characterization tests for the async path and add a disconnect
cancellation test.

Author: AugustoSandim

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
@@ -336,51 +336,38 @@ 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
-        current_state = await asyncio.to_thread(
-            source_chat_graph.get_state,
-            config=RunnableConfig(configurable={"thread_id": session_id}),
+        # 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).
+        await source_chat_graph.aupdate_state(
+            config, {"messages": [HumanMessage(content=message)]}
         )
 
-        # 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 (mirrors the get_state() calls above). While the LLM generates,
-        # emit an SSE comment every KEEPALIVE_INTERVAL_SECONDS so the connection
-        # never goes idle — otherwise proxies (incl. the Next.js rewrite in front of
-        # FastAPI) may drop it and the reply only appears after a refetch.
-        # 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.
         invoke_task = asyncio.create_task(
-            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,
-                        }
-                    ),
-                )
+            source_chat_graph.ainvoke(
+                input={"source_id": source_id, "model_override": model_override},  # type: ignore[call-overload]
+                config=config,
             )
         )
         while True:
@@ -391,6 +378,9 @@ async def stream_source_chat_response(
                 # 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"
 
@@ -424,10 +414,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"),
@@ -453,6 +450,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,

open_notebook/graphs/source_chat.py

@@ -45,7 +45,7 @@ def _source_content_is_available(
     return isinstance(full_text, str) and bool(full_text.strip())
 
 
-def call_model_with_source_context(
+async def call_model_with_source_context(
     state: SourceChatState, config: RunnableConfig
 ) -> dict:
     """
@@ -58,50 +58,26 @@ def call_model_with_source_context(
     4. Tracks context indicators for referenced insights/content
     """
     try:
-        return _call_model_with_source_context_inner(state, config)
+        return await _call_model_with_source_context_inner(state, config)
     except OpenNotebookError:
         raise
     except Exception as e:
         error_class, user_message = classify_error(e)
         raise error_class(user_message) from e
 
 
-def _call_model_with_source_context_inner(
+async def _call_model_with_source_context_inner(
     state: SourceChatState, config: RunnableConfig
 ) -> dict:
     source_id = state.get("source_id")
     if not source_id:
         raise ValueError("source_id is required in state")
 
-    # Build source context using build_source_context (run async code in new loop)
-    def build_context():
-        """Build context in a new event loop"""
-        new_loop = asyncio.new_event_loop()
-        try:
-            asyncio.set_event_loop(new_loop)
-            return new_loop.run_until_complete(
-                build_source_context(
-                    source_id=source_id,
-                    max_tokens=50000,  # Reasonable limit for source context
-                )
-            )
-        finally:
-            new_loop.close()
-            asyncio.set_event_loop(None)
-
-    # Get the built context
-    try:
-        # Try to get the current event loop
-        asyncio.get_running_loop()
-        # If we're in an event loop, run in a thread with a new loop
-        import concurrent.futures
-
-        with concurrent.futures.ThreadPoolExecutor() as executor:
-            future = executor.submit(build_context)
-            context_data = future.result()
-    except RuntimeError:
-        # No event loop running, safe to create a new one
-        context_data = build_context()
+    # Build source context (awaiting directly keeps the node cancellable).
+    context_data = await build_source_context(
+        source_id=source_id,
+        max_tokens=50000,  # Reasonable limit for source context
+    )
 
     # Extract source and insights from context
     source = None
@@ -149,47 +125,16 @@ def build_context():
     )
     payload = [SystemMessage(content=system_prompt)] + state.get("messages", [])
 
-    # Handle async model provisioning from sync context
-    def run_in_new_loop():
-        """Run the async function in a new event loop"""
-        new_loop = asyncio.new_event_loop()
-        try:
-            asyncio.set_event_loop(new_loop)
-            return new_loop.run_until_complete(
-                provision_langchain_model(
-                    str(payload),
-                    config.get("configurable", {}).get("model_id")
-                    or state.get("model_override"),
-                    "chat",
-                    max_tokens=8192,
-                )
-            )
-        finally:
-            new_loop.close()
-            asyncio.set_event_loop(None)
-
-    try:
-        # Try to get the current event loop
-        asyncio.get_running_loop()
-        # If we're in an event loop, run in a thread with a new loop
-        import concurrent.futures
-
-        with concurrent.futures.ThreadPoolExecutor() as executor:
-            future = executor.submit(run_in_new_loop)
-            model = future.result()
-    except RuntimeError:
-        # No event loop running, safe to use asyncio.run()
-        model = asyncio.run(
-            provision_langchain_model(
-                str(payload),
-                config.get("configurable", {}).get("model_id")
-                or state.get("model_override"),
-                "chat",
-                max_tokens=8192,
-            )
-        )
+    # Provision the model asynchronously (cancellable)
+    model = await provision_langchain_model(
+        str(payload),
+        config.get("configurable", {}).get("model_id")
+        or state.get("model_override"),
+        "chat",
+        max_tokens=8192,
+    )
 
-    ai_message = model.invoke(payload)
+    ai_message = await model.ainvoke(payload)
 
     # Clean thinking content from AI response (e.g., <think>...</think> tags)
     content = extract_text_content(ai_message.content)
@@ -211,12 +156,47 @@ def _format_source_context(context_data: Dict) -> str:
     return format_source_context(context_data)
 
 
+class HybridSqliteSaver(SqliteSaver):
+    """Sync SqliteSaver with async delegates for langgraph's async run path.
+
+    The source-chat node is async so the model call can be cancelled when the
+    client disconnects; langgraph's ``ainvoke``/``aupdate_state`` then require
+    an async checkpointer. SqliteSaver is sync-only (its async methods raise
+    NotImplementedError), so delegate each async method to the corresponding
+    sync one on a worker thread. The module-level sync connection — and the
+    sync ``get_state`` callers elsewhere — keep working unchanged.
+    """
+
+    async def aget(self, config):
+        return await asyncio.to_thread(self.get, config)
+
+    async def aget_tuple(self, config):
+        return await asyncio.to_thread(self.get_tuple, config)
+
+    async def alist(self, config, *, filter=None, before=None, limit=None):
+        items = await asyncio.to_thread(
+            lambda: list(self.list(config, filter=filter, before=before, limit=limit))
+        )
+        for item in items:
+            yield item
+
+    async def aput(self, config, checkpoint, metadata, new_versions):
+        return await asyncio.to_thread(
+            self.put, config, checkpoint, metadata, new_versions
+        )
+
+    async def aput_writes(self, config, writes, task_id, task_path=""):
+        return await asyncio.to_thread(
+            self.put_writes, config, writes, task_id, task_path
+        )
+
+
 # Create SQLite checkpointer
 conn = sqlite3.connect(
     LANGGRAPH_CHECKPOINT_FILE,
     check_same_thread=False,
 )
-memory = SqliteSaver(conn)
+memory = HybridSqliteSaver(conn)
 
 # Create the StateGraph
 source_chat_state = StateGraph(SourceChatState)

tests/test_chat_routers_characterization.py

@@ -14,7 +14,7 @@
 tests/test_crud_404.py.
 """
 
-import time
+import asyncio
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -346,21 +346,62 @@ async def test_stream_source_chat_emits_keepalive_while_invoke_runs():
     with patch.object(
         source_chat_router, "KEEPALIVE_INTERVAL_SECONDS", 0.01
     ), patch.object(source_chat_router, "source_chat_graph") as mock_graph:
-        mock_graph.get_state.return_value = _graph_state({"messages": []})
+        mock_graph.aupdate_state = AsyncMock()
 
-        def slow_invoke(*_args, **_kwargs):
-            time.sleep(0.1)
+        async def slow_ainvoke(*_args, **_kwargs):
+            await asyncio.sleep(0.1)
             return {"messages": []}
 
-        mock_graph.invoke.side_effect = slow_invoke
+        mock_graph.ainvoke.side_effect = slow_ainvoke
+
+        request = MagicMock()
+        request.is_disconnected = AsyncMock(return_value=False)
 
         chunks = []
         async for chunk in stream_source_chat_response(
-            "chat_session:abc", "source:xyz", "hello"
+            request, "chat_session:abc", "source:xyz", "hello"
         ):
             chunks.append(chunk)
 
     assert chunks[0].startswith('data: {"type": "user_message"')
     assert chunks[-1].startswith('data: {"type": "complete"')
     # Keepalive comments are emitted between the user_message and completion.
     assert ": ping\n\n" in chunks
+
+
+@pytest.mark.asyncio
+async def test_stream_source_chat_cancels_invoke_on_disconnect():
+    """When the client disconnects mid-generation, generation is cancelled
+    server-side and the stream ends without a completion event."""
+    from api.routers import source_chat as source_chat_router
+    from api.routers.source_chat import stream_source_chat_response
+
+    cancelled = asyncio.Event()
+
+    async def blocking_ainvoke(*_args, **_kwargs):
+        try:
+            await asyncio.Event().wait()  # blocks until cancelled
+        except asyncio.CancelledError:
+            cancelled.set()
+            raise
+
+    with patch.object(
+        source_chat_router, "KEEPALIVE_INTERVAL_SECONDS", 0.01
+    ), patch.object(source_chat_router, "source_chat_graph") as mock_graph:
+        mock_graph.aupdate_state = AsyncMock()
+        mock_graph.ainvoke.side_effect = blocking_ainvoke
+
+        request = MagicMock()
+        request.is_disconnected = AsyncMock(return_value=True)
+
+        chunks = []
+        async for chunk in stream_source_chat_response(
+            request, "chat_session:abc", "source:xyz", "hello"
+        ):
+            chunks.append(chunk)
+
+    # The user message is still yielded up front, but the stream stops without
+    # a completion event and the in-flight invoke task was cancelled.
+    assert chunks[0].startswith('data: {"type": "user_message"')
+    assert not any(c.startswith('data: {"type": "complete"') for c in chunks)
+    assert cancelled.is_set()