From 3a1d8a449cef749f1e779cd0195066c39b2144ce Mon Sep 17 00:00:00 2001
From: geens <niels.geens@gmail.com>
Date: Mon, 11 Nov 2024 13:51:27 +0100
Subject: [PATCH] Fixed stdout

---
 sia/web_agent.py           |  3 ++-
 sia/web_io_buffer.py       | 15 ++++++++++++++-
 sia/web_socket_manager.py  | 11 +++++++++++
 web/src/components/App.jsx | 10 ++++++++--
 web/src/constants.jsx      |  1 +
 5 files changed, 36 insertions(+), 4 deletions(-)

diff --git a/sia/web_agent.py b/sia/web_agent.py
index 6e2e420..898bcae 100644
--- a/sia/web_agent.py
+++ b/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)
diff --git a/sia/web_io_buffer.py b/sia/web_io_buffer.py
index 4399ac0..d919a22 100644
--- a/sia/web_io_buffer.py
+++ b/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 = ""
\ No newline at end of file
+        self._stdout_buffer = ""
+        for handler in self._stdout_change_handlers:
+            handler(self._stdout_buffer)
\ No newline at end of file
diff --git a/sia/web_socket_manager.py b/sia/web_socket_manager.py
index a0b2367..2574a9b 100644
--- a/sia/web_socket_manager.py
+++ b/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."""
diff --git a/web/src/components/App.jsx b/web/src/components/App.jsx
index bc4a283..1353877 100644
--- a/web/src/components/App.jsx
+++ b/web/src/components/App.jsx
@@ -42,7 +42,7 @@ const App = () => {
           setValidationError(message.validation_error);
           break;
         case ServerMessageType.OUTPUT_UPDATE:
-          setOutput(prev => prev + message.data);
+          setOutput(message.output);
           break;
       }
     });
@@ -72,6 +72,12 @@ const App = () => {
     }
   };
 
+  const handleClearOutput = () => {
+    sendMessage({
+      "type": ClientMessageType.CLEAR_OUTPUT
+    });
+  };
+
   useEffect(() => {
     if (agentState === AgentState.APPROVE_RESPONSE) {
       sendMessage({
@@ -151,7 +157,7 @@ const App = () => {
           onApprove={handleApprove}
           onToggleDiff={() => setShowDiff(!showDiff)}
           onSendInput={handleSendInput}
-          onClearOutput={() => setOutput('')}
+          onClearOutput={handleClearOutput}
         />
 
         <FloatingButton onClick={() => setShowSidebar(!showSidebar)} />
diff --git a/web/src/constants.jsx b/web/src/constants.jsx
index 3b2d676..d93c802 100644
--- a/web/src/constants.jsx
+++ b/web/src/constants.jsx
@@ -23,6 +23,7 @@ export const ClientMessageType = {
   MODIFY_RESPONSE: 'MODIFY_RESPONSE',
   APPROVE_RESPONSE: 'APPROVE_RESPONSE',
   SEND_INPUT: 'SEND_INPUT',
+  CLEAR_OUTPUT: 'CLEAR_OUTPUT',
 };
 
 export const ServerMessageType = {