llm/SIA
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Approve context

Browse Source
This commit is contained in:
geens
2024-11-02 19:17:27 +01:00
parent 96d14fc46d
commit 2a15708145
9 changed files with 296 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
sia/__main__.py
View File

@@ -43,9 +43,20 @@ async def run_web_server():
await server.start() await server.start()
try: try:
# Keep running until interrupted
while True: while True:
await asyncio.sleep(1) if server._response_tasks:
done, pending = await asyncio.wait(
server._response_tasks,
timeout=0.1,
return_when=asyncio.FIRST_COMPLETED
)
for task in done:
try:
await task
except Exception as e:
print(f"Task error: {e}")
server._response_tasks.remove(task)
await asyncio.sleep(0.1)
except KeyboardInterrupt: except KeyboardInterrupt:
print("\nShutting down...") print("\nShutting down...")
finally: finally:

79
sia/web_agent.py
View File

@@ -54,7 +54,6 @@ class WebAgent(BaseAgent):
self._command_result: Optional[CommandResult] = None self._command_result: Optional[CommandResult] = None
self.context = self._compile_context() self.context = self._compile_context()
@property @property
def current_state(self) -> WebAgentState: def current_state(self) -> WebAgentState:
"""Get the current state of the agent.""" """Get the current state of the agent."""
@@ -93,43 +92,68 @@ class WebAgent(BaseAgent):
def _notify_response_change(self) -> None: def _notify_response_change(self) -> None:
"""Notify all handlers of response change.""" """Notify all handlers of response change."""
for handler in self._response_change_handlers: for handler in self._response_change_handlers:
handler(self._response) handler(self.response)
def _set_response(self, response: str) -> None: def _set_response(self, response: str) -> None:
"""Set response and notify handlers."""
self.response = response
self._validation_error = self._validator.validate(response)
for handler in self._response_change_handlers:
try:
handler(response)
except Exception as e:
print(f"Error in response handler: {e}")
def _set_state(self, new_state: WebAgentState) -> None:
""" """
Set the current response and notify handlers. Set the current state and notify handlers.
Args: Args:
response: New response new_state: New state to set
""" """
self._response = response self._current_state = new_state
self._validation_error = self._validator.validate(self._response)
self._notify_response_change()
def approve_context(self) -> None:
if self._current_state != WebAgentState.CONTEXT_APPROVAL:
raise Exception("Not in CONTEXT_APPROVAL state")
self._set_response("")
self._current_state = WebAgentState.INFERENCE
self._notify_state_change()
response_token_iter = self._llm.infer(self._system_prompt, self.context)
for token in response_token_iter:
self._set_response(self._response + token)
self._current_state = WebAgentState.RESPONSE_APPROVAL
self._notify_state_change() self._notify_state_change()
async def approve_context(self) -> None:
if self._current_state != WebAgentState.CONTEXT_APPROVAL:
error_msg = f"Not in CONTEXT_APPROVAL state (current: {self._current_state})"
raise Exception(error_msg)
self._set_state(WebAgentState.INFERENCE)
self._set_response("")
response_token_iter = self._llm.infer(self._system_prompt, self.context)
response = ""
try:
for token in response_token_iter:
response += token
await self._set_response(response)
print(f"{token}")
except Exception as e:
print(f"Error during inference: {e}")
self._set_state(WebAgentState.CONTEXT_APPROVAL)
return
await self._set_response(response)
self._set_state(WebAgentState.RESPONSE_APPROVAL)
async def _set_response(self, response: str) -> None:
"""Set response and notify handlers asynchronously."""
self.response = response
self._validation_error = self._validator.validate(response)
for handler in self._response_change_handlers:
try:
await handler(response)
except Exception as e:
print(f"Error in response handler: {e}")
def approve_response(self) -> None: def approve_response(self) -> None:
if self._current_state != WebAgentState.RESPONSE_APPROVAL: if self._current_state != WebAgentState.RESPONSE_APPROVAL:
raise Exception("Not in RESPONSE_APPROVAL state") raise Exception("Not in RESPONSE_APPROVAL state")
self._current_state = WebAgentState.UPDATE self._set_state(WebAgentState.UPDATE)
self._notify_state_change()
parse_result = self._parser.parse(self._response) parse_result = self._parser.parse(self.response)
if isinstance(parse_result, Command): if isinstance(parse_result, Command):
result = parse_result.execute(self._working_memory) result = parse_result.execute(self._working_memory)
self._command_result = result self._command_result = result
@@ -137,7 +161,6 @@ class WebAgent(BaseAgent):
self._working_memory.add_entry(parse_result) self._working_memory.add_entry(parse_result)
self._working_memory.update() self._working_memory.update()
self._context = self._compile_context() self.context = self._compile_context()
self._current_state = WebAgentState.CONTEXT_APPROVAL self._set_state(WebAgentState.CONTEXT_APPROVAL)
self._notify_state_change()

