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

api/routers/source_chat.py

@@ -1,8 +1,9 @@
 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
@@ -28,6 +29,52 @@
 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
+
+# Per-session locks serialize the read-modify-write sequence (snapshot -> append
+# user message -> invoke) in `stream_source_chat_response`. 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 _get_session_lock(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 _release_session_lock(session_id: str, entry: _SessionLock) -> None:
+    entry.lock.release()
+    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)
+
+
 # Request/Response models
 class CreateSourceChatSessionRequest(BaseModel):
     source_id: str = Field(..., description="Source ID to create chat session for")
@@ -76,6 +123,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,47 +385,83 @@ 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}
+    )
+    invoke_task: Optional[asyncio.Task] = None
+    # Serialize snapshot -> append -> invoke per session. Two concurrent requests
+    # for the same thread would otherwise both read the same trailing message and
+    # each start a generation. Held for the whole stream; released in finally.
+    lock_entry = _get_session_lock(session_id)
+    acquired = False
     try:
-        # Get current state
-        # Use sync get_state() in a thread since SqliteSaver doesn't support async
+        await lock_entry.lock.acquire()
+        acquired = True
+        # 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 turn is already the trailing (unanswered) one. 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.
         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 message_id is not None
+                and getattr(last_message, "id", None) == message_id
+            )
+        if not already_pending:
+            await source_chat_graph.aupdate_state(
+                config, {"messages": [HumanMessage(content=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 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,
             )
         )
+        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 outer try/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:
@@ -402,10 +493,28 @@ 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)
+        if acquired:
+            _release_session_lock(session_id, lock_entry)
+        else:
+            # Cancelled while waiting to acquire — decrement the holder count
+            # registered in `_get_session_lock` without releasing an unheld lock.
+            lock_entry.holders -= 1
+            if (
+                lock_entry.holders == 0
+                and _session_locks.get(session_id) is lock_entry
+            ):
+                _session_locks.pop(session_id, None)
 
 
 @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 +540,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={

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 |

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.test.tsx

@@ -107,4 +107,18 @@ describe('ChatPanel composer', () => {
 
     expect(onSendMessage).not.toHaveBeenCalled()
   })
+
+  it('announces generating (not saving) for notebook chat without stop support', () => {
+    render(
+      <ChatPanel
+        messages={[]}
+        isStreaming={true}
+        contextIndicators={null}
+        onSendMessage={vi.fn()}
+      />
+    )
+
+    expect(screen.getByLabelText('chat.generating')).toBeInTheDocument()
+    expect(screen.queryByLabelText('common.saving')).not.toBeInTheDocument()
+  })
 })

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('chat.generating')}
+              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>
   )