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.
150 lines
7.2 KiB
Python
150 lines
7.2 KiB
Python
"""Transport hardening: secrets never reach executors, Windows script quoting, command shapes, atomic state files."""
|
|
import json
|
|
import os
|
|
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.store import Store
|
|
from agentd.transports.local import LocalTransport
|
|
from agentd.transports.warpgate_ssh import PosixWarpgateTransport, WarpgateSshTransport, WindowsWarpgateTransport
|
|
from tests.fakes import config_data
|
|
|
|
|
|
class RecordingWin(WindowsWarpgateTransport):
|
|
"""Windows transport with the relay replaced by an in-memory recorder."""
|
|
def __init__(self, cfg, terminal):
|
|
super().__init__(cfg, terminal)
|
|
self.files, self.launched = {}, []
|
|
|
|
def mkdir(self, path, clean=False):
|
|
pass
|
|
|
|
def write_text(self, path, text):
|
|
self.files[path] = text
|
|
|
|
def launch(self, task, script_path, run_dir):
|
|
self.launched.append((task, script_path, run_dir))
|
|
|
|
|
|
class LocalSecretsTest(unittest.TestCase):
|
|
def test_executor_env_has_no_plane_tokens(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
data = config_data()
|
|
data["agentd"]["state_dir"] = str(Path(tmp) / "state")
|
|
data["terminals"]["oregon"]["repo"] = tmp
|
|
data["projects"]["GWENT"]["webhook_secret_key"] = "HOOK_GW"
|
|
cfg = Config(data)
|
|
self.assertEqual(cfg.secret_keys(), {"TOK_OR", "TOK_W10", "HOOK_GW"})
|
|
t = LocalTransport(cfg, "oregon")
|
|
ex = make_executor(cfg, "oregon", "shell", t)
|
|
script = 'echo "or=${TOK_OR:-none} w10=${TOK_W10:-none} hook=${HOOK_GW:-none} keep=${KEEP_ME:-none}"'
|
|
with mock.patch.dict(os.environ, {"TOK_OR": "s1", "TOK_W10": "s2", "HOOK_GW": "s3", "KEEP_ME": "yes"}):
|
|
h = ex.start(Run("GW-9-1", tmp, "", script=script))
|
|
end = time.time() + 10
|
|
while ex.poll(h) != "finished" and time.time() < end:
|
|
time.sleep(0.05)
|
|
self.assertEqual(ex.collect(h).last_message, "or=none w10=none hook=none keep=yes")
|
|
|
|
|
|
class WindowsScriptTest(unittest.TestCase):
|
|
def setUp(self):
|
|
self.cfg = config.load(ROOT / "config" / "agentd.example.yaml")
|
|
self.t = RecordingWin(self.cfg, "win10")
|
|
self.ex = make_executor(self.cfg, "win10", "codex", self.t)
|
|
|
|
def test_codex_run_and_resume_commands(self):
|
|
run_dir = self.t.join(self.ex.run_root(), "GW-5-1")
|
|
fresh = self.ex.command(Run("GW-5-1", self.t.repo, "p", "gpt-6-astra", "high"), run_dir)
|
|
self.assertTrue(fresh.startswith("& 'C:\\Users\\Administrator\\AppData\\Local\\Programs\\OpenAI\\Codex\\bin\\codex.exe' exec --json"))
|
|
self.assertIn("-C '%s' -s workspace-write" % self.t.repo, fresh)
|
|
self.assertIn("-o '%s\\last.md'" % run_dir, fresh)
|
|
resumed = self.ex.command(Run("GW-5-2", self.t.repo, "p", "gpt-6-astra", "high", thread_id="th-9"), run_dir)
|
|
self.assertIn(" exec resume th-9 --json -m gpt-6-astra", resumed)
|
|
self.assertNotIn(" -C ", resumed) # `codex exec resume` rejects -C/-s; the runner already cd's into the repo
|
|
self.assertNotIn(" -s ", resumed)
|
|
tiered = self.ex.command(Run("GW-5-3", self.t.repo, "p", "gpt-6.1-sol", "medium", tier="fast"), run_dir)
|
|
self.assertIn("-c 'service_tier=\"fast\"'", tiered)
|
|
|
|
def test_run_dirs_are_unique_per_run(self):
|
|
a = self.ex.start(Run("GW-5-1", self.t.repo, "p", "m", "e"))
|
|
b = self.ex.start(Run("GW-5-2", self.t.repo, "p", "m", "e"))
|
|
self.assertNotEqual(a["run_dir"], b["run_dir"])
|
|
self.assertEqual([x[0] for x in self.t.launched], ["agentd-GW-5-1", "agentd-GW-5-2"])
|
|
self.assertTrue(a["run_dir"].endswith("\\runtime\\agentd\\GW-5-1"))
|
|
|
|
def test_runner_quotes_paths_with_spaces_and_apostrophes(self):
|
|
_, script = self.t.runner(r"C:\a b\O'Neil\run", r"C:\my repo\O'Neil", "& 'C:\\x\\codex.exe' exec -", "events.jsonl")
|
|
self.assertIn(r"Set-Location -LiteralPath 'C:\my repo\O''Neil' -ErrorAction Stop", script)
|
|
self.assertIn(r"'C:\a b\O''Neil\run\prompt.md'", script)
|
|
self.assertIn(r"'C:\a b\O''Neil\run\exit.txt' -Value ('exit=97", script) # cd failure still ends the run
|
|
self.assertIn(r"Add-Content -Path 'C:\a b\O''Neil\run\events.jsonl'", script)
|
|
|
|
def test_launch_script_quoting(self):
|
|
s = self.t.launch_script("agentd-GW-1-1", r"C:\My Docs\O'Neil\run.ps1")
|
|
self.assertIn(r"-File \"C:\My Docs\O''Neil\run.ps1\"'", s)
|
|
self.assertIn("/TN 'agentd-GW-1-1'", s)
|
|
|
|
def test_relay_env_has_no_plane_tokens(self):
|
|
seen = {}
|
|
|
|
def fake_run(args, **kw):
|
|
seen.update(kw["env"])
|
|
return mock.Mock(returncode=0, stdout=b"X=ok", stderr=b"")
|
|
with mock.patch.dict(os.environ, {"PLANE_WINDOW10_API_TOKEN": "s"}), \
|
|
mock.patch("agentd.transports.warpgate_ssh.subprocess.run", fake_run):
|
|
WarpgateSshTransport(self.cfg, "win10").ps("'X=ok'")
|
|
self.assertNotIn("PLANE_WINDOW10_API_TOKEN", seen)
|
|
self.assertIn("WIN10_TIMEOUT", seen)
|
|
|
|
|
|
class StoreTest(unittest.TestCase):
|
|
def test_save_is_atomic_and_heartbeat_is_complete_json(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
st = Store(tmp, echo=False)
|
|
st.save({"projects": {"A": {"issues": {}}}})
|
|
with mock.patch("agentd.store.os.replace", side_effect=OSError("disk full")):
|
|
with self.assertRaises(OSError):
|
|
st.save({"projects": {"B": {}}})
|
|
self.assertIn("A", st.load()["projects"]) # old state survives a failed write
|
|
st.heartbeat({"errors": [], "inflight": []})
|
|
hb = json.loads((Path(tmp) / "heartbeat.json").read_text())
|
|
self.assertEqual(hb["errors"], [])
|
|
self.assertIn("ts", hb)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|
|
|
|
|
|
class GuardTest(unittest.TestCase):
|
|
def test_local_clean_refuses_dir_with_pid_file(self):
|
|
from agentd.transports import TransportError
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
t = LocalTransport(Config(config_data()), "oregon")
|
|
(Path(tmp) / "prompt.md").write_text("x")
|
|
t.mkdir(tmp, clean=True)
|
|
self.assertEqual(t.probe(tmp, ("exit.txt", "pid.txt")), [])
|
|
(Path(tmp) / "pid.txt").write_text("1")
|
|
with self.assertRaises(TransportError):
|
|
t.mkdir(tmp, clean=True)
|
|
self.assertEqual(t.probe(tmp, ("exit.txt", "pid.txt")), ["pid.txt"])
|
|
|
|
def test_windows_probe_and_clean_guard(self):
|
|
from agentd.transports import TransportError
|
|
cfg = config.load(ROOT / "config" / "agentd.example.yaml")
|
|
t = WindowsWarpgateTransport(cfg, "win10")
|
|
scripts = []
|
|
t.ps = lambda s, timeout=150: (scripts.append(s), "PROBE=launch.pid,pid.txt\n")[1]
|
|
self.assertEqual(t.probe("C:\\r\\GW-1-1", ("exit.txt", "launch.pid", "pid.txt")), ["launch.pid", "pid.txt"])
|
|
self.assertIn("Join-Path $d $_", scripts[-1])
|
|
t.ps = lambda s, timeout=150: (scripts.append(s), "HASPID\n")[1]
|
|
with self.assertRaises(TransportError):
|
|
t.mkdir("C:\\r\\GW-1-1", clean=True)
|
|
self.assertIn("'pid.txt'", scripts[-1])
|