70
sia/web_io_buffer.py
View File

@@ -1,38 +1,38 @@
from threading import Lock
from .io_buffer import IOBuffer from .io_buffer import IOBuffer
class WebIOBuffer(IOBuffer): class WebIOBuffer(IOBuffer):
""" """
IOBuffer implementation for web interface communication. Thread-safe WebIOBuffer that maintains the synchronous IOBuffer interface.
Uses threading primitives instead of asyncio for synchronization.
This class manages input/output through memory buffers instead of system IO,
allowing for asynchronous communication through a web interface. It maintains
separate buffers for stdin and stdout to simulate bidirectional communication.
""" """
def __init__(self): def __init__(self):
"""Initialize empty input and output buffers."""
self._stdin_buffer = "" self._stdin_buffer = ""
self._stdout_buffer = "" self._stdout_buffer = ""
self._lock = Lock()
def append_stdin(self, content: str) -> None:
"""Thread-safe append to stdin buffer."""
if not content:
return
with self._lock:
self._stdin_buffer += content
def read(self) -> str: def read(self) -> str:
""" """Thread-safe read from stdin buffer."""
Read and clear all available input from the stdin buffer. with self._lock:
content = self._stdin_buffer
Returns: self._stdin_buffer = ""
str: Content from the stdin buffer, or empty string if buffer is empty return content
"""
content = self._stdin_buffer
self._stdin_buffer = ""
return content
def write(self, content: str) -> None: def write(self, content: str) -> None:
""" """Thread-safe write to stdout buffer."""
Append content to the stdout buffer. if not content:
return
Args:
content: String content to append to the buffer with self._lock:
"""
if content:
self._stdout_buffer += content self._stdout_buffer += content
def buffer_length(self) -> int: def buffer_length(self) -> int:
@@ -55,14 +55,16 @@ class WebIOBuffer(IOBuffer):
self._stdin_buffer += content self._stdin_buffer += content
def get_stdout(self) -> str: def get_stdout(self) -> str:
""" """Thread-safe get stdout content."""
Get the current content of the stdout buffer. with self._lock:
return self._stdout_buffer
Returns:
str: Current content of the stdout buffer
"""
return self._stdout_buffer
def clear_stdout(self) -> None: def clear_stdout(self) -> None:
"""Clear the stdout buffer.""" """Thread-safe clear stdout buffer."""
self._stdout_buffer = "" with self._lock:
self._stdout_buffer = ""
def buffer_length(self) -> int:
"""Thread-safe get stdin buffer length."""
with self._lock:
return len(self._stdin_buffer)

104
sia/web_server.py
View File

