From 8206c475e59a9cfa89186cafb820648fb5fccf75 Mon Sep 17 00:00:00 2001
From: AEmotionStudio <163354043+AEmotionStudio@users.noreply.github.com>
Date: Tue, 20 Jan 2026 16:10:23 -0800
Subject: [PATCH] fix: make connect() return bool to verify connection success
 in reconnect loop

---
 bot/comfyui/websocket.py | 30 ++++++++++++++++++++----------
 1 file changed, 20 insertions(+), 10 deletions(-)

diff --git a/bot/comfyui/websocket.py b/bot/comfyui/websocket.py
index edcd58c..4247ecc 100644
--- a/bot/comfyui/websocket.py
+++ b/bot/comfyui/websocket.py
@@ -42,13 +42,17 @@ class ComfyUIWebSocket:
         # Lock to prevent race conditions between connect() and disconnect()
         self._state_lock = asyncio.Lock()
 
-    async def connect(self):
-        """Connect to the WebSocket."""
+    async def connect(self) -> bool:
+        """Connect to the WebSocket.
+        
+        Returns:
+            True if connection was established successfully, False if aborted.
+        """
         async with self._state_lock:
             # Check if disconnect was called - don't proceed if so
             if not self._should_reconnect:
                 logger.info("Connect aborted: disconnect was requested")
-                return
+                return False
                 
             if self.session is None or self.session.closed:
                 self.session = aiohttp.ClientSession()
@@ -62,11 +66,12 @@ class ComfyUIWebSocket:
                         await self.ws.close()
                     if self.session and not self.session.closed:
                         await self.session.close()
-                    return
+                    return False
                 self._running = True
                 self._reconnect_attempts = 0  # Reset on successful connection
                 self._listen_task = asyncio.create_task(self._listen())
                 logger.info(f"Connected to ComfyUI WebSocket at {self.ws_url}")
+                return True
             except Exception as e:
                 logger.error(f"Failed to connect to WebSocket: {e}")
                 if self.session and not self.session.closed:
@@ -183,12 +188,17 @@ class ComfyUIWebSocket:
 
             try:
                 attempts_made = self._reconnect_attempts
-                await self.connect()
-                logger.info(
-                    f"Successfully reconnected after "
-                    f"{attempts_made} attempt(s)."
-                )
-                return
+                connected = await self.connect()
+                if connected:
+                    logger.info(
+                        f"Successfully reconnected after "
+                        f"{attempts_made} attempt(s)."
+                    )
+                    return
+                else:
+                    # connect() returned False - disconnect was called
+                    logger.info("Reconnection aborted: disconnect was requested")
+                    return
             except Exception as e:
                 logger.warning(f"Reconnection attempt failed: {e}")