first commit

layers/presentation/websocket_handler.py

@@ -0,0 +1,97 @@
+"""
+WebSocket connection handler
+Presentation Layer - manages WebSocket connections and real-time communication
+"""
+import asyncio
+from typing import List
+from fastapi import WebSocket, WebSocketDisconnect
+import logging
+
+logger = logging.getLogger(__name__)
+
+class WebSocketManager:
+    """Manages WebSocket connections and broadcasting"""
+    
+    def __init__(self):
+        self.active_connections: List[WebSocket] = []
+    
+    async def connect(self, websocket: WebSocket) -> None:
+        """Accept and store new WebSocket connection"""
+        await websocket.accept()
+        self.active_connections.append(websocket)
+        logger.info(f"New client connected. Total clients: {len(self.active_connections)}")
+    
+    def disconnect(self, websocket: WebSocket) -> None:
+        """Remove WebSocket connection"""
+        if websocket in self.active_connections:
+            self.active_connections.remove(websocket)
+            logger.info(f"Client disconnected. Total clients: {len(self.active_connections)}")
+    
+    async def send_personal_message(self, message: str, websocket: WebSocket) -> None:
+        """Send message to specific WebSocket connection"""
+        try:
+            await websocket.send_text(message)
+        except Exception as e:
+            logger.error(f"Error sending personal message: {e}")
+            self.disconnect(websocket)
+    
+    async def broadcast(self, message: str) -> None:
+        """Broadcast message to all connected clients"""
+        if not self.active_connections:
+            return
+        
+        try:
+            # Send to all connections concurrently
+            tasks = [
+                self._safe_send_message(connection, message) 
+                for connection in self.active_connections.copy()
+            ]
+            
+            # Execute all sends concurrently and handle exceptions
+            results = await asyncio.gather(*tasks, return_exceptions=True)
+            
+            # Remove failed connections
+            failed_connections = []
+            for i, result in enumerate(results):
+                if isinstance(result, Exception):
+                    failed_connections.append(self.active_connections[i])
+            
+            for connection in failed_connections:
+                self.disconnect(connection)
+                
+        except Exception as e:
+            logger.error(f"Error in broadcast: {e}")
+    
+    async def _safe_send_message(self, websocket: WebSocket, message: str) -> None:
+        """Safely send message to WebSocket with error handling"""
+        try:
+            await websocket.send_text(message)
+        except WebSocketDisconnect:
+            # Connection was closed
+            raise
+        except Exception as e:
+            logger.error(f"Error sending message to client: {e}")
+            raise
+    
+    def get_connection_count(self) -> int:
+        """Get number of active connections"""
+        return len(self.active_connections)
+    
+    async def ping_all_connections(self) -> int:
+        """Ping all connections to check health, return number of healthy connections"""
+        if not self.active_connections:
+            return 0
+        
+        healthy_connections = []
+        for connection in self.active_connections.copy():
+            try:
+                await connection.ping()
+                healthy_connections.append(connection)
+            except Exception:
+                logger.debug("Removing unhealthy connection")
+        
+        self.active_connections = healthy_connections
+        return len(healthy_connections)
+
+# Global WebSocket manager instance
+websocket_manager = WebSocketManager()