317 lines
12 KiB
Python
317 lines
12 KiB
Python
from aiohttp import web, WSMsgType
|
|
from pathlib import Path
|
|
from typing import Set, Optional
|
|
import asyncio
|
|
import json
|
|
from dataclasses import dataclass
|
|
import mimetypes
|
|
from functools import partial
|
|
|
|
from .web_agent import WebAgent, WebAgentState
|
|
from .web_io_buffer import WebIOBuffer
|
|
|
|
mimetypes.add_type("application/javascript", ".js")
|
|
mimetypes.add_type("application/javascript", ".jsx")
|
|
mimetypes.add_type("text/javascript", ".js")
|
|
mimetypes.add_type("text/javascript", ".jsx")
|
|
|
|
@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. Serves static files from a web directory.
|
|
"""
|
|
def __init__(self, agent: WebAgent, io_buffer: WebIOBuffer,
|
|
host: str = "0.0.0.0", port: int = 8080,
|
|
static_dir: Optional[Path] = None):
|
|
"""Initialize the web server."""
|
|
self.agent = agent
|
|
self.io_buffer = io_buffer
|
|
self.host = host
|
|
self.port = port
|
|
self.app = web.Application()
|
|
self.static_dir = static_dir or Path(__file__).parent.parent / "static"
|
|
self._init_routes()
|
|
self._clients: Set[web.WebSocketResponse] = set()
|
|
self._lock = asyncio.Lock()
|
|
self._response_tasks: List[asyncio.Task] = [] # Initialize tasks list
|
|
self.runner: Optional[web.AppRunner] = None
|
|
|
|
# Register wrapped agent handlers
|
|
self.agent.add_state_change_handler(self._wrap_async_handler(self._handle_state_change))
|
|
self.agent.add_response_change_handler(self._wrap_async_handler(self._handle_response_change))
|
|
|
|
def _wrap_async_handler(self, coro_func):
|
|
"""Wrap an async handler function to be called synchronously."""
|
|
def wrapper(*args, **kwargs):
|
|
loop = asyncio.get_event_loop()
|
|
task = loop.create_task(coro_func(*args, **kwargs))
|
|
self._response_tasks.append(task)
|
|
self._response_tasks = [t for t in self._response_tasks if not t.done()]
|
|
return task
|
|
return wrapper
|
|
|
|
def _init_routes(self):
|
|
"""Initialize application routes."""
|
|
self.app.middlewares.append(self._cors_middleware)
|
|
self.app.router.add_get("/ws", self._handle_websocket)
|
|
|
|
if self.static_dir.exists():
|
|
self.app.router.add_get("/", self._serve_index)
|
|
self.app.router.add_static("/static/", self.static_dir, show_index=False)
|
|
self.app.router.add_static("/assets/", self.static_dir / "assets", show_index=False)
|
|
self.app.router.add_get("/{path:.*}", self._serve_index)
|
|
|
|
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:
|
|
self.agent.approve_context()
|
|
|
|
elif message_type == ClientMessage.APPROVE_RESPONSE:
|
|
if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL:
|
|
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:
|
|
self.agent._set_response(new_response)
|
|
|
|
elif message_type == ClientMessage.SEND_INPUT:
|
|
input_text = data.get("data")
|
|
if input_text:
|
|
if not input_text.endswith('\n'):
|
|
input_text += '\n'
|
|
|
|
self.io_buffer.append_stdin(input_text)
|
|
|
|
current_length = self.io_buffer.buffer_length()
|
|
return
|
|
else:
|
|
print(f"Agent in wrong state: {self.agent.current_state}")
|
|
|
|
async def _serve_index(self, request: web.Request) -> web.Response:
|
|
"""Serve the React application HTML for any unmatched routes."""
|
|
index_path = self.static_dir / "index.html"
|
|
if not index_path.exists():
|
|
raise web.HTTPNotFound()
|
|
|
|
with open(index_path, "r") as f:
|
|
html_content = f.read()
|
|
|
|
return web.Response(
|
|
text=html_content,
|
|
content_type="text/html"
|
|
)
|
|
|
|
@web.middleware
|
|
async def _cors_middleware(self, request: web.Request, handler):
|
|
"""Handle CORS headers."""
|
|
if request.method == "OPTIONS":
|
|
response = web.Response()
|
|
else:
|
|
response = await handler(request)
|
|
|
|
response.headers.update({
|
|
'Access-Control-Allow-Origin': '*',
|
|
'Access-Control-Allow-Methods': 'GET, POST, OPTIONS',
|
|
'Access-Control-Allow-Headers': 'Content-Type',
|
|
'Access-Control-Max-Age': '86400',
|
|
})
|
|
return response
|
|
|
|
@property
|
|
def clients(self) -> Set[web.WebSocketResponse]:
|
|
"""Get the set of connected clients."""
|
|
return self._clients
|
|
|
|
async def start(self):
|
|
"""Start the web server."""
|
|
# Verify static directory exists
|
|
if not self.static_dir.exists():
|
|
raise FileNotFoundError(f"Static directory not found: {self.static_dir}")
|
|
|
|
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 http://{self.host}:{self.port}")
|
|
print(f"Serving static files from: {self.static_dir}")
|
|
print(f"WebSocket endpoint 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."""
|
|
if not ws.closed:
|
|
try:
|
|
await ws.send_json(message)
|
|
return True
|
|
except Exception as e:
|
|
print(f"Error sending to client: {e}")
|
|
return False
|
|
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((ws, task))
|
|
else:
|
|
disconnected.add(ws)
|
|
|
|
if tasks:
|
|
for ws, task in tasks:
|
|
try:
|
|
success = await task
|
|
if not success:
|
|
disconnected.add(ws)
|
|
except Exception:
|
|
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(heartbeat=30)
|
|
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):
|
|
message_type = data.get("type")
|
|
|
|
if message_type == ClientMessage.APPROVE_CONTEXT:
|
|
if self.agent.current_state == WebAgentState.CONTEXT_APPROVAL:
|
|
await self.agent.approve_context()
|
|
elif message_type == ClientMessage.APPROVE_RESPONSE:
|
|
if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL:
|
|
self.agent.approve_response() # Now synchronous
|
|
|
|
elif message_type == ClientMessage.MODIFY_RESPONSE:
|
|
if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL:
|
|
new_response = data.get("data")
|
|
if new_response:
|
|
self.agent._set_response(new_response)
|
|
|
|
elif message_type == ClientMessage.SEND_INPUT:
|
|
input_text = data.get("data")
|
|
if input_text:
|
|
self.agent.set_input(input_text)
|
|
self.io_buffer.append_stdin(input_text + "\n")
|
|
else:
|
|
print(f"Agent in wrong state: {self.agent.current_state}")
|
|
|
|
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
|
|
)
|
|
|
|
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:
|
|
try:
|
|
await self._send_to_client(ws, message)
|
|
except Exception as e:
|
|
print(f"Error sending to client: {e}")
|
|
disconnected.add(ws)
|
|
else:
|
|
disconnected.add(ws)
|
|
|
|
self._clients -= disconnected |