reconnect on receive

Author: mishushakov

template/server/messaging.py

@@ -83,55 +83,63 @@ async def connect(self):
     async def reconnect(self, max_retries: int = 5, retry_delay: float = 0.1):
         """Reconnect the WebSocket if it's disconnected with retry logic."""
         logger.info(f"Attempting to reconnect WebSocket {self.context_id}")
-        
+
         # Close existing connection if any
         if self._ws is not None:
             try:
                 await self._ws.close()
             except Exception as e:
                 logger.warning(f"Error closing existing WebSocket: {e}")
-        
+
         # Cancel existing receive task if any
         if self._receive_task is not None and not self._receive_task.done():
             self._receive_task.cancel()
             try:
                 await self._receive_task
             except asyncio.CancelledError:
                 pass
-        
+
         # Reset WebSocket and task references
         self._ws = None
         self._receive_task = None
-        
+
         # Attempt to reconnect with fixed delay
         for attempt in range(max_retries):
             try:
                 await self.connect()
-                logger.info(f"Successfully reconnected WebSocket {self.context_id} on attempt {attempt + 1}")
+                logger.info(
+                    f"Successfully reconnected WebSocket {self.context_id} on attempt {attempt + 1}"
+                )
                 return True
             except Exception as e:
                 if attempt < max_retries - 1:
-                    logger.warning(f"Reconnection attempt {attempt + 1} failed: {e}. Retrying in {retry_delay}s...")
+                    logger.warning(
+                        f"Reconnection attempt {attempt + 1} failed: {e}. Retrying in {retry_delay}s..."
+                    )
                     await asyncio.sleep(retry_delay)
                 else:
-                    logger.error(f"Failed to reconnect WebSocket {self.context_id} after {max_retries} attempts: {e}")
+                    logger.error(
+                        f"Failed to reconnect WebSocket {self.context_id} after {max_retries} attempts: {e}"
+                    )
                     return False
-        
+
         return False
 
     def is_connected(self) -> bool:
         """Check if the WebSocket is connected and healthy."""
         return (
-            self._ws is not None 
-            and not self._ws.closed 
-            and self._receive_task is not None 
+            self._ws is not None
+            and not self._ws.closed
+            and self._receive_task is not None
             and not self._receive_task.done()
         )
 
     async def ensure_connected(self):
         """Ensure WebSocket is connected, reconnect if necessary."""
         if not self.is_connected():
-            logger.warning(f"WebSocket {self.context_id} is not connected, attempting to reconnect")
+            logger.warning(
+                f"WebSocket {self.context_id} is not connected, attempting to reconnect"
+            )
             success = await self.reconnect()
             if not success:
                 raise Exception(f"Failed to reconnect WebSocket {self.context_id}")
@@ -302,10 +310,10 @@ async def change_current_directory(
     ):
         message_id = str(uuid.uuid4())
         self._executions[message_id] = Execution(in_background=True)
-        
+
         # Ensure WebSocket is connected before changing directory
         await self.ensure_connected()
-        
+
         if language == "python":
             request = self._get_execute_request(message_id, f"%cd {path}", True)
         elif language == "deno":
@@ -328,7 +336,7 @@ async def change_current_directory(
 
         if self._ws is None:
             raise Exception("WebSocket not connected")
-        
+
         await self._ws.send(request)
 
         async for item in self._wait_for_result(message_id):
@@ -412,6 +420,20 @@ async def _receive_message(self):
                 await self._process_message(json.loads(message))
         except Exception as e:
             logger.error(f"WebSocket received error while receiving messages: {str(e)}")
+
+            # Attempt to reconnect when connection drops
+            logger.info("Attempting to reconnect due to connection loss...")
+            reconnect_success = await self.reconnect()
+
+            if reconnect_success:
+                logger.info("Successfully reconnected after connection loss")
+                # Continue receiving messages with the new connection
+                try:
+                    async for message in self._ws:
+                        await self._process_message(json.loads(message))
+                except Exception as reconnect_e:
+                    logger.error(f"Error in reconnected WebSocket: {str(reconnect_e)}")
+
             # Mark all pending executions as failed due to connection loss
             for execution in self._executions.values():
                 await execution.queue.put(