@@ -5,6 +5,7 @@ import asyncio
import json import json
from dataclasses import dataclass from dataclasses import dataclass
import mimetypes import mimetypes
from functools import partial
from .web_agent import WebAgent, WebAgentState from .web_agent import WebAgent, WebAgentState
from .web_io_buffer import WebIOBuffer from .web_io_buffer import WebIOBuffer
@@ -37,7 +38,7 @@ class WebServer:
integrates with the WebAgent. Serves static files from a web directory. integrates with the WebAgent. Serves static files from a web directory.
""" """
def __init__(self, agent: WebAgent, io_buffer: WebIOBuffer, 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): static_dir: Optional[Path] = None):
"""Initialize the web server.""" """Initialize the web server."""
self.agent = agent self.agent = agent
@@ -49,11 +50,22 @@ class WebServer:
self._init_routes() self._init_routes()
self._clients: Set[web.WebSocketResponse] = set() self._clients: Set[web.WebSocketResponse] = set()
self._lock = asyncio.Lock() self._lock = asyncio.Lock()
self._response_tasks: List[asyncio.Task] = [] # Initialize tasks list
self.runner: Optional[web.AppRunner] = None self.runner: Optional[web.AppRunner] = None
# Register agent handlers # Register wrapped agent handlers
self.agent.add_state_change_handler(self._handle_state_change) self.agent.add_state_change_handler(self._wrap_async_handler(self._handle_state_change))
self.agent.add_response_change_handler(self._handle_response_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): def _init_routes(self):
"""Initialize application routes.""" """Initialize application routes."""
@@ -61,16 +73,44 @@ class WebServer:
self.app.router.add_get("/ws", self._handle_websocket) self.app.router.add_get("/ws", self._handle_websocket)
if self.static_dir.exists(): if self.static_dir.exists():
self.app.router.add_get("/", 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("/static/", self.static_dir, show_index=False)
self.app.router.add_static("/assets/", self.static_dir / "assets", 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("/{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: async def _serve_index(self, request: web.Request) -> web.Response:
""" """Serve the React application HTML for any unmatched routes."""
Serve the React application HTML for any unmatched routes to support
client-side routing.
"""
index_path = self.static_dir / "index.html" index_path = self.static_dir / "index.html"
if not index_path.exists(): if not index_path.exists():
raise web.HTTPNotFound() raise web.HTTPNotFound()
@@ -129,7 +169,8 @@ class WebServer:
try: try:
await ws.send_json(message) await ws.send_json(message)
return True return True
except Exception: except Exception as e:
print(f"Error sending to client: {e}")
return False return False
return False return False
@@ -216,27 +257,28 @@ class WebServer:
return ws return ws
async def _handle_client_message(self, data: dict): async def _handle_client_message(self, data: dict):
"""Handle incoming client message."""
message_type = data.get("type") message_type = data.get("type")
if message_type == ClientMessage.APPROVE_CONTEXT: if message_type == ClientMessage.APPROVE_CONTEXT:
if self.agent.current_state == WebAgentState.CONTEXT_APPROVAL: if self.agent.current_state == WebAgentState.CONTEXT_APPROVAL:
await self.agent.approve_context() await self.agent.approve_context()
elif message_type == ClientMessage.APPROVE_RESPONSE: elif message_type == ClientMessage.APPROVE_RESPONSE:
if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL: 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: elif message_type == ClientMessage.MODIFY_RESPONSE:
if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL: if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL:
new_response = data.get("data") new_response = data.get("data")
if new_response: if new_response:
await self.agent._set_response(new_response) self.agent._set_response(new_response)
elif message_type == ClientMessage.SEND_INPUT: elif message_type == ClientMessage.SEND_INPUT:
input_text = data.get("data") input_text = data.get("data")
if input_text: if input_text:
self.agent.set_input(input_text)
self.io_buffer.append_stdin(input_text + "\n") 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): async def _handle_state_change(self, new_state: WebAgentState):
"""Handle state changes from the WebAgent.""" """Handle state changes from the WebAgent."""
@@ -249,3 +291,27 @@ class WebServer:
response, response,
validation_error=self.agent._validation_error 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

28
test/mock_web_agent.py
View File

@@ -1,4 +1,3 @@
import asyncio
from typing import Optional, List, Callable from typing import Optional, List, Callable
from sia.web_agent import WebAgentState from sia.web_agent import WebAgentState
@@ -21,35 +20,34 @@ class MockWebAgent:
def add_response_change_handler(self, handler): def add_response_change_handler(self, handler):
self._response_handlers.append(handler) self._response_handlers.append(handler)
def set_input(self, input_text):
pass
async def _notify_state_handlers(self, new_state: WebAgentState): def _notify_state_handlers(self, new_state: WebAgentState):
"""Notify all state handlers of a state change.""" """Notify all state handlers of a state change."""
for handler in self._state_handlers: for handler in self._state_handlers:
await handler(new_state) handler(new_state)
async def _notify_response_handlers(self, response: str): def _notify_response_handlers(self, response: str):
"""Notify all response handlers of a response change.""" """Notify all response handlers of a response change."""
for handler in self._response_handlers: for handler in self._response_handlers:
await handler(response) handler(response)
async def _set_response(self, response: str): def _set_response(self, response: str):
"""Set response and notify handlers.""" """Set response and notify handlers."""
self.response = response self.response = response
await self._notify_response_handlers(response) self._notify_response_handlers(response)
async def approve_context(self): def approve_context(self):
"""Handle context approval with immediate state updates.""" """Handle context approval with immediate state updates."""
# Change to INFERENCE state and notify # Change to INFERENCE state and notify
self._current_state = WebAgentState.INFERENCE self._current_state = WebAgentState.INFERENCE
for handler in self._state_handlers: self._notify_state_handlers(WebAgentState.INFERENCE)
await handler(WebAgentState.INFERENCE)
# Set response and notify # Set response and notify
self.response = "<reasoning>test</reasoning>" self._set_response("<reasoning>test</reasoning>")
for handler in self._response_handlers:
await handler(self.response)
# Change to RESPONSE_APPROVAL state and notify # Change to RESPONSE_APPROVAL state and notify
self._current_state = WebAgentState.RESPONSE_APPROVAL self._current_state = WebAgentState.RESPONSE_APPROVAL
for handler in self._state_handlers: self._notify_state_handlers(WebAgentState.RESPONSE_APPROVAL)
await handler(WebAgentState.RESPONSE_APPROVAL)

5
test/standard_io_buffer_test.py
View File

@@ -159,11 +159,6 @@ if __name__ == '__main__':
stdout, stderr, code = self.run_script('read', special_chars) stdout, stderr, code = self.run_script('read', special_chars)
self.assertEqual(code, 0, f"Script failed with stderr: {stderr}") self.assertEqual(code, 0, f"Script failed with stderr: {stderr}")
expected = f'Read: {special_chars}' expected = f'Read: {special_chars}'
# Print actual bytes for debugging
if stdout != expected:
print("\nDebug output:")
print("Expected bytes:", expected.encode('utf-8'))
print("Actual bytes:", stdout.encode('utf-8'))
self.assertEqual(stdout, expected) self.assertEqual(stdout, expected)
def test_empty_input(self): def test_empty_input(self):

26
test/web_agent_test.py
View File

@@ -1,3 +1,4 @@
import asyncio
import unittest import unittest
from datetime import datetime from datetime import datetime
from unittest.mock import Mock, MagicMock, patch, call from unittest.mock import Mock, MagicMock, patch, call
@@ -123,7 +124,7 @@ class WebAgentTest(unittest.TestCase):
self.agent._current_state = WebAgentState.UPDATE self.agent._current_state = WebAgentState.UPDATE
with self.assertRaises(Exception) as context: with self.assertRaises(Exception) as context:
self.agent.approve_context() asyncio.run(self.agent.approve_context())
self.assertIn("Not in CONTEXT_APPROVAL state", str(context.exception)) self.assertIn("Not in CONTEXT_APPROVAL state", str(context.exception))
def test_approve_response_state_error(self): def test_approve_response_state_error(self):
@@ -132,7 +133,7 @@ class WebAgentTest(unittest.TestCase):
self.agent.approve_response() self.agent.approve_response()
self.assertIn("Not in RESPONSE_APPROVAL state", str(context.exception)) self.assertIn("Not in RESPONSE_APPROVAL state", str(context.exception))
def test_context_approval_flow(self): async def test_context_approval_flow(self):
"""Test complete context approval flow with state transitions.""" """Test complete context approval flow with state transitions."""
# Register handlers # Register handlers
self.agent.add_state_change_handler(self.state_change_handler) self.agent.add_state_change_handler(self.state_change_handler)
@@ -140,6 +141,7 @@ class WebAgentTest(unittest.TestCase):
# Setup LLM response # Setup LLM response
def mock_infer(prompt: str, context: str) -> Iterator[str]: def mock_infer(prompt: str, context: str) -> Iterator[str]:
# Return complete response at once for testing
yield "<reasoning>test reasoning</reasoning>" yield "<reasoning>test reasoning</reasoning>"
self.mock_llm.infer.side_effect = mock_infer self.mock_llm.infer.side_effect = mock_infer
@@ -149,6 +151,9 @@ class WebAgentTest(unittest.TestCase):
# Approve context # Approve context
self.agent.approve_context() self.agent.approve_context()
# Wait for response tasks to complete
await asyncio.gather(*self.web_server._response_tasks)
# Verify state transitions # Verify state transitions
expected_states = [ expected_states = [
WebAgentState.INFERENCE, WebAgentState.INFERENCE,
@@ -156,10 +161,11 @@ class WebAgentTest(unittest.TestCase):
] ]
self.assertEqual(self.state_changes, expected_states) self.assertEqual(self.state_changes, expected_states)
# Verify response updates # Updated test: Should now expect streaming updates
self.assertEqual(len(self.response_changes), 2) # Empty + final response self.assertEqual(len(self.response_changes), len(self.received_messages))
self.assertEqual(self.response_changes[0], "") for i, msg in enumerate(self.received_messages):
self.assertEqual(self.response_changes[1], "<reasoning>test reasoning</reasoning>") self.assertEqual(msg["type"], "RESPONSE_UPDATE")
self.assertEqual(msg["data"], self.response_changes[i])
def test_response_approval_flow_command(self): def test_response_approval_flow_command(self):
"""Test response approval flow with command execution.""" """Test response approval flow with command execution."""
@@ -200,7 +206,7 @@ class WebAgentTest(unittest.TestCase):
# Set initial state and response # Set initial state and response
self.agent._current_state = WebAgentState.RESPONSE_APPROVAL self.agent._current_state = WebAgentState.RESPONSE_APPROVAL
self.agent._response = "<reasoning>test reasoning</reasoning>" self.agent.response = "<reasoning>test reasoning</reasoning>"
# Approve response # Approve response
self.agent.approve_response() self.agent.approve_response()
@@ -218,7 +224,7 @@ class WebAgentTest(unittest.TestCase):
] ]
self.assertEqual(self.state_changes, expected_states) self.assertEqual(self.state_changes, expected_states)
def test_response_validation(self): async def test_response_validation(self):
"""Test response validation during response setting.""" """Test response validation during response setting."""
# Register handlers # Register handlers
self.agent.add_response_change_handler(self.response_change_handler) self.agent.add_response_change_handler(self.response_change_handler)
@@ -227,8 +233,8 @@ class WebAgentTest(unittest.TestCase):
error_message = "Invalid XML" error_message = "Invalid XML"
self.mock_validator.validate.return_value = error_message self.mock_validator.validate.return_value = error_message
# Set response # Set response using async method
self.agent._set_response("<invalid>") await self.agent._set_response("<invalid>")
# Verify validation error stored # Verify validation error stored
self.assertEqual(self.agent._validation_error, error_message) self.assertEqual(self.agent._validation_error, error_message)

