f6acef23 · samitouri/QOSami_Agents_Zero-2

@@ -39,25 +39,26 @@ class CodeExecution(Tool):

39 self.allow_running = bool(self.args.get("allow_running", False))

40 reset = bool(self.args.get("reset", False) or runtime_arg == "reset")

41

42 + cfg = _get_config(self.agent)

43 +

44 if runtime_arg == "python":

45 response = await self.execute_python_code(

44 - code=self.args["code"], session=session, reset=reset

46 + cfg, code=self.args["code"], session=session, reset=reset

47 )

48 elif runtime_arg == "nodejs":

49 response = await self.execute_nodejs_code(

48 - code=self.args["code"], session=session, reset=reset

50 + cfg, code=self.args["code"], session=session, reset=reset

51 )

52 elif runtime_arg == "terminal":

53 response = await self.execute_terminal_command(

52 - command=self.args["code"], session=session, reset=reset

54 + cfg, command=self.args["code"], session=session, reset=reset

55 )

56 elif runtime_arg == "output":

55 - cfg = _get_config(self.agent)

57 response = await self.get_terminal_output(

57 - session=session, timeouts=cfg["output_timeouts"]

58 + cfg, session=session, timeouts=cfg["output_timeouts"]

59 )

60 elif runtime_arg == "reset":

60 - response = await self.reset_terminal(session=session)

61 + response = await self.reset_terminal(cfg, session=session)

62 else:

63 response = self.agent.read_prompt(

64 "fw.code.runtime_wrong.md", runtime=runtime_arg

@@ -87,9 +88,8 @@ class CodeExecution(Tool):

88 async def after_execution(self, response, **kwargs):

89 self.agent.hist_add_tool_result(self.name, response.message, **(response.additional or {}))

90

90 - async def prepare_state(self, reset=False, session: int | None = None):

91 + async def prepare_state(self, cfg: dict, reset=False, session: int | None = None):

92 self.state: State | None = self.agent.get_data("_cet_state")

92 - cfg = _get_config(self.agent)

93 ssh_enabled = cfg["ssh_enabled"]

94

95 # always reset state when ssh_enabled changes

@@ -131,43 +131,40 @@ class CodeExecution(Tool):

131 self.agent.set_data("_cet_state", self.state)

132 return self.state

133

134 - async def execute_python_code(self, session: int, code: str, reset: bool = False):

134 + async def execute_python_code(self, cfg: dict, session: int, code: str, reset: bool = False):

135 escaped_code = shlex.quote(code)

136 command = f"ipython -c {escaped_code}"

137 prefix = "python> " + self.format_command_for_output(code) + "\n\n"

138 - return await self.terminal_session(session, command, reset, prefix)

138 + return await self.terminal_session(cfg, session, command, reset, prefix)

139

140 - async def execute_nodejs_code(self, session: int, code: str, reset: bool = False):

140 + async def execute_nodejs_code(self, cfg: dict, session: int, code: str, reset: bool = False):

141 escaped_code = shlex.quote(code)

142 command = f"node /exe/node_eval.js {escaped_code}"

143 prefix = "node> " + self.format_command_for_output(code) + "\n\n"

144 - return await self.terminal_session(session, command, reset, prefix)

144 + return await self.terminal_session(cfg, session, command, reset, prefix)

145

146 async def execute_terminal_command(

147 - self, session: int, command: str, reset: bool = False

147 + self, cfg: dict, session: int, command: str, reset: bool = False

148 ):

149 - cfg = _get_config(self.agent)

149 prefix = (

150 ("bash>" if not runtime.is_windows() or cfg["ssh_enabled"] else "PS>")

151 + self.format_command_for_output(command)

152 + "\n\n"

153 )

155 - return await self.terminal_session(session, command, reset, prefix)

154 + return await self.terminal_session(cfg, session, command, reset, prefix)

155

156 async def terminal_session(

158 - self, session: int, command: str, reset: bool = False, prefix: str = "", timeouts: dict | None = None

157 + self, cfg: dict, session: int, command: str, reset: bool = False, prefix: str = "", timeouts: dict | None = None

158 ):

160 - self.state = await self.prepare_state(reset=reset, session=session)

159 + self.state = await self.prepare_state(cfg, reset=reset, session=session)

160

161 await self.agent.handle_intervention() # wait for intervention and handle it, if paused

162

163 # Check if session is running and handle it

164 if not self.allow_running:

166 - if response := await self.handle_running_session(session):

165 + if response := await self.handle_running_session(cfg, session):

166 return response

167

169 - cfg = _get_config(self.agent)

170 -

168 # try again on lost connection

169 for i in range(2):

170 try:

@@ -188,6 +185,7 @@ class CodeExecution(Tool):

185 background_color="white", font_color="#1B4F72", bold=True

186 ).print(f"{self.agent.agent_name} code execution output{locl}")

