Approve context

This commit is contained in:
2024-11-02 19:17:27 +01:00
parent 96d14fc46d
commit 2a15708145
9 changed files with 296 additions and 149 deletions

View File

@@ -5,6 +5,7 @@ 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
@@ -37,7 +38,7 @@ class WebServer:
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,
host: str = "0.0.0.0", port: int = 8080,
static_dir: Optional[Path] = None):
"""Initialize the web server."""
self.agent = agent
@@ -49,11 +50,22 @@ class WebServer:
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 agent handlers
self.agent.add_state_change_handler(self._handle_state_change)
self.agent.add_response_change_handler(self._handle_response_change)
# 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."""
@@ -61,16 +73,44 @@ class WebServer:
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)
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 to support
client-side routing.
"""
"""Serve the React application HTML for any unmatched routes."""
index_path = self.static_dir / "index.html"
if not index_path.exists():
raise web.HTTPNotFound()
@@ -129,7 +169,8 @@ class WebServer:
try:
await ws.send_json(message)
return True
except Exception:
except Exception as e:
print(f"Error sending to client: {e}")
return False
return False
@@ -216,27 +257,28 @@ class WebServer:
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()
await self.agent.approve_context()
elif message_type == ClientMessage.APPROVE_RESPONSE:
if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL:
await self.agent.approve_response()
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:
await self.agent._set_response(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."""
@@ -249,3 +291,27 @@ class WebServer:
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