fix: persist WebSocket streaming messages to disk (#174)

* fix: persist WebSocket streaming messages to disk

The WebSocket streaming path (get_chatbot_reply_stream) updates the
in-memory session history but never persists it to disk. The REST
endpoint already does this via BackgroundTasks, but the WebSocket
handler was missing the equivalent call.

Added persist_session() after each WebSocket exchange completes,
using asyncio.to_thread() to avoid blocking the event loop.

Fixes #173

* test: add unit test for WebSocket session persistence

---------

Co-authored-by: Bervianto Leo Pratama <[email protected]>

Author: sharma-sugurthi

chatbot-core/api/routes/chatbot.py

@@ -134,6 +134,8 @@ async def chatbot_stream(websocket: WebSocket, session_id: str):
                 json.dumps({"end": True})
             )
 
+            asyncio.create_task(asyncio.to_thread(persist_session, session_id))
+
     except WebSocketDisconnect:
         logger.info(
             "WebSocket disconnected for session %s",

chatbot-core/tests/unit/mocks/test_env.py

@@ -62,6 +62,11 @@ def mock_get_chatbot_reply_stream(mocker):
     """Mock the get_chatbot_reply_stream function."""
     return mocker.patch("api.routes.chatbot.get_chatbot_reply_stream")
 
+@pytest.fixture
+def mock_persist_session(mocker):
+    """Mock the persist_session function."""
+    return mocker.patch("api.routes.chatbot.persist_session")
+
 
 @pytest.fixture
 def mock_process_uploaded_file(mocker):

chatbot-core/tests/unit/routes/test_chatbot.py

@@ -156,3 +156,27 @@ def test_websocket_invalid_session_returns_error_and_closes(
     with client.websocket_connect("/sessions/bad-session/stream") as ws:
         error = ws.receive_json()
         assert error == {"error": "Session not found"}
+
+
+def test_websocket_persist_session_called_after_streaming(
+    client, mock_session_exists, mock_get_chatbot_reply_stream, mock_persist_session
+):
+    """persist_session must be called after a WebSocket streaming exchange completes.
+
+    Before fix: the WebSocket path updated in-memory history but never
+    persisted to disk — messages were lost on server restart.
+    After fix: asyncio.create_task runs persist_session in the background.
+    """
+    mock_session_exists.return_value = True
+
+    async def fake_stream(_session_id, _message):
+        yield "token"
+
+    mock_get_chatbot_reply_stream.side_effect = fake_stream
+
+    with client.websocket_connect("/sessions/test-session-id/stream") as ws:
+        ws.send_json({"message": "Hello"})
+        ws.receive_json()  # token
+        ws.receive_json()  # end marker
+
+    mock_persist_session.assert_called_with("test-session-id")