Add auto_approver
This commit is contained in:
@@ -122,5 +122,5 @@ if __name__ == "__main__":
|
|||||||
loop = asyncio.new_event_loop()
|
loop = asyncio.new_event_loop()
|
||||||
config = Config()
|
config = Config()
|
||||||
main = loop.run_until_complete(Main.create(config))
|
main = loop.run_until_complete(Main.create(config))
|
||||||
print(f"Web server started at http://{config.host}:{config.port}")
|
print(f"Web server started at http://localhost:{config.port}")
|
||||||
web.run_app(main.app, loop=loop, host=config.host, port=config.port)
|
web.run_app(main.app, loop=loop, host=config.host, port=config.port)
|
||||||
162
sia/auto_approver.py
Normal file
162
sia/auto_approver.py
Normal file
@@ -0,0 +1,162 @@
|
|||||||
|
from threading import Thread, Event
|
||||||
|
from typing import Callable, TypeAlias, TypedDict
|
||||||
|
from .web_agent import WebAgent, WebAgentState
|
||||||
|
|
||||||
|
class AutoApproverConfig(TypedDict):
|
||||||
|
context_enabled: bool
|
||||||
|
response_enabled: bool
|
||||||
|
context_timeout: float
|
||||||
|
response_timeout: float
|
||||||
|
|
||||||
|
ConfigChangeHandler: TypeAlias = Callable[[AutoApproverConfig], None]
|
||||||
|
|
||||||
|
class AutoApprover:
|
||||||
|
"""
|
||||||
|
Handles automatic approvals for WebAgent states after configurable timeouts.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, agent: WebAgent):
|
||||||
|
"""
|
||||||
|
Initialize auto approver with a WebAgent instance.
|
||||||
|
"""
|
||||||
|
self.agent = agent
|
||||||
|
self._context_timeout = 5.0
|
||||||
|
self._response_timeout = 10.0
|
||||||
|
self._context_enabled = False
|
||||||
|
self._response_enabled = False
|
||||||
|
self._stop_event = Event()
|
||||||
|
self._context_thread: Thread | None = None
|
||||||
|
self._response_thread: Thread | None = None
|
||||||
|
self._config_change_handlers: list[ConfigChangeHandler] = []
|
||||||
|
|
||||||
|
self.agent.add_state_change_handler(self._handle_state_change)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def config(self) -> AutoApproverConfig:
|
||||||
|
return AutoApproverConfig(
|
||||||
|
context_enabled=self._context_enabled,
|
||||||
|
response_enabled=self._response_enabled,
|
||||||
|
context_timeout=self._context_timeout,
|
||||||
|
response_timeout=self._response_timeout
|
||||||
|
)
|
||||||
|
|
||||||
|
def set_config(self, config: AutoApproverConfig) -> None:
|
||||||
|
self._context_enabled = config['context_enabled']
|
||||||
|
self._response_enabled = config['response_enabled']
|
||||||
|
self._context_timeout = config['context_timeout']
|
||||||
|
self._response_timeout = config['response_timeout']
|
||||||
|
|
||||||
|
if self._context_enabled and self.agent.state == WebAgentState.CONTEXT_APPROVAL:
|
||||||
|
self._start_context_thread()
|
||||||
|
else:
|
||||||
|
self._stop_context_thread()
|
||||||
|
|
||||||
|
if self._response_enabled and self.agent.state == WebAgentState.RESPONSE_APPROVAL:
|
||||||
|
self._start_response_thread()
|
||||||
|
else:
|
||||||
|
self._stop_response_thread()
|
||||||
|
|
||||||
|
self._notify_config_change()
|
||||||
|
|
||||||
|
def add_config_change_handler(self, handler: ConfigChangeHandler) -> None:
|
||||||
|
self._config_change_handlers.append(handler)
|
||||||
|
|
||||||
|
def _notify_config_change(self) -> None:
|
||||||
|
current_config = self.config
|
||||||
|
for handler in self._config_change_handlers:
|
||||||
|
handler(current_config)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def context_timeout(self) -> float:
|
||||||
|
return self._context_timeout
|
||||||
|
|
||||||
|
@context_timeout.setter
|
||||||
|
def context_timeout(self, timeout: float) -> None:
|
||||||
|
if self._context_enabled:
|
||||||
|
raise ValueError("Cannot change timeout while auto-approval is enabled")
|
||||||
|
self._context_timeout = timeout
|
||||||
|
self._notify_config_change()
|
||||||
|
|
||||||
|
@property
|
||||||
|
def response_timeout(self) -> float:
|
||||||
|
return self._response_timeout
|
||||||
|
|
||||||
|
@response_timeout.setter
|
||||||
|
def response_timeout(self, timeout: float) -> None:
|
||||||
|
if self._response_enabled:
|
||||||
|
raise ValueError("Cannot change timeout while auto-approval is enabled")
|
||||||
|
self._response_timeout = timeout
|
||||||
|
self._notify_config_change()
|
||||||
|
|
||||||
|
@property
|
||||||
|
def context_enabled(self) -> bool:
|
||||||
|
return self._context_enabled
|
||||||
|
|
||||||
|
@context_enabled.setter
|
||||||
|
def context_enabled(self, enabled: bool) -> None:
|
||||||
|
if enabled == self._context_enabled:
|
||||||
|
return
|
||||||
|
self._context_enabled = enabled
|
||||||
|
self._stop_context_thread()
|
||||||
|
if enabled and self.agent.state == WebAgentState.CONTEXT_APPROVAL:
|
||||||
|
self._start_context_thread()
|
||||||
|
self._notify_config_change()
|
||||||
|
|
||||||
|
@property
|
||||||
|
def response_enabled(self) -> bool:
|
||||||
|
return self._response_enabled
|
||||||
|
|
||||||
|
@response_enabled.setter
|
||||||
|
def response_enabled(self, enabled: bool) -> None:
|
||||||
|
if enabled == self._response_enabled:
|
||||||
|
return
|
||||||
|
self._response_enabled = enabled
|
||||||
|
self._stop_response_thread()
|
||||||
|
if enabled and self.agent.state == WebAgentState.RESPONSE_APPROVAL:
|
||||||
|
self._start_response_thread()
|
||||||
|
self._notify_config_change()
|
||||||
|
|
||||||
|
def _handle_state_change(self, new_state: WebAgentState) -> None:
|
||||||
|
if new_state == WebAgentState.CONTEXT_APPROVAL:
|
||||||
|
self._start_context_thread()
|
||||||
|
elif self._context_enabled:
|
||||||
|
self._stop_context_thread()
|
||||||
|
|
||||||
|
if new_state == WebAgentState.RESPONSE_APPROVAL:
|
||||||
|
self._start_response_thread()
|
||||||
|
elif self._response_enabled:
|
||||||
|
self._stop_response_thread()
|
||||||
|
|
||||||
|
def _stop_context_thread(self) -> None:
|
||||||
|
if self._context_thread:
|
||||||
|
self._stop_event.set()
|
||||||
|
self._context_thread = None
|
||||||
|
self._stop_event.clear()
|
||||||
|
|
||||||
|
def _stop_response_thread(self) -> None:
|
||||||
|
if self._response_thread:
|
||||||
|
self._stop_event.set()
|
||||||
|
self._response_thread = None
|
||||||
|
self._stop_event.clear()
|
||||||
|
|
||||||
|
def _start_context_thread(self) -> None:
|
||||||
|
self._context_thread = Thread(target=self._context_approval_thread)
|
||||||
|
self._context_thread.start()
|
||||||
|
|
||||||
|
def _start_response_thread(self) -> None:
|
||||||
|
self._response_thread = Thread(target=self._response_approval_thread)
|
||||||
|
self._response_thread.start()
|
||||||
|
|
||||||
|
def _context_approval_thread(self) -> None:
|
||||||
|
if self._stop_event.wait(self._context_timeout):
|
||||||
|
return
|
||||||
|
if (self._context_enabled and
|
||||||
|
self.agent.state == WebAgentState.CONTEXT_APPROVAL):
|
||||||
|
self.agent.approve_context()
|
||||||
|
|
||||||
|
def _response_approval_thread(self) -> None:
|
||||||
|
if self._stop_event.wait(self._response_timeout):
|
||||||
|
return
|
||||||
|
if (self._response_enabled and
|
||||||
|
self.agent.state == WebAgentState.RESPONSE_APPROVAL):
|
||||||
|
self.agent.approve_response()
|
||||||
@@ -6,6 +6,7 @@ from dataclasses import dataclass
|
|||||||
|
|
||||||
from .web_agent import WebAgent, WebAgentState
|
from .web_agent import WebAgent, WebAgentState
|
||||||
from .web_io_buffer import WebIOBuffer
|
from .web_io_buffer import WebIOBuffer
|
||||||
|
from .auto_approver import AutoApprover, AutoApproverConfig
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class ClientMessage:
|
class ClientMessage:
|
||||||
@@ -15,6 +16,7 @@ class ClientMessage:
|
|||||||
MODIFY_RESPONSE = "MODIFY_RESPONSE"
|
MODIFY_RESPONSE = "MODIFY_RESPONSE"
|
||||||
SEND_INPUT = "SEND_INPUT"
|
SEND_INPUT = "SEND_INPUT"
|
||||||
CLEAR_OUTPUT = "CLEAR_OUTPUT"
|
CLEAR_OUTPUT = "CLEAR_OUTPUT"
|
||||||
|
AUTO_APPROVER_CONFIG = "AUTO_APPROVER_CONFIG"
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class ServerMessage:
|
class ServerMessage:
|
||||||
@@ -23,6 +25,7 @@ class ServerMessage:
|
|||||||
CONTEXT_UPDATE = "CONTEXT_UPDATE"
|
CONTEXT_UPDATE = "CONTEXT_UPDATE"
|
||||||
RESPONSE_UPDATE = "RESPONSE_UPDATE"
|
RESPONSE_UPDATE = "RESPONSE_UPDATE"
|
||||||
OUTPUT_UPDATE = "OUTPUT_UPDATE"
|
OUTPUT_UPDATE = "OUTPUT_UPDATE"
|
||||||
|
AUTO_APPROVER_CONFIG = "AUTO_APPROVER_CONFIG"
|
||||||
|
|
||||||
class WebSocketManager:
|
class WebSocketManager:
|
||||||
"""
|
"""
|
||||||
@@ -38,11 +41,13 @@ class WebSocketManager:
|
|||||||
self.agent = agent
|
self.agent = agent
|
||||||
self.io_buffer = io_buffer
|
self.io_buffer = io_buffer
|
||||||
self._clients: Set[web.WebSocketResponse] = set()
|
self._clients: Set[web.WebSocketResponse] = set()
|
||||||
|
self.auto_approver = AutoApprover(agent)
|
||||||
|
|
||||||
self.agent.add_state_change_handler(self._wrap_async(self._handle_state_change))
|
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_context_change_handler(self._wrap_async(self._handle_context_change))
|
||||||
self.agent.add_response_change_handler(self._wrap_async(self._handle_response_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))
|
self.io_buffer.add_stdout_change_handler(self._wrap_async(self._handle_stdout_change))
|
||||||
|
self.auto_approver.add_config_change_handler(self._wrap_async(self._handle_auto_approver_config))
|
||||||
return self
|
return self
|
||||||
|
|
||||||
async def _broadcast_message(self, message: Dict):
|
async def _broadcast_message(self, message: Dict):
|
||||||
@@ -99,6 +104,13 @@ class WebSocketManager:
|
|||||||
"output": output
|
"output": output
|
||||||
})
|
})
|
||||||
|
|
||||||
|
async def _handle_auto_approver_config(self, config: AutoApproverConfig):
|
||||||
|
"""Handle auto approver config changes."""
|
||||||
|
await self._broadcast_message({
|
||||||
|
"type": ServerMessage.AUTO_APPROVER_CONFIG,
|
||||||
|
"config": config
|
||||||
|
})
|
||||||
|
|
||||||
async def _handle_client_message(self, request: web.Request, data: dict[str, Any]):
|
async def _handle_client_message(self, request: web.Request, data: dict[str, Any]):
|
||||||
"""Handle incoming client message."""
|
"""Handle incoming client message."""
|
||||||
message_type = data.get("type")
|
message_type = data.get("type")
|
||||||
@@ -117,9 +129,13 @@ class WebSocketManager:
|
|||||||
if self.agent.state == WebAgentState.RESPONSE_APPROVAL:
|
if self.agent.state == WebAgentState.RESPONSE_APPROVAL:
|
||||||
self.agent.modify_response(data.get("response"))
|
self.agent.modify_response(data.get("response"))
|
||||||
elif message_type == ClientMessage.SEND_INPUT:
|
elif message_type == ClientMessage.SEND_INPUT:
|
||||||
self.io_buffer.append_stdin(data.get("input"))
|
self.io_buffer.append_stdin(data.get("input"))
|
||||||
elif message_type == ClientMessage.CLEAR_OUTPUT:
|
elif message_type == ClientMessage.CLEAR_OUTPUT:
|
||||||
self.io_buffer.clear_stdout()
|
self.io_buffer.clear_stdout()
|
||||||
|
elif message_type == ClientMessage.AUTO_APPROVER_CONFIG:
|
||||||
|
config = data.get("config")
|
||||||
|
if isinstance(config, dict):
|
||||||
|
self.auto_approver.set_config(config)
|
||||||
|
|
||||||
async def handle_websocket(self, request: web.Request) -> web.WebSocketResponse:
|
async def handle_websocket(self, request: web.Request) -> web.WebSocketResponse:
|
||||||
"""Handle new WebSocket connections and messages."""
|
"""Handle new WebSocket connections and messages."""
|
||||||
@@ -142,6 +158,10 @@ class WebSocketManager:
|
|||||||
"response": self.agent.response,
|
"response": self.agent.response,
|
||||||
"validation_error": self.agent._validation_error
|
"validation_error": self.agent._validation_error
|
||||||
})
|
})
|
||||||
|
await ws.send_json({
|
||||||
|
"type": ServerMessage.AUTO_APPROVER_CONFIG,
|
||||||
|
"config": self.auto_approver.config
|
||||||
|
})
|
||||||
async for msg in ws:
|
async for msg in ws:
|
||||||
if msg.type == WSMsgType.TEXT:
|
if msg.type == WSMsgType.TEXT:
|
||||||
data = json.loads(msg.data)
|
data = json.loads(msg.data)
|
||||||
@@ -150,4 +170,4 @@ class WebSocketManager:
|
|||||||
print(f"WebSocket connection closed with error: {ws.exception()}")
|
print(f"WebSocket connection closed with error: {ws.exception()}")
|
||||||
finally:
|
finally:
|
||||||
self._clients.remove(ws)
|
self._clients.remove(ws)
|
||||||
return ws
|
return ws
|
||||||
|
|||||||
87
test/auto_approver_test.py
Normal file
87
test/auto_approver_test.py
Normal file
@@ -0,0 +1,87 @@
|
|||||||
|
import unittest
|
||||||
|
from unittest.mock import Mock, patch
|
||||||
|
from threading import Event
|
||||||
|
import time
|
||||||
|
|
||||||
|
from sia.auto_approver import AutoApprover
|
||||||
|
from sia.web_agent import WebAgent, WebAgentState
|
||||||
|
|
||||||
|
class AutoApproverTest(unittest.TestCase):
|
||||||
|
def setUp(self):
|
||||||
|
"""Create mock agent and approver for each test"""
|
||||||
|
self.mock_agent = Mock(spec=WebAgent)
|
||||||
|
self.mock_agent.state = WebAgentState.UPDATE
|
||||||
|
self.approver = AutoApprover(self.mock_agent)
|
||||||
|
|
||||||
|
def test_timeout_setters(self):
|
||||||
|
"""Test timeout setters block changes when enabled"""
|
||||||
|
self.approver.context_timeout = 10.0
|
||||||
|
self.approver.response_timeout = 15.0
|
||||||
|
|
||||||
|
self.approver.context_enabled = True
|
||||||
|
with self.assertRaises(ValueError):
|
||||||
|
self.approver.context_timeout = 20.0
|
||||||
|
|
||||||
|
self.approver.response_enabled = True
|
||||||
|
with self.assertRaises(ValueError):
|
||||||
|
self.approver.response_timeout = 25.0
|
||||||
|
|
||||||
|
def test_enable_starts_thread_in_correct_state(self):
|
||||||
|
"""Test enabling only starts thread in matching state"""
|
||||||
|
self.mock_agent.state = WebAgentState.CONTEXT_APPROVAL
|
||||||
|
self.approver.context_enabled = True
|
||||||
|
time.sleep(0.1) # Allow thread to start
|
||||||
|
self.assertTrue(self.approver._context_thread.is_alive())
|
||||||
|
self.assertIsNone(self.approver._response_thread)
|
||||||
|
|
||||||
|
def test_state_change_stops_threads(self):
|
||||||
|
"""Test state changes stop irrelevant threads"""
|
||||||
|
# Start in context approval with thread
|
||||||
|
self.mock_agent.state = WebAgentState.CONTEXT_APPROVAL
|
||||||
|
self.approver.context_enabled = True
|
||||||
|
time.sleep(0.1)
|
||||||
|
self.assertTrue(self.approver._context_thread.is_alive())
|
||||||
|
|
||||||
|
# Change state and verify thread stops
|
||||||
|
self.mock_agent.state = WebAgentState.INFERENCE
|
||||||
|
self.approver._handle_state_change(WebAgentState.INFERENCE)
|
||||||
|
time.sleep(0.1)
|
||||||
|
self.assertIsNone(self.approver._context_thread)
|
||||||
|
|
||||||
|
def test_quick_state_cycle(self):
|
||||||
|
"""Test thread stops when state changes before timeout"""
|
||||||
|
self.approver.context_timeout = 5.0
|
||||||
|
self.mock_agent.state = WebAgentState.CONTEXT_APPROVAL
|
||||||
|
self.approver.context_enabled = True
|
||||||
|
|
||||||
|
# Change state quickly
|
||||||
|
time.sleep(0.1)
|
||||||
|
self.mock_agent.state = WebAgentState.INFERENCE
|
||||||
|
self.approver._handle_state_change(WebAgentState.INFERENCE)
|
||||||
|
|
||||||
|
# Verify no approval happened
|
||||||
|
self.mock_agent.approve_context.assert_not_called()
|
||||||
|
|
||||||
|
def test_disable_stops_thread(self):
|
||||||
|
"""Test disabling stops active thread"""
|
||||||
|
self.mock_agent.state = WebAgentState.CONTEXT_APPROVAL
|
||||||
|
self.approver.context_enabled = True
|
||||||
|
time.sleep(0.1)
|
||||||
|
self.assertTrue(self.approver._context_thread.is_alive())
|
||||||
|
|
||||||
|
self.approver.context_enabled = False
|
||||||
|
self.assertIsNone(self.approver._context_thread)
|
||||||
|
|
||||||
|
def test_successful_auto_approval(self):
|
||||||
|
"""Test thread approves after timeout when conditions met"""
|
||||||
|
self.approver.context_timeout = 0.1
|
||||||
|
self.mock_agent.state = WebAgentState.CONTEXT_APPROVAL
|
||||||
|
self.approver.context_enabled = True
|
||||||
|
|
||||||
|
time.sleep(0.2)
|
||||||
|
self.mock_agent.approve_context.assert_called_once()
|
||||||
|
|
||||||
|
def tearDown(self):
|
||||||
|
"""Clean up any running threads"""
|
||||||
|
self.approver.context_enabled = False
|
||||||
|
self.approver.response_enabled = False
|
||||||
@@ -68,6 +68,11 @@ class WebSocketManagerTest(AioHTTPTestCase):
|
|||||||
app.router.add_get('/ws', self.ws_manager.handle_websocket)
|
app.router.add_get('/ws', self.ws_manager.handle_websocket)
|
||||||
return app
|
return app
|
||||||
|
|
||||||
|
async def _skip_initial_messages(self, ws):
|
||||||
|
"""Helper to skip past the initial connection messages"""
|
||||||
|
for _ in range(4): # Now 4 messages including auto approver config
|
||||||
|
await ws.receive_json()
|
||||||
|
|
||||||
@unittest_run_loop
|
@unittest_run_loop
|
||||||
async def test_initial_connection(self):
|
async def test_initial_connection(self):
|
||||||
"""Test initial WebSocket connection and state messages"""
|
"""Test initial WebSocket connection and state messages"""
|
||||||
@@ -88,25 +93,34 @@ class WebSocketManagerTest(AioHTTPTestCase):
|
|||||||
self.assertEqual(msg['response'], '')
|
self.assertEqual(msg['response'], '')
|
||||||
self.assertIsNone(msg['validation_error'])
|
self.assertIsNone(msg['validation_error'])
|
||||||
|
|
||||||
|
# Verify initial auto approver config
|
||||||
|
msg = await ws.receive_json()
|
||||||
|
self.assertEqual(msg['type'], ServerMessage.AUTO_APPROVER_CONFIG)
|
||||||
|
self.assertIn('config', msg)
|
||||||
|
|
||||||
@unittest_run_loop
|
@unittest_run_loop
|
||||||
async def test_approve_context(self):
|
async def test_approve_context(self):
|
||||||
"""Test context approval flow"""
|
"""Test context approval flow"""
|
||||||
async with self.client.ws_connect('/ws') as ws:
|
async with self.client.ws_connect('/ws') as ws:
|
||||||
for _ in range(3):
|
await self._skip_initial_messages(ws)
|
||||||
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
|
|
||||||
await ws.send_json({
|
await ws.send_json({
|
||||||
'type': ClientMessage.APPROVE_CONTEXT
|
'type': ClientMessage.APPROVE_CONTEXT
|
||||||
})
|
})
|
||||||
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
|
|
||||||
|
msg = await ws.receive_json()
|
||||||
self.assertEqual(msg['type'], ServerMessage.STATE_CHANGE)
|
self.assertEqual(msg['type'], ServerMessage.STATE_CHANGE)
|
||||||
self.assertEqual(msg['state'], WebAgentState.INFERENCE.name)
|
self.assertEqual(msg['state'], WebAgentState.INFERENCE.name)
|
||||||
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
|
|
||||||
|
msg = await ws.receive_json()
|
||||||
self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE)
|
self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE)
|
||||||
self.assertEqual(msg['response'], "")
|
self.assertEqual(msg['response'], "")
|
||||||
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
|
|
||||||
|
msg = await ws.receive_json()
|
||||||
self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE)
|
self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE)
|
||||||
self.assertEqual(msg['response'], "<reasoning>test reasoning</reasoning>")
|
self.assertEqual(msg['response'], "<reasoning>test reasoning</reasoning>")
|
||||||
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
|
|
||||||
|
msg = await ws.receive_json()
|
||||||
self.assertEqual(msg['type'], ServerMessage.STATE_CHANGE)
|
self.assertEqual(msg['type'], ServerMessage.STATE_CHANGE)
|
||||||
self.assertEqual(msg['state'], WebAgentState.RESPONSE_APPROVAL.name)
|
self.assertEqual(msg['state'], WebAgentState.RESPONSE_APPROVAL.name)
|
||||||
|
|
||||||
@@ -116,53 +130,51 @@ class WebSocketManagerTest(AioHTTPTestCase):
|
|||||||
test_input = "test input"
|
test_input = "test input"
|
||||||
|
|
||||||
async with self.client.ws_connect('/ws') as ws:
|
async with self.client.ws_connect('/ws') as ws:
|
||||||
for _ in range(3):
|
await self._skip_initial_messages(ws)
|
||||||
await ws.receive_json()
|
|
||||||
|
|
||||||
# Send input
|
|
||||||
await ws.send_json({
|
await ws.send_json({
|
||||||
'type': ClientMessage.SEND_INPUT,
|
'type': ClientMessage.SEND_INPUT,
|
||||||
'input': test_input
|
'input': test_input
|
||||||
})
|
})
|
||||||
|
|
||||||
# Delay for processing
|
|
||||||
await asyncio.sleep(0.5)
|
await asyncio.sleep(0.5)
|
||||||
|
|
||||||
# Verify input was added to buffer
|
|
||||||
self.assertEqual(self.io_buffer.read(), test_input)
|
self.assertEqual(self.io_buffer.read(), test_input)
|
||||||
|
|
||||||
@unittest_run_loop
|
@unittest_run_loop
|
||||||
async def test_modify_response(self):
|
async def test_modify_response(self):
|
||||||
"""Test modifying response in RESPONSE_APPROVAL state"""
|
"""Test modifying response in RESPONSE_APPROVAL state"""
|
||||||
# Set agent to RESPONSE_APPROVAL state
|
|
||||||
async with self.client.ws_connect('/ws') as ws:
|
async with self.client.ws_connect('/ws') as ws:
|
||||||
for _ in range(3):
|
await self._skip_initial_messages(ws)
|
||||||
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
|
|
||||||
await ws.send_json({
|
await ws.send_json({
|
||||||
'type': ClientMessage.APPROVE_CONTEXT
|
'type': ClientMessage.APPROVE_CONTEXT
|
||||||
})
|
})
|
||||||
|
|
||||||
|
# Skip to RESPONSE_APPROVAL state
|
||||||
for _ in range(4):
|
for _ in range(4):
|
||||||
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
|
await ws.receive_json()
|
||||||
|
|
||||||
modified_response = "<reasoning>modified test</reasoning>"
|
modified_response = "<reasoning>modified test</reasoning>"
|
||||||
await ws.send_json({
|
await ws.send_json({
|
||||||
'type': ClientMessage.MODIFY_RESPONSE,
|
'type': ClientMessage.MODIFY_RESPONSE,
|
||||||
'response': modified_response
|
'response': modified_response
|
||||||
})
|
})
|
||||||
msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5)
|
|
||||||
|
msg = await ws.receive_json()
|
||||||
self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE)
|
self.assertEqual(msg['type'], ServerMessage.RESPONSE_UPDATE)
|
||||||
self.assertEqual(msg['response'], modified_response)
|
self.assertEqual(msg['response'], modified_response)
|
||||||
|
|
||||||
@unittest_run_loop
|
@unittest_run_loop
|
||||||
async def test_multiple_clients(self):
|
async def test_multiple_clients(self):
|
||||||
"""Test broadcasting to multiple WebSocket clients"""
|
"""Test broadcasting to multiple WebSocket clients"""
|
||||||
# Connect first client
|
|
||||||
async with self.client.ws_connect('/ws') as ws1, self.client.ws_connect('/ws') as ws2:
|
async with self.client.ws_connect('/ws') as ws1, self.client.ws_connect('/ws') as ws2:
|
||||||
for _ in range(3):
|
await self._skip_initial_messages(ws1)
|
||||||
await asyncio.wait_for(ws1.receive_json(), timeout=0.5)
|
await self._skip_initial_messages(ws2)
|
||||||
await asyncio.wait_for(ws2.receive_json(), timeout=0.5)
|
|
||||||
await ws1.send_json({
|
await ws1.send_json({
|
||||||
'type': ClientMessage.APPROVE_CONTEXT
|
'type': ClientMessage.APPROVE_CONTEXT
|
||||||
})
|
})
|
||||||
|
|
||||||
msg1 = await ws1.receive_json()
|
msg1 = await ws1.receive_json()
|
||||||
msg2 = await ws2.receive_json()
|
msg2 = await ws2.receive_json()
|
||||||
self.assertEqual(msg1['type'], ServerMessage.STATE_CHANGE)
|
self.assertEqual(msg1['type'], ServerMessage.STATE_CHANGE)
|
||||||
@@ -170,6 +182,60 @@ class WebSocketManagerTest(AioHTTPTestCase):
|
|||||||
self.assertEqual(msg2['type'], ServerMessage.STATE_CHANGE)
|
self.assertEqual(msg2['type'], ServerMessage.STATE_CHANGE)
|
||||||
self.assertEqual(msg2['state'], WebAgentState.INFERENCE.name)
|
self.assertEqual(msg2['state'], WebAgentState.INFERENCE.name)
|
||||||
|
|
||||||
|
@unittest_run_loop
|
||||||
|
async def test_auto_approver_config(self):
|
||||||
|
"""Test setting and receiving auto approver config"""
|
||||||
|
async with self.client.ws_connect('/ws') as ws:
|
||||||
|
await self._skip_initial_messages(ws)
|
||||||
|
|
||||||
|
test_config = {
|
||||||
|
'context_enabled': True,
|
||||||
|
'response_enabled': True,
|
||||||
|
'context_timeout': 10.0,
|
||||||
|
'response_timeout': 20.0
|
||||||
|
}
|
||||||
|
|
||||||
|
await ws.send_json({
|
||||||
|
'type': ClientMessage.AUTO_APPROVER_CONFIG,
|
||||||
|
'config': test_config
|
||||||
|
})
|
||||||
|
|
||||||
|
msg = await ws.receive_json()
|
||||||
|
self.assertEqual(msg['type'], ServerMessage.AUTO_APPROVER_CONFIG)
|
||||||
|
self.assertEqual(msg['config'], test_config)
|
||||||
|
|
||||||
|
self.ws_manager.auto_approver.context_enabled = False
|
||||||
|
self.ws_manager.auto_approver.response_enabled = False
|
||||||
|
|
||||||
|
@unittest_run_loop
|
||||||
|
async def test_auto_approver_broadcast(self):
|
||||||
|
"""Test broadcasting auto approver config changes to multiple clients"""
|
||||||
|
async with self.client.ws_connect('/ws') as ws1, self.client.ws_connect('/ws') as ws2:
|
||||||
|
await self._skip_initial_messages(ws1)
|
||||||
|
await self._skip_initial_messages(ws2)
|
||||||
|
|
||||||
|
test_config = {
|
||||||
|
'context_enabled': True,
|
||||||
|
'response_enabled': False,
|
||||||
|
'context_timeout': 15.0,
|
||||||
|
'response_timeout': 25.0
|
||||||
|
}
|
||||||
|
|
||||||
|
await ws1.send_json({
|
||||||
|
'type': ClientMessage.AUTO_APPROVER_CONFIG,
|
||||||
|
'config': test_config
|
||||||
|
})
|
||||||
|
|
||||||
|
msg1 = await ws1.receive_json()
|
||||||
|
msg2 = await ws2.receive_json()
|
||||||
|
self.assertEqual(msg1['type'], ServerMessage.AUTO_APPROVER_CONFIG)
|
||||||
|
self.assertEqual(msg2['type'], ServerMessage.AUTO_APPROVER_CONFIG)
|
||||||
|
self.assertEqual(msg1['config'], test_config)
|
||||||
|
self.assertEqual(msg2['config'], test_config)
|
||||||
|
|
||||||
|
self.ws_manager.auto_approver.context_enabled = False
|
||||||
|
self.ws_manager.auto_approver.response_enabled = False
|
||||||
|
|
||||||
def tearDown(self):
|
def tearDown(self):
|
||||||
"""Clean up resources"""
|
"""Clean up resources"""
|
||||||
super().tearDown()
|
super().tearDown()
|
||||||
|
|||||||
@@ -23,6 +23,12 @@ const App = () => {
|
|||||||
const [showDiff, setShowDiff] = useState(false);
|
const [showDiff, setShowDiff] = useState(false);
|
||||||
const [showSidebar, setShowSidebar] = useState(false);
|
const [showSidebar, setShowSidebar] = useState(false);
|
||||||
const [agentState, setAgentState] = useState(AgentState.CONTEXT_APPROVAL);
|
const [agentState, setAgentState] = useState(AgentState.CONTEXT_APPROVAL);
|
||||||
|
const [autoApproverConfig, setAutoApproverConfig] = useState({
|
||||||
|
context_enabled: false,
|
||||||
|
response_enabled: false,
|
||||||
|
context_timeout: 5.0,
|
||||||
|
response_timeout: 10.0
|
||||||
|
});
|
||||||
|
|
||||||
const { wsState, sendMessage, addMessageHandler } = useWebSocket('ws://localhost:8080/ws');
|
const { wsState, sendMessage, addMessageHandler } = useWebSocket('ws://localhost:8080/ws');
|
||||||
|
|
||||||
@@ -44,10 +50,20 @@ const App = () => {
|
|||||||
case ServerMessageType.OUTPUT_UPDATE:
|
case ServerMessageType.OUTPUT_UPDATE:
|
||||||
setOutput(message.output);
|
setOutput(message.output);
|
||||||
break;
|
break;
|
||||||
|
case ServerMessageType.AUTO_APPROVER_CONFIG:
|
||||||
|
setAutoApproverConfig(message.config);
|
||||||
|
break;
|
||||||
}
|
}
|
||||||
});
|
});
|
||||||
}, [addMessageHandler]);
|
}, [addMessageHandler]);
|
||||||
|
|
||||||
|
const handleAutoApproverConfigChange = (config) => {
|
||||||
|
sendMessage({
|
||||||
|
type: ClientMessageType.AUTO_APPROVER_CONFIG,
|
||||||
|
config: config
|
||||||
|
});
|
||||||
|
};
|
||||||
|
|
||||||
const handleApprove = () => {
|
const handleApprove = () => {
|
||||||
if (agentState === AgentState.CONTEXT_APPROVAL) {
|
if (agentState === AgentState.CONTEXT_APPROVAL) {
|
||||||
sendMessage({
|
sendMessage({
|
||||||
@@ -158,6 +174,8 @@ const App = () => {
|
|||||||
onToggleDiff={() => setShowDiff(!showDiff)}
|
onToggleDiff={() => setShowDiff(!showDiff)}
|
||||||
onSendInput={handleSendInput}
|
onSendInput={handleSendInput}
|
||||||
onClearOutput={handleClearOutput}
|
onClearOutput={handleClearOutput}
|
||||||
|
autoApproverConfig={autoApproverConfig}
|
||||||
|
onAutoApproverConfigChange={handleAutoApproverConfigChange}
|
||||||
/>
|
/>
|
||||||
|
|
||||||
<FloatingButton onClick={() => setShowSidebar(!showSidebar)} />
|
<FloatingButton onClick={() => setShowSidebar(!showSidebar)} />
|
||||||
|
|||||||
@@ -1,5 +1,7 @@
|
|||||||
import React from 'react';
|
import React from 'react';
|
||||||
import { Button } from '@/components/ui/button';
|
import { Button } from '@/components/ui/button';
|
||||||
|
import { Switch } from '@/components/ui/switch';
|
||||||
|
import { Input } from '@/components/ui/input';
|
||||||
import { AgentState, Tabs } from '../constants';
|
import { AgentState, Tabs } from '../constants';
|
||||||
|
|
||||||
export const Sidebar = ({
|
export const Sidebar = ({
|
||||||
@@ -7,10 +9,12 @@ export const Sidebar = ({
|
|||||||
agentState,
|
agentState,
|
||||||
showDiff,
|
showDiff,
|
||||||
input,
|
input,
|
||||||
|
autoApproverConfig,
|
||||||
onApprove,
|
onApprove,
|
||||||
onToggleDiff,
|
onToggleDiff,
|
||||||
onSendInput,
|
onSendInput,
|
||||||
onClearOutput,
|
onClearOutput,
|
||||||
|
onAutoApproverConfigChange,
|
||||||
}) => (
|
}) => (
|
||||||
<aside className={`
|
<aside className={`
|
||||||
fixed top-0 right-0 h-screen w-64 bg-white shadow-lg
|
fixed top-0 right-0 h-screen w-64 bg-white shadow-lg
|
||||||
@@ -19,6 +23,50 @@ export const Sidebar = ({
|
|||||||
${showSidebar ? 'translate-x-0' : 'translate-x-full md:translate-x-0'}
|
${showSidebar ? 'translate-x-0' : 'translate-x-full md:translate-x-0'}
|
||||||
`}>
|
`}>
|
||||||
<div className="p-4 space-y-4 md:pt-0">
|
<div className="p-4 space-y-4 md:pt-0">
|
||||||
|
<div className="space-y-2">
|
||||||
|
<div className="flex items-center justify-between">
|
||||||
|
<label>Auto-approve Context</label>
|
||||||
|
<Switch
|
||||||
|
checked={autoApproverConfig.context_enabled}
|
||||||
|
onCheckedChange={(checked) => onAutoApproverConfigChange({
|
||||||
|
...autoApproverConfig,
|
||||||
|
context_enabled: checked
|
||||||
|
})}
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
<Input
|
||||||
|
type="number"
|
||||||
|
value={autoApproverConfig.context_timeout}
|
||||||
|
onChange={(e) => onAutoApproverConfigChange({
|
||||||
|
...autoApproverConfig,
|
||||||
|
context_timeout: parseFloat(e.target.value)
|
||||||
|
})}
|
||||||
|
disabled={autoApproverConfig.context_enabled}
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
|
||||||
|
<div className="space-y-2">
|
||||||
|
<div className="flex items-center justify-between">
|
||||||
|
<label>Auto-approve Response</label>
|
||||||
|
<Switch
|
||||||
|
checked={autoApproverConfig.response_enabled}
|
||||||
|
onCheckedChange={(checked) => onAutoApproverConfigChange({
|
||||||
|
...autoApproverConfig,
|
||||||
|
response_enabled: checked
|
||||||
|
})}
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
<Input
|
||||||
|
type="number"
|
||||||
|
value={autoApproverConfig.response_timeout}
|
||||||
|
onChange={(e) => onAutoApproverConfigChange({
|
||||||
|
...autoApproverConfig,
|
||||||
|
response_timeout: parseFloat(e.target.value)
|
||||||
|
})}
|
||||||
|
disabled={autoApproverConfig.response_enabled}
|
||||||
|
/>
|
||||||
|
</div>
|
||||||
|
|
||||||
<Button
|
<Button
|
||||||
variant="outline"
|
variant="outline"
|
||||||
onClick={onToggleDiff}
|
onClick={onToggleDiff}
|
||||||
@@ -52,4 +100,4 @@ export const Sidebar = ({
|
|||||||
</Button>
|
</Button>
|
||||||
</div>
|
</div>
|
||||||
</aside>
|
</aside>
|
||||||
);
|
);
|
||||||
|
|||||||
21
web/src/components/ui/switch.jsx
Normal file
21
web/src/components/ui/switch.jsx
Normal file
@@ -0,0 +1,21 @@
|
|||||||
|
import React from 'react';
|
||||||
|
|
||||||
|
export const Switch = ({ checked, onCheckedChange, disabled }) => (
|
||||||
|
<label className="relative inline-flex items-center cursor-pointer">
|
||||||
|
<input
|
||||||
|
type="checkbox"
|
||||||
|
className="sr-only peer"
|
||||||
|
checked={checked}
|
||||||
|
onChange={(e) => onCheckedChange(e.target.checked)}
|
||||||
|
disabled={disabled}
|
||||||
|
/>
|
||||||
|
<div className={`
|
||||||
|
w-11 h-6 bg-gray-200 rounded-full peer
|
||||||
|
peer-checked:after:translate-x-full peer-checked:bg-blue-600
|
||||||
|
after:content-[''] after:absolute after:top-[2px] after:left-[2px]
|
||||||
|
after:bg-white after:rounded-full after:h-5 after:w-5
|
||||||
|
after:transition-all
|
||||||
|
${disabled ? 'opacity-50 cursor-not-allowed' : ''}
|
||||||
|
`}></div>
|
||||||
|
</label>
|
||||||
|
);
|
||||||
@@ -24,6 +24,7 @@ export const ClientMessageType = {
|
|||||||
APPROVE_RESPONSE: 'APPROVE_RESPONSE',
|
APPROVE_RESPONSE: 'APPROVE_RESPONSE',
|
||||||
SEND_INPUT: 'SEND_INPUT',
|
SEND_INPUT: 'SEND_INPUT',
|
||||||
CLEAR_OUTPUT: 'CLEAR_OUTPUT',
|
CLEAR_OUTPUT: 'CLEAR_OUTPUT',
|
||||||
|
AUTO_APPROVER_CONFIG: 'AUTO_APPROVER_CONFIG',
|
||||||
};
|
};
|
||||||
|
|
||||||
export const ServerMessageType = {
|
export const ServerMessageType = {
|
||||||
@@ -31,4 +32,5 @@ export const ServerMessageType = {
|
|||||||
CONTEXT_UPDATE: 'CONTEXT_UPDATE',
|
CONTEXT_UPDATE: 'CONTEXT_UPDATE',
|
||||||
RESPONSE_UPDATE: 'RESPONSE_UPDATE',
|
RESPONSE_UPDATE: 'RESPONSE_UPDATE',
|
||||||
OUTPUT_UPDATE: 'OUTPUT_UPDATE',
|
OUTPUT_UPDATE: 'OUTPUT_UPDATE',
|
||||||
|
AUTO_APPROVER_CONFIG: 'AUTO_APPROVER_CONFIG',
|
||||||
};
|
};
|
||||||
Reference in New Issue
Block a user