jj.wang
4f45d368ad
Before the launch call the planned run (run id, terminal, handle, bookkeeping) is saved as a launching marker. If the launch call fails or agentd dies before recording the run, the next tick probes the run directory (exit.txt / launch.pid / pid.txt): evidence -> adopt the run; nothing -> relaunch with the same run id in <run_id>-r2 (-r3 ...); probe unreachable -> keep the marker, which holds the terminal slot. Transports refuse to clean a run directory holding a pid file. A terminal whose relay fails is skipped for the rest of the tick (heartbeat.down) while other terminals carry on.
220 lines
11 KiB
Python
220 lines
11 KiB
Python
"""Remote POSIX transport (warpgate-ssh, os linux/darwin): the generated bash is executed by a local bash in
|
|
place of the relay, so quoting, chunked uploads, detached launch and process-group kill are exercised for real."""
|
|
import base64
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import tempfile
|
|
import time
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
from agentd import ROOT, config
|
|
from agentd.config import Config
|
|
from agentd.executors import Run, make_executor
|
|
from agentd.transports import make_transport
|
|
from agentd.transports.warpgate_ssh import PosixWarpgateTransport, relay_argv
|
|
from tests.fakes import config_data
|
|
|
|
FAKE_CLAUDE = """#!/usr/bin/env bash
|
|
prompt="$(cat)"
|
|
echo "args: $*" >&2
|
|
printf '{"type":"result","subtype":"success","is_error":false,"result":"[review] got %s","session_id":"sess-7"}' "${#prompt}"
|
|
"""
|
|
|
|
|
|
def remote_data(root):
|
|
d = config_data()
|
|
d["vault"] = {"client": "/opt/vault/agent-access.py", "config": "https://vault.example/identity.json"}
|
|
d["terminals"]["mac"] = {"transport": "warpgate-ssh", "target": "mac-target", "os": "linux", "launch": "nohup",
|
|
"repo": str(root / "my repo's"), "run_root": str(root / "run root"), "plane_user": "oregon",
|
|
"token_key": "TOK_MAC", "concurrency": 1, "executors": ["claude-code", "shell"],
|
|
"bins": {"claude": str(root / "bin dir" / "claude")}}
|
|
return d
|
|
|
|
|
|
class LocalBash(PosixWarpgateTransport):
|
|
"""The relay replaced by `bash -c` on this machine; records every script (one entry = one relay call)."""
|
|
def __init__(self, cfg, terminal):
|
|
super().__init__(cfg, terminal)
|
|
self.calls = []
|
|
|
|
def sh(self, script, timeout=150, marker=None):
|
|
self.calls.append(script)
|
|
out = subprocess.run(["bash", "-c", script], capture_output=True, text=True, timeout=timeout,
|
|
env=self.cfg.scrubbed_env())
|
|
if marker and marker not in out.stdout:
|
|
raise AssertionError("marker %s missing: %s %s" % (marker, out.stdout, out.stderr))
|
|
return out.stdout
|
|
|
|
|
|
def wait(ex, h, seconds=10):
|
|
end = time.time() + seconds
|
|
while time.time() < end:
|
|
if ex.poll(h) == "finished":
|
|
return True
|
|
time.sleep(0.1)
|
|
return False
|
|
|
|
|
|
class PosixRemoteTest(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.root = Path(self.tmp.name)
|
|
(self.root / "my repo's").mkdir()
|
|
(self.root / "bin dir").mkdir()
|
|
fake = self.root / "bin dir" / "claude"
|
|
fake.write_text(FAKE_CLAUDE)
|
|
fake.chmod(0o755)
|
|
self.cfg = Config(remote_data(self.root))
|
|
self.t = LocalBash(self.cfg, "mac")
|
|
|
|
def tearDown(self):
|
|
self.tmp.cleanup()
|
|
|
|
def test_factory_picks_posix_for_linux_and_config_is_valid(self):
|
|
self.assertEqual(config.validate(remote_data(self.root)), [])
|
|
self.assertIsInstance(make_transport(self.cfg, "mac"), PosixWarpgateTransport)
|
|
self.assertEqual(make_transport(self.cfg, "win10").os, "windows")
|
|
|
|
def test_live_config_runs_oregon_remotely(self):
|
|
cfg = config.load(ROOT / "config" / "agentd.yaml")
|
|
t = make_transport(cfg, "oregon")
|
|
self.assertIsInstance(t, PosixWarpgateTransport)
|
|
self.assertEqual((t.target, t.os), ("oregon", "linux"))
|
|
self.assertEqual(cfg.executors["codex"]["tier_defaults"], {"gpt-6.1-sol": "fast"})
|
|
|
|
def test_relay_argv_and_scrubbed_env(self):
|
|
argv = relay_argv(self.cfg, "mac-target", ["bash", "-c", "echo hi"], 30)
|
|
self.assertEqual(argv[1:], ["-I", "/opt/vault/agent-access.py", "--config", "https://vault.example/identity.json",
|
|
"run", "--target", "mac-target", "--timeout", "30", "--", "bash", "-c", "echo hi"])
|
|
seen = {}
|
|
|
|
def fake_run(args, **kw):
|
|
seen.update(args=args, env=kw["env"])
|
|
return mock.Mock(returncode=0, stdout=b"MKDIR=ok\n", stderr=b"")
|
|
with mock.patch.dict(os.environ, {"TOK_MAC": "s", "TOK_W10": "s"}), \
|
|
mock.patch("agentd.transports.warpgate_ssh.subprocess.run", fake_run):
|
|
PosixWarpgateTransport(self.cfg, "mac").mkdir("/x")
|
|
self.assertNotIn("TOK_MAC", seen["env"])
|
|
self.assertNotIn("TOK_W10", seen["env"])
|
|
self.assertEqual(seen["args"][-3:-1], ["bash", "-c"])
|
|
|
|
def test_claude_round_is_one_call_and_resume_uses_session_id(self):
|
|
ex = make_executor(self.cfg, "mac", "claude-code", self.t)
|
|
prompt = "it's \"quoted\" $HOME `x` \\ 中文\n"
|
|
h = ex.start(Run("GW-4-1", self.cfg.terminals["mac"]["repo"], prompt, model="sonnet"))
|
|
self.assertEqual(len(self.t.calls), 1) # mkdir + prompt + run.sh + launch in one relay call
|
|
self.assertIn("setsid nohup bash 'run.sh' </dev/null >/dev/null 2>&1 &", self.t.calls[0])
|
|
self.assertTrue(wait(ex, h))
|
|
rd = Path(h["run_dir"])
|
|
self.assertEqual((rd / "prompt.md").read_text(encoding="utf-8"), prompt)
|
|
self.assertEqual((rd / "pid.txt").read_text().strip(), (rd / "launch.pid").read_text().strip())
|
|
res = ex.collect(h)
|
|
self.assertEqual((res.last_message, res.exit_code, res.thread_id),
|
|
("[review] got %d" % len(prompt.rstrip("\n")), 0, "sess-7"))
|
|
cmd = ex.command(Run("GW-4-2", "r", "p", model="sonnet", thread_id="sess-7"), h["run_dir"])
|
|
self.assertEqual(cmd, "'%s' -p --resume sess-7" % (self.root / "bin dir" / "claude"))
|
|
|
|
def test_runner_quotes_paths_and_cd_failure_writes_exit(self):
|
|
_, script = self.t.runner("/a b/O'Neil/run", "/no such/repo's", "cmd", "out.json")
|
|
self.assertIn("cd '/no such/repo'\\''s' ||", script)
|
|
self.assertIn("< '/a b/O'\\''Neil/run'/prompt.md > '/a b/O'\\''Neil/run'/out.json", script)
|
|
ex = make_executor(self.cfg, "mac", "shell", self.t)
|
|
h = ex.start(Run("GW-4-3", str(self.root / "missing"), "", script="true"))
|
|
self.assertTrue(wait(ex, h))
|
|
self.assertEqual(ex.exit_code(h), 97)
|
|
|
|
def test_chunked_upload_and_reads(self):
|
|
self.t.CHUNK = 64
|
|
path = str(self.root / "run root" / "x y" / "f'.txt")
|
|
text = "line1 'q' \"d\" $(nope)\n" * 20 + "last ünïcode"
|
|
self.t.mkdir(os.path.dirname(path))
|
|
n = len(self.t.calls)
|
|
self.t.write_text(path, text)
|
|
self.assertGreater(len(self.t.calls) - n, 2)
|
|
self.assertEqual(Path(path).read_text(encoding="utf-8"), text)
|
|
self.assertFalse(os.path.exists(path + ".b64"))
|
|
self.assertEqual(self.t.read_text(path), text)
|
|
self.assertIsNone(self.t.read_text(path + ".missing"))
|
|
self.assertEqual(self.t.grep_first(path, "$(nope)"), "line1 'q' \"d\" $(nope)")
|
|
self.assertIsNone(self.t.grep_first(path, "absent"))
|
|
self.assertEqual(self.t.tail(path, 2), "line1 'q' \"d\" $(nope)\nlast ünïcode")
|
|
self.assertEqual(self.t.tail(path + ".missing", 2), "")
|
|
|
|
def test_big_prompt_goes_up_before_the_launch_call(self):
|
|
self.t.CHUNK = 100
|
|
ex = make_executor(self.cfg, "mac", "shell", self.t)
|
|
h = ex.start(Run("GW-4-4", self.cfg.terminals["mac"]["repo"], "z" * 500, script="wc -c"))
|
|
self.assertGreater(len(self.t.calls), 2)
|
|
self.assertTrue(wait(ex, h))
|
|
self.assertEqual(ex.collect(h).last_message, "500")
|
|
|
|
@unittest.skipUnless(shutil.which("setsid") and shutil.which("pgrep"), "needs setsid/pgrep")
|
|
def test_kill_ends_the_whole_process_group(self):
|
|
ex = make_executor(self.cfg, "mac", "shell", self.t)
|
|
marker = "agentd-test-%d" % os.getpid()
|
|
h = ex.start(Run("GW-4-5", self.cfg.terminals["mac"]["repo"], "",
|
|
script="bash -c 'sleep 300; : %s' & sleep 300; : %s" % (marker, marker)))
|
|
end = time.time() + 5
|
|
while time.time() < end and not subprocess.run(["pgrep", "-f", marker], capture_output=True).stdout:
|
|
time.sleep(0.1)
|
|
self.assertTrue(subprocess.run(["pgrep", "-f", marker], capture_output=True).stdout)
|
|
ex.kill(h)
|
|
self.assertIn("KILLED=ok left=0", subprocess.run(["bash", "-c", self.t.kill_script(h["run_dir"])],
|
|
capture_output=True, text=True).stdout)
|
|
self.assertFalse(subprocess.run(["pgrep", "-f", marker], capture_output=True).stdout)
|
|
self.assertEqual(ex.poll(h), "running") # killed runners write no exit.txt; core treats it as cancelled
|
|
|
|
def test_kill_after_finish_does_not_signal_a_reused_pid(self):
|
|
ex = make_executor(self.cfg, "mac", "shell", self.t)
|
|
h = ex.start(Run("GW-4-6", self.cfg.terminals["mac"]["repo"], "", script="true"))
|
|
self.assertTrue(wait(ex, h))
|
|
(Path(h["run_dir"]) / "pid.txt").write_text("1\n")
|
|
self.assertIn("KILLED=finished", subprocess.run(["bash", "-c", self.t.kill_script(h["run_dir"])],
|
|
capture_output=True, text=True).stdout)
|
|
|
|
def test_darwin_detaches_without_setsid_binary(self):
|
|
d = remote_data(self.root)
|
|
d["terminals"]["mac"]["os"] = "darwin"
|
|
t = PosixWarpgateTransport(Config(d), "mac")
|
|
s = t.launch_script("/r", "run.sh")
|
|
self.assertIn("nohup perl -MPOSIX -e 'POSIX::setsid() or die", s)
|
|
self.assertNotIn("setsid nohup", s)
|
|
self.assertIn("echo $! > launch.pid", s)
|
|
|
|
def test_semantic_checks(self):
|
|
d = remote_data(self.root)
|
|
d["terminals"]["mac"]["launch"] = "schtasks"
|
|
d["terminals"]["mac"]["run_root"] = "~/runs"
|
|
del d["vault"]["config"]
|
|
errs = config.validate(d)
|
|
for frag in ("needs launch=nohup", "mac.run_root: must be an absolute path", "need vault.config"):
|
|
self.assertTrue(any(frag in e for e in errs), (frag, errs))
|
|
d = remote_data(self.root)
|
|
d["terminals"]["win10"]["target"] = "win11"
|
|
self.assertTrue(any("only reach target 'win10'" in e for e in config.validate(d)))
|
|
|
|
def test_upload_payload_is_plain_base64(self):
|
|
s = self.t.put_script("/r/p.md", base64.b64encode("x'y".encode()).decode())
|
|
self.assertIn("printf %s 'eCd5' | base64 -d > '/r/p.md.part' && mv -f '/r/p.md.part' '/r/p.md'", s)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|
|
|
|
def test_probe_dir_name_and_pid_guard(self):
|
|
ex = make_executor(self.cfg, "mac", "shell", self.t)
|
|
run = Run("GW-5-1", self.cfg.terminals["mac"]["repo"], "", script="sleep 0.2", dir_name="GW-5-1-r2")
|
|
self.assertEqual(ex.probe(ex.handle(run)), [])
|
|
h = ex.start(run)
|
|
self.assertEqual((os.path.basename(h["run_dir"]), h["task"]), ("GW-5-1-r2", "agentd-GW-5-1-r2"))
|
|
self.assertTrue(wait(ex, h))
|
|
self.assertEqual(ex.probe(h), ["exit.txt", "launch.pid", "pid.txt"])
|
|
n = len(self.t.calls)
|
|
with self.assertRaises(AssertionError): # LocalBash: marker LAUNCHED= missing (HASPID)
|
|
ex.start(run)
|
|
self.assertIn("HASPID", self.t.calls[n])
|
|
self.assertTrue(os.path.exists(os.path.join(h["run_dir"], "exit.txt"))) # nothing wiped
|