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 source_chat_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,94 @@ 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}
+    )
     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}),
-        )
-
-        # 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}
-                ),
+        # Serializes the turn for this session and persists the pending user
+        # message before generation starts. Held for the whole stream.
+        async with source_chat_turn(
+            graph=source_chat_graph,
+            session_id=session_id,
+            config=config,
+            message=message,
+            message_id=message_id,
+        ):
+            # Send user message event
+            user_event = {"type": "user_message", "content": message, "timestamp": None}
+            yield f"data: {json.dumps(user_event)}\n\n"
+
+            # 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,
+            try:
+                last_keepalive = time.monotonic()
+                while True:
+                    done, _ = await asyncio.wait(
+                        {invoke_task}, timeout=DISCONNECT_POLL_INTERVAL_SECONDS
+                    )
+                    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
+                    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
+
+                # 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(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"
+                    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
@@ -406,6 +448,7 @@ async def stream_source_chat_response(
 
 @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 +474,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={

api/source_chat_service.py

@@ -0,0 +1,133 @@
+"""Turn coordination for source chat.
+
+Owns the policy a source-chat turn needs before generation can start: the
+per-session lock that serializes the snapshot -> append -> invoke sequence, and
+the decision to persist the pending human turn. The graph is passed in so the
+caller's module attribute stays authoritative and this module is usable
+standalone.
+"""
+
+import asyncio
+from contextlib import asynccontextmanager
+from typing import Any, AsyncIterator, Optional, Protocol
+
+from langchain_core.messages import HumanMessage
+from langchain_core.runnables import RunnableConfig
+
+
+class CheckpointGraph(Protocol):
+    """The slice of a compiled LangGraph a turn needs."""
+
+    def get_state(self, config: RunnableConfig) -> Any: ...
+
+    async def aupdate_state(self, config: RunnableConfig, values: Any) -> Any: ...
+
+
+# Per-session locks serialize the read-modify-write sequence (snapshot -> append
+# user message -> invoke). Without them, two concurrent requests for the same
+# thread could both read the same trailing message and each start a generation.
+# Created lazily; refcounted so the entry is evicted once the last holder
+# releases — a long-lived process must not keep one lock per session it has ever
+# seen.
+class _SessionLock:
+    __slots__ = ("lock", "holders")
+
+    def __init__(self) -> None:
+        self.lock = asyncio.Lock()
+        self.holders = 0
+
+
+_session_locks: dict[str, _SessionLock] = {}
+
+
+def _register_holder(session_id: str) -> _SessionLock:
+    # No `await` between these dict ops, so on a single event loop the
+    # read/create/increment is atomic. Registering the caller as a holder before
+    # it awaits `acquire` keeps the entry alive until it releases.
+    entry = _session_locks.get(session_id)
+    if entry is None:
+        entry = _SessionLock()
+        _session_locks[session_id] = entry
+    entry.holders += 1
+    return entry
+
+
+def _drop_holder(session_id: str, entry: _SessionLock) -> None:
+    entry.holders -= 1
+    # The `is entry` guard is defensive: the entry is only evicted when this was
+    # the last holder, so `session_id` must still map to this same entry.
+    if entry.holders == 0 and _session_locks.get(session_id) is entry:
+        _session_locks.pop(session_id, None)
+
+
+@asynccontextmanager
+async def session_turn_lock(session_id: str) -> AsyncIterator[None]:
+    """Hold the session's turn lock for the duration of the block."""
+    entry = _register_holder(session_id)
+    acquired = False
+    try:
+        await entry.lock.acquire()
+        acquired = True
+        yield
+    finally:
+        # Cancelled while waiting to acquire: drop the holder registered above
+        # without releasing an unheld lock.
+        if acquired:
+            entry.lock.release()
+        _drop_holder(session_id, entry)
+
+
+async def persist_pending_human_turn(
+    graph: CheckpointGraph,
+    config: RunnableConfig,
+    message: str,
+    message_id: Optional[str] = None,
+) -> bool:
+    """Append the user message to the checkpoint unless it is already pending.
+
+    Persisting up front makes the message survive a mid-generation disconnect
+    (the frontend refetches the checkpoint on cancel/complete and would
+    otherwise drop it). The guard keys on the client message id, not content: a
+    retry that reuses the same id is deduplicated, while two distinct identical
+    messages get distinct ids and are both kept. A completed exchange always
+    ends with an AI message, so a trailing human turn is necessarily still
+    pending.
+
+    Returns whether the message was appended.
+    """
+    # SqliteSaver has no async read, so snapshot off the event loop.
+    current_state = await asyncio.to_thread(graph.get_state, config=config)
+    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 message_id is not None
+            and getattr(last_message, "id", None) == message_id
+        )
+    if already_pending:
+        return False
+    await graph.aupdate_state(
+        config, {"messages": [HumanMessage(content=message, id=message_id)]}
+    )
+    return True
+
+
+@asynccontextmanager
+async def source_chat_turn(
+    graph: CheckpointGraph,
+    session_id: str,
+    config: RunnableConfig,
+    message: str,
+    message_id: Optional[str] = None,
+) -> AsyncIterator[None]:
+    """Serialize and persist one source-chat turn, then run the caller's block.
+
+    The lock is held for the whole block so a concurrent request for the same
+    thread cannot snapshot the same trailing message and start a second
+    generation.
+    """
+    async with session_turn_lock(session_id):
+        await persist_pending_human_turn(graph, config, message, message_id)
+        yield

docs/7-DEVELOPMENT/decisions/ADR-009-hybrid-sqlite-checkpointer-bridge.md

@@ -0,0 +1,29 @@
+# ADR-009: Async checkpointer via a sync SqliteSaver thread bridge
+
+- **Status**: Accepted
+- **Date**: 2026-09
+- **Related**: #1264 (source-chat SSE keepalive + server-side cancellation), [ADR-004](ADR-004-background-workers.md) (worker-isolated long-running work)
+
+## Context
+
+To cancel source-chat generation when a client disconnects (#1264), the graph had to run on langgraph's async path (`ainvoke`/`aupdate_state`), which requires an async checkpointer. The existing checkpointer is `SqliteSaver`, which is sync-only: its async methods raise `NotImplementedError`. Switching to a genuinely async SQLite checkpointer would add a dependency and a second persistence format to reason about, while a second kind of checkpointer for one graph invites drift between the chat and source-chat graphs.
+
+## Decision
+
+**Wrap the existing sync `SqliteSaver` in a thin `HybridSqliteSaver` that delegates each async method to the corresponding sync method on a worker thread (`asyncio.to_thread`), and keep the module-level `sqlite3.connect(..., check_same_thread=False)` connection and all sync `get_state` callers unchanged.**
+
+Concurrency safety comes from `SqliteSaver` itself, not from this bridge: it holds an internal `threading.Lock` and funnels every read/write through a `cursor()` context manager that acquires that lock. So even though the `to_thread` delegates can run on different worker threads at once, only one thread ever touches the single shared SQLite connection at a time — the serialization invariant the async path needs.
+
+The async/event-loop and thread boundaries stay clean: the model node runs as a cancellable `asyncio` task on the loop; only the checkpoint I/O crosses to a worker thread, and the per-session `asyncio.Lock` in the router serializes the snapshot → append → invoke sequence per conversation.
+
+## Alternatives considered
+
+- **Dedicated async checkpointer (`AsyncSqliteSaver` / `aio-sqlite`)** — rejected: adds a dependency, a second checkpoint format, and a second persistence path to test, for no benefit over the sync saver since the DB work is already lock-serialized and short.
+- **Run the whole graph synchronously in a thread** — the pre-existing approach — rejected: it made the model call uncancellable, which is exactly the gap #1264's follow-up needed to close.
+- **A single global `asyncio.Lock` around all checkpointer calls** — rejected as redundant: `SqliteSaver` already serializes access internally; a second lock in front would add contention without adding correctness. The per-session lock in the router is a different concern (read-modify-write atomicity across `get_state` + `aupdate_state` + `ainvoke`), and it lives there rather than in the saver.
+
+## Consequences
+
+- One checkpointer type and one SQLite file serve both the chat and source-chat graphs; the async path is a small delegation layer, not a second persistence implementation.
+- Any future checkpointer API surface must keep the `aget`/`aget_tuple`/`aput`/`aput_writes`/`alist` delegation in sync with `SqliteSaver`'s sync surface; if langgraph adds new async methods, the bridge raises `NotImplementedError` until extended.
+- Correctness depends on `SqliteSaver`'s internal `threading.Lock` continuing to guard every connection access. If that upstream guarantee ever changed, this bridge would need its own lock or a real async saver.

docs/7-DEVELOPMENT/decisions/README.md

@@ -50,5 +50,6 @@ What this makes easier, what it makes harder, what to watch. (bullets)
 | [ADR-006](ADR-006-migration-granularity.md) | Migration granularity follows merge granularity, not release granularity | Accepted |
 | [ADR-007](ADR-007-optin-runtimes.md) | Heavy extraction runtimes (Docling, Crawl4AI local) are opt-in, installed at startup | Accepted |
 | [ADR-008](ADR-008-notebook-scoped-search.md) | Notebook scope is an optional filter on the existing search functions | Accepted |
+| [ADR-009](ADR-009-hybrid-sqlite-checkpointer-bridge.md) | Async checkpointer via a sync SqliteSaver thread bridge | Accepted |
 | [PDR-001](PDR-001-single-user-first.md) | Single-user first; don't preclude multi-user | Accepted |
 | [PDR-002](PDR-002-provider-agnostic-core.md) | Provider-agnostic core by default | Accepted |