Fix assistant losing answers when the chat panel is closed or reconnects

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({

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

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"]))