Fixed stdout

sia/web_agent.py

@@ -153,8 +153,9 @@ class WebAgent(BaseAgent):
         for token in response_token_iter:
             response += token
             self._set_response(response)
-            print(f"{token}")
+            print(f"{token}", end='')
         self._set_state(WebAgentState.RESPONSE_APPROVAL)
+        print()
     
     def _approve_response_thread(self) -> None:
         parse_result = self._parser.parse(self._response)

sia/web_io_buffer.py

@@ -1,4 +1,5 @@
 from threading import Lock
+from typing import Callable, List
 
 from .io_buffer import IOBuffer
 
@@ -10,6 +11,14 @@ class WebIOBuffer(IOBuffer):
     def __init__(self):
         self._stdin_buffer = ""
         self._stdout_buffer = ""
+        self._stdout_change_handlers: List[Callable[[str], None]] = []
+    
+    def add_stdout_change_handler(self, handler: Callable[[str], None]) -> None:
+        """
+        Add a callback for stdout changes.
+        """
+        if handler not in self._stdout_change_handlers:
+            self._stdout_change_handlers.append(handler)
     
     def read(self) -> str:
         """Thread-safe read from stdin buffer."""
@@ -20,6 +29,8 @@ class WebIOBuffer(IOBuffer):
     def write(self, content: str) -> None:
         """Thread-safe write to stdout buffer."""
         self._stdout_buffer += content
+        for handler in self._stdout_change_handlers:
+            handler(self._stdout_buffer)
             
     def buffer_length(self) -> int:
         """
@@ -45,4 +56,6 @@ class WebIOBuffer(IOBuffer):
             
     def clear_stdout(self) -> None:
         """Thread-safe clear stdout buffer."""
-        self._stdout_buffer = ""
+        self._stdout_buffer = ""
+        for handler in self._stdout_change_handlers:
+            handler(self._stdout_buffer)

sia/web_socket_manager.py

@@ -14,6 +14,7 @@ class ClientMessage:
     APPROVE_RESPONSE = "APPROVE_RESPONSE" 
     MODIFY_RESPONSE = "MODIFY_RESPONSE"
     SEND_INPUT = "SEND_INPUT"
+    CLEAR_OUTPUT = "CLEAR_OUTPUT"
 
 @dataclass
 class ServerMessage:
@@ -41,6 +42,7 @@ class WebSocketManager:
         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))
         return self
 
     async def _broadcast_message(self, message: Dict):
@@ -89,6 +91,13 @@ class WebSocketManager:
             "response": response,
             "validation_error": validation_error
         })
+    
+    async def _handle_stdout_change(self, output: str):
+        """Handle output changes from the WebAgent."""
+        await self._broadcast_message({
+            "type": ServerMessage.OUTPUT_UPDATE,
+            "output": output
+        })
 
     async def _handle_client_message(self, request: web.Request, data: dict[str, Any]):
         """Handle incoming client message."""
@@ -109,6 +118,8 @@ class WebSocketManager:
                 self.agent.modify_response(data.get("response"))
         elif message_type == ClientMessage.SEND_INPUT:
                 self.io_buffer.append_stdin(data.get("input"))
+        elif message_type == ClientMessage.CLEAR_OUTPUT:
+            self.io_buffer.clear_stdout()
 
     async def handle_websocket(self, request: web.Request) -> web.WebSocketResponse:
         """Handle new WebSocket connections and messages."""