fix: handle malformed JSON in WebSocket handler gracefully (fixes #257) (#259)

* fix: handle malformed JSON in WebSocket handler gracefully

Wrap json.loads() in chatbot_stream with a try/except for
json.JSONDecodeError. On malformed input, the handler now:

1. Logs a warning with the session ID
2. Sends a JSON error back: {"error": "Invalid JSON format."}
3. Continues listening for the next message

Previously, a single malformed message would propagate to the
outer except Exception block, terminating the entire WebSocket
connection.

Fixes #257

* test: add WebSocket unit tests for malformed JSON handling

---------

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

Author: sharma-sugurthi

chatbot-core/api/routes/chatbot.py

@@ -105,7 +105,18 @@ async def chatbot_stream(websocket: WebSocket, session_id: str):
     try:
         while True:
             data = await websocket.receive_text()
-            message_data = json.loads(data)
+
+            try:
+                message_data = json.loads(data)
+            except json.JSONDecodeError:
+                logger.warning(
+                    "Malformed JSON from session %s", session_id
+                )
+                await websocket.send_text(
+                    json.dumps({"error": "Invalid JSON format."})
+                )
+                continue
+
             user_message = message_data.get("message", "")
 
             if not user_message:

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

@@ -57,6 +57,12 @@ def mock_get_chatbot_reply(mocker):
     """Mock the get_chatbot_reply function."""
     return mocker.patch("api.routes.chatbot.get_chatbot_reply")
 
+@pytest.fixture
+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_process_uploaded_file(mocker):
     """Mock the process_uploaded_file function."""

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

@@ -64,3 +64,95 @@ def test_delete_chat_not_found(client, mock_delete_session):
 
     assert response.status_code == 404
     assert response.json() == {"detail": "Session not found."}
+
+
+# =========================
+# WebSocket Tests
+# =========================
+
+def test_websocket_malformed_json_returns_error_and_stays_alive(
+    client, mock_session_exists, mock_get_chatbot_reply_stream
+):
+    """Malformed JSON must return an error message without closing the connection.
+
+    Before fix: a single bad message like 'hello' would propagate to the
+    outer except block and terminate the WebSocket connection entirely.
+    After fix: the handler catches json.JSONDecodeError, sends an error
+    JSON, and continues listening.
+    """
+    mock_session_exists.return_value = True
+
+    async def fake_stream(_session_id, _message):
+        yield "reply-token"
+
+    mock_get_chatbot_reply_stream.side_effect = fake_stream
+
+    with client.websocket_connect("/sessions/test-session-id/stream") as ws:
+        # Send malformed JSON
+        ws.send_text("not valid json")
+        error_response = ws.receive_json()
+        assert error_response == {"error": "Invalid JSON format."}
+
+        # Connection is still alive — send a valid message
+        ws.send_json({"message": "Hello"})
+        token_response = ws.receive_json()
+        assert "token" in token_response
+        end_response = ws.receive_json()
+        assert end_response == {"end": True}
+
+
+def test_websocket_valid_json_streams_response(
+    client, mock_session_exists, mock_get_chatbot_reply_stream
+):
+    """A valid JSON message must produce streaming tokens followed by end marker."""
+    mock_session_exists.return_value = True
+
+    async def fake_stream(_session_id, _message):
+        yield "Hello"
+        yield " world"
+
+    mock_get_chatbot_reply_stream.side_effect = fake_stream
+
+    with client.websocket_connect("/sessions/test-session-id/stream") as ws:
+        ws.send_json({"message": "What is Jenkins?"})
+
+        token1 = ws.receive_json()
+        assert token1 == {"token": "Hello"}
+        token2 = ws.receive_json()
+        assert token2 == {"token": " world"}
+        end = ws.receive_json()
+        assert end == {"end": True}
+
+
+def test_websocket_empty_message_is_skipped(
+    client, mock_session_exists, mock_get_chatbot_reply_stream
+):
+    """An empty message field must be silently skipped without error."""
+    mock_session_exists.return_value = True
+
+    async def fake_stream(_session_id, _message):
+        yield "response"
+
+    mock_get_chatbot_reply_stream.side_effect = fake_stream
+
+    with client.websocket_connect("/sessions/test-session-id/stream") as ws:
+        # Send valid JSON with empty message — should be skipped
+        ws.send_json({"message": ""})
+
+        # Send a real message to prove connection is still alive
+        ws.send_json({"message": "Real question"})
+        token = ws.receive_json()
+        assert token == {"token": "response"}
+        end = ws.receive_json()
+        assert end == {"end": True}
+
+
+def test_websocket_invalid_session_returns_error_and_closes(
+    client, mock_session_exists
+):
+    """Connecting with a non-existent session ID must return error and close."""
+    mock_session_exists.return_value = False
+
+    with client.websocket_connect("/sessions/bad-session/stream") as ws:
+        error = ws.receive_json()
+        assert error == {"error": "Session not found"}