From 9dfd5ddc396f6ccd08c49a7761d0b733b1506f83 Mon Sep 17 00:00:00 2001 From: Niels Geens Date: Fri, 15 Nov 2024 17:09:07 +0100 Subject: [PATCH] Add auto_approver --- sia/__main__.py | 2 +- sia/auto_approver.py | 162 +++++++++++++++++++++++++++++++ sia/web_socket_manager.py | 24 ++++- test/auto_approver_test.py | 87 +++++++++++++++++ test/web_socket_manager_test.py | 108 +++++++++++++++++---- web/src/components/App.jsx | 18 ++++ web/src/components/Sidebar.jsx | 50 +++++++++- web/src/components/ui/switch.jsx | 21 ++++ web/src/constants.jsx | 2 + 9 files changed, 449 insertions(+), 25 deletions(-) create mode 100644 sia/auto_approver.py create mode 100644 test/auto_approver_test.py create mode 100644 web/src/components/ui/switch.jsx diff --git a/sia/__main__.py b/sia/__main__.py index 5541aa7..2a80273 100644 --- a/sia/__main__.py +++ b/sia/__main__.py @@ -122,5 +122,5 @@ if __name__ == "__main__": loop = asyncio.new_event_loop() config = 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) \ No newline at end of file diff --git a/sia/auto_approver.py b/sia/auto_approver.py new file mode 100644 index 0000000..dd3b4d8 --- /dev/null +++ b/sia/auto_approver.py @@ -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() diff --git a/sia/web_socket_manager.py b/sia/web_socket_manager.py index 2574a9b..3b2800c 100644 --- a/sia/web_socket_manager.py +++ b/sia/web_socket_manager.py @@ -6,6 +6,7 @@ from dataclasses import dataclass from .web_agent import WebAgent, WebAgentState from .web_io_buffer import WebIOBuffer +from .auto_approver import AutoApprover, AutoApproverConfig @dataclass class ClientMessage: @@ -15,6 +16,7 @@ class ClientMessage: MODIFY_RESPONSE = "MODIFY_RESPONSE" SEND_INPUT = "SEND_INPUT" CLEAR_OUTPUT = "CLEAR_OUTPUT" + AUTO_APPROVER_CONFIG = "AUTO_APPROVER_CONFIG" @dataclass class ServerMessage: @@ -23,6 +25,7 @@ class ServerMessage: CONTEXT_UPDATE = "CONTEXT_UPDATE" RESPONSE_UPDATE = "RESPONSE_UPDATE" OUTPUT_UPDATE = "OUTPUT_UPDATE" + AUTO_APPROVER_CONFIG = "AUTO_APPROVER_CONFIG" class WebSocketManager: """ @@ -38,11 +41,13 @@ class WebSocketManager: self.agent = agent self.io_buffer = io_buffer 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_context_change_handler(self._wrap_async(self._handle_context_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.auto_approver.add_config_change_handler(self._wrap_async(self._handle_auto_approver_config)) return self async def _broadcast_message(self, message: Dict): @@ -99,6 +104,13 @@ class WebSocketManager: "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]): """Handle incoming client message.""" message_type = data.get("type") @@ -117,9 +129,13 @@ class WebSocketManager: if self.agent.state == WebAgentState.RESPONSE_APPROVAL: self.agent.modify_response(data.get("response")) 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: 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: """Handle new WebSocket connections and messages.""" @@ -142,6 +158,10 @@ class WebSocketManager: "response": self.agent.response, "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: if msg.type == WSMsgType.TEXT: data = json.loads(msg.data) @@ -150,4 +170,4 @@ class WebSocketManager: print(f"WebSocket connection closed with error: {ws.exception()}") finally: self._clients.remove(ws) - return ws \ No newline at end of file + return ws diff --git a/test/auto_approver_test.py b/test/auto_approver_test.py new file mode 100644 index 0000000..5a27ddb --- /dev/null +++ b/test/auto_approver_test.py @@ -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 diff --git a/test/web_socket_manager_test.py b/test/web_socket_manager_test.py index 9b3cc96..2ac3909 100644 --- a/test/web_socket_manager_test.py +++ b/test/web_socket_manager_test.py @@ -68,6 +68,11 @@ class WebSocketManagerTest(AioHTTPTestCase): app.router.add_get('/ws', self.ws_manager.handle_websocket) 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 async def test_initial_connection(self): """Test initial WebSocket connection and state messages""" @@ -88,25 +93,34 @@ class WebSocketManagerTest(AioHTTPTestCase): self.assertEqual(msg['response'], '') 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 async def test_approve_context(self): """Test context approval flow""" async with self.client.ws_connect('/ws') as ws: - for _ in range(3): - msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5) + await self._skip_initial_messages(ws) + await ws.send_json({ '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['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['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['response'], "test 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['state'], WebAgentState.RESPONSE_APPROVAL.name) @@ -116,53 +130,51 @@ class WebSocketManagerTest(AioHTTPTestCase): test_input = "test input" async with self.client.ws_connect('/ws') as ws: - for _ in range(3): - await ws.receive_json() + await self._skip_initial_messages(ws) - # Send input await ws.send_json({ 'type': ClientMessage.SEND_INPUT, 'input': test_input }) - # Delay for processing await asyncio.sleep(0.5) - - # Verify input was added to buffer self.assertEqual(self.io_buffer.read(), test_input) @unittest_run_loop async def test_modify_response(self): """Test modifying response in RESPONSE_APPROVAL state""" - # Set agent to RESPONSE_APPROVAL state async with self.client.ws_connect('/ws') as ws: - for _ in range(3): - msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5) + await self._skip_initial_messages(ws) + await ws.send_json({ 'type': ClientMessage.APPROVE_CONTEXT }) + + # Skip to RESPONSE_APPROVAL state for _ in range(4): - msg = await asyncio.wait_for(ws.receive_json(), timeout=0.5) + await ws.receive_json() + modified_response = "modified test" await ws.send_json({ 'type': ClientMessage.MODIFY_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['response'], modified_response) @unittest_run_loop async def test_multiple_clients(self): """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: - for _ in range(3): - await asyncio.wait_for(ws1.receive_json(), timeout=0.5) - await asyncio.wait_for(ws2.receive_json(), timeout=0.5) + await self._skip_initial_messages(ws1) + await self._skip_initial_messages(ws2) + await ws1.send_json({ 'type': ClientMessage.APPROVE_CONTEXT }) + msg1 = await ws1.receive_json() msg2 = await ws2.receive_json() self.assertEqual(msg1['type'], ServerMessage.STATE_CHANGE) @@ -170,6 +182,60 @@ class WebSocketManagerTest(AioHTTPTestCase): self.assertEqual(msg2['type'], ServerMessage.STATE_CHANGE) 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): """Clean up resources""" super().tearDown() diff --git a/web/src/components/App.jsx b/web/src/components/App.jsx index 1353877..a0e391e 100644 --- a/web/src/components/App.jsx +++ b/web/src/components/App.jsx @@ -23,6 +23,12 @@ const App = () => { const [showDiff, setShowDiff] = useState(false); const [showSidebar, setShowSidebar] = useState(false); 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'); @@ -44,10 +50,20 @@ const App = () => { case ServerMessageType.OUTPUT_UPDATE: setOutput(message.output); break; + case ServerMessageType.AUTO_APPROVER_CONFIG: + setAutoApproverConfig(message.config); + break; } }); }, [addMessageHandler]); + const handleAutoApproverConfigChange = (config) => { + sendMessage({ + type: ClientMessageType.AUTO_APPROVER_CONFIG, + config: config + }); + }; + const handleApprove = () => { if (agentState === AgentState.CONTEXT_APPROVAL) { sendMessage({ @@ -158,6 +174,8 @@ const App = () => { onToggleDiff={() => setShowDiff(!showDiff)} onSendInput={handleSendInput} onClearOutput={handleClearOutput} + autoApproverConfig={autoApproverConfig} + onAutoApproverConfigChange={handleAutoApproverConfigChange} /> setShowSidebar(!showSidebar)} /> diff --git a/web/src/components/Sidebar.jsx b/web/src/components/Sidebar.jsx index 7562427..c69d67d 100644 --- a/web/src/components/Sidebar.jsx +++ b/web/src/components/Sidebar.jsx @@ -1,5 +1,7 @@ import React from 'react'; import { Button } from '@/components/ui/button'; +import { Switch } from '@/components/ui/switch'; +import { Input } from '@/components/ui/input'; import { AgentState, Tabs } from '../constants'; export const Sidebar = ({ @@ -7,10 +9,12 @@ export const Sidebar = ({ agentState, showDiff, input, + autoApproverConfig, onApprove, onToggleDiff, onSendInput, onClearOutput, + onAutoApproverConfigChange, }) => ( -); \ No newline at end of file +); diff --git a/web/src/components/ui/switch.jsx b/web/src/components/ui/switch.jsx new file mode 100644 index 0000000..7bf86ab --- /dev/null +++ b/web/src/components/ui/switch.jsx @@ -0,0 +1,21 @@ +import React from 'react'; + +export const Switch = ({ checked, onCheckedChange, disabled }) => ( + +); diff --git a/web/src/constants.jsx b/web/src/constants.jsx index d93c802..ff9d4a9 100644 --- a/web/src/constants.jsx +++ b/web/src/constants.jsx @@ -24,6 +24,7 @@ export const ClientMessageType = { APPROVE_RESPONSE: 'APPROVE_RESPONSE', SEND_INPUT: 'SEND_INPUT', CLEAR_OUTPUT: 'CLEAR_OUTPUT', + AUTO_APPROVER_CONFIG: 'AUTO_APPROVER_CONFIG', }; export const ServerMessageType = { @@ -31,4 +32,5 @@ export const ServerMessageType = { CONTEXT_UPDATE: 'CONTEXT_UPDATE', RESPONSE_UPDATE: 'RESPONSE_UPDATE', OUTPUT_UPDATE: 'OUTPUT_UPDATE', + AUTO_APPROVER_CONFIG: 'AUTO_APPROVER_CONFIG', }; \ No newline at end of file