No output from command tool

[zokahn]

Hi,

Thank you so much for this spectacular tool. After a while, the session can't continue to a lack of output from calling commands in:
python/tools/code_execution_tool.py

there is no specific logging or debugging built in. I am depending on what agent-zero can fix itself by looking at its code. Please help by identifying debug options or the root cause for this.
I will add to it when I find more. Checked it on main and development versions. It is keeping me from progressing in projects. Thanks!

[zokahn]

working workaround:

y--- a/python/tools/code_execution_tool.py
+++ b/python/tools/code_execution_tool.py
@@ -1,13 +1,12 @@
 from dataclasses import dataclass
-import os, json, contextlib, subprocess, ast, shlex
+import os, json, contextlib, subprocess, ast, shlex, asyncio
 from io import StringIO
 import time
 from typing import Literal
 from python.helpers import files, messages
 from agent import Agent
 from python.helpers.tool import Tool, Response
-from python.helpers import files
 from python.helpers.print_style import PrintStyle
 from python.helpers.shell_local import LocalInteractiveSession
 from python.helpers.shell_ssh import SSHInteractiveSession
 from python.helpers.docker import DockerContainerManager
@@ -15,33 +14,34 @@
 @dataclass
 class State:
     shell: LocalInteractiveSession | SSHInteractiveSession
     docker: DockerContainerManager | None
-        
 
 class CodeExecution(Tool):
 
-    def execute(self,**kwargs):
-
-        if self.agent.handle_intervention(): return Response(message="", break_loop=False)  # wait for intervention and handle it, if paused
+    def execute(self, **kwargs):
+        if self.agent.handle_intervention():
+            return Response(message="", break_loop=False)
         
         self.prepare_state()
-
-        # os.chdir(files.get_abs_path("./work_dir")) #change CWD to work_dir
         
         runtime = self.args["runtime"].lower().strip()
-        if runtime == "python":
-            response = self.execute_python_code(self.args["code"])
-        elif runtime == "nodejs":
-            response = self.execute_nodejs_code(self.args["code"])
-        elif runtime == "terminal":
-            response = self.execute_terminal_command(self.args["code"])
-        elif runtime == "output":
-            response = self.get_terminal_output()
-        else:
-            response = files.read_file("./prompts/fw.code_runtime_wrong.md", runtime=runtime)
+        try:
+            if runtime == "python":
+                response = self.execute_python_code(self.args["code"])
+            elif runtime == "nodejs":
+                response = self.execute_nodejs_code(self.args["code"])
+            elif runtime == "terminal":
+                response = self.execute_terminal_command(self.args["code"])
+            elif runtime == "output":
+                response = self.get_terminal_output()
+            else:
+                response = files.read_file("./prompts/fw.code_runtime_wrong.md", runtime=runtime)
+        except Exception as e:
+            response = f"Error during execution: {str(e)}"
 
-        if not response: response = files.read_file("./prompts/fw.code_no_output.md")
+        if not response:
+            response = files.read_file("./prompts/fw.code_no_output.md")
         return Response(message=response, break_loop=False)
 
     def after_execution(self, response, **kwargs):
         msg_response = files.read_file("./prompts/fw.tool_response.md", tool_name=self.name, tool_response=response.message)
@@ -49,57 +49,75 @@
 
     def prepare_state(self):
         self.state = self.agent.get_data("cot_state")
         if not self.state:
-
-            #initialize docker container if execution in docker is configured
             if self.agent.config.code_exec_docker_enabled:
                 docker = DockerContainerManager(name=self.agent.config.code_exec_docker_name, image=self.agent.config.code_exec_docker_image, ports=self.agent.config.code_exec_docker_ports, volumes=self.agent.config.code_exec_docker_volumes)
                 docker.start_container()
-            else: docker = None
+            else:
+                docker = None
 
-            #initialize local or remote interactive shell insterface
             if self.agent.config.code_exec_ssh_enabled:
-                shell = SSHInteractiveSession(self.agent.config.code_exec_ssh_addr,self.agent.config.code_exec_ssh_port,self.agent.config.code_exec_ssh_user,self.agent.config.code_exec_ssh_pass)
-            else: shell = LocalInteractiveSession()
+                shell = SSHInteractiveSession(self.agent.config.code_exec_ssh_addr, self.agent.config.code_exec_ssh_port, self.agent.config.code_exec_ssh_user, self.agent.config.code_exec_ssh_pass)
+            else:
+                shell = LocalInteractiveSession()
                 
-            self.state = State(shell=shell,docker=docker)
+            self.state = State(shell=shell, docker=docker)
             shell.connect()
         self.agent.set_data("cot_state", self.state)
     
+    def sanitize_command(self, command):
+        # Implement more robust command sanitization here
+        return shlex.quote(command)
+
     def execute_python_code(self, code):
-        escaped_code = shlex.quote(code)
-        command = f'python3 -c {escaped_code}'
+        sanitized_code = self.sanitize_command(code)
+        command = f'python3 -c {sanitized_code}'
         return self.terminal_session(command)
 
     def execute_nodejs_code(self, code):
-        escaped_code = shlex.quote(code)
-        command = f'node -e {escaped_code}'
+        sanitized_code = self.sanitize_command(code)
+        command = f'node -e {sanitized_code}'
         return self.terminal_session(command)
 
     def execute_terminal_command(self, command):
-        return self.terminal_session(command)
+        sanitized_command = self.sanitize_command(command)
+        # Redirect stderr to stdout
+        wrapped_command = f"{sanitized_command} 2>&1"
+        return self.terminal_session(wrapped_command)
 
-    def terminal_session(self, command):
+    async def terminal_session_async(self, command):
+        if self.agent.handle_intervention():
+            return ""
 
-        if self.agent.handle_intervention(): return ""  # wait for intervention and handle it, if paused
-       
         self.state.shell.send_command(command)
 
-        PrintStyle(background_color="white",font_color="#1B4F72",bold=True).print(f"{self.agent.agent_name} code execution output:")
-        return self.get_terminal_output()
+        PrintStyle(background_color="white", font_color="#1B4F72", bold=True).print(f"{self.agent.agent_name} code execution output:")
+        
+        full_output = ""
+        async for output in self.read_output_async():
+            if self.agent.handle_intervention():
+                return full_output
+            
+            PrintStyle(font_color="#85C1E9").stream(output)
+            full_output += output
+
+        return full_output.strip()
+
+    async def read_output_async(self):
+        timeout = 10  # 10 seconds timeout
+        start_time = time.time()
+        while True:
+            if time.time() - start_time > timeout:
+                yield "Command execution timed out."
+                break
+            output = await asyncio.to_thread(self.state.shell.read_output)
+            if not output[1]:  # If no partial output
+                if output[0]:  # If there's full output
+                    yield output[0]
+                break
+            yield output[1]
+            await asyncio.sleep(0.1)  # Short delay to prevent busy-waiting
+
+    def terminal_session(self, command):
+        return asyncio.run(self.terminal_session_async(command))
 
     def get_terminal_output(self):
-        idle=0
-        while True:       
-            time.sleep(0.1)  # Wait for some output to be generated
-            full_output, partial_output = self.state.shell.read_output()
-
-            if self.agent.handle_intervention(): return full_output  # wait for intervention and handle it, if paused
-        
-            if partial_output:
-                PrintStyle(font_color="#85C1E9").stream(partial_output)
-                idle=0    
-            else:
-                idle+=1
-                if ( full_output and idle > 30 ) or ( not full_output and idle > 100 ): return full_output