Add context info in ui

sia/web/context_info_websocket.py

@@ -0,0 +1,48 @@
+from aiohttp import web, WSMsgType
+from typing import Dict, Set
+
+from .util import wrap_async
+from ..web_agent import WebAgent
+
+class ContextInfoWebSocket:
+    """
+    WebSocket handler for context info updates.
+    Broadcasts context info to all connected clients when it changes.
+    """
+
+    def __init__(self, agent: WebAgent):
+        self._agent = agent
+        self._clients: Set[web.WebSocketResponse] = set()
+        self._agent.add_context_change_handler(wrap_async(self._handle_context_change))
+
+    async def _broadcast_message(self, message: Dict):
+        """Broadcast message to all connected clients."""
+        disconnected = set()
+        for ws in self._clients:
+            try:
+                await ws.send_json(message)
+            except ConnectionResetError:
+                disconnected.add(ws)
+        self._clients -= disconnected
+
+    async def _handle_context_change(self, context_info: Dict):
+        """Handle context changes from the agent."""
+        await self._broadcast_message(context_info)
+
+    async def handle_connection(self, request: web.Request) -> web.WebSocketResponse:
+        """Handle new WebSocket connections."""
+        ws = web.WebSocketResponse(heartbeat=30)
+        await ws.prepare(request)
+
+        try:
+            # Send initial context info
+            await ws.send_json(self._agent.context_info)
+            self._clients.add(ws)
+
+            async for msg in ws:
+                if msg.type == WSMsgType.ERROR:
+                    print(f"WebSocket connection closed with error: {ws.exception()}")
+        finally:
+            self._clients.discard(ws)
+
+        return ws