jj.wang
222 lines
11 KiB
Python
222 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") # PAD-12: direct-ssh over Tailscale, Warpgate relay as fallback
|
|
self.assertIsInstance(t.primary, PosixWarpgateTransport)
|
|
self.assertIsInstance(t.fallback, PosixWarpgateTransport)
|
|
self.assertEqual((t.primary.channel.stdin, t.fallback.channel.stdin), (True, False))
|
|
self.assertEqual((t.fallback.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
|