111 lines
4.2 KiB
Python
111 lines
4.2 KiB
Python
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:
|
|
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
|
|
|
|
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")) |