refactor(tools): unify remote result assembly, compact execute_code docs and comments

This commit is contained in:
Teknium
2026-09-02 22:33:03 -07:00
parent 89881ad022
commit 02f92ec707
5 changed files with 241 additions and 559 deletions
+61 -149
View File
@@ -1,28 +1,18 @@
"""Session-persistent Python kernels for execute_code.
"""Session-persistent Python kernels for execute_code: one child per (owner, mode,
interpreter, cwd, tool-set), one code cell per call, state survives across calls.
One Python child stays alive per (owner, mode, interpreter, cwd, tool-set) and
runs one code cell per call, so variables/imports/data survive across calls.
Constraints, in order: (1) SAME security envelope as per-call (``_build_child_env``
scrubbing, ``_rpc_server_loop`` token + per-cell tool budget, ANSI strip + secret
redaction) — only lifetime widens. (2) A wedged kernel dies, never hangs the agent:
timeout/interrupt kills the process tree and drops the registry entry; state loss is
deliberate (a cell cannot be interrupted in place safely). (3) Env frozen at spawn:
later passthrough is invisible until ``reset=true`` (the result names the kernel).
Design constraints, in order: (1) the SAME security envelope as per-call —
``_build_child_env`` scrubbing, ``_rpc_server_loop`` with the same token and
per-cell tool budget, the same ANSI strip + secret redaction; nothing here widens
what a script can reach, only how long it lives. (2) A wedged kernel dies, never
hangs the agent: timeout or interrupt kills the whole process tree and drops the
registry entry; losing state is deliberate since one cell cannot be interrupted
in place without leaving the interpreter unknown. (3) The env is frozen at spawn:
passthrough registered later is invisible until ``reset=true`` (the result names
the kernel so this is diagnosable).
Wire protocol (host <-> child): requests are one JSON object per stdin line
``{"id", "code"}``; responses are framed on stdout as
``<SENTINEL> <byte-length>\\n<json>`` with a per-kernel random SENTINEL from the
environment. Bytes outside frames are raw fd-level output (subprocesses inherit
the real stdout), attributed to the running cell — calls are serialized per
kernel. A script forging a frame can only fake its own cell result (same trust
position as a per-call script printing a forged success message).
Also hosts what ``tools.code_kernel_remote`` shares: owner resolution, the
registry lifecycle, and the runner's cell-exec core.
Wire protocol: one JSON request per stdin line ``{"id", "code"}``; replies framed on
stdout as ``<SENTINEL> <byte-length>\\n<json>`` with a per-kernel random SENTINEL from
the env. Bytes outside frames are raw fd output attributed to the running cell (calls
are serialized per kernel). A forged frame can only fake its own cell result.
Also hosts what ``tools.code_kernel_remote`` shares: owner resolution, registry, cell core.
"""
from __future__ import annotations
@@ -46,13 +36,11 @@ logger = logging.getLogger(__name__)
_IS_WINDOWS = sys.platform == "win32"
# Runner-side cap on captured python-level output; the host applies its own
# MAX_STDOUT truncation again.
# Runner-side cap on captured python-level output; the host re-applies its own MAX_STDOUT cap.
_RUNNER_CAPTURE_BYTES = 1_000_000
# Shared by both generated runners (which define _CAPTURE_LIMIT first): exec one
# request in the persistent GLOBALS namespace and build the response payload.
# `__name__` is `__main__` so scripts behave like the per-call path.
# Shared by both generated runners (which define _CAPTURE_LIMIT first): exec one request in the
# persistent GLOBALS namespace, build the payload. `__name__` is `__main__` as on the per-call path.
RUNNER_CELL_SOURCE = '''\
GLOBALS = {"__name__": "__main__", "__builtins__": __builtins__}
@@ -150,17 +138,15 @@ if __name__ == "__main__":
class CellAuthority:
"""The approval/context identity of exactly one execute_code cell.
Interpreter state persists across cells; RPC authority must not. Each cell
installs a fresh authority — captured from the CALLING thread at cell start,
exactly what ``propagate_context_to_thread`` would capture for a per-call RPC
thread — and retires it when the cell settles, so a late tool call (a
background thread the cell left behind, a raced client write) is refused
instead of running under a stale approval/session/turn identity.
Interpreter state persists across cells; RPC authority must not. Each cell installs a
fresh authority captured from the CALLING thread at cell start (what
``propagate_context_to_thread`` captures for a per-call RPC thread) and retires it when
the cell settles, so a late tool call (leaked background thread, raced client write) is
refused instead of running under a stale approval/session/turn identity.
"""
def __init__(self, task_id: str):
import contextvars
self.task_id = task_id
self.ctx = contextvars.copy_context()
self.active = True
@@ -168,12 +154,10 @@ class CellAuthority:
self._callbacks = (None, None)
try:
from tools.thread_context import _callback_api
self._api = _callback_api()
self._callbacks = (self._api[0](), self._api[1]())
except Exception:
# Fail-closed, mirroring propagate_context_to_thread: with no
# callbacks installed, dangerous approvals deny.
# Fail-closed like propagate_context_to_thread: no callbacks → dangerous approvals deny.
self._api = None
def retire(self) -> None:
@@ -182,17 +166,13 @@ class CellAuthority:
def dispatch(self, tool_name: str, tool_args: dict) -> str:
"""Run one tool call under THIS cell's context and callbacks."""
from tools.code_execution_tool import tool_error
if not self.active:
return tool_error(
"No active execute_code cell: the cell this kernel call "
"belonged to has settled, so its tool authority is retired."
)
return tool_error("No active execute_code cell: the cell this kernel call "
"belonged to has settled, so its tool authority is retired.")
return self.ctx.run(self._invoke, tool_name, tool_args)
def _invoke(self, tool_name: str, tool_args: dict) -> str:
from model_tools import handle_function_call
previous = None
if self._api is not None:
get_approval, get_sudo, set_approval, set_sudo = self._api
@@ -259,7 +239,6 @@ class SessionKernel:
self.stop_event.set()
if self.alive():
from tools.code_execution_tool import _kill_process_group
_kill_process_group(self.proc, escalate=True)
sock, self.server_sock = self.server_sock, None
try:
@@ -271,16 +250,12 @@ class SessionKernel:
pass
if self.tmpdir:
import shutil
shutil.rmtree(self.tmpdir, ignore_errors=True)
class KernelRegistry:
"""Key -> kernel map plus its lock (shared with the remote registry).
Kernels are popped under the lock and torn down outside it — teardown
may block on the child process or the remote transport.
"""
"""Key -> kernel map plus its lock (shared with the remote registry). Kernels are popped
under the lock and torn down outside it — teardown may block on the child or the transport."""
def __init__(self, teardown: Callable[[Any], None]):
self.kernels: Dict[Tuple, Any] = {}
@@ -305,57 +280,45 @@ class KernelRegistry:
_REGISTRY = KernelRegistry(lambda kernel: kernel.teardown())
_KERNELS: Dict[Tuple, SessionKernel] = _REGISTRY.kernels
# Bounded lifecycle defaults (config: code_execution.max_session_kernels /
# code_execution.kernel_idle_timeout). A long-lived gateway must never
# accumulate one live child per finished conversation: stable owner id,
# owner-teardown disposal, idle reaping, max-live bound.
# Bounded lifecycle defaults (config: code_execution.max_session_kernels / kernel_idle_timeout).
# A long-lived gateway must never accumulate one live child per finished conversation:
# stable owner id, owner-teardown disposal, idle reaping, max-live bound.
DEFAULT_MAX_SESSION_KERNELS = 4
DEFAULT_KERNEL_IDLE_TIMEOUT = 1800
def _lifecycle_limits() -> Tuple[int, int]:
from tools.code_execution_tool import _load_config
config = _load_config()
def limit(key: str, default: int) -> int:
try:
return max(1, int(config.get(key, default)))
except (TypeError, ValueError):
return default
return (limit("max_session_kernels", DEFAULT_MAX_SESSION_KERNELS),
limit("kernel_idle_timeout", DEFAULT_KERNEL_IDLE_TIMEOUT))
def _resolve_owner(task_id: str) -> str:
"""The stable identity a session kernel belongs to.
"""The stable identity a session kernel belongs to: the conversation's approval session key
(context-propagated, stable across turns, distinct per session). ``run_agent`` mints a fresh
task id per turn, so a task-keyed kernel would neither survive the next turn nor be torn down
with anything; the task id is only the last-resort owner (embeds/tests without a session).
The conversation's approval session key: context-propagated, stable across
turns, distinct per session. ``run_agent`` mints a fresh task id per
top-level turn, so a task-keyed kernel would neither survive the next turn
nor ever be torn down with anything; the task id is only the last-resort
owner for embeds and tests with no session context.
Delegated children run in a copy of the parent's context and INHERIT its
approval session key — without the ``::child::`` qualifier a child's
execute_code would attach to the parent's kernel and read its in-memory
state (verified live, both directions). Children get their own kernels,
keyed by their delegation session id.
Delegated children INHERIT the parent's approval session key — without the ``::child::``
qualifier a child's execute_code would attach to the parent's kernel and read its state
(verified live, both directions). Children get their own kernels keyed by delegation session id.
"""
try:
from tools.approval import get_current_session_key
session_key = get_current_session_key(default="")
except Exception:
session_key = ""
owner = session_key or (task_id or "")
try:
from agent.delegation_context import is_delegated_child_context
if is_delegated_child_context():
from gateway.session_context import get_session_env
child_id = get_session_env("HERMES_SESSION_ID", "") or (task_id or "")
owner = f"{owner}::child::{child_id}"
except Exception:
@@ -380,25 +343,16 @@ atexit.register(shutdown_all_kernels)
def _rpc_forever(kernel: SessionKernel, max_tool_calls: int,
sandbox_tools: frozenset) -> None:
"""Serve tool RPC for the kernel's whole life.
``_rpc_server_loop`` serves one connection and returns on disconnect or its
300s idle timeout; a kernel legitimately idles longer between cells, so
re-accept until teardown (the client stub reconnects: HERMES_RPC_PERSISTENT).
The serving thread carries NO frozen authority: every dispatch routes
through the CURRENT cell's ``CellAuthority``, so a later cell's tool calls
run under that cell's context, not whatever the first cell captured.
"""
"""Serve tool RPC for the kernel's whole life: ``_rpc_server_loop`` returns on disconnect or
its 300s idle timeout, and a kernel idles longer between cells, so re-accept until teardown
(the client stub reconnects: HERMES_RPC_PERSISTENT). The serving thread carries NO frozen
authority — every dispatch routes through the CURRENT cell's ``CellAuthority``."""
from tools.code_execution_tool import _rpc_server_loop, tool_error
def _dispatch(tool_name: str, tool_args: dict) -> str:
authority = kernel.cell_authority
if authority is None:
return tool_error("No active execute_code cell: this kernel has no cell "
"authority installed.")
return tool_error("No active execute_code cell: this kernel has no cell authority installed.")
return authority.dispatch(tool_name, tool_args)
while not kernel.stop_event.is_set():
_rpc_server_loop(kernel.server_sock, "", kernel.tool_call_log, kernel.tool_call_counter,
max_tool_calls, sandbox_tools, kernel.stop_event, kernel.rpc_token,
@@ -408,19 +362,15 @@ def _rpc_forever(kernel: SessionKernel, max_tool_calls: int,
def _stdout_reader(kernel: SessionKernel) -> None:
"""Split the child's stdout into protocol frames and raw passthrough."""
from tools.code_execution_tool import MAX_STDOUT_BYTES
assert kernel.proc is not None and kernel.proc.stdout is not None
stream = kernel.proc.stdout
marker = ("\n" + kernel.sentinel + " ").encode("utf-8")
def raw(data: bytes) -> None:
kernel.raw.append(data, MAX_STDOUT_BYTES)
buf = b""
while True:
# read1: return as soon as any bytes arrive. A plain read(n) on a
# BufferedReader blocks until n bytes or EOF, which would sit on a
# complete frame smaller than the buffer forever.
# read1 returns as soon as any bytes arrive; a plain read(n) on a BufferedReader
# blocks until n bytes or EOF and would sit on a complete small frame forever.
chunk = stream.read1(4096)
if not chunk:
if buf:
@@ -431,8 +381,7 @@ def _stdout_reader(kernel: SessionKernel) -> None:
while True:
index = buf.find(marker)
if index < 0:
# Keep a marker-sized tail in case the marker is split
# across reads; everything before it is raw output.
# Keep a marker-sized tail (marker may be split across reads); the rest is raw.
spill = buf[: -len(marker)] if len(buf) > len(marker) else b""
if spill:
raw(spill)
@@ -448,8 +397,7 @@ def _stdout_reader(kernel: SessionKernel) -> None:
try:
length = int(rest[:newline])
except ValueError:
# Not a real frame header (user output that happens to
# contain the marker bytes); treat the marker as raw.
# Not a real frame header (user output containing the marker bytes): raw.
raw(marker)
buf = rest
continue
@@ -469,7 +417,6 @@ def _stdout_reader(kernel: SessionKernel) -> None:
def _stderr_reader(kernel: SessionKernel) -> None:
from tools.code_execution_tool import MAX_STDERR_BYTES
assert kernel.proc is not None and kernel.proc.stderr is not None
while True:
chunk = kernel.proc.stderr.read1(4096)
@@ -501,52 +448,40 @@ def _bind_rpc_socket(kernel: SessionKernel) -> str:
def _spawn(kernel: SessionKernel, *, child_python: str, child_cwd: str,
sandbox_tools: frozenset, max_tool_calls: int) -> None:
from tools.code_execution_tool import _build_child_env, generate_hermes_tools_module
kernel.tmpdir = tempfile.mkdtemp(prefix="hermes_kernel_")
kernel.rpc_token = secrets.token_urlsafe(32)
kernel.sentinel = "@@HERMES-KERNEL-" + secrets.token_urlsafe(16) + "@@"
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)
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
# Cells clip stdout to the inline cap; the full text spills to the kernel's
# own tmpdir so the agent can read_file the middle instead of re-running.
# Full clipped stdout spills to the kernel's tmpdir so the agent can read_file the middle.
child_env["HERMES_KERNEL_SPILL_DIR"] = kernel.tmpdir
# Tell the generated client to reconnect after the RPC server's idle
# timeout — a kernel outlives the 300s window between cells.
# Generated client reconnects after the RPC server's 300s idle timeout between cells.
child_env["HERMES_RPC_PERSISTENT"] = "1"
kernel.proc = subprocess.Popen(
[child_python, os.path.join(kernel.tmpdir, "hermes_kernel_runner.py")],
# Strict mode resolves an empty cwd: the kernel's own staging dir
# then plays the per-call tmpdir's role.
# 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,
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 context/callbacks into the server thread for the
# kernel's whole life. Authority is rebound per cell via CellAuthority.
# Deliberately NOT propagate_context_to_thread: that would freeze the spawning cell's
# context/callbacks into the server thread for life. Authority is rebound per cell.
for target, args in ((_rpc_forever, (kernel, max_tool_calls, sandbox_tools)),
(_stdout_reader, (kernel,)), (_stderr_reader, (kernel,))):
threading.Thread(target=target, args=args, daemon=True).start()
def _acquire_kernel(key: Tuple, reset: bool) -> Tuple[SessionKernel, bool]:
"""Look up or register the kernel for *key*; returns (kernel, state_reset).
Every entry also sweeps idle-expired kernels and enforces the process-wide
LRU cap, so a long-lived host stays bounded even for owners that never
toggle or reset. Doomed kernels are popped under the lock, torn down outside it.
"""
"""Look up or register the kernel for *key*; returns (kernel, state_reset). Every entry also
sweeps idle-expired kernels and enforces the process-wide LRU cap (doomed kernels are popped
under the lock, torn down outside it), so a long-lived host stays bounded."""
cap, idle_timeout = _lifecycle_limits()
with _REGISTRY.lock:
now = time.monotonic()
@@ -594,23 +529,17 @@ def _cell_result(kernel: SessionKernel, key: Tuple, status: str, payload: Dict[s
from tools.code_execution_tool import _sandbox_failure_hint, _truncate_stdout_text
from agent.redact import redact_sensitive_text
from tools.ansi_strip import strip_ansi
def clean(text: str) -> str:
return redact_sensitive_text(strip_ansi(text), code_file=True)
if status in ("timeout", "interrupted"):
# No safe way to interrupt one cell in place: kill the kernel,
# report the state loss, let the next call respawn.
# No safe way to interrupt one cell in place: kill the kernel, report the loss, respawn next call.
_REGISTRY.discard(key, kernel)
duration = round(time.monotonic() - exec_start, 2)
kernel.execution_count = int(payload.get("execution_count", kernel.execution_count + 1))
stderr_raw = kernel.stderr.drain()
stdout_text = clean(str(payload.get("stdout", "")) + kernel.raw.drain())
cell_stderr = clean(str(payload.get("stderr", "")) + stderr_raw)
stdout_text, stdout_metadata = _truncate_stdout_text(stdout_text)
cell_status = payload.get("status", "")
result: Dict[str, Any] = {
"status": status, "output": stdout_text, "exit_code": 0,
@@ -619,9 +548,7 @@ def _cell_result(kernel: SessionKernel, key: Tuple, status: str, payload: Dict[s
"execution_count": kernel.execution_count, "state_reset": state_reset},
}
result.update(stdout_metadata)
# Cell-side spill (runner clipped before replying): surface the full-output
# path with the same read_file recipe as the host-side spill.
# Cell-side spill (runner clipped before replying): same read_file recipe as the host-side spill.
cell_spill = str(payload.get("stdout_spill_path", "") or "")
if cell_spill and payload.get("stdout_clipped"):
result["stdout_spill_path"] = cell_spill
@@ -630,7 +557,6 @@ def _cell_result(kernel: SessionKernel, key: Tuple, status: str, payload: Dict[s
f'— page it with read_file(path="{cell_spill}", offset=...) instead of re-running. '
"(Kernel state persists: printing a narrower slice next call is often cheaper.)"
)
if status == "timeout":
message = (f"Cell timed out after {timeout}s; the session kernel was killed and its "
"state was lost. The next execute_code call starts a fresh kernel.")
@@ -638,7 +564,6 @@ def _cell_result(kernel: SessionKernel, key: Tuple, status: str, payload: Dict[s
output=(stdout_text + "\n\n⏰ " + message) if stdout_text else ("⏰ " + message))
elif status == "interrupted":
from tools.code_execution_tool import _format_interrupted_output
result.update(exit_code=-1, output=_format_interrupted_output(stdout_text),
error="Interrupted; the session kernel was killed and its state was lost.")
elif cell_status == "error":
@@ -667,41 +592,29 @@ def execute_in_session_kernel(
code: str, *, task_id: str, mode: str, child_python: str, child_cwd: str,
sandbox_tools: frozenset, timeout: int, max_tool_calls: int, reset: bool, is_interrupted,
) -> str:
"""Run one cell in the (owner, mode, python, cwd, tools) session kernel.
The owner is the conversation's session key (``_resolve_owner``), not the
per-turn task id, so state survives across user turns of one conversation
and dies with the session.
"""
"""Run one cell in the (owner, mode, python, cwd, tools) session kernel. The owner is the
session key (``_resolve_owner``), not the per-turn task id, so state survives across turns."""
key = (_resolve_owner(task_id) or "", mode, child_python, child_cwd, tuple(sorted(sandbox_tools)))
exec_start = time.monotonic()
kernel, state_reset = _acquire_kernel(key, reset)
reused = kernel.proc is not None
# Captured on the calling thread BEFORE the cell runs — the same snapshot a
# per-call RPC thread would have received — and installed atomically on the
# kernel so the serving thread dispatches this cell's tool calls under this
# cell's approval/session/turn identity.
# Captured on the calling thread BEFORE the cell runs (the snapshot a per-call RPC thread
# would get) and installed on the kernel so RPC dispatches under THIS cell's identity.
authority = CellAuthority(task_id)
with kernel.lock:
try:
if kernel.proc is None:
_spawn(kernel, child_python=child_python, child_cwd=child_cwd,
sandbox_tools=sandbox_tools, max_tool_calls=max_tool_calls)
assert kernel.proc is not None and kernel.proc.stdin is not None
# Per-cell tool budget: the RPC loop enforces counter < max, so a
# fresh cell starts from zero without restarting the server.
# Per-cell tool budget: the RPC loop enforces counter < max; reset without restarting.
kernel.tool_call_counter[0] = 0
# Anything raw that leaked between cells belongs to no cell.
kernel.raw.drain()
kernel.stderr.drain()
kernel.cell_authority = authority
kernel.proc.stdin.write((json.dumps({"id": uuid.uuid4().hex, "code": code}) + "\n").encode("utf-8"))
kernel.proc.stdin.flush()
status, payload = _await_cell(kernel, timeout, is_interrupted)
result = _cell_result(
kernel, key, status, payload,
@@ -718,7 +631,6 @@ def execute_in_session_kernel(
"duration_seconds": round(time.monotonic() - exec_start, 2),
}, ensure_ascii=False)
finally:
# The cell has settled on every path (success, exception, timeout,
# exit, kernel death): its tool authority retires with it, so
# The cell has settled on every path: its tool authority retires with it, so
# nothing the cell left running can dispatch under it.
authority.retire()