plane-agentd/tests/test_transports.py

"""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["webhook"] = {"secret_key": "HOOK_GW", "tick_command": ["true"]}
            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])