Agent-System/scripts/tests/test_agent_worktree.py

"""Worktree lifecycle checks use real Git and the real agentctl registry."""
import fcntl
import importlib.util
import hashlib
import json
import os
from pathlib import Path
import shutil
import subprocess
import sys
import tempfile
import time
import unittest
from unittest import mock


SOURCE = Path(__file__).resolve().parents[1]


class WorktreeTest(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.addCleanup(self.tmp.cleanup)
        self.root = Path(self.tmp.name) / "repo"
        self.root.mkdir()
        (self.root / "scripts").mkdir()
        for name in ("agent-worktree.sh", "agent_worktree.py", "agentctl.py", "agentctl", "agentctl_w4.py", "agentctl_catalog.py", "agentctl_visibility.py"):
            shutil.copy2(SOURCE / name, self.root / "scripts" / name)
        (self.root / "configs").mkdir(parents=True, exist_ok=True)
        shutil.copy2(SOURCE.parent / "configs" / "agent-governance.json",
                     self.root / "configs" / "agent-governance.json")
        (self.root / ".gitignore").write_text(".worktrees/\n.locks/\nignored-evidence/\n")
        (self.root / "source.txt").write_text("base\n")
        self.env = {**os.environ, "PYTHONDONTWRITEBYTECODE": "1"}
        for name in ("AGENTCTL_RUN_ID", "AGENTCTL_TASK_ID"):
            self.env.pop(name, None)
        self.git("init", "-b", "dev")
        self.git("config", "user.email", "test@example.invalid")
        self.git("config", "user.name", "Test")
        self.git("add", ".")
        self.git("commit", "-qm", "fixture")
        self.base = self.git("rev-parse", "HEAD").stdout.strip()

    def run_cmd(self, args, cwd=None, env=None):
        return subprocess.run(args, cwd=cwd or self.root, env=env or self.env,
                              text=True, capture_output=True, timeout=20)

    def git(self, *args, cwd=None):
        result = self.run_cmd(["git", *args], cwd=cwd)
        self.assertEqual(result.returncode, 0, result.stderr)
        return result

    def ctl(self, *args, cwd=None):
        return self.run_cmd([sys.executable, str(self.root / "scripts/agentctl.py"), *args], cwd=cwd)

    def init(self):
        result = self.ctl("init")
        self.assertEqual(result.returncode, 0, result.stderr)

    def helper(self, *args, script_root=None, cwd=None, env=None):
        script = (script_root or self.root) / "scripts/agent-worktree.sh"
        return self.run_cmd(["bash", str(script), *args], cwd=cwd, env=env)

    def create(self, branch="feature/task"):
        result = self.helper("create", branch)
        self.assertEqual(result.returncode, 0, result.stderr)
        return Path(json.loads(result.stdout)["worktree"])

    def remove_with_visibility_fixture(self, *args):
        """Fixture only process cwd enumeration; Git, registry and lock checks stay real.

        This host may hide unrelated same-user daemon cwds. Positive removal tests
        model a fully visible process namespace containing the helper itself.
        Live cwd rejection below still runs the actual host process enumeration.
        """
        code = '''
import importlib.util, os, sys
from pathlib import Path
from unittest import mock
spec = importlib.util.spec_from_file_location("helper", sys.argv[1])
helper = importlib.util.module_from_spec(spec)
spec.loader.exec_module(helper)
def visible_cwds():
    yield str(os.getpid()), Path.cwd()
sys.argv = sys.argv[1:]
with mock.patch.object(helper, "iter_same_user_cwds", visible_cwds):
    sys.exit(helper.main())
'''
        return self.run_cmd([sys.executable, "-c", code, str(self.root / "scripts/agent_worktree.py"), "remove", *args])

    def test_list_is_read_only_without_init(self):
        result = self.helper("list")
        self.assertEqual(result.returncode, 0, result.stderr)
        self.assertIn(str(self.root), result.stdout)
        self.assertFalse((self.root / ".locks").exists())
        self.assertFalse((self.root / ".worktrees").exists())

    def test_write_requires_explicit_init(self):
        result = self.helper("create", "feature/task")
        self.assertNotEqual(result.returncode, 0)
        self.assertIn("NOT_INITIALIZED", result.stderr)
        self.assertFalse((self.root / ".worktrees").exists())
        self.assertFalse((self.root / ".locks").exists())

    def test_create_from_feature_uses_current_dev_and_canonical_main(self):
        self.init()
        first = self.create()
        self.git("commit", "--allow-empty", "-qm", "feature-only", cwd=first)
        result = self.helper("create", "fix/second", script_root=first, cwd=first / "scripts")
        self.assertEqual(result.returncode, 0, result.stderr)
        data = json.loads(result.stdout)
        self.assertEqual(data["base_dev_sha"], self.base)
        self.assertEqual(Path(data["worktree"]).resolve(), (self.root / ".worktrees/second").resolve())
        self.assertEqual(self.git("rev-parse", "fix/second").stdout.strip(), self.base)
        self.assertFalse((first / ".worktrees").exists())
        self.assertFalse((first / ".locks").exists())

    def test_create_with_main_on_unrelated_branch_uses_dev(self):
        self.init()
        self.git("checkout", "-b", "feature/legacy")
        self.git("commit", "--allow-empty", "-qm", "legacy-only")
        self.create()
        self.assertEqual(self.git("rev-parse", "feature/task").stdout.strip(), self.base)
        self.assertEqual(self.git("branch", "--show-current").stdout.strip(), "feature/legacy")

    def test_rejects_traversal_nested_and_protected_branches(self):
        self.init()
        for branch in ("dev", "main", "test", "feature/../escape", "feature/a/b", "fix/.hidden", "feature/-x"):
            with self.subTest(branch=branch):
                self.assertNotEqual(self.helper("create", branch).returncode, 0)
        self.assertFalse((self.root / ".worktrees").exists())

    def test_same_task_id_and_existing_branch_do_not_collide(self):
        self.init()
        first = self.create()
        self.assertNotEqual(self.helper("create", "fix/task").returncode, 0)
        self.assertNotEqual(self.helper("create", "feature/task").returncode, 0)
        self.assertEqual(self.git("branch", "--show-current", cwd=first).stdout.strip(), "feature/task")

    def test_symlink_parent_or_target_is_refused(self):
        self.init()
        outside = Path(self.tmp.name) / "outside"
        outside.mkdir()
        parent = self.root / ".worktrees"
        parent.symlink_to(outside, target_is_directory=True)
        self.assertNotEqual(self.helper("create", "feature/task").returncode, 0)
        self.assertEqual(list(outside.iterdir()), [])
        parent.unlink()
        parent.mkdir()
        (parent / "task").symlink_to(outside, target_is_directory=True)
        self.assertNotEqual(self.helper("create", "feature/task").returncode, 0)
        self.assertEqual(list(outside.iterdir()), [])

    def test_stable_lock_refuses_concurrent_writer(self):
        self.init()
        lock = self.root / ".locks/os/worktrees.lock"
        with lock.open("a+") as holder:
            inode = lock.stat().st_ino
            fcntl.flock(holder.fileno(), fcntl.LOCK_EX)
            result = self.helper("create", "feature/task")
            self.assertNotEqual(result.returncode, 0)
            self.assertIn("busy", result.stderr)
        self.create()
        self.assertEqual(lock.stat().st_ino, inode)

    def test_forged_inherited_lease_is_refused(self):
        self.init()
        env = {**self.env, "AGENTCTL_RUN_ID": "forged", "AGENTCTL_TASK_ID": "task"}
        self.assertNotEqual(self.helper("create", "feature/task", env=env).returncode, 0)
        self.assertFalse((self.root / ".worktrees/task").exists())

    def test_cleanup_legacy_mapping_and_nonforced_branch_delete_with_visibility_fixture(self):
        self.init()
        legacy = self.root / ".worktrees/feature-task"
        self.git("worktree", "add", "-b", "feature/task", str(legacy), self.base)
        result = self.remove_with_visibility_fixture("feature/task", "--delete-branch")
        self.assertEqual(result.returncode, 0, result.stderr)
        self.assertFalse(legacy.exists())
        self.assertNotIn("refs/heads/feature/task", self.git("show-ref").stdout)

    def test_cleanup_refuses_unmerged_commit(self):
        self.init()
        task = self.create()
        self.git("commit", "--allow-empty", "-qm", "unique", cwd=task)
        result = self.helper("remove", "feature/task", "--delete-branch")
        self.assertNotEqual(result.returncode, 0)
        self.assertIn("ancestor", result.stderr)
        self.assertTrue(task.exists())

    def test_branch_delete_refusal_preserves_branch_without_force_with_visibility_fixture(self):
        self.init()
        self.git("branch", "feature/legacy", self.base)
        self.git("commit", "--allow-empty", "-qm", "dev-next")
        task = self.create()
        self.git("checkout", "feature/legacy")
        result = self.remove_with_visibility_fixture("feature/task", "--delete-branch")
        self.assertNotEqual(result.returncode, 0)
        self.assertIn("branch retained", result.stderr)
        self.assertFalse(task.exists())
        self.assertIn("refs/heads/feature/task", self.git("show-ref").stdout)

    def test_cleanup_refuses_tracked_untracked_and_ignored_evidence(self):
        self.init()
        task = self.create()
        for relative in ("source.txt", "new-evidence.txt", "ignored-evidence/only-copy.log"):
            with self.subTest(relative=relative):
                file = task / relative
                original = file.read_text() if file.exists() else None
                file.parent.mkdir(exist_ok=True)
                file.write_text("unique evidence\n")
                result = self.helper("remove", "feature/task", "--delete-branch")
                self.assertNotEqual(result.returncode, 0)
                self.assertTrue(file.exists())
                self.assertEqual(file.read_text(), "unique evidence\n")
                if original is not None:
                    file.write_text(original)
                else:
                    file.unlink()

    def test_cleanup_refuses_live_process_cwd(self):
        self.init()
        task = self.create()
        process = subprocess.Popen([sys.executable, "-c", "import time; time.sleep(30)"], cwd=task)
        try:
            result = self.helper("remove", "feature/task")
            self.assertNotEqual(result.returncode, 0)
            self.assertIn(f"process {process.pid} still uses", result.stderr)
            self.assertTrue(task.exists())
        finally:
            process.terminate()
            process.wait(timeout=5)

    def test_cleanup_refuses_unfinished_git_operation(self):
        self.init()
        task = self.create()
        marker = Path(self.git("rev-parse", "--path-format=absolute", "--git-path", "MERGE_HEAD", cwd=task).stdout.strip())
        marker.write_text(self.base + "\n")
        result = self.helper("remove", "feature/task")
        self.assertNotEqual(result.returncode, 0)
        self.assertIn("unfinished", result.stderr)
        self.assertTrue(task.exists())

    def test_cleanup_actual_branch_mapping_with_visibility_fixture(self):
        self.init()
        first = self.create()
        other = self.root / ".worktrees/legacy-other"
        self.git("worktree", "add", "-b", "fix/task", str(other), self.base)
        result = self.remove_with_visibility_fixture("fix/task")
        self.assertEqual(result.returncode, 0, result.stderr)
        self.assertTrue(first.exists())
        self.assertFalse(other.exists())

    def test_inherited_worktree_lease_does_not_double_lock(self):
        self.init()
        task = self.create()
        result = self.ctl("task", "register", "--id", "task", "--owner", "test", "--scope", "scripts",
                          "--capability", "git:worktrees", cwd=task)
        self.assertEqual(result.returncode, 0, result.stderr)
        result = self.ctl("with", "--task", "task", "--resource", "git:worktrees", "--cwd", str(self.root),
                          "--", "bash", str(self.root / "scripts/agent-worktree.sh"), "create", "fix/child", cwd=task)
        self.assertEqual(result.returncode, 0, result.stderr)
        self.assertTrue((self.root / ".worktrees/child").exists())

    def test_cleanup_lifecycle_lock_prevents_a_new_run(self):
        self.init()
        task = self.create()
        result = self.ctl("task", "register", "--id", "task", "--owner", "test", "--scope", "scripts", cwd=task)
        self.assertEqual(result.returncode, 0, result.stderr)
        key = hashlib.sha256(str(task.resolve()).encode()).hexdigest()
        lock = self.root / ".locks/os" / f"worktree-{key}.lock"
        marker = Path(self.tmp.name) / "should-not-exist"
        with lock.open("a+") as holder:
            fcntl.flock(holder.fileno(), fcntl.LOCK_EX)
            result = self.ctl("with", "--task", "task", "--resource", "path:scripts", "--wait", "0", "--",
                              sys.executable, "-c", "from pathlib import Path; import sys; Path(sys.argv[1]).touch()",
                              str(marker))
        self.assertNotEqual(result.returncode, 0)
        self.assertFalse(marker.exists())

    def test_cleanup_refuses_active_registry_operation_even_with_other_cwd(self):
        self.init()
        task = self.create()
        result = self.ctl("task", "register", "--id", "task", "--owner", "test", "--scope", "scripts", cwd=task)
        self.assertEqual(result.returncode, 0, result.stderr)
        ready = Path(self.tmp.name) / "ready"
        child_code = "from pathlib import Path; import os,time,sys; os.chdir(sys.argv[2]); Path(sys.argv[1]).touch(); time.sleep(30)"
        proc = subprocess.Popen([sys.executable, str(self.root / "scripts/agentctl.py"), "with", "--task", "task",
                                 "--resource", "path:scripts", "--", sys.executable,
                                 "-c", child_code, str(ready), str(self.root)], cwd=self.root, env=self.env,
                                stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        try:
            deadline = time.monotonic() + 5
            while not ready.exists() and proc.poll() is None and time.monotonic() < deadline:
                time.sleep(0.02)
            self.assertTrue(ready.exists(), "supervised operation did not start")
            result = self.helper("remove", "feature/task")
            self.assertNotEqual(result.returncode, 0)
            self.assertTrue(task.exists())
            self.assertNotIn("still uses", result.stderr)
            self.assertRegex(result.stderr, "active registered operation|live or unresolved operation")
        finally:
            proc.terminate()
            proc.communicate(timeout=8)


class ProcessVisibilityTest(unittest.TestCase):
    def test_unknown_same_uid_inaccessible_cwd_is_not_skipped(self):
        spec = importlib.util.spec_from_file_location("agent_worktree_under_test", SOURCE / "agent_worktree.py")
        helper = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(helper)

        def unreadable():
            raise helper.Refused("cannot verify cwd of same-user process 987654321; cleanup_pending")

        with tempfile.TemporaryDirectory() as temp:
            with mock.patch.object(helper, "iter_same_user_cwds", side_effect=unreadable):
                with self.assertRaisesRegex(helper.Refused, "cannot verify cwd"):
                    helper.assert_no_process_cwd(Path(temp) / "task")

    def test_linux_permission_error_path_remains_fail_closed(self):
        spec = importlib.util.spec_from_file_location("agent_worktree_under_test", SOURCE / "agent_worktree.py")
        helper = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(helper)
        with tempfile.TemporaryDirectory() as temp:
            proc = Path(temp) / "987654321"
            proc.mkdir()
            (proc / "comm").write_text("coding-agent\n")
            with mock.patch.object(helper.sys, "platform", "linux"), \
                    mock.patch.object(helper.Path, "iterdir", return_value=iter([proc])), \
                    mock.patch.object(helper.os, "getuid", return_value=0), \
                    mock.patch.object(helper.os, "readlink", side_effect=PermissionError), \
                    mock.patch.object(Path, "stat", lambda self: type("S", (), {"st_uid": 0})()):
                with self.assertRaisesRegex(helper.Refused, "cannot verify cwd"):
                    helper.assert_no_process_cwd(Path(temp) / "task")


if __name__ == "__main__":
    unittest.main()