From e1db50da4a49548629494278b0e58bd73916acbe Mon Sep 17 00:00:00 2001
From: Manas <manas.gooner11@gmail.com>
Date: Sat, 4 Jan 2025 03:08:09 +0530
Subject: [PATCH] fix(websockets): improve error logging for edge-cases

---
 src/diart/websockets.py | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/src/diart/websockets.py b/src/diart/websockets.py
index 834d074..67e665d 100644
--- a/src/diart/websockets.py
+++ b/src/diart/websockets.py
@@ -212,10 +212,21 @@ def _on_message_received(
             WebSocket server instance
         message : AnyStr
             Received message data
+
+        Raises
+        ------
+        Warning
+            If client not found in self._clients. Common cases:
+            - Message received before client initialization complete
+            - Message received after client cleanup started
+            - Race condition in client connection/disconnection
+        Error
+            If message processing fails. Error will be logged.
         """
         client_id = client["id"]
 
         if client_id not in self._clients:
+            logger.warning(f"Received message from unregistered client {client_id}")
             return
 
         try:
@@ -240,12 +251,23 @@ def send(self, client_id: Text, message: AnyStr) -> None:
             Target client identifier
         message : AnyStr
             Message to send
+
+        Raises
+        ------
+        Warning
+            If client is not found in server.clients. Common cases:
+            - Client disconnected but cleanup is not complete
+            - Network issues caused client to drop
+            - Race condition between disconnect and message send
+        Error
+            If message sending fails. Error will be logged and re-raised.
         """
         if not message:
             return
 
         client = next((c for c in self.server.clients if c["id"] == client_id), None)
         if client is None:
+            logger.warning(f"Failed to send message to client {client_id}: client not found")
             return
 
         try:
@@ -261,8 +283,19 @@ def close(self, client_id: Text) -> None:
         ----------
         client_id : Text
             Client identifier to close
+
+        Raises
+        ------
+        Warning
+            If client not found in self._clients. Common cases:
+            - Already cleaned up by another error handler
+            - Multiple error handlers trying to cleanup same client
+            - Cleanup triggered for client that never fully connected
+        Error
+            If cleanup fails. Error will be logged and client will be force-removed.
         """
         if client_id not in self._clients:
+            logger.warning(f"Attempted to close non-existent client {client_id}")
             return
 
         try: