From 1f123c77e75a810682d968332501ae354d6730e2 Mon Sep 17 00:00:00 2001
From: algizn97 <algiz_n@mail.ru>
Date: Wed, 19 Nov 2025 16:50:18 +0500
Subject: [PATCH] Fixed the websocket

---
 app/bybit/telegram_message_handler.py |  4 +++-
 app/bybit/web_socket.py               | 18 +++++++++++++-----
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/app/bybit/telegram_message_handler.py b/app/bybit/telegram_message_handler.py
index a17ac58..e6251df 100644
--- a/app/bybit/telegram_message_handler.py
+++ b/app/bybit/telegram_message_handler.py
@@ -1,6 +1,6 @@
 import logging.config
 import math
-# import json
+import json
 import app.telegram.keyboards.inline as kbi
 import database.request as rq
 from app.bybit.get_functions.get_instruments_info import get_instruments_info
@@ -373,5 +373,7 @@ class TelegramMessageHandler:
                         )
                         logger.info("Stop trading for symbol: %s, create_type: %s, stop_order_type: %s: %s",
                                     symbol, create_type, stop_order_type, tg_id)
+            else:
+                logger.info("Execution update: %s", json.dumps(message))
         except Exception as e:
             logger.error("Error in telegram_message_handler: %s", e, exc_info=True)
diff --git a/app/bybit/web_socket.py b/app/bybit/web_socket.py
index f69e85a..225cae5 100644
--- a/app/bybit/web_socket.py
+++ b/app/bybit/web_socket.py
@@ -11,6 +11,13 @@ logging.config.dictConfig(LOGGING_CONFIG)
 logger = logging.getLogger("web_socket")
 
 
+class CustomWebSocket(WebSocket):
+    def _on_error(self, error):
+        logger.error(f"WebSocket error: {error}")
+        # Здесь можно добавить собственную логику, например уведомления, метрики и т.д.
+        return super()._on_error(error)
+
+
 class WebSocketBot:
     """
     Class to handle WebSocket connections and messages.
@@ -43,9 +50,9 @@ class WebSocketBot:
                     continue
 
                 if tg_id in self.user_sockets:
-                    self.user_sockets.clear()
-                    self.user_messages.clear()
-                    self.user_keys.clear()
+                    self.user_sockets.pop(tg_id, None)
+                    self.user_messages.pop(tg_id, None)
+                    self.user_keys.pop(tg_id, None)
                     logger.info(
                         "Closed old websocket for user %s due to key change", tg_id
                     )
@@ -73,12 +80,13 @@ class WebSocketBot:
     async def try_connect_user(self, api_key, api_secret, tg_id):
         """Try to connect a user to the WebSocket."""
         try:
-            self.ws_private = WebSocket(
+            self.ws_private = CustomWebSocket(
                 demo=True,
                 testnet=False,
                 channel_type="private",
                 api_key=api_key,
                 api_secret=api_secret,
+                ping_interval=20,
             )
 
             self.user_sockets[tg_id] = self.ws_private
@@ -98,7 +106,7 @@ class WebSocketBot:
             return True
         except Exception as e:
             logger.error("Error connecting user %s: %s", tg_id, e)
-            return False
+            await asyncio.sleep(5)
 
     async def handle_order_update(self, message, tg_id):
         """Handle order updates."""