empathy-framework 3.2.3__py3-none-any.whl → 3.8.2__py3-none-any.whl

hot_reload/websocket.py

@@ -0,0 +1,176 @@
+"""WebSocket notification system for hot-reload events.
+
+Notifies connected clients when wizards are reloaded.
+
+Copyright 2025 Smart AI Memory, LLC
+Licensed under Fair Source 0.9
+"""
+
+import asyncio
+import logging
+from typing import Any
+
+from fastapi import WebSocket
+
+logger = logging.getLogger(__name__)
+
+
+class ReloadNotificationManager:
+    """Manages WebSocket connections for hot-reload notifications.
+
+    Broadcasts reload events to all connected clients.
+    """
+
+    def __init__(self):
+        """Initialize notification manager."""
+        self._connections: list[WebSocket] = []
+        self._lock = asyncio.Lock()
+
+    async def connect(self, websocket: WebSocket) -> None:
+        """Register a new WebSocket connection.
+
+        Args:
+            websocket: WebSocket connection to register
+
+        """
+        await websocket.accept()
+
+        async with self._lock:
+            self._connections.append(websocket)
+
+        logger.info(f"WebSocket connected (total: {len(self._connections)})")
+
+        # Send welcome message
+        await self._send_to_client(
+            websocket,
+            {
+                "event": "connected",
+                "message": "Hot-reload notifications enabled",
+                "active_connections": len(self._connections),
+            },
+        )
+
+    async def disconnect(self, websocket: WebSocket) -> None:
+        """Unregister a WebSocket connection.
+
+        Args:
+            websocket: WebSocket connection to unregister
+
+        """
+        async with self._lock:
+            if websocket in self._connections:
+                self._connections.remove(websocket)
+
+        logger.info(f"WebSocket disconnected (remaining: {len(self._connections)})")
+
+    async def broadcast(self, message: dict[str, Any]) -> None:
+        """Broadcast message to all connected clients.
+
+        Args:
+            message: Message dictionary to broadcast
+
+        """
+        if not self._connections:
+            logger.debug("No active connections, skipping broadcast")
+            return
+
+        async with self._lock:
+            # Create a copy to avoid modification during iteration
+            connections = self._connections.copy()
+
+        # Broadcast to all connections
+        disconnected = []
+        for websocket in connections:
+            try:
+                await self._send_to_client(websocket, message)
+            except Exception as e:
+                logger.error(f"Error sending to client: {e}")
+                disconnected.append(websocket)
+
+        # Remove disconnected clients
+        if disconnected:
+            async with self._lock:
+                for websocket in disconnected:
+                    if websocket in self._connections:
+                        self._connections.remove(websocket)
+
+            logger.info(f"Removed {len(disconnected)} disconnected clients")
+
+    async def _send_to_client(
+        self,
+        websocket: WebSocket,
+        message: dict[str, Any],
+    ) -> None:
+        """Send message to a specific client.
+
+        Args:
+            websocket: WebSocket connection
+            message: Message to send
+
+        """
+        try:
+            await websocket.send_json(message)
+        except Exception as e:
+            logger.error(f"Failed to send message to client: {e}")
+            raise
+
+    def get_connection_count(self) -> int:
+        """Get number of active connections.
+
+        Returns:
+            Number of connected clients
+
+        """
+        return len(self._connections)
+
+
+# Global notification manager instance
+_notification_manager: ReloadNotificationManager | None = None
+
+
+def get_notification_manager() -> ReloadNotificationManager:
+    """Get the global notification manager.
+
+    Returns:
+        Global ReloadNotificationManager instance
+
+    """
+    global _notification_manager
+    if _notification_manager is None:
+        _notification_manager = ReloadNotificationManager()
+    return _notification_manager
+
+
+def create_notification_callback() -> callable:
+    """Create a callback function for the wizard reloader.
+
+    Returns:
+        Callback function that broadcasts notifications
+
+    """
+
+    def notification_callback(message: dict[str, Any]) -> None:
+        """Callback to broadcast reload notifications.
+
+        Args:
+            message: Notification message
+
+        """
+        manager = get_notification_manager()
+
+        # Run async broadcast in event loop
+        try:
+            loop = asyncio.get_event_loop()
+            if loop.is_running():
+                # Create a task if loop is running
+                asyncio.create_task(manager.broadcast(message))
+            else:
+                # Run directly if no loop running
+                asyncio.run(manager.broadcast(message))
+        except RuntimeError:
+            # No event loop, log warning
+            logger.warning("No event loop available for notification")
+        except Exception as e:
+            logger.error(f"Error in notification callback: {e}")
+
+    return notification_callback