36
test/web_server_test.py
View File

@@ -6,6 +6,7 @@ from sia.web_agent import WebAgentState
from unittest.mock import patch from unittest.mock import patch
import asyncio import asyncio
import asyncio import asyncio
import time
import json import json
import json import json
import shutil import shutil
@@ -29,6 +30,7 @@ class WebServerTest(AioHTTPTestCase):
self.static_dir.mkdir(exist_ok=True) self.static_dir.mkdir(exist_ok=True)
# Create web server with temp static directory # Create web server with temp static directory
self.agent._current_state = WebAgentState.CONTEXT_APPROVAL
self.web_server = WebServer( self.web_server = WebServer(
self.agent, self.agent,
self.io_buffer, self.io_buffer,
@@ -54,7 +56,6 @@ class WebServerTest(AioHTTPTestCase):
elif msg.type == WSMsgType.BINARY: elif msg.type == WSMsgType.BINARY:
messages.append(json.loads(msg.data.decode())) messages.append(json.loads(msg.data.decode()))
elif msg.type == WSMsgType.CLOSED: elif msg.type == WSMsgType.CLOSED:
print("WebSocket closed normally")
break break
elif msg.type == WSMsgType.ERROR: elif msg.type == WSMsgType.ERROR:
self.fail(f"WebSocket error while waiting for messages: {msg.data}") self.fail(f"WebSocket error while waiting for messages: {msg.data}")
@@ -108,6 +109,7 @@ class WebServerTest(AioHTTPTestCase):
await ws.send_json({ await ws.send_json({
"type": ClientMessage.APPROVE_CONTEXT, "type": ClientMessage.APPROVE_CONTEXT,
}) })
await asyncio.sleep(0.1) # Let state changes propagate
# Should receive 3 messages: INFERENCE state, response, RESPONSE_APPROVAL state # Should receive 3 messages: INFERENCE state, response, RESPONSE_APPROVAL state
messages = await self.wait_for_messages(ws, count=3) messages = await self.wait_for_messages(ws, count=3)
@@ -124,7 +126,7 @@ class WebServerTest(AioHTTPTestCase):
self.assertEqual(messages[2]["type"], ServerMessage.STATE_CHANGE) self.assertEqual(messages[2]["type"], ServerMessage.STATE_CHANGE)
self.assertEqual(messages[2]["data"], WebAgentState.RESPONSE_APPROVAL.name) self.assertEqual(messages[2]["data"], WebAgentState.RESPONSE_APPROVAL.name)
@unittest_run_loop @unittest_run_loop
async def test_validation_error(self): async def test_validation_error(self):
"""Test handling of validation errors.""" """Test handling of validation errors."""
async with self.client.ws_connect("/ws") as ws: async with self.client.ws_connect("/ws") as ws:
@@ -134,7 +136,7 @@ class WebServerTest(AioHTTPTestCase):
# Set validation error and response # Set validation error and response
self.agent._validation_error = "Invalid XML" self.agent._validation_error = "Invalid XML"
await self.agent._set_response("<invalid>") self.agent._set_response("<invalid>") # Now synchronous
await asyncio.sleep(0.1) # Let changes propagate await asyncio.sleep(0.1) # Let changes propagate
# Should receive response update with error # Should receive response update with error
@@ -204,24 +206,34 @@ class WebServerTest(AioHTTPTestCase):
await ws1.close() await ws1.close()
await ws2.close() await ws2.close()
@unittest_run_loop async def wait_for_buffer_content(self, timeout=1.0):
start = time.time()
while time.time() - start < timeout:
if self.io_buffer.buffer_length() > 0:
return True
await asyncio.sleep(0.01)
return False
async def test_send_input(self): async def test_send_input(self):
"""Test sending input to agent."""
async with self.client.ws_connect("/ws") as ws: async with self.client.ws_connect("/ws") as ws:
# Drain initial messages
await self.drain_messages(ws) await self.drain_messages(ws)
await asyncio.sleep(0.1) # Let connection settle await asyncio.sleep(0.1)
# Send input
test_input = "test input" test_input = "test input"
await ws.send_json({ await ws.send_json({
"type": ClientMessage.SEND_INPUT, "type": ClientMessage.SEND_INPUT,
"data": test_input "data": test_input
}) })
await asyncio.sleep(0.1) # Let input process
# Wait for buffer to have content
self.assertTrue(await self.wait_for_buffer_content())
# Verify input added to buffer with newline content = self.io_buffer.read()
self.assertEqual(self.io_buffer.read(), test_input + "\n") self.assertEqual(
content,
test_input + "\n",
f"Expected '{test_input}\\n' but got '{content}'"
)
@unittest_run_loop @unittest_run_loop
async def test_invalid_json(self): async def test_invalid_json(self):

82
web/src/components/App.jsx
View File

@@ -1,12 +1,10 @@
import React, { useState, useEffect, useRef } from 'react'; import React, { useState, useEffect, useRef } from 'react';
import Editor from '@monaco-editor/react'; import Editor from '@monaco-editor/react';
import { Card, CardContent, CardHeader, CardTitle } from '../components/ui/card.jsx'; import { Card, CardContent, CardHeader, CardTitle } from '@/components/ui/card';
import { Button } from '../components/ui/button.jsx'; import { Button } from '@/components/ui/button';
import { ScrollArea } from '../components/ui/scroll-area.jsx'; import { ScrollArea } from '@/components/ui/scroll-area';
import { Alert, AlertDescription } from '../components/ui/alert.jsx'; import { Alert, AlertDescription } from '@/components/ui/alert';
import { Tabs, TabsContent, TabsList, TabsTrigger } from '../components/ui/tabs.jsx'; import { Tabs, TabsContent, TabsList, TabsTrigger } from '@/components/ui/tabs';
import { Textarea } from '../components/ui/textarea.jsx';
import { Input } from '../components/ui/input.jsx';
const WebSocketState = { const WebSocketState = {
CONNECTING: 'CONNECTING', CONNECTING: 'CONNECTING',
@@ -29,6 +27,13 @@ const MessageType = {
VALIDATION_ERROR: 'VALIDATION_ERROR', VALIDATION_ERROR: 'VALIDATION_ERROR',
}; };
const ClientMessageType = {
APPROVE_CONTEXT: 'APPROVE_CONTEXT',
APPROVE_RESPONSE: 'APPROVE_RESPONSE',
MODIFY_RESPONSE: 'MODIFY_RESPONSE',
SEND_INPUT: 'SEND_INPUT',
};
const SIAInterface = () => { const SIAInterface = () => {
// Editor content state // Editor content state
const [context, setContext] = useState(''); const [context, setContext] = useState('');
@@ -40,7 +45,7 @@ const SIAInterface = () => {
// UI state // UI state
const [activeTab, setActiveTab] = useState('input'); const [activeTab, setActiveTab] = useState('input');
const [agentState, setAgentState] = useState(AgentState.UPDATE); const [agentState, setAgentState] = useState(AgentState.CONTEXT_APPROVAL);
const [wsState, setWsState] = useState(WebSocketState.DISCONNECTED); const [wsState, setWsState] = useState(WebSocketState.DISCONNECTED);
// WebSocket ref // WebSocket ref
@@ -58,10 +63,7 @@ const SIAInterface = () => {
}; };
useEffect(() => { useEffect(() => {
// Connect to WebSocket
connectWebSocket(); connectWebSocket();
// Cleanup on unmount
return () => { return () => {
if (wsRef.current) { if (wsRef.current) {
wsRef.current.close(); wsRef.current.close();
@@ -80,10 +82,13 @@ const SIAInterface = () => {
ws.onclose = () => { ws.onclose = () => {
setWsState(WebSocketState.DISCONNECTED); setWsState(WebSocketState.DISCONNECTED);
// Attempt to reconnect after delay
setTimeout(connectWebSocket, 2000); setTimeout(connectWebSocket, 2000);
}; };
ws.onerror = (error) => {
console.error('WebSocket error:', error);
};
ws.onmessage = (event) => { ws.onmessage = (event) => {
try { try {
const message = JSON.parse(event.data); const message = JSON.parse(event.data);
@@ -113,21 +118,26 @@ const SIAInterface = () => {
} }
}; };
const sendWebSocketMessage = (type, data) => { const sendWebSocketMessage = (type, data = null) => {
if (wsRef.current?.readyState === WebSocket.OPEN) { if (wsRef.current?.readyState === WebSocket.OPEN) {
wsRef.current.send(JSON.stringify({ type, data })); const message = {
type: type, // Ensure exact string match
data: data || undefined // Only include if not null
};
console.log('Sending message:', JSON.stringify(message)); // Log exact message being sent
wsRef.current.send(JSON.stringify(message));
} else {
console.error('WebSocket not connected, state:', wsRef.current?.readyState);
} }
}; };
const handleSendInput = (withStdin = false) => { const handleSendInput = (withStdin = false) => {
if (!input.trim()) return; if (!input.trim()) return;
// Send input to backend sendWebSocketMessage(ClientMessageType.SEND_INPUT, input);
sendWebSocketMessage('SEND_INPUT', input);
setInput(''); setInput('');
if (withStdin) { if (withStdin) {
// Prepare read_stdin response
const readStdinXml = `<read_stdin> const readStdinXml = `<read_stdin>
<!-- Content will be populated with stdin --> <!-- Content will be populated with stdin -->
</read_stdin>`; </read_stdin>`;
@@ -139,12 +149,26 @@ const SIAInterface = () => {
}; };
const handleApproveContext = () => { const handleApproveContext = () => {
sendWebSocketMessage('APPROVE_CONTEXT'); console.log('Current state:', agentState);
console.log('WebSocket state:', wsRef.current?.readyState);
if (agentState === AgentState.CONTEXT_APPROVAL &&
wsRef.current?.readyState === WebSocket.OPEN) {
console.log('Sending approve context message');
sendWebSocketMessage(ClientMessageType.APPROVE_CONTEXT);
} else {
console.warn(
'Cannot approve context.',
'Agent state:', agentState,
'WebSocket state:', wsRef.current?.readyState
);
}
}; };
const handleApproveResponse = (modified = false) => { const handleApproveResponse = (modified = false) => {
const response = modified ? modifiedResponse : originalResponse; if (agentState === AgentState.RESPONSE_APPROVAL) {
sendWebSocketMessage('APPROVE_RESPONSE', response); const response = modified ? modifiedResponse : originalResponse;
sendWebSocketMessage(ClientMessageType.APPROVE_RESPONSE, response);
}
}; };
const handleResetModified = () => { const handleResetModified = () => {
@@ -190,7 +214,7 @@ const SIAInterface = () => {
<div className="mt-4"> <div className="mt-4">
<Button <Button
onClick={handleApproveContext} onClick={handleApproveContext}
disabled={agentState !== AgentState.CONTEXT_APPROVAL} disabled={agentState !== AgentState.CONTEXT_APPROVAL || wsState !== WebSocketState.CONNECTED}
className="w-full" className="w-full"
> >
Approve Context Approve Context
@@ -223,6 +247,7 @@ const SIAInterface = () => {
<div className="flex space-x-2"> <div className="flex space-x-2">
<Button <Button
onClick={() => handleSendInput(false)} onClick={() => handleSendInput(false)}
disabled={wsState !== WebSocketState.CONNECTED}
className="flex-1" className="flex-1"
> >
Send Input Send Input
@@ -230,6 +255,7 @@ const SIAInterface = () => {
<Button <Button
onClick={() => handleSendInput(true)} onClick={() => handleSendInput(true)}
variant="secondary" variant="secondary"
disabled={wsState !== WebSocketState.CONNECTED}
className="flex-1" className="flex-1"
> >
Send & Read Stdin Send & Read Stdin
@@ -259,7 +285,11 @@ const SIAInterface = () => {
)} )}
<Button <Button
onClick={() => handleApproveResponse(false)} onClick={() => handleApproveResponse(false)}
disabled={agentState !== AgentState.RESPONSE_APPROVAL || validationError} disabled={
agentState !== AgentState.RESPONSE_APPROVAL ||
validationError ||
wsState !== WebSocketState.CONNECTED
}
className="w-full" className="w-full"
> >
Approve Original Response Approve Original Response
@@ -297,7 +327,11 @@ const SIAInterface = () => {
</Button> </Button>
<Button <Button
onClick={() => handleApproveResponse(true)} onClick={() => handleApproveResponse(true)}
disabled={agentState !== AgentState.RESPONSE_APPROVAL || validationError} disabled={
agentState !== AgentState.RESPONSE_APPROVAL ||
validationError ||
wsState !== WebSocketState.CONNECTED
}
className="flex-1" className="flex-1"
> >
Approve Modified Response Approve Modified Response
@@ -334,4 +368,4 @@ const SIAInterface = () => {
); );
}; };
export default SIAInterface; export default SIAInterface;