@samitouri / QOSami_Agents_Zero / commits / 97575a2d

local TTY implementation

frdel committed Aug 7, 2025 at 13:32 UTC 97575a2d8b1b318df759639f885e3079a8983fca
11 files changed +234 -219
agent.py
-12
@@ -221,18 +221,6 @@ class AgentConfig:
221 profile: str = ""
222 memory_subdir: str = ""
223 knowledge_subdirs: list[str] = field(default_factory=lambda: ["default", "custom"])
224 - code_exec_docker_enabled: bool = False
225 - code_exec_docker_name: str = "A0-dev"
226 - code_exec_docker_image: str = "agent0ai/agent-zero-run:development"
227 - code_exec_docker_ports: dict[str, int] = field(
228 - default_factory=lambda: {"22/tcp": 55022, "80/tcp": 55080}
229 - )
230 - code_exec_docker_volumes: dict[str, dict[str, str]] = field(
231 - default_factory=lambda: {
232 - files.get_base_dir(): {"bind": "/a0", "mode": "rw"},
233 - files.get_abs_path("work_dir"): {"bind": "/root", "mode": "rw"},
234 - }
235 - )
224 code_exec_ssh_enabled: bool = True
225 code_exec_ssh_addr: str = "localhost"
226 code_exec_ssh_port: int = 55022
docker/run/fs/exe/run_A0.sh
+2 -3
@@ -10,9 +10,8 @@ echo "Starting A0..."
10 exec python /a0/run_ui.py \
11 --dockerized=true \
12 --port=80 \
13 - --host="0.0.0.0" \
14 - --code_exec_docker_enabled=false \
15 - --code_exec_ssh_enabled=true \
13 + --host="0.0.0.0"
14 + # --code_exec_ssh_enabled=true \
15 # --code_exec_ssh_addr="localhost" \
16 # --code_exec_ssh_port=22 \
17 # --code_exec_ssh_user="root" \
initialize.py
+1 -29
@@ -79,19 +79,7 @@ def initialize_agent():
79 memory_subdir=current_settings["agent_memory_subdir"],
80 knowledge_subdirs=[current_settings["agent_knowledge_subdir"], "default"],
81 mcp_servers=current_settings["mcp_servers"],
82 - code_exec_docker_enabled=False,
83 - # code_exec_docker_name = "A0-dev",
84 - # code_exec_docker_image = "agent0ai/agent-zero:development",
85 - # code_exec_docker_ports = { "22/tcp": 55022, "80/tcp": 55080 }
86 - # code_exec_docker_volumes = {
87 - # files.get_base_dir(): {"bind": "/a0", "mode": "rw"},
88 - # files.get_abs_path("work_dir"): {"bind": "/root", "mode": "rw"},
89 - # },
90 - # code_exec_ssh_enabled = True,
91 - # code_exec_ssh_addr = "localhost",
92 - # code_exec_ssh_port = 55022,
93 - # code_exec_ssh_user = "root",
94 - # code_exec_ssh_pass = "",
82 + # code_exec params get initialized in _set_runtime_config
83 # additional = {},
84 )
85
@@ -176,19 +164,3 @@ def _set_runtime_config(config: AgentConfig, set: settings.Settings):
164 for key, value in ssh_conf.items():
165 if hasattr(config, key):
166 setattr(config, key, value)
179 -
180 - # if config.code_exec_docker_enabled:
181 - # config.code_exec_docker_ports["22/tcp"] = ssh_conf["code_exec_ssh_port"]
182 - # config.code_exec_docker_ports["80/tcp"] = ssh_conf["code_exec_http_port"]
183 - # config.code_exec_docker_name = f"{config.code_exec_docker_name}-{ssh_conf['code_exec_ssh_port']}-{ssh_conf['code_exec_http_port']}"
184 -
185 - # dman = docker.DockerContainerManager(
186 - # logger=log.Log(),
187 - # name=config.code_exec_docker_name,
188 - # image=config.code_exec_docker_image,
189 - # ports=config.code_exec_docker_ports,
190 - # volumes=config.code_exec_docker_volumes,
191 - # )
192 - # dman.start_container()
193 -
194 - # config.code_exec_ssh_pass = asyncio.run(rfc_exchange.get_root_password())
python/api/rfc.py
+4
@@ -8,6 +8,10 @@ class RFC(ApiHandler):
8 def requires_csrf(cls) -> bool:
9 return False
10
11 + @classmethod
12 + def requires_auth(cls) -> bool:
13 + return False
14 +
15 async def process(self, input: dict, request: Request) -> dict | Response:
16 result = await runtime.handle_rfc(input) # type: ignore
17 return result
python/helpers/settings.py
+58 -11
@@ -71,7 +71,6 @@ class Settings(TypedDict):
71 memory_memorize_consolidation: bool
72 memory_memorize_replace_threshold: float
73
74 -
74 api_keys: dict[str, str]
75
76 auth_login: str
@@ -113,7 +112,15 @@ class SettingsField(TypedDict, total=False):
112 title: str
113 description: str
114 type: Literal[
116 - "text", "number", "select", "range", "textarea", "password", "switch", "button", "html"
115 + "text",
116 + "number",
117 + "select",
118 + "range",
119 + "textarea",
120 + "password",
121 + "switch",
122 + "button",
123 + "html",
124 ]
125 value: Any
126 min: float
@@ -141,7 +148,6 @@ SETTINGS_FILE = files.get_abs_path("tmp/settings.json")
148 _settings: Settings | None = None
149
150
144 -
151 def convert_out(settings: Settings) -> SettingsOutput:
152 default_settings = get_default_settings()
153
@@ -478,7 +484,7 @@ def convert_out(settings: Settings) -> SettingsOutput:
484 "type": "number",
485 "value": settings["browser_model_rl_output"],
486 }
481 - )
487 + )
488
489 browser_model_fields.append(
490 {
@@ -498,7 +504,6 @@ def convert_out(settings: Settings) -> SettingsOutput:
504 "tab": "agent",
505 }
506
501 -
507 # basic auth section
508 auth_fields: list[SettingsField] = []
509
@@ -599,7 +604,8 @@ def convert_out(settings: Settings) -> SettingsOutput:
604 "value": settings["agent_profile"],
605 "options": [
606 {"value": subdir, "label": subdir}
602 - for subdir in files.get_subdirectories("agents") if subdir != "_example"
607 + for subdir in files.get_subdirectories("agents")
608 + if subdir != "_example"
609 ],
610 }
611 )
@@ -626,7 +632,6 @@ def convert_out(settings: Settings) -> SettingsOutput:
632 "tab": "agent",
633 }
634
629 -
635 memory_fields: list[SettingsField] = []
636
637 memory_fields.append(
@@ -856,6 +861,46 @@ def convert_out(settings: Settings) -> SettingsOutput:
861 "tab": "developer",
862 }
863
864 + # code_exec_fields: list[SettingsField] = []
865 +
866 + # code_exec_fields.append(
867 + # {
868 + # "id": "code_exec_ssh_enabled",
869 + # "title": "Use SSH for code execution",
870 + # "description": "Code execution will use SSH to connect to the terminal. When disabled, a local python terminal interface is used instead. SSH should only be used in development environment or when encountering issues with the local python terminal interface.",
871 + # "type": "switch",
872 + # "value": settings["code_exec_ssh_enabled"],
873 + # }
874 + # )
875 +
876 + # code_exec_fields.append(
877 + # {
878 + # "id": "code_exec_ssh_addr",
879 + # "title": "Code execution SSH address",
880 + # "description": "Address of the SSH server for code execution. Only applies when SSH is enabled.",
881 + # "type": "text",
882 + # "value": settings["code_exec_ssh_addr"],
883 + # }
884 + # )
885 +
886 + # code_exec_fields.append(
887 + # {
888 + # "id": "code_exec_ssh_port",
889 + # "title": "Code execution SSH port",
890 + # "description": "Port of the SSH server for code execution. Only applies when SSH is enabled.",
891 + # "type": "text",
892 + # "value": settings["code_exec_ssh_port"],
893 + # }
894 + # )
895 +
896 + # code_exec_section: SettingsSection = {
897 + # "id": "code_exec",
898 + # "title": "Code execution",
899 + # "description": "Configuration of code execution by the agent.",
900 + # "fields": code_exec_fields,
901 + # "tab": "developer",
902 + # }
903 +
904 # Speech to text section
905 stt_fields: list[SettingsField] = []
906
@@ -1085,6 +1130,7 @@ def convert_out(settings: Settings) -> SettingsOutput:
1130 mcp_server_section,
1131 backup_section,
1132 dev_section,
1133 + # code_exec_section,
1134 ]
1135 }
1136 return result
@@ -1147,7 +1193,7 @@ def normalize_settings(settings: Settings) -> Settings:
1193 # adjust settings values to match current version if needed
1194 if "version" not in copy or copy["version"] != default["version"]:
1195 _adjust_to_version(copy, default)
1150 - copy["version"] = default["version"] # sync version
1196 + copy["version"] = default["version"] # sync version
1197
1198 # remove keys that are not in default
1199 keys_to_remove = [key for key in copy if key not in default]
@@ -1162,7 +1208,7 @@ def normalize_settings(settings: Settings) -> Settings:
1208 try:
1209 copy[key] = type(value)(copy[key]) # type: ignore
1210 if isinstance(copy[key], str):
1165 - copy[key] = copy[key].strip() # strip strings
1211 + copy[key] = copy[key].strip() # strip strings
1212 except (ValueError, TypeError):
1213 copy[key] = value # make default instead
1214
@@ -1179,6 +1225,7 @@ def _adjust_to_version(settings: Settings, default: Settings):
1225 if "agent_profile" not in settings or settings["agent_profile"] == "default":
1226 settings["agent_profile"] = "agent0"
1227
1228 +
1229 def _read_settings_file() -> Settings | None:
1230 if os.path.exists(SETTINGS_FILE):
1231 content = files.read_file(SETTINGS_FILE)
@@ -1441,9 +1488,9 @@ def set_root_password(password: str):
1488 def get_runtime_config(set: Settings):
1489 if runtime.is_dockerized():
1490 return {
1491 + "code_exec_ssh_enabled": False,
1492 "code_exec_ssh_addr": "localhost",
1493 "code_exec_ssh_port": 22,
1446 - "code_exec_http_port": 80,
1494 "code_exec_ssh_user": "root",
1495 }
1496 else:
@@ -1455,9 +1502,9 @@ def get_runtime_config(set: Settings):
1502 if host.endswith("/"):
1503 host = host[:-1]
1504 return {
1505 + "code_exec_ssh_enabled": True,
1506 "code_exec_ssh_addr": host,
1507 "code_exec_ssh_port": set["rfc_port_ssh"],
1460 - "code_exec_http_port": set["rfc_port_http"],
1508 "code_exec_ssh_user": "root",
1509 }
1510
python/helpers/shell_local.py
+27 -49
@@ -3,70 +3,48 @@ import subprocess
3 import time
4 import sys
5 from typing import Optional, Tuple
6 +from python.helpers import tty_session
7 +from python.helpers.shell_ssh import clean_string
8
9 class LocalInteractiveSession:
10 def __init__(self):
9 - self.process = None
11 + self.session: tty_session.TTYSession|None = None
12 self.full_output = ''
13
14 async def connect(self):
13 - # Start a new subprocess with the appropriate shell for the OS
14 - if sys.platform.startswith('win'):
15 - # Windows
16 - self.process = subprocess.Popen(
17 - ['cmd.exe'],
18 - stdin=subprocess.PIPE,
19 - stdout=subprocess.PIPE,
20 - stderr=subprocess.PIPE,
21 - text=True,
22 - bufsize=1
23 - )
24 - else:
25 - # macOS and Linux
26 - self.process = subprocess.Popen(
27 - ['/bin/bash'],
28 - stdin=subprocess.PIPE,
29 - stdout=subprocess.PIPE,
30 - stderr=subprocess.PIPE,
31 - text=True,
32 - bufsize=1
33 - )
15 + self.session = tty_session.TTYSession("/bin/bash")
16 + await self.session.start()
17 + await self.session.read_full_until_idle(idle_timeout=1, total_timeout=1)
18
35 - def close(self):
36 - if self.process:
37 - self.process.terminate()
38 - self.process.wait()
19 + async def close(self):
20 + if self.session:
21 + self.session.kill()
22 + # self.session.wait()
23
40 - def send_command(self, command: str):
41 - if not self.process:
24 + async def send_command(self, command: str):
25 + if not self.session:
26 raise Exception("Shell not connected")
27 self.full_output = ""
44 - self.process.stdin.write(command + '\n') # type: ignore
45 - self.process.stdin.flush() # type: ignore
28 + await self.session.sendline(command)
29
30 async def read_output(self, timeout: float = 0, reset_full_output: bool = False) -> Tuple[str, Optional[str]]:
48 - if not self.process:
31 + if not self.session:
32 raise Exception("Shell not connected")
33
34 if reset_full_output:
35 self.full_output = ""
53 - partial_output = ''
54 - start_time = time.time()
55 -
56 - while (timeout <= 0 or time.time() - start_time < timeout):
57 - rlist, _, _ = select.select([self.process.stdout], [], [], 0.1)
58 - if rlist:
59 - line = self.process.stdout.readline() # type: ignore
60 - if line:
61 - partial_output += line
62 - self.full_output += line
63 - time.sleep(0.1)
64 - else:
65 - break # No more output
66 - else:
67 - break # No data available
36 +
37 + # get output from terminal
38 + partial_output = await self.session.read_full_until_idle(idle_timeout=0.01, total_timeout=timeout)
39 + self.full_output += partial_output
40 +
41 + # clean output
42 + partial_output = clean_string(partial_output)
43 + clean_full_output = clean_string(self.full_output)
44 +
45 + print("\n\n"+self.full_output.encode("unicode_escape").decode("ascii")+"\n\n")
46 +
47
48 if not partial_output:
70 - return self.full_output, None
71 -
72 - return self.full_output, partial_output
\ No newline at end of file
49 + return clean_full_output, None
50 + return clean_full_output, partial_output
\ No newline at end of file
python/helpers/shell_ssh.py
+52 -33
@@ -28,11 +28,20 @@ class SSHInteractiveSession:
28 self.last_command = b""
29 self.trimmed_command_length = 0 # Initialize trimmed_command_length
30
31 - async def connect(self):
32 - # try 3 times with wait and then except
31 + async def connect(self, keepalive_interval: int = 5):
32 + """
33 + Establish the SSH connection and start an interactive shell.
34 +
35 + Parameters
36 + ----------
37 + keepalive_interval : int
38 + Interval in **seconds** between keep-alive packets sent by Paramiko.
39 + A value ≤ 0 disables Paramiko’s keep-alive feature.
40 + """
41 errors = 0
42 while True:
43 try:
44 + # --- establish TCP/SSH session ---------------------------------
45 self.client.connect(
46 self.hostname,
47 self.port,
@@ -41,16 +50,25 @@ class SSHInteractiveSession:
50 allow_agent=False,
51 look_for_keys=False,
52 )
53 +
54 + # --------- NEW: enable transport-level keep-alives -------------
55 + transport = self.client.get_transport()
56 + if transport and keepalive_interval > 0:
57 + # sends an SSH_MSG_IGNORE every <keepalive_interval> seconds
58 + transport.set_keepalive(keepalive_interval)
59 + # ----------------------------------------------------------------
60 +
61 + # invoke interactive shell
62 self.shell = self.client.invoke_shell(width=100, height=50)
45 - # self.shell.send(f'PS1="{SSHInteractiveSession.ps1_label}"'.encode())
46 - # return
47 - self.shell.send("stty -echo\n".encode()) # disable shell echo
63 + self.shell.send("stty -echo\n".encode()) # disable local echo
64
49 - while True: # wait for end of initial output
65 + # wait for initial prompt/output to settle
66 + while True:
67 full, part = await self.read_output()
68 if full and not part:
69 return
70 time.sleep(0.1)
71 +
72 except Exception as e:
73 errors += 1
74 if errors < 3:
@@ -60,18 +78,17 @@ class SSHInteractiveSession:
78 content=f"SSH Connection attempt {errors}...",
79 temp=True,
80 )
63 -
81 time.sleep(5)
82 else:
83 raise e
84
68 - def close(self):
85 + async def close(self):
86 if self.shell:
87 self.shell.close()
88 if self.client:
89 self.client.close()
90
74 - def send_command(self, command: str):
91 + async def send_command(self, command: str):
92 if not self.shell:
93 raise Exception("Shell not connected")
94 self.full_output = b""
@@ -137,8 +154,8 @@ class SSHInteractiveSession:
154 decoded_partial_output = partial_output.decode("utf-8", errors="replace")
155 decoded_full_output = self.full_output.decode("utf-8", errors="replace")
156
140 - decoded_partial_output = self.clean_string(decoded_partial_output)
141 - decoded_full_output = self.clean_string(decoded_full_output)
157 + decoded_partial_output = clean_string(decoded_partial_output)
158 + decoded_full_output = clean_string(decoded_full_output)
159
160 return decoded_full_output, decoded_partial_output
161
@@ -190,32 +207,34 @@ class SSHInteractiveSession:
207
208 return data
209
193 - def clean_string(self, input_string):
194 - # Remove ANSI escape codes
195 - ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
196 - cleaned = ansi_escape.sub("", input_string)
210 +def clean_string(input_string):
211 + # Remove ANSI escape codes
212 + ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
213 + cleaned = ansi_escape.sub("", input_string)
214
198 - # remove null bytes
199 - cleaned = cleaned.replace("\x00", "")
215 + # remove null bytes
216 + cleaned = cleaned.replace("\x00", "")
217
201 - # remove ipython \r\r\n> sequences from the start
202 - cleaned = re.sub(r'^[ \r]*(?:\r*\n>[ \r]*)*', '', cleaned)
218 + # remove ipython \r\r\n> sequences from the start
219 + cleaned = re.sub(r'^[ \r]*(?:\r*\n>[ \r]*)*', '', cleaned)
220 + # also remove any amount of '> ' sequences from the start
221 + cleaned = re.sub(r'^(>\s*)+', '', cleaned)
222
204 - # Replace '\r\n' with '\n'
205 - cleaned = cleaned.replace("\r\n", "\n")
223 + # Replace '\r\n' with '\n'
224 + cleaned = cleaned.replace("\r\n", "\n")
225
207 - # remove leading \r and spaces
208 - cleaned = cleaned.lstrip("\r ")
226 + # remove leading \r and spaces
227 + cleaned = cleaned.lstrip("\r ")
228
210 - # Split the string by newline characters to process each segment separately
211 - lines = cleaned.split("\n")
229 + # Split the string by newline characters to process each segment separately
230 + lines = cleaned.split("\n")
231
213 - for i in range(len(lines)):
214 - # Handle carriage returns '\r' by splitting and taking the last part
215 - parts = [part for part in lines[i].split("\r") if part.strip()]
216 - if parts:
217 - lines[i] = parts[
218 - -1
219 - ].rstrip() # Overwrite with the last part after the last '\r'
232 + for i in range(len(lines)):
233 + # Handle carriage returns '\r' by splitting and taking the last part
234 + parts = [part for part in lines[i].split("\r") if part.strip()]
235 + if parts:
236 + lines[i] = parts[
237 + -1
238 + ].rstrip() # Overwrite with the last part after the last '\r'
239
221 - return "\n".join(lines)
240 + return "\n".join(lines)
python/helpers/tty_session.py
+28 -4
@@ -15,9 +15,7 @@ sys.stdout.reconfigure(errors="replace") # type: ignore
15
16
17 class TTYSession:
18 - def __init__(
19 - self, cmd, *, cwd=None, env=None, encoding="utf-8", echo=False
20 - ): # ← NEW kw-arg `echo`
18 + def __init__(self, cmd, *, cwd=None, env=None, encoding="utf-8", echo=False):
19 self.cmd = cmd if isinstance(cmd, str) else " ".join(cmd)
20 self.cwd = cwd
21 self.env = env or os.environ.copy()
@@ -26,6 +24,17 @@ class TTYSession:
24 self._proc = None
25 self._buf = asyncio.Queue()
26
27 + def __del__(self):
28 + # Simple cleanup on object destruction
29 + import nest_asyncio
30 +
31 + nest_asyncio.apply()
32 + if hasattr(self, "close"):
33 + try:
34 + asyncio.run(self.close())
35 + except Exception:
36 + pass
37 +
38 # ── user-facing coroutines ────────────────────────────────────────
39 async def start(self):
40 if _IS_WIN:
@@ -36,7 +45,22 @@ class TTYSession:
45 self._proc = await _spawn_posix_pty(
46 self.cmd, self.cwd, self.env, self.echo
47 ) # ← pass echo
39 - asyncio.create_task(self._pump_stdout())
48 + self._pump_task = asyncio.create_task(self._pump_stdout())
49 +
50 + async def close(self):
51 + # Cancel the pump task if it exists
52 + if hasattr(self, "_pump_task") and self._pump_task:
53 + self._pump_task.cancel()
54 + try:
55 + await self._pump_task
56 + except asyncio.CancelledError:
57 + pass
58 + # Terminate the process if it exists
59 + if self._proc:
60 + self._proc.terminate()
61 + await self._proc.wait()
62 + self._proc = None
63 + self._pump_task = None
64
65 async def send(self, data: str | bytes):
66 if self._proc is None:
python/tools/code_execution_tool.py
+56 -78
@@ -16,7 +16,6 @@ import re
16 @dataclass
17 class State:
18 shells: dict[int, LocalInteractiveSession | SSHInteractiveSession]
19 - docker: DockerContainerManager | None
19
20
21 class CodeExecution(Tool):
@@ -25,10 +24,6 @@ class CodeExecution(Tool):
24
25 await self.agent.handle_intervention() # wait for intervention and handle it, if paused
26
28 - await self.prepare_state()
29 -
30 - # os.chdir(files.get_abs_path("./work_dir")) #change CWD to work_dir
31 -
27 runtime = self.args.get("runtime", "").lower().strip()
28 session = int(self.args.get("session", 0))
29
@@ -80,59 +75,48 @@ class CodeExecution(Tool):
75 async def after_execution(self, response, **kwargs):
76 self.agent.hist_add_tool_result(self.name, response.message)
77
83 - async def prepare_state(self, reset=False, session=None):
84 - self.state = self.agent.get_data("_cet_state")
85 - if not self.state or reset:
86 -
87 - # initialize docker container if execution in docker is configured
88 - if not self.state and self.agent.config.code_exec_docker_enabled:
89 - docker = DockerContainerManager(
90 - logger=self.agent.context.log,
91 - name=self.agent.config.code_exec_docker_name,
92 - image=self.agent.config.code_exec_docker_image,
93 - ports=self.agent.config.code_exec_docker_ports,
94 - volumes=self.agent.config.code_exec_docker_volumes,
78 + async def prepare_state(self, reset=False, session: int | None = None):
79 + self.state: State | None = self.agent.get_data("_cet_state")
80 + if not self.state:
81 + # initialize shells dictionary if not exists
82 + shells: dict[int, LocalInteractiveSession | SSHInteractiveSession] = {}
83 + else:
84 + shells = self.state.shells.copy()
85 +
86 + # Only reset the specified session if provided
87 + if reset and session is not None and session in shells:
88 + await shells[session].close()
89 + del shells[session]
90 + elif reset and not session:
91 + # Close all sessions if full reset requested
92 + for s in list(shells.keys()):
93 + await shells[s].close()
94 + shells = {}
95 +
96 + # initialize local or remote interactive shell interface for session 0 if needed
97 + if session is not None and session not in shells:
98 + if self.agent.config.code_exec_ssh_enabled:
99 + pswd = (
100 + self.agent.config.code_exec_ssh_pass
101 + if self.agent.config.code_exec_ssh_pass
102 + else await rfc_exchange.get_root_password()
103 + )
104 + shell = SSHInteractiveSession(
105 + self.agent.context.log,
106 + self.agent.config.code_exec_ssh_addr,
107 + self.agent.config.code_exec_ssh_port,
108 + self.agent.config.code_exec_ssh_user,
109 + pswd,
110 )
96 - docker.start_container()
111 else:
98 - docker = self.state.docker if self.state else None
99 -
100 - # initialize shells dictionary if not exists
101 - shells = {} if not self.state else self.state.shells.copy()
102 -
103 - # Only reset the specified session if provided
104 - if session is not None and session in shells:
105 - shells[session].close()
106 - del shells[session]
107 - elif reset and not session:
108 - # Close all sessions if full reset requested
109 - for s in list(shells.keys()):
110 - shells[s].close()
111 - shells = {}
112 -
113 - # initialize local or remote interactive shell interface for session 0 if needed
114 - if 0 not in shells:
115 - if self.agent.config.code_exec_ssh_enabled:
116 - pswd = (
117 - self.agent.config.code_exec_ssh_pass
118 - if self.agent.config.code_exec_ssh_pass
119 - else await rfc_exchange.get_root_password()
120 - )
121 - shell = SSHInteractiveSession(
122 - self.agent.context.log,
123 - self.agent.config.code_exec_ssh_addr,
124 - self.agent.config.code_exec_ssh_port,
125 - self.agent.config.code_exec_ssh_user,
126 - pswd,
127 - )
128 - else:
129 - shell = LocalInteractiveSession()
112 + shell = LocalInteractiveSession()
113
131 - shells[0] = shell
132 - await shell.connect()
114 + shells[session] = shell
115 + await shell.connect()
116
134 - self.state = State(shells=shells, docker=docker)
117 + self.state = State(shells=shells)
118 self.agent.set_data("_cet_state", self.state)
119 + return self.state
120
121 async def execute_python_code(self, session: int, code: str, reset: bool = False):
122 escaped_code = shlex.quote(code)
@@ -156,45 +140,35 @@ class CodeExecution(Tool):
140 self, session: int, command: str, reset: bool = False, prefix: str = ""
141 ):
142
143 + self.state = await self.prepare_state(reset=reset, session=session)
144 +
145 await self.agent.handle_intervention() # wait for intervention and handle it, if paused
146 # try again on lost connection
147 for i in range(2):
148 try:
149
164 - if reset:
165 - await self.reset_terminal()
166 -
167 - if session not in self.state.shells:
168 - if self.agent.config.code_exec_ssh_enabled:
169 - pswd = (
170 - self.agent.config.code_exec_ssh_pass
171 - if self.agent.config.code_exec_ssh_pass
172 - else await rfc_exchange.get_root_password()
173 - )
174 - shell = SSHInteractiveSession(
175 - self.agent.context.log,
176 - self.agent.config.code_exec_ssh_addr,
177 - self.agent.config.code_exec_ssh_port,
178 - self.agent.config.code_exec_ssh_user,
179 - pswd,
180 - )
181 - else:
182 - shell = LocalInteractiveSession()
183 - self.state.shells[session] = shell
184 - await shell.connect()
185 -
186 - self.state.shells[session].send_command(command)
150 + await self.state.shells[session].send_command(command)
151 +
152 + locl = (
153 + " (local)"
154 + if isinstance(self.state.shells[session], LocalInteractiveSession)
155 + else (
156 + " (remote)"
157 + if isinstance(self.state.shells[session], SSHInteractiveSession)
158 + else " (unknown)"
159 + )
160 + )
161
162 PrintStyle(
163 background_color="white", font_color="#1B4F72", bold=True
190 - ).print(f"{self.agent.agent_name} code execution output")
164 + ).print(f"{self.agent.agent_name} code execution output{locl}")
165 return await self.get_terminal_output(session=session, prefix=prefix)
166
167 except Exception as e:
168 if i == 1:
169 # try again on lost connection
170 PrintStyle.error(str(e))
197 - await self.prepare_state(reset=True)
171 + await self.prepare_state(reset=True, session=session)
172 continue
173 else:
174 raise e
@@ -221,6 +195,10 @@ class CodeExecution(Tool):
195 sleep_time=0.1,
196 prefix="",
197 ):
198 +
199 + # if not self.state:
200 + self.state = await self.prepare_state(session=session)
201 +
202 # Common shell prompt regex patterns (add more as needed)
203 prompt_patterns = [
204 re.compile(r"\\(venv\\).+[$#] ?$"), # (venv) ...$ or (venv) ...#
webui/components/messages/action-buttons/simple-action-buttons.js
+5
@@ -9,6 +9,11 @@ function getTextContent(element) {
9 for (const child of element.children) {
10 // Skip action buttons
11 if (child.classList.contains("action-buttons")) continue;
12 + // If the child is an image, copy its src URL
13 + if (child.tagName && child.tagName.toLowerCase() === "img") {
14 + if (child.src) textParts.push(child.src);
15 + continue;
16 + }
17 // Get text content from the child
18 const text = child.innerText || "";
19 if (text.trim()) {
webui/public/code_exec.svg new
+1
@@ -0,0 +1 @@
1 +<svg xmlns="http://www.w3.org/2000/svg" height="24px" viewBox="0 -960 960 960" width="24px" fill="#e8eaed"><path d="M192-212q-26 0-43-17t-17-43v-416q0-26 17-43t43-17h576q26 0 43 17t17 43v416q0 26-17 43t-43 17H192Zm0-28h576q12 0 22-10t10-22v-368H160v368q0 12 10 22t22 10Zm108-77-19-19 103-104-104-104 20-19 123 123-123 123Zm206 11v-28h188v28H506Z"/></svg>
\ No newline at end of file