From 2ab915184cd640238e07defffe9c5767f995336b Mon Sep 17 00:00:00 2001 From: mohamed zwanski Date: Thu, 22 Oct 2026 21:42:52 +0100 Subject: [PATCH 1/2] Persist assistant responses when the chat WebSocket disconnects (#133) - Keep consuming Claude's stream after the client goes away so the full answer is stored in the conversation history - Persist partial responses if the stream is aborted - Reuse the live session when reconnecting to the same conversation and wait for any in-flight response before signalling the client Co-Authored-By: Claude Opus 5.5 --- server/routers/assistant_chat.py | 43 ++++-- server/services/assistant_chat_session.py | 64 ++++++++- test_assistant_chat_session.py | 161 ++++++++++++++++++++++ 3 files changed, 252 insertions(+), 16 deletions(-) create mode 100644 test_assistant_chat_session.py diff --git a/server/routers/assistant_chat.py b/server/routers/assistant_chat.py index 1c3ece5c8..e3f569464 100644 --- a/server/routers/assistant_chat.py +++ b/server/routers/assistant_chat.py @@ -7,7 +7,7 @@ import json import logging -from typing import Optional +from typing import AsyncIterator, Optional from fastapi import APIRouter, HTTPException, WebSocket, WebSocketDisconnect from pydantic import BaseModel @@ -197,6 +197,29 @@ async def close_session(project_name: str): # WebSocket Endpoint # ============================================================================ +async def _forward_chunks(websocket: WebSocket, chunks: AsyncIterator[dict]) -> bool: + """ + Forward streamed chunks to the WebSocket. + + If the client goes away mid-response (panel closed, tab hidden, network + blip), keep consuming the stream so Claude's answer still completes and is + persisted to the conversation history. + + Returns: + True if the WebSocket is still usable, False if it disconnected. + """ + connected = True + async for chunk in chunks: + if not connected: + continue + try: + await websocket.send_json(chunk) + except Exception: + connected = False + logger.info("Assistant WebSocket closed mid-response; finishing response in background") + return connected + + @router.websocket("/ws/{project_name}") async def assistant_chat_websocket(websocket: WebSocket, project_name: str): """ @@ -270,11 +293,11 @@ async def assistant_chat_websocket(websocket: WebSocket, project_name: str): ) logger.debug("Session created, starting...") - # Stream the initial greeting - async for chunk in session.start(): - if logger.isEnabledFor(logging.DEBUG): - logger.debug(f"Sending chunk: {chunk.get('type')}") - await websocket.send_json(chunk) + # Re-attach to a live session for the same conversation, or + # start a fresh one (streams the greeting for new chats) + chunks = session.resume() if session.is_ready else session.start() + if not await _forward_chunks(websocket, chunks): + break logger.debug("Session start complete") except Exception as e: logger.exception(f"Error starting assistant session for {project_name}") @@ -302,8 +325,8 @@ async def assistant_chat_websocket(websocket: WebSocket, project_name: str): continue # Stream Claude's response - async for chunk in session.send_message(user_content): - await websocket.send_json(chunk) + if not await _forward_chunks(websocket, session.send_message(user_content)): + break elif msg_type == "answer": # User answered a structured question @@ -330,8 +353,8 @@ async def assistant_chat_websocket(websocket: WebSocket, project_name: str): user_response = str(answers) # Stream Claude's response - async for chunk in session.send_message(user_response): - await websocket.send_json(chunk) + if not await _forward_chunks(websocket, session.send_message(user_response)): + break else: await websocket.send_json({ diff --git a/server/services/assistant_chat_session.py b/server/services/assistant_chat_session.py index 1a16606d5..eabe7775d 100755 --- a/server/services/assistant_chat_session.py +++ b/server/services/assistant_chat_session.py @@ -7,12 +7,14 @@ but cannot modify any files. """ +import asyncio import json import logging import os import shutil import sys import threading +from contextlib import aclosing from datetime import datetime from pathlib import Path from typing import AsyncGenerator, Optional @@ -208,6 +210,19 @@ def __init__(self, project_name: str, project_dir: Path, conversation_id: Option self._client_entered: bool = False self.created_at = datetime.now() self._history_loaded: bool = False # Track if we've loaded history for resumed conversations + # Serializes queries so a response that is still streaming (e.g. after the + # browser disconnected) finishes and is persisted before the next one starts. + self._query_lock = asyncio.Lock() + + @property + def is_ready(self) -> bool: + """Whether the Claude client is initialized and can accept messages.""" + return self.client is not None and self._client_entered + + async def wait_until_idle(self) -> None: + """Wait for any in-flight response to finish streaming and be persisted.""" + async with self._query_lock: + pass async def close(self) -> None: """Clean up resources and close the Claude client.""" @@ -349,6 +364,17 @@ async def start(self) -> AsyncGenerator[dict, None]: # _history_loaded stays False so send_message() will include history yield {"type": "response_done"} + async def resume(self) -> AsyncGenerator[dict, None]: + """ + Re-attach a client to an already running session. + + Waits for any response still being generated (e.g. one that was in + flight when the browser disconnected) so it is persisted before the + client reloads the conversation history. + """ + await self.wait_until_idle() + yield {"type": "response_done"} + async def send_message(self, user_message: str) -> AsyncGenerator[dict, None]: """ Send user message and stream Claude's response. @@ -371,6 +397,15 @@ async def send_message(self, user_message: str) -> AsyncGenerator[dict, None]: yield {"type": "error", "content": "No conversation ID set."} return + # aclosing() propagates an early close so partial responses get persisted + async with self._query_lock, aclosing(self._send_message_locked(user_message)) as chunks: + async for chunk in chunks: + yield chunk + + async def _send_message_locked(self, user_message: str) -> AsyncGenerator[dict, None]: + """Body of send_message(); caller must hold self._query_lock.""" + assert self.conversation_id is not None + # Store user message in database add_message(self.project_dir, self.conversation_id, "user", user_message) @@ -399,8 +434,9 @@ async def send_message(self, user_message: str) -> AsyncGenerator[dict, None]: logger.info(f"Loaded {len(history)} messages from conversation history") try: - async for chunk in self._query_claude(message_to_send): - yield chunk + async with aclosing(self._query_claude(message_to_send)) as chunks: + async for chunk in chunks: + yield chunk yield {"type": "response_done"} except Exception as e: logger.exception("Error during Claude query") @@ -461,10 +497,11 @@ async def _query_claude(self, message: str) -> AsyncGenerator[dict, None]: yield {"type": "error", "content": "Rate limited. Please try again later."} return raise - - # Store the complete response in the database - if full_response and self.conversation_id: - add_message(self.project_dir, self.conversation_id, "assistant", full_response) + finally: + # Persist the response even if streaming was interrupted (client + # disconnected, session closed) so it shows up when history reloads. + if full_response and self.conversation_id: + add_message(self.project_dir, self.conversation_id, "assistant", full_response) def get_conversation_id(self) -> Optional[int]: """Get the current conversation ID.""" @@ -490,6 +527,10 @@ async def create_session( """ Create a new session for a project, closing any existing one. + If the existing session is already serving the requested conversation it is + returned as-is (check ``session.is_ready`` and use ``resume()`` instead of + ``start()``). + Args: project_name: Name of the project project_dir: Absolute path to the project directory @@ -498,6 +539,17 @@ async def create_session( old_session: Optional[AssistantChatSession] = None with _sessions_lock: + existing = _sessions.get(project_name) + if ( + existing is not None + and conversation_id is not None + and existing.conversation_id == conversation_id + and existing.project_dir == project_dir + and existing.is_ready + ): + # Reconnecting to the same conversation (e.g. the panel was closed and + # reopened): keep the live session so an in-flight response isn't lost. + return existing old_session = _sessions.pop(project_name, None) session = AssistantChatSession(project_name, project_dir, conversation_id) _sessions[project_name] = session diff --git a/test_assistant_chat_session.py b/test_assistant_chat_session.py new file mode 100644 index 000000000..b1454f186 --- /dev/null +++ b/test_assistant_chat_session.py @@ -0,0 +1,161 @@ +""" +Assistant Chat Session Tests +============================ + +Regression tests for issue #133: assistant answers were lost when the chat +panel was closed (WebSocket disconnected) while a response was streaming. + +Run with: python -m pytest test_assistant_chat_session.py +""" + +import asyncio +from pathlib import Path +from unittest.mock import patch + +import pytest + +from server.routers import assistant_chat as router +from server.services import assistant_chat_session as svc + + +class TextBlock: + def __init__(self, text: str): + self.text = text + + +class AssistantMessage: + def __init__(self, text: str): + self.content = [TextBlock(text)] + + +class FakeClient: + async def query(self, message: str) -> None: + self.last_query = message + + async def __aexit__(self, *args) -> None: + pass + + +def _fake_receive(parts: list[str], gate: asyncio.Event | None = None): + async def receive(client, logger): + for i, part in enumerate(parts): + if gate is not None and i == 1: + await gate.wait() + yield AssistantMessage(part) + return receive + + +def _ready_session(conversation_id: int = 1) -> svc.AssistantChatSession: + session = svc.AssistantChatSession("proj", Path("/tmp/proj"), conversation_id) + session.client = FakeClient() # type: ignore[assignment] + session._client_entered = True + session._history_loaded = True + return session + + +class ClosingWebSocket: + """WebSocket stub that fails on the Nth send, like a closed browser tab.""" + + def __init__(self, fail_after: int): + self.sent: list[dict] = [] + self.fail_after = fail_after + + async def send_json(self, data: dict) -> None: + if len(self.sent) >= self.fail_after: + raise RuntimeError("WebSocket is closed") + self.sent.append(data) + + +def test_response_persisted_when_websocket_disconnects_mid_stream(): + stored: list[tuple[str, str]] = [] + + async def run(): + session = _ready_session() + ws = ClosingWebSocket(fail_after=1) + with ( + patch.object(svc, "add_message", lambda _d, _c, role, content: stored.append((role, content))), + patch.object(svc, "safe_receive_response", _fake_receive(["Hello ", "world", "!"])), + ): + connected = await router._forward_chunks(ws, session.send_message("hi")) + return connected, ws.sent + + connected, sent = asyncio.run(run()) + + assert connected is False + assert sent == [{"type": "text", "content": "Hello "}] + # Full answer still stored even though the client only saw the first chunk + assert stored == [("user", "hi"), ("assistant", "Hello world!")] + + +def test_partial_response_persisted_when_stream_is_aborted(): + stored: list[tuple[str, str]] = [] + + async def run(): + session = _ready_session() + with ( + patch.object(svc, "add_message", lambda _d, _c, role, content: stored.append((role, content))), + patch.object(svc, "safe_receive_response", _fake_receive(["partial", " rest"])), + ): + gen = session.send_message("hi") + await gen.__anext__() + await gen.aclose() + + asyncio.run(run()) + assert stored == [("user", "hi"), ("assistant", "partial")] + + +def test_create_session_reuses_live_session_for_same_conversation(): + async def run(): + with patch.dict(svc._sessions, clear=True): + existing = _ready_session(conversation_id=7) + svc._sessions["proj"] = existing + + same = await svc.create_session("proj", Path("/tmp/proj"), conversation_id=7) + assert same is existing + + other = await svc.create_session("proj", Path("/tmp/proj"), conversation_id=8) + assert other is not existing + assert svc._sessions["proj"] is other + + asyncio.run(run()) + + +def test_resume_waits_for_in_flight_response(): + stored: list[tuple[str, str]] = [] + order: list[str] = [] + + async def run(): + session = _ready_session() + gate = asyncio.Event() + + def add(_d, _c, role, content): + stored.append((role, content)) + order.append(f"stored:{role}") + + async def consume_message(): + async for _ in session.send_message("hi"): + pass + + async def consume_resume(): + async for chunk in session.resume(): + order.append(chunk["type"]) + + with ( + patch.object(svc, "add_message", add), + patch.object(svc, "safe_receive_response", _fake_receive(["a", "b"], gate)), + ): + msg_task = asyncio.create_task(consume_message()) + await asyncio.sleep(0) + resume_task = asyncio.create_task(consume_resume()) + await asyncio.sleep(0.01) + assert "response_done" not in order + gate.set() + await asyncio.gather(msg_task, resume_task) + + asyncio.run(run()) + assert stored[-1] == ("assistant", "ab") + assert order.index("stored:assistant") < order.index("response_done") + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) From f5dfda6e660db1a6b6947fb1f306bb57efd6049d Mon Sep 17 00:00:00 2001 From: mohamed zwanski Date: Thu, 22 Oct 2026 21:47:00 +0100 Subject: [PATCH 2/2] Keep assistant chat alive when the panel closes and resync on reconnect (#133) - Keep AssistantChat mounted after the panel is first opened (per project) so closing it no longer drops the WebSocket mid-response; mark the hidden panel inert and scroll the message list directly so it can't steal focus or scroll the page - On resume/reconnect, wait for the server to finish any in-flight reply, then reload the stored history instead of showing stale cached messages - Stop reconnecting after the chat unmounts Co-Authored-By: Claude Opus 5.5 --- ui/src/components/AssistantChat.tsx | 13 ++++---- ui/src/components/AssistantPanel.tsx | 12 ++++++- ui/src/hooks/useAssistantChat.ts | 48 ++++++++++++++++++++++++++-- 3 files changed, 64 insertions(+), 9 deletions(-) diff --git a/ui/src/components/AssistantChat.tsx b/ui/src/components/AssistantChat.tsx index 0592644e6..be689f539 100644 --- a/ui/src/components/AssistantChat.tsx +++ b/ui/src/components/AssistantChat.tsx @@ -38,7 +38,7 @@ export function AssistantChat({ }: AssistantChatProps) { const [inputValue, setInputValue] = useState('') const [showHistory, setShowHistory] = useState(false) - const messagesEndRef = useRef(null) + const messagesContainerRef = useRef(null) const inputRef = useRef(null) const hasStartedRef = useRef(false) const lastConversationIdRef = useRef(undefined) @@ -77,10 +77,12 @@ export function AssistantChat({ previousActiveConversationIdRef.current = activeConversationId }, [activeConversationId, onConversationCreated]) - // Auto-scroll to bottom on new messages + // Auto-scroll to bottom on new messages. Scroll the container directly rather + // than using scrollIntoView, which can scroll the page while the panel is hidden. useEffect(() => { - messagesEndRef.current?.scrollIntoView({ behavior: 'smooth' }) - }, [messages]) + const container = messagesContainerRef.current + container?.scrollTo({ top: container.scrollHeight, behavior: 'smooth' }) + }, [messages, initialMessages]) // Start or resume the chat session when component mounts or conversationId changes useEffect(() => { @@ -228,7 +230,7 @@ export function AssistantChat({ {/* Messages area */} -
+
{isLoadingConversation ? (
@@ -252,7 +254,6 @@ export function AssistantChat({ {displayMessages.map((message) => ( ))} -
)}
diff --git a/ui/src/components/AssistantPanel.tsx b/ui/src/components/AssistantPanel.tsx index ae9168123..3ca240b1b 100644 --- a/ui/src/components/AssistantPanel.tsx +++ b/ui/src/components/AssistantPanel.tsx @@ -48,6 +48,14 @@ function setStoredConversationId(projectName: string, conversationId: number | n } export function AssistantPanel({ projectName, isOpen, onClose }: AssistantPanelProps) { + // Keep the chat mounted after it has been opened for this project so closing + // the panel doesn't drop the WebSocket (and any response still streaming in). + const [openedForProject, setOpenedForProject] = useState(isOpen ? projectName : null) + if (isOpen && openedForProject !== projectName) { + setOpenedForProject(projectName) + } + const keepChatMounted = openedForProject === projectName + // Load initial conversation ID from localStorage const [conversationId, setConversationId] = useState(() => getStoredConversationId(projectName) @@ -172,6 +180,7 @@ export function AssistantPanel({ projectName, isOpen, onClose }: AssistantPanelP role="dialog" aria-label="Project Assistant" aria-hidden={!isOpen} + inert={!isOpen} > {/* Resize handle */}
- {isOpen && ( + {keepChatMounted && ( (null); const [currentQuestions, setCurrentQuestions] = useState(null); + const queryClient = useQueryClient(); const wsRef = useRef(null); const currentAssistantMessageRef = useRef(null); const reconnectAttempts = useRef(0); @@ -47,10 +49,19 @@ export function useAssistantChat({ const pingIntervalRef = useRef(null); const reconnectTimeoutRef = useRef(null); const checkAndSendTimeoutRef = useRef(null); + // Mirrors conversationId so WebSocket callbacks can read the latest value + const conversationIdRef = useRef(null); + // Conversation being resumed; set while waiting for the server to confirm + const resumingConversationIdRef = useRef(null); + + useEffect(() => { + conversationIdRef.current = conversationId; + }, [conversationId]); // Clean up on unmount useEffect(() => { return () => { + reconnectAttempts.current = maxReconnectAttempts; // Prevent reconnection after unmount if (pingIntervalRef.current) { clearInterval(pingIntervalRef.current); } @@ -87,8 +98,20 @@ export function useAssistantChat({ ws.onopen = () => { setConnectionStatus("connected"); + const isReconnect = reconnectAttempts.current > 0; reconnectAttempts.current = 0; + // After an unexpected disconnect, re-attach to the conversation. The + // server finishes any in-flight response, then we reload history so a + // reply that streamed while we were offline isn't lost. + if (isReconnect && conversationIdRef.current) { + resumingConversationIdRef.current = conversationIdRef.current; + setIsLoading(true); + ws.send( + JSON.stringify({ type: "start", conversation_id: conversationIdRef.current }), + ); + } + // Start ping interval to keep connection alive pingIntervalRef.current = window.setInterval(() => { if (ws.readyState === WebSocket.OPEN) { @@ -232,9 +255,26 @@ export function useAssistantChat({ } case "response_done": { - setIsLoading(false); currentAssistantMessageRef.current = null; + const resumedId = resumingConversationIdRef.current; + if (resumedId !== null) { + // The stored history is now complete (including any reply that + // finished while disconnected). Reload it and let it replace the + // live messages to avoid duplicates. + resumingConversationIdRef.current = null; + queryClient + .invalidateQueries({ queryKey: ["conversation", projectName, resumedId] }) + .finally(() => { + setMessages([]); + setIsLoading(false); + }); + queryClient.invalidateQueries({ queryKey: ["conversations", projectName] }); + break; + } + + setIsLoading(false); + // Mark current message as done streaming setMessages((prev) => { const lastMessage = prev[prev.length - 1]; @@ -253,6 +293,7 @@ export function useAssistantChat({ } case "error": { + resumingConversationIdRef.current = null; setIsLoading(false); onError?.(data.content); @@ -278,7 +319,7 @@ export function useAssistantChat({ console.error("Failed to parse WebSocket message:", e); } }; - }, [projectName, onError]); + }, [projectName, onError, queryClient]); const start = useCallback( (existingConversationId?: number | null) => { @@ -301,6 +342,9 @@ export function useAssistantChat({ if (existingConversationId) { payload.conversation_id = existingConversationId; setConversationId(existingConversationId); + resumingConversationIdRef.current = existingConversationId; + } else { + resumingConversationIdRef.current = null; } if (import.meta.env.DEV) { console.debug('[useAssistantChat] Sending start message:', payload);