Ignore WS heartbeats and reject empty chat messages (#32)
Deploy Beta / unit-tests (push) Successful in 10s
Unit Tests / test (push) Successful in 9s
Deploy Beta / docker (push) Successful in 25s
Deploy Beta / deploy-beta (push) Successful in 56s

## 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: #32
This commit was merged in pull request #32.
This commit is contained in:
2026-07-28 05:12:13 -07:00
parent 5d5b448868
commit ee3d47c8c3
8 changed files with 170 additions and 4 deletions
+18 -1
View File
@@ -27,6 +27,7 @@ from .services.title_generator import title_generator
from .services.moderation_classifier import moderation_classifier, ModerationLabel from .services.moderation_classifier import moderation_classifier, ModerationLabel
from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType
from .services.data_analysis_service import AsyncDataAnalysisService from .services.data_analysis_service import AsyncDataAnalysisService
from .utils import has_usable_user_prompt, is_heartbeat_payload, normalize_user_message
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -238,13 +239,29 @@ class ChatConsumerAgain(AsyncWebsocketConsumer):
logger.debug(f"Bytes Data: {bytes_data}") logger.debug(f"Bytes Data: {bytes_data}")
if text_data: if text_data:
data = json.loads(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) conversation_id = data.get("conversation_id", None)
email = data.get("email", None) email = data.get("email", None)
file = data.get("file", None) file = data.get("file", None)
file_type = data.get("fileType", "") file_type = data.get("fileType", "")
model = data.get("modelName", "Turbo") 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: if not conversation_id:
# we need to create a new conversation # we need to create a new conversation
# we will generate a name for it too # we will generate a name for it too
+19 -2
View File
@@ -22,6 +22,7 @@ from .services.title_generator import title_generator
from .services.moderation_classifier import moderation_classifier, ModerationLabel from .services.moderation_classifier import moderation_classifier, ModerationLabel
from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType
from .services.data_analysis_service import AsyncDataAnalysisService from .services.data_analysis_service import AsyncDataAnalysisService
from .utils import has_usable_user_prompt, is_heartbeat_payload, normalize_user_message
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -288,13 +289,29 @@ class ChatConsumerGraph(AsyncWebsocketConsumer):
print("Text Data: ", text_data) print("Text Data: ", text_data)
if text_data: if text_data:
data = json.loads(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") 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) conversation_id = data.get("conversation_id", None)
email = data.get("email", None) email = data.get("email", None)
file = data.get("file", None) file = data.get("file", None)
file_type = data.get("fileType", "") 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: if not conversation_id:
title = await title_generator.generate_async(message) title = await title_generator.generate_async(message)
conversation_id = await create_conversation(message, email, title) conversation_id = await create_conversation(message, email, title)
+5
View File
@@ -122,6 +122,11 @@ class PromptSerializer(serializers.ModelSerializer):
"id", "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 BasicUserSerializer(serializers.ModelSerializer):
class Meta: class Meta:
@@ -1,4 +1,5 @@
from unittest import mock from unittest import mock
import json
from asgiref.sync import sync_to_async from asgiref.sync import sync_to_async
from channels.testing import WebsocketCommunicator from channels.testing import WebsocketCommunicator
@@ -412,3 +413,61 @@ class WebSocketRoutingTestCase(TransactionTestCase):
[str(route.pattern) for route in websocket_urlpatterns], [str(route.pattern) for route in websocket_urlpatterns],
["ws/chat_again/$", "ws/conditional_chat/$"], ["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()
@@ -58,6 +58,22 @@ class PromptSerializerTestCase(TestCase):
self.assertFalse(serializer.is_valid()) self.assertFalse(serializer.is_valid())
self.assertIn("message", serializer.errors) 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): def test_user_created_is_required(self):
serializer = PromptSerializer(data={"message": "hi"}) serializer = PromptSerializer(data={"message": "hi"})
+27 -1
View File
@@ -11,7 +11,13 @@ from chat_backend.ollama_config import (
ollama_llm_kwargs, ollama_llm_kwargs,
ollama_model, 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): class ExtractTokenUsageTestCase(SimpleTestCase):
@@ -134,3 +140,23 @@ class OllamaConfigFallbackTestCase(SimpleTestCase):
with override_settings(OLLAMA_MODEL="llama3.2"): with override_settings(OLLAMA_MODEL="llama3.2"):
del settings.OLLAMA_EMBED_MODEL del settings.OLLAMA_EMBED_MODEL
self.assertEqual(ollama_embed_model(), "llama3.2") 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 "))
+20
View File
@@ -1,6 +1,26 @@
import datetime 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): def last_day_of_month(any_day):
# The day 28 exists in every month. 4 days later, it's always next month # 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) next_month = any_day.replace(day=28) + datetime.timedelta(days=4)
+6
View File
@@ -469,6 +469,12 @@ class ConversationDetailView(APIView):
# make sure that our model exists and it is running # make sure that our model exists and it is running
prompt = request.data.get("prompt") 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") conversation_id = request.data.get("conversation_id")
is_user = bool(request.data.get("is_user")) is_user = bool(request.data.get("is_user"))