WI-3 (#830): compaction observers → revoke + D4 same-PID generation auto-revoke (#842)

packages/mosaic/src/mutator-gate/runtime_tools_unittest.py

@@ -9,6 +9,7 @@ import json
 import os
 import runpy
 import socket
+import stat
 import subprocess
 import sys
 import tempfile
@@ -16,10 +17,13 @@ import threading
 import unittest
 from contextlib import redirect_stderr
 from pathlib import Path
+from types import SimpleNamespace
 from unittest.mock import patch
 
 
 TOOLS_DIR = Path(__file__).parents[2] / "framework/tools/lease-broker"
+if str(TOOLS_DIR) not in sys.path:
+    sys.path.insert(0, str(TOOLS_DIR))
 
 
 def load_tool(module_name: str, filename: str):
@@ -78,6 +82,9 @@ class LaunchRuntimeTest(unittest.TestCase):
         def execute(command: str, argv: list[str], environment: dict[str, str]) -> None:
             calls["execute"] = (command, argv, environment)
 
+        def initialize_generation(path: Path, generation: int) -> None:
+            calls["generation"] = (path, generation)
+
         result = LAUNCHER.main(
             ["--runtime", "claude", "--", "claude", "--print", "hello"],
             environ={
@@ -87,6 +94,7 @@ class LaunchRuntimeTest(unittest.TestCase):
             },
             request=request,
             execute=execute,
+            initialize_generation=initialize_generation,
         )
 
         self.assertEqual(result, 0)
@@ -101,6 +109,14 @@ class LaunchRuntimeTest(unittest.TestCase):
         self.assertEqual(environment["MOSAIC_LEASE_SESSION_ID"], session_id)
         self.assertEqual(environment["MOSAIC_RUNTIME_GENERATION"], "7")
         self.assertEqual(environment["MOSAIC_LEASE_RUNTIME"], "claude")
+        self.assertEqual(
+            environment["MOSAIC_LEASE_GENERATION_FILE"],
+            f"/run/test/generation-{session_id}.state",
+        )
+        self.assertEqual(
+            calls["generation"],
+            (Path(f"/run/test/generation-{session_id}.state"), 7),
+        )
         self.assertEqual(environment["PRESERVED"], "yes")
 
     def test_dangerous_claude_mode_is_owned_and_injected_by_the_wrapper(self) -> None:
@@ -110,6 +126,7 @@ class LaunchRuntimeTest(unittest.TestCase):
             environ={"MOSAIC_LEASE_BROKER_SOCKET": "/broker"},
             request=lambda *_args: {"ok": True, "session_id": "e" * 64},
             execute=lambda *args: executed.append(args),
+            initialize_generation=lambda *_args: None,
         )
         self.assertEqual(result, 0)
         self.assertEqual(
@@ -135,6 +152,7 @@ class LaunchRuntimeTest(unittest.TestCase):
             environ={"MOSAIC_LEASE_BROKER_SOCKET": "/broker"},
             request=lambda *_args: {"ok": True, "session_id": "f" * 64},
             execute=lambda *args: executed.append(args),
+            initialize_generation=lambda *_args: None,
         )
         self.assertEqual(result, 0)
         self.assertEqual(executed[0][0:2], ("pi", ["pi", "--print", "hello"]))
@@ -174,6 +192,21 @@ class LaunchRuntimeTest(unittest.TestCase):
                 self.assertEqual(result, 1)
                 self.assertEqual(executed, [])
 
+    def test_generation_initialization_failure_denies_before_exec(self) -> None:
+        with redirect_stderr(io.StringIO()):
+            self.assertEqual(
+                LAUNCHER.main(
+                    ["--runtime", "pi", "--", "pi"],
+                    environ={"MOSAIC_LEASE_BROKER_SOCKET": "/broker"},
+                    request=lambda *_args: {"ok": True, "session_id": "a" * 64},
+                    execute=lambda *_args: self.fail("must not execute"),
+                    initialize_generation=lambda *_args: (_ for _ in ()).throw(
+                        OSError("unsafe state")
+                    ),
+                ),
+                1,
+            )
+
     def test_registration_exceptions_fail_closed(self) -> None:
         failures = [ValueError("bad"), OSError("down"), json.JSONDecodeError("bad", "x", 0)]
         for failure in failures:
@@ -199,6 +232,7 @@ class LaunchRuntimeTest(unittest.TestCase):
                     environ={"MOSAIC_LEASE_BROKER_SOCKET": "/x"},
                     request=lambda *_args: {"ok": True, "session_id": "c" * 64},
                     execute=lambda *_args: (_ for _ in ()).throw(OSError("missing")),
+                    initialize_generation=lambda *_args: None,
                 ),
                 1,
             )
@@ -284,6 +318,23 @@ class ExecutableEntrypointTest(unittest.TestCase):
             runpy.run_path(str(TOOLS_DIR / "launch-runtime.py"), run_name="__main__")
         self.assertEqual(raised.exception.code, 64)
 
+    def test_revoker_entrypoint_denies_when_identity_environment_is_absent(self) -> None:
+        with patch.object(
+            sys,
+            "argv",
+            [
+                str(TOOLS_DIR / "revoke-lease.py"),
+                "--runtime",
+                "claude",
+                "--reason",
+                "pre-compact",
+            ],
+        ), patch.dict(os.environ, {}, clear=True), redirect_stderr(
+            io.StringIO()
+        ), self.assertRaises(SystemExit) as raised:
+            runpy.run_path(str(TOOLS_DIR / "revoke-lease.py"), run_name="__main__")
+        self.assertEqual(raised.exception.code, 2)
+
     def test_gate_entrypoint_denies_when_identity_environment_is_absent(self) -> None:
         class Stdin:
             buffer = io.BytesIO(b'{"tool_name":"Bash"}')
@@ -325,6 +376,19 @@ class MutatorGateTest(unittest.TestCase):
             )
         return result, stderr.getvalue(), calls
 
+    def test_generation_file_is_the_effective_generation_authority(self) -> None:
+        calls: list[dict[str, object]] = []
+        result = GATE.main(
+            ["--runtime", "pi"],
+            environ=self.environment(),
+            stream=io.BytesIO(b'{"tool_name":"bash"}'),
+            request=lambda _path, payload: calls.append(payload)
+            or {"ok": True, "decision": "allow"},
+            resolve_generation=lambda _environment: 9,
+        )
+        self.assertEqual(result, 0)
+        self.assertEqual(calls[0]["runtime_generation"], 9)
+
     def test_allow_and_denial_decisions(self) -> None:
         allowed, allowed_stderr, calls = self.run_main()
         self.assertEqual(allowed, 0)
@@ -432,5 +496,252 @@ class MutatorGateTest(unittest.TestCase):
                         self.assertEqual(fake.shutdown_how, socket.SHUT_WR)
 
 
