Separate and tested web_socket_manager

2024-11-03 22:26:12 +01:00
parent 4bd497e1ca
commit e243c6ac3b
11 changed files with 433 additions and 776 deletions
--- a/sia/web_socket_manager.py
+++ b/sia/web_socket_manager.py
@@ -0,0 +1,114 @@
+from aiohttp import web, WSMsgType
+from typing import Any, Dict, Set
+import asyncio
+import json
+from dataclasses import dataclass
+
+from .web_agent import WebAgent, WebAgentState
+from .web_io_buffer import WebIOBuffer
+
+@dataclass
+class ClientMessage:
+    """Messages that can be received from clients."""
+    APPROVE_CONTEXT = "APPROVE_CONTEXT"
+    APPROVE_RESPONSE = "APPROVE_RESPONSE" 
+    MODIFY_RESPONSE = "MODIFY_RESPONSE"
+    SEND_INPUT = "SEND_INPUT"
+
+@dataclass
+class ServerMessage:
+    """Messages that can be sent to clients."""
+    STATE_CHANGE = "STATE_CHANGE"
+    CONTEXT_UPDATE = "CONTEXT_UPDATE"
+    RESPONSE_UPDATE = "RESPONSE_UPDATE"
+    OUTPUT_UPDATE = "OUTPUT_UPDATE"
+
+class WebSocketManager:
+    """
+    Manages WebSocket connections and integrates with the WebAgent
+    """
+    def __init__(self,
+                agent: WebAgent,
+                io_buffer: WebIOBuffer):
+        """Initialize the web server."""
+        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)
+
+    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
+        
+    def _handle_state_change(self, new_state: WebAgentState):
+        """Handle state changes from the WebAgent."""
+        asyncio.run(self.broadcast_message({
+            "type": ServerMessage.STATE_CHANGE,
+            "state": new_state.name,
+        }))
+
+    def _handle_response_change(self, response: str, validation_error: str):
+        """Handle response changes from the WebAgent."""
+        asyncio.run(self.broadcast_message({
+            "type": ServerMessage.RESPONSE_UPDATE,
+            "response": response,
+            "validation_error": validation_error
+        }))
+
+    async def handle_websocket(self, request: web.Request) -> web.WebSocketResponse:
+        """Handle new WebSocket connections and messages."""
+        ws = web.WebSocketResponse(heartbeat=30)
+        await ws.prepare(request)
+
+        self._clients.add(ws)
+        
+        try:
+            await ws.send_json({
+                "type": ServerMessage.STATE_CHANGE,
+                "state": self.agent.current_state.name
+            })
+            await ws.send_json({
+                "type": ServerMessage.CONTEXT_UPDATE,
+                "context": self.agent.context
+            })
+            await ws.send_json({
+                "type": ServerMessage.RESPONSE_UPDATE,
+                "response": self.agent.response,
+                "validation_error": self.agent._validation_error
+            })
+            async for msg in ws:
+                if msg.type == WSMsgType.TEXT:
+                    #try:
+                    data = json.loads(msg.data)
+                    await self._handle_client_message(request, data)
+                    #except Exception as e:
+                    #    print(f"Error handling incoming websocket message: {msg}\n{e}")
+                elif msg.type == WSMsgType.ERROR:
+                    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"))