Files
aiturk-hermes-ide/tools/code_kernel_remote.py
T

577 lines
21 KiB
Python

"""Session-persistent kernels for REMOTE terminal backends (docker/ssh/modal).
Closes the gap tracked in hermes-agent#96873: local execute_code holds a
persistent kernel child (tools/code_kernel.py); remote backends previously
re-shipped and re-ran a fresh script per call, losing all interpreter state.
The remote transport offers exactly one primitive — ``env.execute(cmd)``,
run-to-completion — so the three things the local kernel gets from owning a
child process are rebuilt on top of it:
1. **A process that outlives one env.execute():** the kernel runner is
started detached (``nohup ... &``) and its PID recorded; each later cell
first probes liveness with ``kill -0``.
2. **A conversation channel:** a file-based CELL protocol in the kernel dir
(``cell_req_NNNNNN.json`` / ``cell_res_NNNNNN.json``), sibling to the
existing file-based TOOL-RPC protocol (req_/res_ files) which is reused
unchanged — the host-side ``_rpc_poll_loop`` is started per cell with the
calling thread's context, which is what gives per-cell tool authority.
3. **Death detection:** a failed liveness probe (transport drop, container
restart, OOM-killed runner) reads as *kernel died: state lost*; the next
call respawns fresh and says so — never a hung poll loop, because every
wait is bounded by the cell timeout.
Same invariants as local: owner = approval session key with the
``::child::{id}`` qualifier for delegated children (imported from
tools.code_kernel — one resolver, cannot drift), same generated tool stubs,
same output post-processing in the caller. ``reset=true`` kills and
respawns. Spawn failure fails OPEN to the per-call path with a note, so a
degraded remote host never blocks execution entirely.
"""
from __future__ import annotations
import atexit
import base64
import json
import logging
import shlex
import threading
import time
import uuid
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple
logger = logging.getLogger(__name__)
# One lock guards the registry; teardown runs outside it (mirrors code_kernel).
_REMOTE_KERNELS: Dict[Tuple, "RemoteKernel"] = {}
_REMOTE_KERNELS_LOCK = threading.Lock()
# How often the host polls the remote for a cell result file. Each poll is
# one env.execute round-trip (typically 0.1-0.4s on ssh/docker), so this is
# a floor, not a rate.
_CELL_POLL_INTERVAL = 0.5
# The remote runner: a tiny forever-loop that polls for cell request files,
# execs them in one persistent namespace, and writes response files. It is
# deliberately transport-agnostic (pure files) and stdlib-only. Cells and
# tool-RPC share the kernel dir but use distinct prefixes.
REMOTE_KERNEL_RUNNER_SOURCE = '''\
"""Auto-generated Hermes REMOTE session-kernel runner (file cell protocol)."""
import contextlib
import io
import json
import os
import sys
import time
import traceback
KDIR = os.environ["HERMES_KERNEL_DIR"]
CELLS = os.path.join(KDIR, "cells")
CAPTURE_LIMIT = {capture_limit}
IDLE_EXIT_SECONDS = {idle_exit}
GLOBALS = {{"__name__": "__main__", "__builtins__": __builtins__}}
def _bounded(text):
if len(text) <= CAPTURE_LIMIT:
return text, False
return text[:CAPTURE_LIMIT], True
def main():
execution_count = 0
last_activity = time.time()
while True:
pending = sorted(
f for f in os.listdir(CELLS)
if f.startswith("cell_req_") and f.endswith(".json")
)
if not pending:
if time.time() - last_activity > IDLE_EXIT_SECONDS:
return # self-reap: nobody is talking to us anymore
time.sleep(0.2)
continue
for name in pending:
req_path = os.path.join(CELLS, name)
try:
with open(req_path, "r", encoding="utf-8") as f:
request = json.load(f)
except Exception:
# Partially-written request (ship in progress): retry next tick.
continue
os.remove(req_path)
last_activity = time.time()
execution_count += 1
out, err = io.StringIO(), io.StringIO()
status = "ok"
trace = ""
try:
with contextlib.redirect_stdout(out), contextlib.redirect_stderr(err):
exec(compile(request["code"], "<cell>", "exec"), GLOBALS)
except SystemExit as exc:
status = "exit"
trace = "SystemExit: " + repr(exc.code)
except BaseException:
status = "error"
trace = traceback.format_exc()
stdout_text, stdout_clipped = _bounded(out.getvalue())
stderr_text, stderr_clipped = _bounded(err.getvalue())
payload = {{
"id": request.get("id", ""),
"status": status,
"stdout": stdout_text,
"stderr": stderr_text,
"stdout_clipped": stdout_clipped,
"stderr_clipped": stderr_clipped,
"traceback": trace,
"execution_count": execution_count,
}}
res_name = name.replace("cell_req_", "cell_res_")
tmp = os.path.join(CELLS, res_name + ".tmp")
with open(tmp, "w", encoding="utf-8") as f:
json.dump(payload, f, ensure_ascii=False)
os.replace(tmp, os.path.join(CELLS, res_name))
if status == "exit":
return
if __name__ == "__main__":
main()
'''
@dataclass
class RemoteKernel:
"""Host-side record of one detached remote kernel process."""
env: Any
env_type: str
kernel_dir: str
pid: str
rpc_token: str
owner: str
created: float = field(default_factory=time.monotonic)
last_used: float = field(default_factory=time.monotonic)
execution_count: int = 0
cell_seq: int = 0
# Cells currently running on this kernel. Reap/evict skip attached
# kernels: killing one mid-cell tears the runner out from under a live
# poll loop (same guard as tools.code_kernel, hermes-agent#101861).
attached: int = 0
def _kernel_key(owner: str, env_type: str, task_env_id: str) -> Tuple:
return (owner, "remote", env_type, task_env_id)
def _is_alive(kernel: RemoteKernel) -> bool:
"""Bounded liveness probe: kill -0 through the transport.
Any transport failure counts as dead — the caller respawns. This is the
"death detection" leg: a dropped ssh connection and a dead runner are
indistinguishable from here, and both have the same correct answer.
"""
try:
probe = kernel.env.execute(
f"kill -0 {shlex.quote(kernel.pid)} 2>/dev/null && echo ALIVE",
cwd="/", timeout=15,
)
return "ALIVE" in (probe.get("output", "") or "")
except Exception:
return False
def _kill(kernel: RemoteKernel) -> None:
"""Best-effort kill of the runner and its subprocesses, then rm -rf."""
try:
kernel.env.execute(
# Kill the runner's process group if the shell gave it one,
# falling back to the single PID.
f"pkill -TERM -P {shlex.quote(kernel.pid)} 2>/dev/null; "
f"kill {shlex.quote(kernel.pid)} 2>/dev/null; true",
cwd="/", timeout=15,
)
except Exception:
logger.debug("remote kernel kill failed (transport?)", exc_info=True)
try:
kernel.env.execute(
f"rm -rf {shlex.quote(kernel.kernel_dir)}", cwd="/", timeout=15,
)
except Exception:
logger.debug("remote kernel dir cleanup failed", exc_info=True)
def shutdown_all_remote_kernels() -> None:
with _REMOTE_KERNELS_LOCK:
kernels = list(_REMOTE_KERNELS.values())
_REMOTE_KERNELS.clear()
for kernel in kernels:
_kill(kernel)
def shutdown_remote_kernels_for_owner(owner: str) -> None:
"""Session-boundary disposal — wired to the same clear_session hook as
local kernels, so /new and session close reap both kinds."""
if not owner:
return
with _REMOTE_KERNELS_LOCK:
doomed = [k for k in _REMOTE_KERNELS if k[0] == owner]
kernels = [_REMOTE_KERNELS.pop(k) for k in doomed]
for kernel in kernels:
_kill(kernel)
def _reap_unlocked(idle_timeout: int) -> List["RemoteKernel"]:
"""Pop idle-expired remote kernels; caller tears them down outside the lock.
Mirrors tools.code_kernel._reap_unlocked. The remote runner itself
self-exits after the same idle window (REMOTE_KERNEL_RUNNER_SOURCE's
IDLE_EXIT_SECONDS), so this only needs to clear the HOST-side
bookkeeping entry — without it, _REMOTE_KERNELS grows one entry per
distinct (owner, env_type, task_env_id) that is never revisited, for
the life of the gateway process.
"""
now = time.monotonic()
doomed = [
key
for key, kernel in _REMOTE_KERNELS.items()
if kernel.attached == 0 and now - kernel.last_used > idle_timeout
]
return [_REMOTE_KERNELS.pop(key) for key in doomed]
def _evict_over_cap_unlocked(keep: Tuple) -> List["RemoteKernel"]:
"""Pop least-recently-used remote kernels beyond the process-wide cap.
Mirrors tools.code_kernel._evict_over_cap_unlocked, reusing the same
max_session_kernels config as an independent bound on _REMOTE_KERNELS.
"""
from tools.code_kernel import _lifecycle_limits
cap, _ = _lifecycle_limits()
if len(_REMOTE_KERNELS) <= cap:
return []
by_age = sorted(
(key for key in _REMOTE_KERNELS if key != keep and _REMOTE_KERNELS[key].attached == 0),
key=lambda key: _REMOTE_KERNELS[key].last_used,
)
doomed = by_age[: len(_REMOTE_KERNELS) - cap]
return [_REMOTE_KERNELS.pop(key) for key in doomed]
atexit.register(shutdown_all_remote_kernels)
def _spawn_remote_kernel(env, env_type: str, owner: str, task_env_id: str,
sandbox_tools: frozenset, *,
idle_exit: int) -> Optional[RemoteKernel]:
"""Start a detached kernel runner on the remote. None on failure."""
from tools.code_execution_tool import (
MAX_STDOUT_BYTES,
_ship_file_to_remote,
_env_temp_dir,
generate_hermes_tools_module,
)
import secrets as _secrets
kernel_dir = f"{_env_temp_dir(env)}/hermes_rkernel_{uuid.uuid4().hex[:12]}"
q_dir = shlex.quote(kernel_dir)
try:
env.execute(f"mkdir -p {q_dir}/cells {q_dir}/rpc", cwd="/", timeout=15)
rpc_token = _secrets.token_urlsafe(32)
runner_src = REMOTE_KERNEL_RUNNER_SOURCE.format(
capture_limit=MAX_STDOUT_BYTES,
idle_exit=idle_exit,
)
_ship_file_to_remote(env, f"{kernel_dir}/kernel_runner.py", runner_src)
tools_src = generate_hermes_tools_module(
list(sandbox_tools), transport="file",
)
_ship_file_to_remote(env, f"{kernel_dir}/hermes_tools.py", tools_src)
env_prefix = (
f"HERMES_KERNEL_DIR={q_dir} "
f"HERMES_RPC_DIR={shlex.quote(kernel_dir + '/rpc')} "
f"HERMES_RPC_TOKEN={shlex.quote(rpc_token)} "
f"PYTHONDONTWRITEBYTECODE=1 PYTHONPATH={q_dir}"
)
started = env.execute(
f"cd {q_dir} && nohup env {env_prefix} python3 kernel_runner.py "
f"> {q_dir}/runner.log 2>&1 & echo PID:$!",
cwd="/", timeout=20,
)
pid = ""
for line in (started.get("output", "") or "").splitlines():
if line.strip().startswith("PID:"):
pid = line.strip()[4:].strip()
break
if not pid.isdigit():
logger.warning("remote kernel spawn returned no PID: %r",
started.get("output", ""))
env.execute(f"rm -rf {q_dir}", cwd="/", timeout=15)
return None
kernel = RemoteKernel(
env=env, env_type=env_type, kernel_dir=kernel_dir,
pid=pid, rpc_token=rpc_token, owner=owner,
)
if not _is_alive(kernel):
# Died instantly (missing python3 was pre-checked by the caller,
# so this is unexpected) — surface the runner log at debug.
try:
log = env.execute(f"cat {q_dir}/runner.log", cwd="/", timeout=10)
logger.warning("remote kernel died at spawn: %s",
(log.get("output", "") or "")[:500])
except Exception:
pass
env.execute(f"rm -rf {q_dir}", cwd="/", timeout=15)
return None
return kernel
except Exception:
logger.warning("remote kernel spawn failed", exc_info=True)
try:
env.execute(f"rm -rf {q_dir}", cwd="/", timeout=15)
except Exception:
pass
return None
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 does output post-processing),
or ``None`` when no kernel could be spawned — the caller falls open to
the per-call path. ``state_lost`` / ``state_reset`` / ``reused`` ride in
the ``kernel`` sub-dict, matching the local kernel's result shape.
"""
from tools.code_kernel import _resolve_owner
owner = _resolve_owner(task_env_id)
key = _kernel_key(owner, env_type, task_env_id)
state_lost = False
state_reset = False
with _REMOTE_KERNELS_LOCK:
expired = _reap_unlocked(idle_exit)
kernel = _REMOTE_KERNELS.get(key)
for doomed in expired:
_kill(doomed)
if kernel is not None and reset:
with _REMOTE_KERNELS_LOCK:
_REMOTE_KERNELS.pop(key, None)
_kill(kernel)
kernel = None
state_reset = True
if kernel is not None and not _is_alive(kernel):
# Transport drop, container restart, self-reaped on idle, OOM — all
# the same answer: report the loss, respawn fresh.
with _REMOTE_KERNELS_LOCK:
_REMOTE_KERNELS.pop(key, None)
_kill(kernel) # best-effort dir cleanup; process is already gone
kernel = None
state_lost = True
reused = kernel is not None
if kernel is None:
kernel = _spawn_remote_kernel(
env, env_type, owner, task_env_id, sandbox_tools,
idle_exit=idle_exit,
)
if kernel is None:
return None # fail open to per-call
with _REMOTE_KERNELS_LOCK:
_REMOTE_KERNELS[key] = kernel
kernel.last_used = time.monotonic()
with _REMOTE_KERNELS_LOCK:
kernel.attached += 1
evicted = _evict_over_cap_unlocked(keep=key)
for doomed in evicted:
_kill(doomed)
try:
return _run_remote_cell(
kernel, key, code, env=env, task_env_id=task_env_id,
sandbox_tools=sandbox_tools, timeout=timeout,
max_tool_calls=max_tool_calls, reused=reused,
state_reset=state_reset, state_lost=state_lost,
)
finally:
with _REMOTE_KERNELS_LOCK:
kernel.attached -= 1
kernel.last_used = time.monotonic()
def _run_remote_cell(
kernel: RemoteKernel,
key: Tuple,
code: str,
*,
env,
task_env_id: str,
sandbox_tools: frozenset,
timeout: int,
max_tool_calls: int,
reused: bool,
state_reset: bool,
state_lost: bool,
) -> Dict[str, Any]:
from tools.code_execution_tool import (
_rpc_poll_loop,
_ship_file_to_remote,
)
from tools.thread_context import propagate_context_to_thread
kernel.cell_seq += 1
seq = f"{kernel.cell_seq:06d}"
q_cells = shlex.quote(f"{kernel.kernel_dir}/cells")
# Clean stale tool-RPC requests from a previous cell before arming this
# cell's poll loop, so a background thread the last cell leaked cannot
# smuggle a call into this cell's authority window.
try:
env.execute(
f"rm -f {shlex.quote(kernel.kernel_dir + '/rpc')}/req_* "
f"{shlex.quote(kernel.kernel_dir + '/rpc')}/res_*",
cwd="/", timeout=10,
)
except Exception:
pass
tool_call_log: list = []
tool_call_counter = [0]
stop_event = 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,
)
rpc_thread.start()
cell_status = "no-result"
cell_payload: Dict[str, Any] = {}
try:
request = json.dumps({"id": seq, "code": code}, ensure_ascii=False)
_ship_file_to_remote(
env, f"{kernel.kernel_dir}/cells/cell_req_{seq}.json.tmp", request,
)
env.execute(
f"mv {q_cells}/cell_req_{seq}.json.tmp {q_cells}/cell_req_{seq}.json",
cwd="/", timeout=10,
)
deadline = time.monotonic() + timeout
res_name = f"cell_res_{seq}.json"
while time.monotonic() < deadline:
try:
probe = env.execute(
f"cat {q_cells}/{shlex.quote(res_name)} 2>/dev/null",
cwd="/", timeout=20,
)
except Exception:
# One flaky round-trip is not kernel death; liveness decides.
time.sleep(_CELL_POLL_INTERVAL)
continue
body = (probe.get("output", "") or "").strip()
if body:
try:
cell_payload = json.loads(body)
cell_status = cell_payload.get("status", "error")
except ValueError:
cell_status = "protocol-error"
env.execute(
f"rm -f {q_cells}/{shlex.quote(res_name)}",
cwd="/", timeout=10,
)
break
time.sleep(_CELL_POLL_INTERVAL)
else:
cell_status = "timeout"
finally:
stop_event.set()
rpc_thread.join(timeout=5)
if cell_status in ("timeout", "protocol-error", "no-result"):
# No safe way to interrupt one cell in place (same contract as
# local): kill the kernel, report the loss, respawn next call.
with _REMOTE_KERNELS_LOCK:
_REMOTE_KERNELS.pop(key, None)
_kill(kernel)
return {
"status": "timeout" if cell_status == "timeout" else "error",
"stdout": "",
"stderr": "",
"traceback": "",
"tool_calls_made": tool_call_counter[0],
"kernel": {
"reused": reused,
"remote": True,
"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."
),
},
}
if cell_status == "exit":
with _REMOTE_KERNELS_LOCK:
_REMOTE_KERNELS.pop(key, None)
_kill(kernel)
kernel.execution_count = int(cell_payload.get("execution_count", 0) or 0)
result: Dict[str, Any] = {
"status": "success" if cell_status in ("ok", "exit") else "error",
"stdout": cell_payload.get("stdout", ""),
"stderr": cell_payload.get("stderr", ""),
"traceback": cell_payload.get("traceback", ""),
"stdout_clipped": bool(cell_payload.get("stdout_clipped")),
"stderr_clipped": bool(cell_payload.get("stderr_clipped")),
"tool_calls_made": tool_call_counter[0],
"kernel": {
"reused": reused,
"remote": True,
"execution_count": kernel.execution_count,
},
}
if cell_status == "exit":
result["kernel"]["ended"] = True
if state_reset:
result["kernel"]["state_reset"] = True
if state_lost:
result["kernel"]["state_lost"] = True
result["kernel"]["note"] = (
"The previous remote kernel was gone (transport drop, container "
"restart, or idle self-exit); state from earlier calls was lost "
"and a fresh kernel was started."
)
if cell_status == "error" and result["traceback"]:
result["error"] = result["traceback"].strip().splitlines()[-1]
return result