Basic web loop works

sia/web_socket_manager.py

@@ -27,18 +27,23 @@ class WebSocketManager:
     """
     Manages WebSocket connections and integrates with the WebAgent
     """
-    def __init__(self,
-                agent: WebAgent,
-                io_buffer: WebIOBuffer):
+    @classmethod
+    async def create(
+            cls,
+            agent: WebAgent,
+            io_buffer: WebIOBuffer):
         """Initialize the web server."""
+        self = cls()
         self.agent = agent
         self.io_buffer = io_buffer
         self._clients: Set[web.WebSocketResponse] = set()
 
-        self.agent.add_state_change_handler(self._handle_state_change)
-        self.agent.add_response_change_handler(self._handle_response_change)
+        self.agent.add_state_change_handler(self._wrap_async(self._handle_state_change))
+        self.agent.add_context_change_handler(self._wrap_async(self._handle_context_change))
+        self.agent.add_response_change_handler(self._wrap_async(self._handle_response_change))
+        return self
 
-    async def broadcast_message(self, message: Dict):
+    async def _broadcast_message(self, message: Dict):
         """Broadcast message to all connected clients."""
         disconnected = set()
         for ws in self._clients:
@@ -47,21 +52,63 @@ class WebSocketManager:
             except ConnectionResetError:
                 disconnected.add(ws)
         self._clients -= disconnected
+
+    def _wrap_async(self, coro_func):
+        """
+        Wraps an async callback to be safely called from another thread.
         
-    def _handle_state_change(self, new_state: WebAgentState):
+        Args:
+            coro_func: The async function to wrap
+            
+        Returns:
+            A sync function that schedules the coroutine on the event loop
+        """
+        loop = asyncio.get_event_loop()
+        def wrapper(*args, **kwargs):
+            loop.call_soon_threadsafe(lambda: asyncio.create_task(coro_func(*args, **kwargs)))
+        return wrapper
+        
+    async def _handle_state_change(self, new_state: WebAgentState):
         """Handle state changes from the WebAgent."""
-        asyncio.run(self.broadcast_message({
+        await self._broadcast_message({
             "type": ServerMessage.STATE_CHANGE,
             "state": new_state.name,
-        }))
+        })
+    
+    async def _handle_context_change(self, context: str):
+        """Handle context changes from the WebAgent."""
+        await self._broadcast_message({
+            "type": ServerMessage.CONTEXT_UPDATE,
+            "context": context
+        })
 
-    def _handle_response_change(self, response: str, validation_error: str):
+    async def _handle_response_change(self, response: str, validation_error: str):
         """Handle response changes from the WebAgent."""
-        asyncio.run(self.broadcast_message({
+        await self._broadcast_message({
             "type": ServerMessage.RESPONSE_UPDATE,
             "response": response,
             "validation_error": validation_error
-        }))
+        })
+
+    async def _handle_client_message(self, request: web.Request, data: dict[str, Any]):
+        """Handle incoming client message."""
+        message_type = data.get("type")
+        
+        if message_type == ClientMessage.APPROVE_CONTEXT:
+            if self.agent.state == WebAgentState.CONTEXT_APPROVAL:
+                if data.get("context") is not None:
+                    self.agent.modify_context(data.get("context"))
+                self.agent.approve_context()
+        elif message_type == ClientMessage.APPROVE_RESPONSE:
+            if self.agent.state == WebAgentState.RESPONSE_APPROVAL:
+                if data.get("response") is not None:
+                    self.agent.modify_response(data.get("response"))
+                self.agent.approve_response()
+        elif message_type == ClientMessage.MODIFY_RESPONSE:
+            if self.agent.state == WebAgentState.RESPONSE_APPROVAL:
+                self.agent.modify_response(data.get("response"))
+        elif message_type == ClientMessage.SEND_INPUT:
+                self.io_buffer.append_stdin(data.get("input"))
 
     async def handle_websocket(self, request: web.Request) -> web.WebSocketResponse:
         """Handle new WebSocket connections and messages."""
@@ -73,7 +120,7 @@ class WebSocketManager:
         try:
             await ws.send_json({
                 "type": ServerMessage.STATE_CHANGE,
-                "state": self.agent.current_state.name
+                "state": self.agent.state.name
             })
             await ws.send_json({
                 "type": ServerMessage.CONTEXT_UPDATE,
@@ -92,20 +139,4 @@ class WebSocketManager:
                     print(f"WebSocket connection closed with error: {ws.exception()}")
         finally:
             self._clients.remove(ws)
-        return ws
-
-    async def _handle_client_message(self, request: web.Request, data: dict[str, Any]):
-        """Handle incoming client message."""
-        message_type = data.get("type")
-        
-        if message_type == ClientMessage.APPROVE_CONTEXT:
-            if self.agent.current_state == WebAgentState.CONTEXT_APPROVAL:
-                await asyncio.to_thread(self.agent.approve_context)
-        elif message_type == ClientMessage.APPROVE_RESPONSE:
-            if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL:
-                await asyncio.to_thread(self.agent.approve_response)
-        elif message_type == ClientMessage.MODIFY_RESPONSE:
-            if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL:
-                await asyncio.to_thread(self.agent.set_response, data.get("response"))
-        elif message_type == ClientMessage.SEND_INPUT:
-                self.io_buffer.append_stdin(data.get("input"))
+        return ws