From 303af7a0143083af64429cedacddd23a76dbde81 Mon Sep 17 00:00:00 2001 From: Teknium <127238744+teknium1@users.noreply.github.com> Date: Wed, 2 Sep 2026 23:00:08 -0700 Subject: [PATCH] refactor(tools): tighten remote-kernel result assembly and per-call staging --- tools/code_execution_tool.py | 89 +++++++++++++----------------------- tools/code_kernel.py | 10 ++-- tools/code_kernel_remote.py | 38 ++++++--------- 3 files changed, 49 insertions(+), 88 deletions(-) diff --git a/tools/code_execution_tool.py b/tools/code_execution_tool.py index 95926d3f7b..9891682207 100644 --- a/tools/code_execution_tool.py +++ b/tools/code_execution_tool.py @@ -28,8 +28,7 @@ from typing import Any, Dict, List, Optional, Tuple from tools.thread_context import propagate_context_to_thread from tools.registry import registry, tool_error -# Env/interpreter resolution and RPC servers live in sibling modules; re-exported -# here so `from tools.code_execution_tool import X` / patch() targets keep working. +# Sibling-module symbols re-exported so `from tools.code_execution_tool import X` / patch() keep working. from tools.code_execution_env import ( # noqa: F401 _SAFE_ENV_PREFIXES, _SECRET_SUBSTRINGS, _HERMES_CHILD_ALLOWED, _WINDOWS_ESSENTIAL_ENV_VARS, _scrub_child_env, _build_child_env, _PROBE_CACHE_MAX, _usable_python_cache, _python_prefix_cache, @@ -53,18 +52,14 @@ DEFAULT_TIMEOUT = 300 # 5 minutes DEFAULT_MAX_TOOL_CALLS = 50 MAX_STDOUT_BYTES = 50_000 # 50 KB MAX_STDERR_BYTES = 10_000 # 10 KB -# Hard ceiling on the spilled file, mirroring web_tools' MAX_STORED_TEXT_CHARS -# rationale: a runaway print loop must not write unbounded bytes to disk. +# Hard ceiling on the spilled file (as web_tools' MAX_STORED_TEXT_CHARS): a runaway print loop must not fill the disk. MAX_SPILLED_STDOUT_BYTES = 5_000_000 def _truncate_stdout_text(stdout_text: str) -> Tuple[str, Dict[str, Any]]: - """Cap stdout by bytes (40% head / 60% tail) with explicit truncation metadata. - - Byte counts ride alongside the textual marker because a client layer can miss - or re-truncate the marker. The omitted middle is spilled to cache/exec and the - result carries the path (recover-don't-rerun, as web_extract's cache/web). - """ + """Cap stdout by bytes (40% head / 60% tail) with explicit truncation metadata: byte counts + ride alongside the textual marker because a client layer can miss or re-truncate it. The + omitted middle is spilled to cache/exec and the result carries the path (recover-don't-rerun).""" stdout_bytes = stdout_text.encode("utf-8", errors="replace") total = len(stdout_bytes) captured = min(total, MAX_STDOUT_BYTES) @@ -123,9 +118,7 @@ def check_sandbox_requirements() -> bool: return config.get("env_type") != "vercel_sandbox" or _check_vercel_sandbox_requirements(config) -# --------------------------------------------------------------------------- -# hermes_tools.py code generator -# --------------------------------------------------------------------------- +# ---- hermes_tools.py code generator ---- # Per-tool stub templates: (signature, docstring, args_dict_expr — the JSON payload sent over RPC). _TOOL_STUBS = { @@ -404,9 +397,14 @@ def _call(tool_name, args): ''' -# --------------------------------------------------------------------------- -# Remote execution support (file-based RPC via terminal backend) -# --------------------------------------------------------------------------- +# ---- Remote execution support (file-based RPC via terminal backend) ---- + +# execute_code's container_config keys (a subset of terminal_tool's; the create path fills the rest). +_CONTAINER_CONFIG_DEFAULTS = ( + ("container_cpu", 1), ("container_memory", 5120), ("container_disk", 51200), ("container_persistent", True), + ("vercel_runtime", ""), ("docker_volumes", []), ("docker_run_as_host_user", False), ("docker_network", True), +) + def _get_or_create_env(task_id: str): """``(env, env_type)`` — the environment the terminal/file tools share for *task_id*, created on @@ -438,14 +436,7 @@ def _get_or_create_env(task_id: str): overrides = _task_env_overrides.get(effective_task_id, {}) container_config = None if _is_container_backend(env_type): - container_config = { - "container_cpu": config.get("container_cpu", 1), "container_memory": config.get("container_memory", 5120), - "container_disk": config.get("container_disk", 51200), - "container_persistent": config.get("container_persistent", True), - "vercel_runtime": config.get("vercel_runtime", ""), "docker_volumes": config.get("docker_volumes", []), - "docker_run_as_host_user": config.get("docker_run_as_host_user", False), - "docker_network": config.get("docker_network", True), - } + container_config = {key: config.get(key, default) for key, default in _CONTAINER_CONFIG_DEFAULTS} logger.info("Creating new %s environment for execute_code task %s...", env_type, effective_task_id[:8]) env = _create_environment( @@ -578,10 +569,7 @@ def _run_remote_per_call(env, env_type: str, code: str, effective_task_id: str, sandbox_dir = f"{_env_temp_dir(env)}/hermes_exec_{uuid.uuid4().hex[:12]}" quoted_sandbox_dir = shlex.quote(sandbox_dir) quoted_rpc_dir = shlex.quote(f"{sandbox_dir}/rpc") - tool_call_log: list = [] - tool_call_counter = [0] - stop_event = threading.Event() - rpc_thread = None + tool_call_counter, stop_event, rpc_thread = [0], threading.Event(), None try: env.execute(f"mkdir -p {quoted_rpc_dir}", cwd="/", timeout=10) rpc_token = secrets.token_urlsafe(32) @@ -591,15 +579,12 @@ def _run_remote_per_call(env, env_type: str, code: str, effective_task_id: str, # Wrapped so the thread inherits the turn's approval context + callbacks # (tools.thread_context) — else sandbox RPC tool calls lose approval routing. rpc_thread = threading.Thread( - target=propagate_context_to_thread(_rpc_poll_loop), - args=(env, f"{sandbox_dir}/rpc", effective_task_id, tool_call_log, tool_call_counter, - max_tool_calls, sandbox_tools, stop_event, rpc_token), - daemon=True, - ) + target=propagate_context_to_thread(_rpc_poll_loop), daemon=True, + args=(env, f"{sandbox_dir}/rpc", effective_task_id, [], tool_call_counter, + max_tool_calls, sandbox_tools, stop_event, rpc_token)) rpc_thread.start() - env_prefix = (f"HERMES_RPC_DIR={quoted_rpc_dir} " - f"HERMES_RPC_TOKEN={shlex.quote(rpc_token)} " - f"PYTHONDONTWRITEBYTECODE=1") + env_prefix = (f"HERMES_RPC_DIR={quoted_rpc_dir} HERMES_RPC_TOKEN={shlex.quote(rpc_token)} " + "PYTHONDONTWRITEBYTECODE=1") tz = os.getenv("HERMES_TIMEZONE", "").strip() if tz: env_prefix += f" TZ={shlex.quote(tz)}" @@ -640,10 +625,8 @@ def _execute_remote(code: str, task_id: Optional[str], enabled_tools: Optional[L (tools/code_kernel_remote.py) first, else the per-call script ship — the fail-open route when a kernel cannot be spawned and the only route for hosts that cannot sustain a background process.""" _cfg = _load_config() - timeout = _cfg.get("timeout", DEFAULT_TIMEOUT) - max_tool_calls = _cfg.get("max_tool_calls", DEFAULT_MAX_TOOL_CALLS) - sandbox_tools = _sandbox_tools_for(enabled_tools) - effective_task_id = task_id or "default" + timeout, max_tool_calls = _cfg.get("timeout", DEFAULT_TIMEOUT), _cfg.get("max_tool_calls", DEFAULT_MAX_TOOL_CALLS) + sandbox_tools, effective_task_id = _sandbox_tools_for(enabled_tools), task_id or "default" env, env_type = _get_or_create_env(effective_task_id) exec_start = time.monotonic() try: @@ -674,9 +657,7 @@ def _execute_remote(code: str, task_id: Optional[str], enabled_tools: Optional[L timeout=timeout, max_tool_calls=max_tool_calls, exec_start=exec_start) -# --------------------------------------------------------------------------- -# Main entry point -# --------------------------------------------------------------------------- +# ---- Main entry point ---- def execute_code( @@ -686,13 +667,10 @@ def execute_code( reset: bool = False, ) -> str: """Run Python in the session's persistent kernel (local) or on the remote terminal backend, - with RPC access to a subset of Hermes tools; returns the JSON result string. - - "Sandbox" means the security envelope (env scrubbing, tool whitelist + call budget, output - redaction), not an isolation jail: default `project` mode runs in the session's cwd with the - project venv. ``enabled_tools`` ∩ SANDBOX_ALLOWED_TOOLS; ``reset`` kills the existing kernel - first (ignored on per-call paths). - """ + with RPC access to a subset of Hermes tools; returns the JSON result string. "Sandbox" means + the security envelope (env scrubbing, tool whitelist + call budget, output redaction), not an + isolation jail: default `project` mode runs in the session's cwd with the project venv. + ``enabled_tools`` ∩ SANDBOX_ALLOWED_TOOLS; ``reset`` kills the existing kernel first.""" if not SANDBOX_AVAILABLE: return tool_error("execute_code sandbox is unavailable in this environment. " "Use normal tool calls (terminal, read_file, write_file, ...) instead.") @@ -787,9 +765,7 @@ def _load_config() -> dict: return {} -# --------------------------------------------------------------------------- -# Execution mode resolution (strict vs project) -# --------------------------------------------------------------------------- +# ---- Execution mode resolution (strict vs project) ---- # Canonical code_execution.mode values (referenced by tests and the config layer). Session # kernels are the only local execution model; a leftover kernel_mode config key is ignored. @@ -810,9 +786,7 @@ def _get_execution_mode() -> str: return DEFAULT_EXECUTION_MODE -# --------------------------------------------------------------------------- -# OpenAI Function-Calling Schema -# --------------------------------------------------------------------------- +# ---- OpenAI Function-Calling Schema ---- # Per-tool documentation lines for the execute_code description, in canonical display order. _TOOL_DOC_LINES = [ @@ -843,8 +817,7 @@ def build_execute_code_schema(enabled_sandbox_tools: set = None, mode = _get_execution_mode() tool_lines = "\n".join(doc for name, doc in _TOOL_DOC_LINES if name in enabled_sandbox_tools) import_examples = [n for n in ("web_search", "terminal") if n in enabled_sandbox_tools] - if not import_examples: - import_examples = sorted(enabled_sandbox_tools)[:2] + import_examples = import_examples or sorted(enabled_sandbox_tools)[:2] import_str = ", ".join(import_examples) + ", ..." if import_examples else "..." if mode == "strict": cwd_note = ( diff --git a/tools/code_kernel.py b/tools/code_kernel.py index acd94bdd96..1a9d412354 100644 --- a/tools/code_kernel.py +++ b/tools/code_kernel.py @@ -30,6 +30,7 @@ import tempfile import threading import time import uuid +from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Tuple logger = logging.getLogger(__name__) @@ -430,11 +431,10 @@ def _bind_rpc_socket(kernel: SessionKernel) -> str: rpc_endpoint = f"tcp://{host}:{port}" else: sock_tmpdir = "/tmp" if sys.platform == "darwin" else tempfile.gettempdir() - kernel.sock_path = os.path.join(sock_tmpdir, f"hermes_rpc_{uuid.uuid4().hex}.sock") + rpc_endpoint = kernel.sock_path = os.path.join(sock_tmpdir, f"hermes_rpc_{uuid.uuid4().hex}.sock") server_sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) server_sock.bind(kernel.sock_path) os.chmod(kernel.sock_path, 0o600) - rpc_endpoint = kernel.sock_path server_sock.listen(1) kernel.server_sock = server_sock return rpc_endpoint @@ -449,8 +449,7 @@ def _spawn(kernel: SessionKernel, *, child_python: str, child_cwd: str, rpc_endpoint = _bind_rpc_socket(kernel) for name, src in (("hermes_tools.py", generate_hermes_tools_module(list(sandbox_tools))), ("hermes_kernel_runner.py", KERNEL_RUNNER_SOURCE)): - with open(os.path.join(kernel.tmpdir, name), "w", encoding="utf-8") as f: - f.write(src) + Path(kernel.tmpdir, name).write_text(src, encoding="utf-8") child_env = _build_child_env(rpc_endpoint=rpc_endpoint, rpc_token=kernel.rpc_token, tmpdir=kernel.tmpdir, child_python=child_python) child_env["HERMES_KERNEL_SENTINEL"] = kernel.sentinel @@ -461,9 +460,8 @@ def _spawn(kernel: SessionKernel, *, child_python: str, child_cwd: str, kernel.proc = subprocess.Popen( [child_python, os.path.join(kernel.tmpdir, "hermes_kernel_runner.py")], # Strict mode passes an empty cwd: the kernel's staging dir plays the per-call tmpdir's role. - cwd=child_cwd or kernel.tmpdir, env=child_env, + cwd=child_cwd or kernel.tmpdir, env=child_env, start_new_session=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE, - start_new_session=True, creationflags=subprocess.CREATE_NO_WINDOW if _IS_WINDOWS else 0, ) # Deliberately NOT propagate_context_to_thread: that would freeze the spawning cell's diff --git a/tools/code_kernel_remote.py b/tools/code_kernel_remote.py index f885c77068..d34ad00b11 100644 --- a/tools/code_kernel_remote.py +++ b/tools/code_kernel_remote.py @@ -243,8 +243,7 @@ def _run_remote_cell(kernel: RemoteKernel, code: str, timeout: int) -> Tuple[str from tools.code_execution_tool import _ship_file_to_remote kernel.cell_seq += 1 seq = f"{kernel.cell_seq:06d}" - q_cells = shlex.quote(f"{kernel.kernel_dir}/cells") - q_res = shlex.quote(f"cell_res_{seq}.json") + q_cells, q_res = shlex.quote(f"{kernel.kernel_dir}/cells"), shlex.quote(f"cell_res_{seq}.json") _ship_file_to_remote(kernel.env, f"{kernel.kernel_dir}/cells/cell_req_{seq}.json.tmp", json.dumps({"id": seq, "code": code}, ensure_ascii=False)) kernel.sh(f"mv {q_cells}/cell_req_{seq}.json.tmp {q_cells}/cell_req_{seq}.json", timeout=10) @@ -271,12 +270,9 @@ def execute_in_remote_kernel( code: str, *, env, env_type: str, task_env_id: str, sandbox_tools: frozenset, timeout: int, max_tool_calls: int, reset: bool, idle_exit: int = 1800, ) -> Optional[Dict[str, Any]]: - """Run one cell in the owner's remote kernel. - - Returns the raw cell result dict (caller post-processes output), or ``None`` when no - kernel could be spawned (caller falls open to per-call). ``state_lost`` / ``state_reset`` / - ``reused`` ride in the ``kernel`` sub-dict, matching the local kernel's result shape. - """ + """Run one cell in the owner's remote kernel. Returns the raw cell result dict (caller + post-processes output), or ``None`` when no kernel could be spawned (caller falls open to + per-call). ``state_lost``/``state_reset``/``reused`` ride in the ``kernel`` sub-dict.""" from tools.code_kernel import _resolve_owner from tools.code_execution_tool import _rpc_poll_loop from tools.thread_context import propagate_context_to_thread @@ -294,16 +290,13 @@ def execute_in_remote_kernel( kernel.sh(f"rm -f {q_rpc}/req_* {q_rpc}/res_*", timeout=10) except Exception: pass - tool_call_log: list = [] tool_call_counter, stop_event = [0], threading.Event() # Per-cell RPC thread carrying THIS call's approval/session context — the remote analogue # of CellAuthority: authority lives exactly as long as the cell's poll loop. rpc_thread = threading.Thread( - target=propagate_context_to_thread(_rpc_poll_loop), - args=(env, f"{kernel.kernel_dir}/rpc", task_env_id, tool_call_log, tool_call_counter, - max_tool_calls, sandbox_tools, stop_event, kernel.rpc_token), - daemon=True, - ) + target=propagate_context_to_thread(_rpc_poll_loop), daemon=True, + args=(env, f"{kernel.kernel_dir}/rpc", task_env_id, [], tool_call_counter, + max_tool_calls, sandbox_tools, stop_event, kernel.rpc_token)) rpc_thread.start() cell_status, cell_payload = "no-result", {} try: @@ -313,26 +306,23 @@ def execute_in_remote_kernel( rpc_thread.join(timeout=5) kernel_info: Dict[str, Any] = {"reused": reused, "remote": True} result: Dict[str, Any] = { - "status": "error", "stdout": cell_payload.get("stdout", ""), - "stderr": cell_payload.get("stderr", ""), "traceback": cell_payload.get("traceback", ""), - "tool_calls_made": tool_call_counter[0], "kernel": kernel_info, + "status": "error", "stdout": cell_payload.get("stdout", ""), "stderr": cell_payload.get("stderr", ""), + "traceback": cell_payload.get("traceback", ""), "tool_calls_made": tool_call_counter[0], "kernel": kernel_info, } if cell_status in ("timeout", "protocol-error", "no-result"): # No safe way to interrupt one cell in place (same contract as local): kill, report, respawn. _REGISTRY.discard(key, kernel) if cell_status == "timeout": result["status"] = "timeout" - note = ("Cell timed out; the remote session kernel was killed and " - "its state was lost. The next call starts a fresh kernel.") - else: - note = "Remote kernel protocol failure; kernel killed, state lost." - kernel_info.update(ended=True, state_lost=True, note=note) + kernel_info.update(ended=True, state_lost=True, note=( + "Cell timed out; the remote session kernel was killed and its state was lost. The next call " + "starts a fresh kernel." if cell_status == "timeout" + else "Remote kernel protocol failure; kernel killed, state lost.")) return result if cell_status == "exit": _REGISTRY.discard(key, kernel) kernel_info["ended"] = True - kernel.execution_count = int(cell_payload.get("execution_count", 0) or 0) - kernel_info["execution_count"] = kernel.execution_count + kernel.execution_count = kernel_info["execution_count"] = int(cell_payload.get("execution_count", 0) or 0) if cell_status in ("ok", "exit"): result["status"] = "success" result["stdout_clipped"] = bool(cell_payload.get("stdout_clipped"))