From ee3d47c8c31d825b914329d3cf8b1ede759ac8dc Mon Sep 17 00:00:00 2001 From: Ryan Westfall Date: Tue, 28 Jul 2026 05:12:13 -0700 Subject: [PATCH] Ignore WS heartbeats and reject empty chat messages (#32) ## Summary - Closes #31 - Ignore WebSocket `type: ping` heartbeats so keepalives no longer create conversations or hit title/LLM pipelines - Reject empty/whitespace user messages in both chat consumers, `PromptSerializer`, and REST conversation prompt POST ## Test plan - [x] `UserPromptGuardTestCase`, `PromptSerializerTestCase` blank/whitespace cases - [x] `WebSocketReceiveGuardTestCase` ping ignore + empty message rejection (both WS routes) - [ ] Deploy to beta; leave idle tab open and confirm no new rogue conversations - [ ] Confirm normal chat send still works Related FE: https://git.aimloperations.com/ai_ml_operations/chat_web_app/issues/51Reviewed-on: https://git.aimloperations.com/ai_ml_operations/chat_backend/pulls/32 --- llm_be/chat_backend/consumers.py | 19 +++++- llm_be/chat_backend/consumers_graph.py | 21 ++++++- llm_be/chat_backend/serializers.py | 5 ++ llm_be/chat_backend/tests/test_consumers.py | 59 +++++++++++++++++++ llm_be/chat_backend/tests/test_serializers.py | 16 +++++ llm_be/chat_backend/tests/test_utils.py | 28 ++++++++- llm_be/chat_backend/utils.py | 20 +++++++ llm_be/chat_backend/views.py | 6 ++ 8 files changed, 170 insertions(+), 4 deletions(-) diff --git a/llm_be/chat_backend/consumers.py b/llm_be/chat_backend/consumers.py index b1d3d59..c1bd1c4 100644 --- a/llm_be/chat_backend/consumers.py +++ b/llm_be/chat_backend/consumers.py @@ -27,6 +27,7 @@ from .services.title_generator import title_generator from .services.moderation_classifier import moderation_classifier, ModerationLabel from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType from .services.data_analysis_service import AsyncDataAnalysisService +from .utils import has_usable_user_prompt, is_heartbeat_payload, normalize_user_message logger = logging.getLogger(__name__) @@ -238,13 +239,29 @@ class ChatConsumerAgain(AsyncWebsocketConsumer): logger.debug(f"Bytes Data: {bytes_data}") if text_data: data = json.loads(text_data) - message = data.get("message", None) + # Keepalive frames must not create conversations or hit the LLM. + if is_heartbeat_payload(data): + return + + message = normalize_user_message(data.get("message", None)) conversation_id = data.get("conversation_id", None) email = data.get("email", None) file = data.get("file", None) file_type = data.get("fileType", "") model = data.get("modelName", "Turbo") + if not has_usable_user_prompt(message, file): + logger.info("Ignoring websocket payload with empty message") + await self.send_json_message( + json.dumps( + { + "type": "error", + "content": "Message text cannot be empty.", + } + ) + ) + return + if not conversation_id: # we need to create a new conversation # we will generate a name for it too diff --git a/llm_be/chat_backend/consumers_graph.py b/llm_be/chat_backend/consumers_graph.py index 2dafff4..4b44bb8 100644 --- a/llm_be/chat_backend/consumers_graph.py +++ b/llm_be/chat_backend/consumers_graph.py @@ -22,6 +22,7 @@ from .services.title_generator import title_generator from .services.moderation_classifier import moderation_classifier, ModerationLabel from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType from .services.data_analysis_service import AsyncDataAnalysisService +from .utils import has_usable_user_prompt, is_heartbeat_payload, normalize_user_message logger = logging.getLogger(__name__) @@ -288,13 +289,29 @@ class ChatConsumerGraph(AsyncWebsocketConsumer): print("Text Data: ", text_data) if text_data: data = json.loads(text_data) + # Keepalive frames must not create conversations or hit the LLM. + if is_heartbeat_payload(data): + return + model = data.get("modelName", "Turbo") - message = data.get("message", None) + message = normalize_user_message(data.get("message", None)) conversation_id = data.get("conversation_id", None) email = data.get("email", None) file = data.get("file", None) file_type = data.get("fileType", "") - + + if not has_usable_user_prompt(message, file): + logger.info("Ignoring websocket payload with empty message") + await self.send_json_message( + json.dumps( + { + "type": "error", + "content": "Message text cannot be empty.", + } + ) + ) + return + if not conversation_id: title = await title_generator.generate_async(message) conversation_id = await create_conversation(message, email, title) diff --git a/llm_be/chat_backend/serializers.py b/llm_be/chat_backend/serializers.py index 9a6947a..646db4e 100644 --- a/llm_be/chat_backend/serializers.py +++ b/llm_be/chat_backend/serializers.py @@ -122,6 +122,11 @@ class PromptSerializer(serializers.ModelSerializer): "id", ) + def validate_message(self, value: str) -> str: + if value is None or not str(value).strip(): + raise serializers.ValidationError("Message text cannot be empty.") + return str(value).strip() + class BasicUserSerializer(serializers.ModelSerializer): class Meta: diff --git a/llm_be/chat_backend/tests/test_consumers.py b/llm_be/chat_backend/tests/test_consumers.py index f440e8e..52e23f4 100644 --- a/llm_be/chat_backend/tests/test_consumers.py +++ b/llm_be/chat_backend/tests/test_consumers.py @@ -1,4 +1,5 @@ from unittest import mock +import json from asgiref.sync import sync_to_async from channels.testing import WebsocketCommunicator @@ -412,3 +413,61 @@ class WebSocketRoutingTestCase(TransactionTestCase): [str(route.pattern) for route in websocket_urlpatterns], ["ws/chat_again/$", "ws/conditional_chat/$"], ) + + +class WebSocketReceiveGuardTestCase(TransactionTestCase): + """Heartbeats / empty prompts must not spawn conversations or LLM work.""" + + @parameterized.expand( + [("chat", "/ws/chat_again/"), ("conditional_chat", "/ws/conditional_chat/")] + ) + async def test_ping_heartbeat_is_ignored(self, _name, path): + communicator = WebsocketCommunicator(application, path) + connected, _ = await communicator.connect() + self.assertTrue(connected) + + with mock.patch( + "chat_backend.consumers.title_generator.generate_async", + new_callable=mock.AsyncMock, + ) as title_chat, mock.patch( + "chat_backend.consumers_graph.title_generator.generate_async", + new_callable=mock.AsyncMock, + ) as title_graph: + await communicator.send_json_to({"type": "ping", "email": "a@b.com"}) + # No reply expected; give the event loop a tick. + self.assertTrue(await communicator.receive_nothing(timeout=0.2)) + title_chat.assert_not_called() + title_graph.assert_not_called() + + count = await sync_to_async(Conversation.objects.count)() + self.assertEqual(count, 0) + await communicator.disconnect() + + @parameterized.expand( + [("chat", "/ws/chat_again/"), ("conditional_chat", "/ws/conditional_chat/")] + ) + async def test_empty_message_is_rejected(self, _name, path): + communicator = WebsocketCommunicator(application, path) + connected, _ = await communicator.connect() + self.assertTrue(connected) + + with mock.patch( + "chat_backend.consumers.title_generator.generate_async", + new_callable=mock.AsyncMock, + ) as title_chat, mock.patch( + "chat_backend.consumers_graph.title_generator.generate_async", + new_callable=mock.AsyncMock, + ) as title_graph: + await communicator.send_json_to( + {"message": " ", "email": "a@b.com", "conversation_id": None} + ) + response = await communicator.receive_from(timeout=1) + payload = json.loads(response) + self.assertEqual(payload["type"], "error") + self.assertIn("empty", payload["content"].lower()) + title_chat.assert_not_called() + title_graph.assert_not_called() + + count = await sync_to_async(Conversation.objects.count)() + self.assertEqual(count, 0) + await communicator.disconnect() diff --git a/llm_be/chat_backend/tests/test_serializers.py b/llm_be/chat_backend/tests/test_serializers.py index 0ef05ac..0a92eca 100644 --- a/llm_be/chat_backend/tests/test_serializers.py +++ b/llm_be/chat_backend/tests/test_serializers.py @@ -58,6 +58,22 @@ class PromptSerializerTestCase(TestCase): self.assertFalse(serializer.is_valid()) self.assertIn("message", serializer.errors) + def test_message_rejects_blank_and_whitespace(self): + for payload in ("", " ", "\n\t"): + serializer = PromptSerializer( + data={"message": payload, "user_created": True} + ) + self.assertFalse(serializer.is_valid(), payload) + self.assertIn("message", serializer.errors) + + def test_message_is_stripped(self): + serializer = PromptSerializer( + data={"message": " hello ", "user_created": True} + ) + + self.assertTrue(serializer.is_valid(), serializer.errors) + self.assertEqual(serializer.validated_data["message"], "hello") + def test_user_created_is_required(self): serializer = PromptSerializer(data={"message": "hi"}) diff --git a/llm_be/chat_backend/tests/test_utils.py b/llm_be/chat_backend/tests/test_utils.py index 6b3f612..a26f7cf 100644 --- a/llm_be/chat_backend/tests/test_utils.py +++ b/llm_be/chat_backend/tests/test_utils.py @@ -11,7 +11,13 @@ from chat_backend.ollama_config import ( ollama_llm_kwargs, ollama_model, ) -from chat_backend.utils import extract_token_usage, last_day_of_month +from chat_backend.utils import ( + extract_token_usage, + has_usable_user_prompt, + is_heartbeat_payload, + last_day_of_month, + normalize_user_message, +) class ExtractTokenUsageTestCase(SimpleTestCase): @@ -134,3 +140,23 @@ class OllamaConfigFallbackTestCase(SimpleTestCase): with override_settings(OLLAMA_MODEL="llama3.2"): del settings.OLLAMA_EMBED_MODEL self.assertEqual(ollama_embed_model(), "llama3.2") + + +class UserPromptGuardTestCase(SimpleTestCase): + def test_heartbeat_payload_detected(self): + self.assertTrue(is_heartbeat_payload({"type": "ping", "email": "a@b.com"})) + self.assertFalse(is_heartbeat_payload({"message": "hi"})) + self.assertFalse(is_heartbeat_payload(None)) + + def test_normalize_user_message(self): + self.assertIsNone(normalize_user_message(None)) + self.assertIsNone(normalize_user_message("")) + self.assertIsNone(normalize_user_message(" \n\t")) + self.assertEqual(normalize_user_message(" hello "), "hello") + + def test_has_usable_user_prompt(self): + self.assertFalse(has_usable_user_prompt(None)) + self.assertFalse(has_usable_user_prompt(" ")) + self.assertFalse(has_usable_user_prompt("", file="base64")) + self.assertTrue(has_usable_user_prompt("hi")) + self.assertTrue(has_usable_user_prompt(" hi ")) diff --git a/llm_be/chat_backend/utils.py b/llm_be/chat_backend/utils.py index 9972693..687b1cf 100644 --- a/llm_be/chat_backend/utils.py +++ b/llm_be/chat_backend/utils.py @@ -1,6 +1,26 @@ import datetime +def is_heartbeat_payload(data) -> bool: + """True for app-level WS keepalive frames (see FE buildHeartbeatPayload).""" + return isinstance(data, dict) and data.get("type") == "ping" + + +def normalize_user_message(message): + """Return stripped message text, or None if missing/blank.""" + if message is None: + return None + if not isinstance(message, str): + message = str(message) + stripped = message.strip() + return stripped or None + + +def has_usable_user_prompt(message, file=None) -> bool: + """Reject empty/whitespace chat text. ``file`` kept for call-site clarity.""" + return normalize_user_message(message) is not None + + def last_day_of_month(any_day): # The day 28 exists in every month. 4 days later, it's always next month next_month = any_day.replace(day=28) + datetime.timedelta(days=4) diff --git a/llm_be/chat_backend/views.py b/llm_be/chat_backend/views.py index d41299a..3813dd8 100644 --- a/llm_be/chat_backend/views.py +++ b/llm_be/chat_backend/views.py @@ -469,6 +469,12 @@ class ConversationDetailView(APIView): # make sure that our model exists and it is running prompt = request.data.get("prompt") + if not isinstance(prompt, str) or not prompt.strip(): + return Response( + {"detail": "Message text cannot be empty."}, + status=status.HTTP_400_BAD_REQUEST, + ) + prompt = prompt.strip() conversation_id = request.data.get("conversation_id") is_user = bool(request.data.get("is_user"))