187 return await self.get_terminal_output(

188 + cfg,

189 session=session,

190 prefix=prefix,

191 timeouts=(timeouts or cfg["code_exec_timeouts"]),

@@ -196,7 +194,7 @@ class CodeExecution(Tool):

194 except Exception as e:

195 if i == 1:

196 PrintStyle.error(str(e))

199 - await self.prepare_state(reset=True, session=session)

197 + await self.prepare_state(cfg, reset=True, session=session)

198 continue

199 else:

200 raise e

@@ -209,6 +207,7 @@ class CodeExecution(Tool):

207

208 async def get_terminal_output(

209 self,

210 + cfg: dict,

211 session=0,

212 reset_full_output=True,

213 first_output_timeout=30,

@@ -219,7 +218,7 @@ class CodeExecution(Tool):

218 prefix="",

219 timeouts: dict | None = None,

220 ):

222 - self.state = await self.prepare_state(session=session)

221 + self.state = await self.prepare_state(cfg, session=session)

222

223 # Override timeouts if a dict is provided

224 if timeouts:

@@ -228,7 +227,6 @@ class CodeExecution(Tool):

227 dialog_timeout = timeouts.get("dialog_timeout", dialog_timeout)

228 max_exec_timeout = timeouts.get("max_exec_timeout", max_exec_timeout)

229

231 - cfg = _get_config(self.agent)

230 prompt_patterns = cfg["prompt_patterns"]

231 dialog_patterns = cfg["dialog_patterns"]

232

@@ -348,6 +346,7 @@ class CodeExecution(Tool):

346

347 async def handle_running_session(

348 self,

349 + cfg: dict,

350 session=0,

351 reset_full_output=True,

352 prefix=""

@@ -357,7 +356,6 @@ class CodeExecution(Tool):

356 if not self.state.shells[session].running:

357 return None

358

360 - cfg = _get_config(self.agent)

359 prompt_patterns = cfg["prompt_patterns"]

360 dialog_patterns = cfg["dialog_patterns"]

361

@@ -372,7 +370,7 @@ class CodeExecution(Tool):

370 truncated_output.splitlines()[-3:] if truncated_output else []

371 )

372 last_lines.reverse()

375 - for _, line in enumerate(last_lines):

373 + for line in last_lines:

374 for pat in prompt_patterns:

375 if pat.search(line.strip()):

376 PrintStyle.info(

@@ -406,7 +404,7 @@ class CodeExecution(Tool):

404 if self.state and session in self.state.shells:

405 self.state.shells[session].running = False

406

409 - async def reset_terminal(self, session=0, reason: str | None = None):

407 + async def reset_terminal(self, cfg: dict, session=0, reason: str | None = None):

408 if reason:

409 PrintStyle(font_color="#FFA500", bold=True).print(

410 f"Resetting terminal session {session}... Reason: {reason}"

@@ -416,7 +414,7 @@ class CodeExecution(Tool):

414 f"Resetting terminal session {session}..."

415 )

416

419 - await self.prepare_state(reset=True, session=session)

417 + await self.prepare_state(cfg, reset=True, session=session)

418 response = self.agent.read_prompt(

419 "fw.code.info.md", info=self.agent.read_prompt("fw.code.reset.md")

420 )

@@ -464,7 +462,6 @@ class CodeExecution(Tool):

462 # ------------------------------------------------------------------

463

464 def _resolve_ssh_enabled(raw_value) -> bool:

467 - """Resolve ssh_enabled: 'auto' detects based on dockerized state."""

465 val = str(raw_value).strip().lower()

466 if val == "auto":

467 return not runtime.is_dockerized()

@@ -472,12 +469,10 @@ def _resolve_ssh_enabled(raw_value) -> bool:

469

470

471 def _resolve_ssh_addr(cfg_addr: str) -> str:

475 - """Resolve SSH address: fall back to rfc_url from settings when empty."""

472 if cfg_addr:

473 return cfg_addr

474 set = settings.get_settings()

475 host = set.get("rfc_url", "localhost")

480 - # Strip protocol and port from URL

476 if "//" in host:

477 host = host.split("//")[1]

478 if ":" in host:

@@ -488,83 +483,39 @@ def _resolve_ssh_addr(cfg_addr: str) -> str:

483

484

485 async def _resolve_ssh_pass(cfg_pass: str) -> str:

491 - """Resolve SSH password: fall back to root_password via RFC when empty."""

486 if cfg_pass:

487 return cfg_pass

488 return await rfc_exchange.get_root_password()

489

490

497 -def _get_config(agent) -> dict:

498 - cfg = plugins.get_plugin_config("code_execution", agent=agent) or {}

491 +def _parse_patterns(raw, flags=0) -> list[re.Pattern]:

492 + lines = [str(p) for p in raw] if isinstance(raw, list) else str(raw).splitlines()

493 + return [re.compile(p.strip(), flags) for p in lines if p.strip()]

494 +

495 +

496 +_TIMEOUT_KEYS = ("first_output_timeout", "between_output_timeout", "max_exec_timeout", "dialog_timeout")

497

500 - # SSH / TTY switch (supports auto/true/false)

501 - ssh_enabled = _resolve_ssh_enabled(cfg.get("ssh_enabled", "auto"))

502 -

503 - # SSH credentials (with RFC fallbacks)

504 - ssh_addr = _resolve_ssh_addr(str(cfg.get("ssh_addr", "")))

505 - ssh_port = int(cfg.get("ssh_port", 55022))

506 - ssh_user = str(cfg.get("ssh_user", "root"))

507 - ssh_pass = str(cfg.get("ssh_pass", ""))

508 -

509 - # Timeouts for python/nodejs/terminal runtimes

510 - code_exec_timeouts = {

511 - "first_output_timeout": int(cfg.get("code_exec_first_output_timeout", 30)),

512 - "between_output_timeout": int(cfg.get("code_exec_between_output_timeout", 15)),

513 - "max_exec_timeout": int(cfg.get("code_exec_max_exec_timeout", 180)),

514 - "dialog_timeout": int(cfg.get("code_exec_dialog_timeout", 5)),

515 - }

498

517 - # Timeouts for "output" runtime

518 - output_timeouts = {

519 - "first_output_timeout": int(cfg.get("output_first_output_timeout", 90)),

520 - "between_output_timeout": int(cfg.get("output_between_output_timeout", 45)),

521 - "max_exec_timeout": int(cfg.get("output_max_exec_timeout", 300)),

522 - "dialog_timeout": int(cfg.get("output_dialog_timeout", 5)),

499 +def _parse_timeouts(cfg: dict, prefix: str, defaults: tuple[int, ...]) -> dict:

500 + return {

501 + key: int(cfg.get(f"{prefix}_{key}", default))

502 + for key, default in zip(_TIMEOUT_KEYS, defaults)

503 }

504

525 - # Prompt patterns (one regex per line, or a list)

526 - prompt_patterns_raw = cfg.get(

527 - "prompt_patterns",

528 - r"($venv$).+[$#] ?$" + "\n"

529 - + r"root@[^:]+:[^#]+# ?$" + "\n"

530 - + r"[a-zA-Z0-9_.-]+@[^:]+:[^$#]+[$#] ?$" + "\n"

531 - + r"$?.*$?\s*PS\s+[^>]+> ?$",

532 - )

533 - if isinstance(prompt_patterns_raw, list):

534 - prompt_lines = [str(p) for p in prompt_patterns_raw]

535 - else:

536 - prompt_lines = str(prompt_patterns_raw).splitlines()

537 - prompt_patterns = [

538 - re.compile(p.strip())

539 - for p in prompt_lines

540 - if p.strip()

541 - ]

542 -

543 - # Dialog patterns (one regex per line, or a list)

544 - dialog_patterns_raw = cfg.get(

545 - "dialog_patterns",

546 - "Y/N\nyes/no\n:\\s*$\n\\?\\s*$",

547 - )

548 - if isinstance(dialog_patterns_raw, list):

549 - dialog_lines = [str(p) for p in dialog_patterns_raw]

550 - else:

551 - dialog_lines = str(dialog_patterns_raw).splitlines()

552 - dialog_patterns = [

553 - re.compile(p.strip(), re.IGNORECASE)

554 - for p in dialog_lines

555 - if p.strip()

556 - ]

505 +

506 +def _get_config(agent) -> dict:

507 + cfg = plugins.get_plugin_config("code_execution", agent=agent) or {}

508

509 return {

559 - "ssh_enabled": ssh_enabled,

560 - "ssh_addr": ssh_addr,

561 - "ssh_port": ssh_port,

562 - "ssh_user": ssh_user,

563 - "ssh_pass": ssh_pass,

564 - "code_exec_timeouts": code_exec_timeouts,

565 - "output_timeouts": output_timeouts,

566 - "prompt_patterns": prompt_patterns,

567 - "dialog_patterns": dialog_patterns,

510 + "ssh_enabled": _resolve_ssh_enabled(cfg.get("ssh_enabled", "auto")),

511 + "ssh_addr": _resolve_ssh_addr(str(cfg.get("ssh_addr", ""))),

512 + "ssh_port": int(cfg.get("ssh_port", 55022)),

513 + "ssh_user": str(cfg.get("ssh_user", "root")),

514 + "ssh_pass": str(cfg.get("ssh_pass", "")),

515 + "code_exec_timeouts": _parse_timeouts(cfg, "code_exec", (30, 15, 180, 5)),

516 + "output_timeouts": _parse_timeouts(cfg, "output", (90, 45, 300, 5)),

517 + "prompt_patterns": _parse_patterns(cfg.get("prompt_patterns", "")),

518 + "dialog_patterns": _parse_patterns(cfg.get("dialog_patterns", ""), re.IGNORECASE),

519 }

520

521

refactor: resolve config once per execution and deduplicate helpers