+class LeaseRevocationTest(unittest.TestCase):
+    def load_modules(self):
+        generation_path = TOOLS_DIR / "lease_generation.py"
+        revoker_path = TOOLS_DIR / "revoke-lease.py"
+        self.assertTrue(generation_path.is_file(), "lease_generation.py must be shipped")
+        self.assertTrue(revoker_path.is_file(), "revoke-lease.py must be shipped")
+        return (
+            load_tool("lease_generation_test", "lease_generation.py"),
+            load_tool("lease_revoker_test", "revoke-lease.py"),
+        )
+
+    def test_generation_parser_and_descriptor_security_rejections(self) -> None:
+        generation, _ = self.load_modules()
+        for value in (None, "", "-1", "no", "é", str(generation.MAX_GENERATION + 1)):
+            with self.subTest(value=value), self.assertRaises(ValueError):
+                generation.parse_generation(value)
+        for value in (-1, generation.MAX_GENERATION + 1):
+            with self.subTest(initialize=value), tempfile.TemporaryDirectory() as directory, self.assertRaises(ValueError):
+                generation.initialize_runtime_generation(Path(directory) / "state", value)
+
+        with patch.object(
+            generation.os,
+            "fstat",
+            return_value=SimpleNamespace(st_mode=stat.S_IFDIR | 0o700, st_uid=os.geteuid(), st_size=0),
+        ), self.assertRaises(ValueError):
+            generation._validate_descriptor(4)
+        with patch.object(
+            generation.os,
+            "fstat",
+            return_value=SimpleNamespace(st_mode=stat.S_IFREG | 0o600, st_uid=os.geteuid() + 1, st_size=0),
+        ), self.assertRaises(ValueError):
+            generation._validate_descriptor(4)
+        with patch.object(
+            generation.os,
+            "fstat",
+            return_value=SimpleNamespace(st_mode=stat.S_IFREG | 0o644, st_uid=os.geteuid(), st_size=0),
+        ), self.assertRaises(ValueError):
+            generation._validate_descriptor(4)
+        with patch.object(
+            generation.os,
+            "fstat",
+            return_value=SimpleNamespace(
+                st_mode=stat.S_IFREG | 0o600,
+                st_uid=os.geteuid(),
+                st_size=generation.MAX_GENERATION_BYTES + 1,
+            ),
+        ), self.assertRaises(ValueError):
+            generation._validate_descriptor(4)
+
+    def test_generation_file_is_private_monotonic_and_rejects_unsafe_state(self) -> None:
+        generation, _ = self.load_modules()
+        with tempfile.TemporaryDirectory() as directory:
+            path = Path(directory) / "runtime.generation"
+            generation.initialize_runtime_generation(path, 4)
+            self.assertEqual(generation.read_runtime_generation({
+                "MOSAIC_RUNTIME_GENERATION": "1",
+                "MOSAIC_LEASE_GENERATION_FILE": str(path),
+            }), 4)
+            self.assertEqual(generation.bump_runtime_generation({
+                "MOSAIC_RUNTIME_GENERATION": "1",
+                "MOSAIC_LEASE_GENERATION_FILE": str(path),
+            }), 5)
+            self.assertEqual(path.read_text(), "5\n")
+            self.assertEqual(stat.S_IMODE(path.stat().st_mode), 0o600)
+
+            path.write_text("bad\n")
+            with self.assertRaises(ValueError):
+                generation.read_runtime_generation({
+                    "MOSAIC_RUNTIME_GENERATION": "1",
+                    "MOSAIC_LEASE_GENERATION_FILE": str(path),
+                })
+            path.write_bytes(b"\xff\n")
+            with self.assertRaises(ValueError):
+                generation.read_runtime_generation({
+                    "MOSAIC_LEASE_GENERATION_FILE": str(path),
+                })
+            path.write_text(f"{generation.MAX_GENERATION}\n")
+            with self.assertRaises(ValueError):
+                generation.bump_runtime_generation({
+                    "MOSAIC_LEASE_GENERATION_FILE": str(path),
+                })
+            with self.assertRaises(ValueError):
+                generation.bump_runtime_generation({"MOSAIC_RUNTIME_GENERATION": "1"})
+
+    def test_generation_write_must_make_progress(self) -> None:
+        generation, _ = self.load_modules()
+        with tempfile.TemporaryDirectory() as directory, patch.object(
+            generation.os, "write", return_value=0
+        ), self.assertRaises(OSError):
+            generation.initialize_runtime_generation(Path(directory) / "state", 1)
+
+    def test_revoker_reuses_broker_revoke_and_optional_generation_bump(self) -> None:
+        _, revoker = self.load_modules()
+        with tempfile.TemporaryDirectory() as directory:
+            generation_file = Path(directory) / "runtime.generation"
+            generation_file.write_text("2\n")
+            generation_file.chmod(0o600)
+            environment = {
+                "MOSAIC_LEASE_BROKER_SOCKET": "/broker",
+                "MOSAIC_LEASE_SESSION_ID": "a" * 64,
+                "MOSAIC_RUNTIME_GENERATION": "2",
+                "MOSAIC_LEASE_GENERATION_FILE": str(generation_file),
+            }
+            calls: list[tuple[Path, dict[str, object]]] = []
+            result = revoker.main(
+                ["--runtime", "pi", "--reason", "session-start-resume", "--bump-generation"],
+                environ=environment,
+                request=lambda path, payload: calls.append((path, payload))
+                or {"ok": True, "state": "UNVERIFIED"},
+            )
+            self.assertEqual(result, 0)
+            self.assertEqual(generation_file.read_text(), "3\n")
+            self.assertEqual(calls[0][0], Path("/broker"))
+            self.assertEqual(calls[0][1], {
+                "action": "revoke_lease",
+                "session_id": "a" * 64,
+                "runtime_generation": 3,
+                "reason": "session-start-resume",
+                "runtime": "pi",
+            })
+
+            calls.clear()
+            self.assertEqual(
+                revoker.main(
+                    ["--runtime", "pi", "--reason", "pi-context-after-compact"],
+                    environ=environment,
+                    request=lambda path, payload: calls.append((path, payload))
+                    or {"ok": True, "state": "UNVERIFIED"},
+                ),
+                0,
+            )
+            self.assertEqual(calls[0][1]["runtime_generation"], 3)
+
+    def test_revoker_broker_framing_and_shape_validation(self) -> None:
+        _, revoker = self.load_modules()
+        with self.assertRaises(ValueError):
+            revoker.broker_request(Path("/broker"), {"reason": "x" * revoker.MAX_FRAME})
+        replies = [
+            (b'{"ok":true,"state":"UNVERIFIED"}\n', {"ok": True, "state": "UNVERIFIED"}),
+            (b'{"ok":true}', ValueError),
+            (b'[]\n', ValueError),
+            (b'x' * (revoker.MAX_FRAME + 1), ValueError),
+        ]
+        for wire_reply, expected in replies:
+            with self.subTest(size=len(wire_reply)):
+                fake = FakeSocket(wire_reply)
+                with patch.object(revoker.socket, "socket", return_value=fake):
+                    if isinstance(expected, type) and issubclass(expected, Exception):
+                        with self.assertRaises(expected):
+                            revoker.broker_request(Path("/broker"), {"action": "revoke_lease"})
+                    else:
+                        self.assertEqual(
+                            revoker.broker_request(Path("/broker"), {"action": "revoke_lease"}),
+                            expected,
+                        )
+                        self.assertEqual(fake.timeout, revoker.BROKER_TIMEOUT_SECONDS)
+                        self.assertEqual(fake.connected, "/broker")
+                        self.assertEqual(fake.shutdown_how, socket.SHUT_WR)
+
+    def test_failed_observer_revocation_advances_the_local_generation_fence(self) -> None:
+        _, revoker = self.load_modules()
+        with tempfile.TemporaryDirectory() as directory:
+            generation_file = Path(directory) / "runtime.generation"
+            generation_file.write_text("8\n")
+            generation_file.chmod(0o600)
+            environment = {
+                "MOSAIC_LEASE_BROKER_SOCKET": "/broker",
+                "MOSAIC_LEASE_SESSION_ID": "a" * 64,
+                "MOSAIC_RUNTIME_GENERATION": "8",
+                "MOSAIC_LEASE_GENERATION_FILE": str(generation_file),
+            }
+            with redirect_stderr(io.StringIO()):
+                result = revoker.main(
+                    ["--runtime", "claude", "--reason", "session-start-compact"],
+                    environ=environment,
+                    request=lambda *_args: (_ for _ in ()).throw(OSError("down")),
+                )
+            self.assertEqual(result, 2)
+            self.assertEqual(generation_file.read_text(), "9\n")
+
+    def test_revoker_fails_closed_on_identity_reply_and_transport_errors(self) -> None:
+        _, revoker = self.load_modules()
+        good = {
+            "MOSAIC_LEASE_BROKER_SOCKET": "/broker",
+            "MOSAIC_LEASE_SESSION_ID": "a" * 64,
+            "MOSAIC_RUNTIME_GENERATION": "1",
+        }
+        malformed_session = {**good, "MOSAIC_LEASE_SESSION_ID": "not-a-session"}
+        cases = [
+            ({}, lambda *_args: {"ok": True, "state": "UNVERIFIED"}),
+            (malformed_session, lambda *_args: {"ok": True, "state": "UNVERIFIED"}),
+            (good, lambda *_args: {"ok": False, "state": "UNVERIFIED"}),
+            (good, lambda *_args: {"ok": True, "state": "VERIFIED"}),
+            (good, lambda *_args: (_ for _ in ()).throw(OSError("down"))),
+            (
+                good,
+                lambda *_args: (_ for _ in ()).throw(
+                    json.JSONDecodeError("bad", "x", 0)
+                ),
+            ),
+        ]
+        for environment, request in cases:
+            with self.subTest(environment=environment), redirect_stderr(io.StringIO()):
+                self.assertEqual(
+                    revoker.main(
+                        ["--runtime", "claude", "--reason", "pre-compact"],
+                        environ=environment,
+                        request=request,
+                    ),
+                    2,
+                )
+
+        with redirect_stderr(io.StringIO()):
+            self.assertEqual(
+                revoker.main(
+                    ["--runtime", "claude", "--reason", "x" * 129],
+                    environ=good,
+                    request=lambda *_args: self.fail("invalid reason reached broker"),
+                ),
+                2,
+            )
+
+        with tempfile.TemporaryDirectory() as directory:
+            generation_file = Path(directory) / "runtime.generation"
+            generation_file.write_text("1\n")
+            generation_file.chmod(0o600)
+            with redirect_stderr(io.StringIO()):
+                self.assertEqual(
+                    revoker.main(
+                        [
+                            "--runtime",
+                            "pi",
+                            "--reason",
+                            "session-start-resume",
+                            "--bump-generation",
+                        ],
+                        environ={
+                            **good,
+                            "MOSAIC_LEASE_GENERATION_FILE": str(generation_file),
+                        },
+                        request=lambda *_args: (_ for _ in ()).throw(OSError("down")),
+                    ),
+                    2,
+                )
+            self.assertEqual(generation_file.read_text(), "2\n")
+
+
 if __name__ == "__main__":
     unittest.main()