SIA/sia/web_socket_manager.py

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"
    CLEAR_OUTPUT = "CLEAR_OUTPUT"

@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
    """
    @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._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))
        self.io_buffer.add_stdout_change_handler(self._wrap_async(self._handle_stdout_change))
        return self

    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 _wrap_async(self, coro_func):
        """
        Wraps an async callback to be safely called from another thread.
        
        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."""
        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
        })

    async def _handle_response_change(self, response: str, validation_error: str):
        """Handle response changes from the WebAgent."""
        await self._broadcast_message({
            "type": ServerMessage.RESPONSE_UPDATE,
            "response": response,
            "validation_error": validation_error
        })
    
    async def _handle_stdout_change(self, output: str):
        """Handle output changes from the WebAgent."""
        await self._broadcast_message({
            "type": ServerMessage.OUTPUT_UPDATE,
            "output": output
        })

    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"))
        elif message_type == ClientMessage.CLEAR_OUTPUT:
            self.io_buffer.clear_stdout()

    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.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:
                    data = json.loads(msg.data)
                    await self._handle_client_message(request, data)
                elif msg.type == WSMsgType.ERROR:
                    print(f"WebSocket connection closed with error: {ws.exception()}")
        finally:
            self._clients.remove(ws)
        return ws