SIA/sia/web_server.py

import asyncio
import json
from typing import Set, Optional
from aiohttp import web, WSMsgType
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"
    VALIDATION_ERROR = "VALIDATION_ERROR"

class WebServer:
    """
    AIOHTTP-based web server that manages WebSocket connections and 
    integrates with the WebAgent.
    """
    
    def __init__(self, agent: WebAgent, io_buffer: WebIOBuffer, host: str = "localhost", port: int = 8080):
        """Initialize the web server."""
        self.agent = agent
        self.io_buffer = io_buffer
        self.host = host
        self.port = port
        self.app = web.Application()
        self._init_routes()
        self._clients: Set[web.WebSocketResponse] = set()
        self._lock = asyncio.Lock()
        self.runner: Optional[web.AppRunner] = None
        
        # Register agent handlers
        self.agent.add_state_change_handler(self._handle_state_change)
        self.agent.add_response_change_handler(self._handle_response_change)

    def _init_routes(self):
        """Initialize application routes."""
        self.app.router.add_get("/ws", self._handle_websocket)
        
    @property
    def clients(self) -> Set[web.WebSocketResponse]:
        """Get the set of connected clients."""
        return self._clients

    async def start(self):
        """Start the web server."""
        self.runner = web.AppRunner(self.app)
        await self.runner.setup()
        site = web.TCPSite(self.runner, self.host, self.port)
        await site.start()
        print(f"Web server started at ws://{self.host}:{self.port}/ws")

    async def stop(self):
        """Stop the web server and clean up resources."""
        if self.runner:
            await self.runner.cleanup()

    async def _send_to_client(self, ws: web.WebSocketResponse, message: dict) -> bool:
        """Send message to client, return True if successful."""
        try:
            await ws.send_json(message)
            return True
        except Exception:
            return False
            
    async def broadcast(self, message_type: str, data: str, **kwargs):
        """Broadcast a message to all connected clients."""
        async with self._lock:
            message = {
                "type": message_type,
                "data": data,
                **kwargs
            }
            
            tasks = []
            disconnected = set()
            
            for ws in self._clients:
                if not ws.closed:
                    task = asyncio.create_task(self._send_to_client(ws, message))
                    tasks.append(task)
                else:
                    disconnected.add(ws)
                    
            if tasks:
                results = await asyncio.gather(*tasks, return_exceptions=True)
                for ws, success in zip(self._clients - disconnected, results):
                    if not success:
                        disconnected.add(ws)
                    
            self._clients -= disconnected

    async def _handle_websocket(self, request: web.Request) -> web.WebSocketResponse:
        """Handle new WebSocket connections and messages."""
        ws = web.WebSocketResponse()
        await ws.prepare(request)
        
        async with self._lock:
            self._clients.add(ws)
        
        try:
            # Send initial state
            await self._send_to_client(ws, {
                "type": ServerMessage.STATE_CHANGE,
                "data": self.agent.current_state.name if self.agent.current_state else "INITIAL"
            })
            
            # Send current context if available
            if self.agent.context:
                await self._send_to_client(ws, {
                    "type": ServerMessage.CONTEXT_UPDATE,
                    "data": self.agent.context
                })
                
            # Send current response if available
            if self.agent.response:
                await self._send_to_client(ws, {
                    "type": ServerMessage.RESPONSE_UPDATE,
                    "data": self.agent.response,
                    "validation_error": self.agent._validation_error
                })
            
            # Handle incoming messages
            async for msg in ws:
                if msg.type == WSMsgType.TEXT:
                    try:
                        data = json.loads(msg.data)
                        await self._handle_client_message(data)
                                
                    except json.JSONDecodeError:
                        print(f"Invalid JSON message: {msg.data}")
                    except Exception as e:
                        print(f"Error handling message: {e}")
                        
                elif msg.type == WSMsgType.ERROR:
                    print(f"WebSocket connection closed with error: {ws.exception()}")
                    break
                    
        finally:
            async with self._lock:
                self._clients.remove(ws)
                
        return ws

    async def _handle_client_message(self, data: dict):
        """Handle incoming client message."""
        message_type = data.get("type")
        
        if message_type == ClientMessage.APPROVE_CONTEXT:
            if self.agent.current_state == WebAgentState.CONTEXT_APPROVAL:
                await self.agent.approve_context()  # Make sure this is awaited
                
        elif message_type == ClientMessage.APPROVE_RESPONSE:
            if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL:
                await self.agent.approve_response()
                
        elif message_type == ClientMessage.MODIFY_RESPONSE:
            if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL:
                new_response = data.get("data")
                if new_response:
                    await self.agent._set_response(new_response)
                    
        elif message_type == ClientMessage.SEND_INPUT:
            input_text = data.get("data")
            if input_text:
                self.io_buffer.append_stdin(input_text + "\n")
        
    async def _handle_state_change(self, new_state: WebAgentState):
        """Handle state changes from the WebAgent."""
        await self.broadcast(ServerMessage.STATE_CHANGE, new_state.name)
        
    async def _handle_response_change(self, response: str):
        """Handle response changes from the WebAgent."""
        await self.broadcast(
            ServerMessage.RESPONSE_UPDATE,
            response,
            validation_error=self.agent._validation_error
        )