f49c9d738f
The lockdown tests pinned literal substrings ("( set -a", "exec python3",
"umask 077", "chmod 700"), which is change-detection. Replay the recorded
commands through a real bash under a temp root instead: dirs come out
0700 and the archive 0600, the sourced env reaches the child, the
child's exit code is the command's, and the token does not leak into the
outer shell (the session-snapshot class). The token-never-in-argv
assertions stay. No new test functions; the replay is skipped on
Windows, where the stdin/argv invariants are still checked.
393 lines
16 KiB
Python
393 lines
16 KiB
Python
"""Remote session kernels (tools/code_kernel_remote.py) — hermes-agent#96873.
|
|
|
|
These tests drive execute_in_remote_kernel against a scripted fake env that
|
|
implements the same contract as docker/ssh/modal envs (run-to-completion
|
|
execute()), with canned outputs for the spawn/liveness/cell round-trips.
|
|
The REAL end-to-end behavior (actual detached processes, real files, real
|
|
kill) was verified live on Windows against a bash-backed env; these tests
|
|
pin the host-side protocol logic: spawn parsing, liveness handling,
|
|
state_lost/state_reset reporting, fail-open, and owner isolation.
|
|
"""
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
|
|
|
|
from tools.code_kernel_remote import (
|
|
_REMOTE_KERNELS,
|
|
RemoteKernel,
|
|
execute_in_remote_kernel,
|
|
shutdown_all_remote_kernels,
|
|
shutdown_remote_kernels_for_owner,
|
|
)
|
|
|
|
|
|
class ScriptedEnv:
|
|
"""Contract-faithful fake: answers env.execute() from a script table.
|
|
|
|
Handlers are (substring, callable) pairs checked in order; the callable
|
|
receives the command and returns the result dict.
|
|
"""
|
|
|
|
def __init__(self, handlers):
|
|
self.handlers = handlers
|
|
self.commands = []
|
|
self.stdin_payloads = []
|
|
|
|
def get_temp_dir(self):
|
|
return "/tmp"
|
|
|
|
def execute(self, command, cwd=None, timeout=None, stdin_data=None):
|
|
self.commands.append(command)
|
|
self.stdin_payloads.append(stdin_data)
|
|
for needle, handler in self.handlers:
|
|
if needle in command:
|
|
return handler(command)
|
|
return {"output": "", "returncode": 0}
|
|
|
|
|
|
def _spawn_ok_handlers(cell_results):
|
|
"""Handlers for a healthy kernel: spawn returns PID, liveness ALIVE,
|
|
cat of a cell result file returns the next canned payload."""
|
|
results = list(cell_results)
|
|
|
|
def cat_handler(command):
|
|
if results:
|
|
return {"output": json.dumps(results.pop(0)), "returncode": 0}
|
|
return {"output": "", "returncode": 0}
|
|
|
|
return [
|
|
("nohup", lambda c: {"output": "PID:4242\n", "returncode": 0}),
|
|
("kill -0", lambda c: {"output": "ALIVE\n", "returncode": 0}),
|
|
("cat ", cat_handler),
|
|
]
|
|
|
|
|
|
def _cell(status="ok", stdout="", execution_count=1, **kw):
|
|
payload = {
|
|
"id": "000001", "status": status, "stdout": stdout, "stderr": "",
|
|
"stdout_clipped": False, "stderr_clipped": False, "traceback": "",
|
|
"execution_count": execution_count,
|
|
}
|
|
payload.update(kw)
|
|
return payload
|
|
|
|
|
|
def _run(env, code="print(1)", *, task="t1", reset=False, timeout=10,
|
|
tools=frozenset({"read_file"})):
|
|
return execute_in_remote_kernel(
|
|
code, env=env, env_type="ssh", task_env_id=task,
|
|
sandbox_tools=tools, timeout=timeout,
|
|
max_tool_calls=5, reset=reset,
|
|
)
|
|
|
|
|
|
class RemoteKernelBase(unittest.TestCase):
|
|
def setUp(self):
|
|
shutdown_all_remote_kernels()
|
|
# No approval session key in tests → owner falls back to task id,
|
|
# which is exactly the isolation-by-key behavior under test.
|
|
self._ship = patch(
|
|
"tools.code_execution_tool._ship_file_to_remote",
|
|
)
|
|
self._ship.start()
|
|
self._poll = patch(
|
|
"tools.code_execution_tool._rpc_poll_loop",
|
|
)
|
|
self._poll.start()
|
|
|
|
def tearDown(self):
|
|
self._ship.stop()
|
|
self._poll.stop()
|
|
shutdown_all_remote_kernels()
|
|
|
|
|
|
class TestSpawnAndReuse(RemoteKernelBase):
|
|
def test_first_call_spawns_second_reuses(self):
|
|
env = ScriptedEnv(_spawn_ok_handlers(
|
|
[_cell(stdout="one\n"), _cell(stdout="two\n", execution_count=2)],
|
|
))
|
|
first = _run(env)
|
|
self.assertEqual(first["status"], "success", first)
|
|
self.assertFalse(first["kernel"]["reused"])
|
|
second = _run(env)
|
|
self.assertTrue(second["kernel"]["reused"])
|
|
self.assertEqual(second["kernel"]["execution_count"], 2)
|
|
# Exactly one spawn happened.
|
|
self.assertEqual(
|
|
sum(1 for c in env.commands if "nohup" in c), 1,
|
|
)
|
|
|
|
def test_spawn_failure_fails_open(self):
|
|
env = ScriptedEnv([
|
|
("nohup", lambda c: {"output": "sh: cannot fork\n", "returncode": 1}),
|
|
])
|
|
self.assertIsNone(_run(env))
|
|
self.assertEqual(len(_REMOTE_KERNELS), 0)
|
|
|
|
def test_reset_kills_and_respawns(self):
|
|
env = ScriptedEnv(_spawn_ok_handlers([_cell(), _cell()]))
|
|
_run(env)
|
|
result = _run(env, reset=True)
|
|
self.assertTrue(result["kernel"].get("state_reset"))
|
|
self.assertFalse(result["kernel"]["reused"])
|
|
self.assertEqual(sum(1 for c in env.commands if "nohup" in c), 2)
|
|
|
|
|
|
class TestDeathDetection(RemoteKernelBase):
|
|
def test_dead_kernel_is_reported_and_respawned(self):
|
|
env = ScriptedEnv(_spawn_ok_handlers([_cell(), _cell()]))
|
|
_run(env)
|
|
# Flip liveness to dead for the next probe only.
|
|
original = env.handlers
|
|
env.handlers = [("kill -0", lambda c: {"output": "", "returncode": 1})] \
|
|
+ [h for h in original if h[0] != "kill -0"]
|
|
# Restore ALIVE after the respawn's own probe would run: the spawn
|
|
# path probes liveness once — make the dead answer one-shot.
|
|
state = {"dead_probes": 0}
|
|
|
|
def flaky_liveness(command):
|
|
state["dead_probes"] += 1
|
|
if state["dead_probes"] == 1:
|
|
return {"output": "", "returncode": 1}
|
|
return {"output": "ALIVE\n", "returncode": 0}
|
|
|
|
env.handlers = [("kill -0", flaky_liveness)] + \
|
|
[h for h in original if h[0] != "kill -0"]
|
|
result = _run(env)
|
|
self.assertEqual(result["status"], "success", result)
|
|
self.assertTrue(result["kernel"].get("state_lost"))
|
|
self.assertIn("state from earlier calls was lost",
|
|
result["kernel"].get("note", ""))
|
|
|
|
def test_cell_timeout_kills_kernel_and_reports(self):
|
|
# cat never returns a result file → cell deadline expires.
|
|
env = ScriptedEnv([
|
|
("nohup", lambda c: {"output": "PID:77\n", "returncode": 0}),
|
|
("kill -0", lambda c: {"output": "ALIVE\n", "returncode": 0}),
|
|
("cat ", lambda c: {"output": "", "returncode": 0}),
|
|
])
|
|
result = _run(env, timeout=2)
|
|
self.assertEqual(result["status"], "timeout")
|
|
self.assertTrue(result["kernel"]["state_lost"])
|
|
self.assertEqual(len(_REMOTE_KERNELS), 0)
|
|
# The kernel was actually killed on the remote.
|
|
self.assertTrue(any("kill " in c for c in env.commands))
|
|
|
|
|
|
class TestOwnershipIsolation(RemoteKernelBase):
|
|
def test_changed_tool_set_spawns_kernel_with_fresh_stubs(self):
|
|
env = ScriptedEnv(_spawn_ok_handlers([_cell(), _cell()]))
|
|
_run(env, tools=frozenset({"read_file"}))
|
|
_run(env, tools=frozenset({"web_search"}))
|
|
|
|
self.assertEqual(len(_REMOTE_KERNELS), 2)
|
|
self.assertEqual(sum(1 for c in env.commands if "nohup" in c), 2)
|
|
keyed_tool_sets = {key[-1] for key in _REMOTE_KERNELS}
|
|
self.assertEqual(
|
|
keyed_tool_sets,
|
|
{("read_file",), ("web_search",)},
|
|
)
|
|
|
|
def test_delegated_children_get_their_own_remote_kernels(self):
|
|
"""Same invariant as local (#94647 review fix): the child context
|
|
qualifier must key a DIFFERENT remote kernel."""
|
|
from agent.delegation_context import delegated_child_context
|
|
|
|
env = ScriptedEnv(_spawn_ok_handlers([_cell(), _cell()]))
|
|
_run(env, task="conv")
|
|
with delegated_child_context("child-9"):
|
|
_run(env, task="conv")
|
|
# Two distinct kernels, two spawns.
|
|
self.assertEqual(len(_REMOTE_KERNELS), 2)
|
|
self.assertEqual(sum(1 for c in env.commands if "nohup" in c), 2)
|
|
|
|
def test_owner_disposal_reaps_only_that_owner(self):
|
|
env = ScriptedEnv(_spawn_ok_handlers([_cell(), _cell()]))
|
|
_run(env, task="owner-a")
|
|
_run(env, task="owner-b")
|
|
self.assertEqual(len(_REMOTE_KERNELS), 2)
|
|
shutdown_remote_kernels_for_owner("owner-a")
|
|
self.assertEqual(len(_REMOTE_KERNELS), 1)
|
|
remaining_owner = next(iter(_REMOTE_KERNELS))[0]
|
|
self.assertEqual(remaining_owner, "owner-b")
|
|
|
|
|
|
class TestIdleReapAndCapEviction(RemoteKernelBase):
|
|
"""Unlike local session kernels, remote kernels had no idle-reap or
|
|
process-wide cap: _REMOTE_KERNELS grew one entry per distinct
|
|
(owner, env_type, task_env_id) that was never revisited, for the life
|
|
of the gateway process."""
|
|
|
|
def test_idle_expired_kernel_is_reaped_on_next_call(self):
|
|
env = ScriptedEnv(_spawn_ok_handlers([_cell(), _cell()]))
|
|
execute_in_remote_kernel(
|
|
"print(1)", env=env, env_type="ssh", task_env_id="stale",
|
|
sandbox_tools=frozenset(), timeout=10, max_tool_calls=5,
|
|
reset=False, idle_exit=1800,
|
|
)
|
|
self.assertEqual(len(_REMOTE_KERNELS), 1)
|
|
# Backdate the kernel's last_used past the idle window — simulates
|
|
# a key that is never revisited again.
|
|
for kernel in _REMOTE_KERNELS.values():
|
|
kernel.last_used -= 2000
|
|
# A call for a DIFFERENT key must reap the stale entry on entry,
|
|
# without ever touching or reviving it.
|
|
execute_in_remote_kernel(
|
|
"print(1)", env=env, env_type="ssh", task_env_id="fresh",
|
|
sandbox_tools=frozenset(), timeout=10, max_tool_calls=5,
|
|
reset=False, idle_exit=1800,
|
|
)
|
|
owners = {key[0] for key in _REMOTE_KERNELS}
|
|
self.assertNotIn("stale", owners)
|
|
self.assertIn("fresh", owners)
|
|
|
|
def test_over_cap_evicts_least_recently_used(self):
|
|
with patch("tools.code_kernel._lifecycle_limits", return_value=(2, 1800)):
|
|
env = ScriptedEnv(_spawn_ok_handlers([_cell() for _ in range(10)]))
|
|
for i in range(3):
|
|
execute_in_remote_kernel(
|
|
"print(1)", env=env, env_type="ssh", task_env_id=f"owner-{i}",
|
|
sandbox_tools=frozenset(), timeout=10, max_tool_calls=5,
|
|
reset=False, idle_exit=1800,
|
|
)
|
|
self.assertEqual(len(_REMOTE_KERNELS), 2)
|
|
owners = {key[0] for key in _REMOTE_KERNELS}
|
|
self.assertNotIn("owner-0", owners)
|
|
self.assertIn("owner-1", owners)
|
|
self.assertIn("owner-2", owners)
|
|
|
|
def test_eviction_skips_kernels_with_a_running_cell(self):
|
|
"""Cap eviction must never kill a kernel mid-cell (the local-kernel
|
|
race from hermes-agent#101861): a busy kernel stays put and a
|
|
settled one goes instead, even if the busy one is older."""
|
|
import threading
|
|
|
|
gate = threading.Event()
|
|
|
|
def slow_cat(command):
|
|
gate.wait(10)
|
|
return {"output": json.dumps(_cell()), "returncode": 0}
|
|
|
|
busy_env = ScriptedEnv([
|
|
("nohup", lambda c: {"output": "PID:4242\n", "returncode": 0}),
|
|
("kill -0", lambda c: {"output": "ALIVE\n", "returncode": 0}),
|
|
("cat ", slow_cat),
|
|
])
|
|
with patch("tools.code_kernel._lifecycle_limits", return_value=(1, 1800)):
|
|
worker = threading.Thread(target=_run, args=(busy_env,), kwargs={"task": "busy"})
|
|
worker.start()
|
|
# Snapshot: the worker thread inserts into the registry concurrently and a live
|
|
# dict iteration raises "dictionary changed size during iteration".
|
|
while not any(k.attached for k in list(_REMOTE_KERNELS.values())):
|
|
time.sleep(0.005)
|
|
env = ScriptedEnv(_spawn_ok_handlers([_cell()]))
|
|
_run(env, task="settled")
|
|
owners = {key[0] for key in _REMOTE_KERNELS}
|
|
self.assertIn("busy", owners)
|
|
gate.set()
|
|
worker.join(10)
|
|
self.assertFalse(any("kill 4242" in c for c in busy_env.commands))
|
|
|
|
|
|
class TestDispatchIntegration(unittest.TestCase):
|
|
"""_execute_remote prefers the kernel and falls open to per-call."""
|
|
|
|
def test_execute_remote_uses_kernel_result(self):
|
|
from tools.code_execution_tool import _execute_remote
|
|
|
|
fake = {
|
|
"status": "success", "stdout": "kernel says hi\n", "stderr": "",
|
|
"traceback": "", "tool_calls_made": 0,
|
|
"kernel": {"reused": True, "remote": True, "execution_count": 3},
|
|
}
|
|
env = ScriptedEnv([
|
|
("command -v python3", lambda c: {"output": "OK\n", "returncode": 0}),
|
|
])
|
|
with patch("tools.code_execution_tool._load_config",
|
|
return_value={"timeout": 30, "max_tool_calls": 5}), \
|
|
patch("tools.code_execution_tool._get_or_create_env",
|
|
return_value=(env, "ssh")), \
|
|
patch("tools.code_kernel_remote.execute_in_remote_kernel",
|
|
return_value=fake):
|
|
result = json.loads(_execute_remote("print()", "t", ["read_file"]))
|
|
self.assertEqual(result["status"], "success")
|
|
self.assertIn("kernel says hi", result["output"])
|
|
self.assertEqual(result["kernel"]["execution_count"], 3)
|
|
|
|
def test_execute_remote_falls_open_to_per_call(self):
|
|
from tools.code_execution_tool import _execute_remote
|
|
from unittest.mock import MagicMock
|
|
|
|
env = ScriptedEnv([
|
|
("command -v python3", lambda c: {"output": "OK\n", "returncode": 0}),
|
|
("python3 script.py", lambda c: {"output": "per-call ran\n",
|
|
"returncode": 0}),
|
|
])
|
|
with patch("tools.code_execution_tool._load_config",
|
|
return_value={"timeout": 30, "max_tool_calls": 5}), \
|
|
patch("tools.code_execution_tool._get_or_create_env",
|
|
return_value=(env, "ssh")), \
|
|
patch("tools.code_kernel_remote.execute_in_remote_kernel",
|
|
return_value=None), \
|
|
patch("tools.code_execution_tool._ship_file_to_remote"), \
|
|
patch("tools.code_execution_tool.threading.Thread",
|
|
return_value=MagicMock()):
|
|
result = json.loads(_execute_remote("print()", "t", ["read_file"]))
|
|
self.assertEqual(result["status"], "success")
|
|
self.assertIn("per-call ran", result["output"])
|
|
|
|
|
|
class TestSharedHostLockdown(RemoteKernelBase):
|
|
"""Shared-host hardening: the kernel dir lives under a shared temp dir, so
|
|
every dir must be owner-only, every remote write owner-only, and the RPC
|
|
token must travel in a sourced env file rather than a ps-visible argv."""
|
|
|
|
def test_spawn_locks_down_dirs_and_hides_token(self):
|
|
self._ship.stop() # let the real ship commands reach env.commands
|
|
env = ScriptedEnv(_spawn_ok_handlers([_cell(stdout="hi\n")]))
|
|
result = _run(env)
|
|
self.assertEqual(result["status"], "success", result)
|
|
kernel = next(iter(_REMOTE_KERNELS.values()))
|
|
# The token never rides a command line: a remote shell's argv is
|
|
# world-readable via ps for the command's whole lifetime.
|
|
self.assertFalse(
|
|
any(kernel.rpc_token in c for c in env.commands),
|
|
"rpc token appeared in a remote command line")
|
|
spawn_cmd = next(c for c in env.commands if "nohup" in c)
|
|
self.assertNotIn("HERMES_RPC_TOKEN=", spawn_cmd)
|
|
mkdir_cmd = next(c for c in env.commands if "mkdir -p" in c)
|
|
if sys.platform != "win32":
|
|
# Behaviour, not command text: replay the recorded dir setup
|
|
# through a real shell under a private temp root; every dir under
|
|
# the shared temp dir must come out owner-only.
|
|
import shutil
|
|
import subprocess
|
|
import tempfile
|
|
root = tempfile.mkdtemp()
|
|
self.addCleanup(shutil.rmtree, root, True)
|
|
local = root + kernel.kernel_dir
|
|
r = subprocess.run(["bash", "-c", mkdir_cmd.replace(kernel.kernel_dir, local)],
|
|
capture_output=True, text=True)
|
|
self.assertEqual(r.returncode, 0, r.stderr)
|
|
for d in (local, f"{local}/cells", f"{local}/rpc"):
|
|
self.assertEqual(os.stat(d).st_mode & 0o777, 0o700, d)
|
|
# Ships write owner-only; on a pipe-capable backend the payload rides
|
|
# stdin, so the base64 (which decodes to the token for kernel.env)
|
|
# never enters argv either.
|
|
ship_cmds = [c for c in env.commands if "base64 -d" in c]
|
|
self.assertTrue(any("kernel.env" in c for c in ship_cmds))
|
|
import base64
|
|
env_ship = next(p for p, c in zip(env.stdin_payloads, env.commands)
|
|
if p and "kernel.env" in c)
|
|
env_content = base64.b64decode(env_ship).decode()
|
|
self.assertIn(f"HERMES_RPC_TOKEN={kernel.rpc_token}", env_content)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|