#!/usr/bin/env python3
"""Branch-focused tests for the lease-gated runtime executables."""

from __future__ import annotations

import importlib.util
import io
import json
import os
import runpy
import socket
import stat
import subprocess
import sys
import tempfile
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):
    spec = importlib.util.spec_from_file_location(module_name, TOOLS_DIR / filename)
    if spec is None or spec.loader is None:
        raise RuntimeError(f"unable to load {filename}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


LAUNCHER = load_tool("lease_runtime_launcher", "launch-runtime.py")
GATE = load_tool("lease_mutator_gate", "mutator-gate.py")


def matching_activation_probe(*_args: object, **_kwargs: object) -> dict[str, object]:
    """Fake activation-capability probe matching what enforcement expects
    (C4, #869 Point-1). Injected into `LAUNCHER.main()` calls below that are
    exercising OTHER branches (registration, exec, generation init, ...) so
    the new version-coupling gate — which runs before those — never blocks
    on host state (no real `mosaic` CLI on PATH in a test sandbox). The
    version-coupling gate's OWN behavior (match/mismatch/absent) is covered
    by its dedicated red-first tests in `version_coupling_unittest.py`."""

    return dict(LAUNCHER.EXPECTED_ACTIVATION_CAPABILITY)


class FakeSocket:
    def __init__(self, *chunks: bytes):
        self.chunks = list(chunks)
        self.timeout = None
        self.connected = None
        self.sent = b""
        self.shutdown_how = None

    def __enter__(self):
        return self

    def __exit__(self, *_args):
        return False

    def settimeout(self, value: float) -> None:
        self.timeout = value

    def connect(self, value: str) -> None:
        self.connected = value

    def sendall(self, value: bytes) -> None:
        self.sent += value

    def shutdown(self, how: int) -> None:
        self.shutdown_how = how

    def recv(self, _size: int) -> bytes:
        return self.chunks.pop(0) if self.chunks else b""


class LaunchRuntimeTest(unittest.TestCase):
    def test_success_registers_then_injects_session_before_exec(self) -> None:
        calls: dict[str, object] = {}
        session_id = "a" * 64

        def request(path: Path, payload: dict[str, object]) -> dict[str, object]:
            calls["path"] = path
            calls["request"] = payload
            return {"ok": True, "session_id": session_id}

        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={
                "MOSAIC_LEASE_BROKER_SOCKET": "/run/test/broker.sock",
                "MOSAIC_RUNTIME_GENERATION": "7",
                "PRESERVED": "yes",
            },
            request=request,
            execute=execute,
            initialize_generation=initialize_generation,
            probe_activation_capability=matching_activation_probe,
        )

        self.assertEqual(result, 0)
        self.assertEqual(calls["path"], Path("/run/test/broker.sock"))
        self.assertEqual(
            calls["request"],
            {"action": "register_anchor", "runtime_generation": 7},
        )
        command, argv, environment = calls["execute"]
        self.assertEqual(command, "claude")
        self.assertEqual(argv, ["claude", "--print", "hello"])
        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:
        executed: list[tuple[str, list[str], dict[str, str]]] = []
        result = LAUNCHER.main(
            ["--runtime", "claude", "--dangerous", "--", "claude", "-p", "hello"],
            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,
            probe_activation_capability=matching_activation_probe,
        )
        self.assertEqual(result, 0)
        self.assertEqual(
            executed[0][1],
            ["claude", "--dangerously-skip-permissions", "-p", "hello"],
        )

        with redirect_stderr(io.StringIO()):
            self.assertEqual(
                LAUNCHER.main(
                    ["--runtime", "pi", "--dangerous", "--", "pi"],
                    environ={"MOSAIC_LEASE_BROKER_SOCKET": "/broker"},
                    request=lambda *_args: {"ok": True, "session_id": "e" * 64},
                    execute=lambda *_args: self.fail("invalid dangerous runtime executed"),
                ),
                64,
            )

    def test_command_without_separator_is_forwarded_unchanged(self) -> None:
        executed: list[tuple[str, list[str], dict[str, str]]] = []
        result = LAUNCHER.main(
            ["--runtime", "pi", "pi", "--print", "hello"],
            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,
            probe_activation_capability=matching_activation_probe,
        )
        self.assertEqual(result, 0)
        self.assertEqual(executed[0][0:2], ("pi", ["pi", "--print", "hello"]))

    def test_missing_command_is_usage_error(self) -> None:
        with redirect_stderr(io.StringIO()):
            self.assertEqual(
                LAUNCHER.main(
                    ["--runtime", "pi", "--"],
                    environ={},
                    request=lambda *_args: {},
                    execute=lambda *_args: None,
                ),
                64,
            )

    def test_registration_validation_and_environment_fail_closed(self) -> None:
        good_session = "b" * 64
        cases = [
            ({}, {"ok": True, "session_id": good_session}),
            ({"MOSAIC_LEASE_BROKER_SOCKET": "/x", "MOSAIC_RUNTIME_GENERATION": "bad"}, {}),
            ({"MOSAIC_LEASE_BROKER_SOCKET": "/x", "MOSAIC_RUNTIME_GENERATION": "-1"}, {}),
            ({"MOSAIC_LEASE_BROKER_SOCKET": "/x"}, {"ok": False, "session_id": good_session}),
            ({"MOSAIC_LEASE_BROKER_SOCKET": "/x"}, {"ok": True, "session_id": 4}),
            ({"MOSAIC_LEASE_BROKER_SOCKET": "/x"}, {"ok": True, "session_id": "b" * 63}),
            ({"MOSAIC_LEASE_BROKER_SOCKET": "/x"}, {"ok": True, "session_id": "z" * 64}),
        ]
        for environment, reply in cases:
            with self.subTest(environment=environment, reply=reply), redirect_stderr(io.StringIO()):
                executed: list[object] = []
                result = LAUNCHER.main(
                    ["--runtime", "pi", "--", "pi"],
                    environ=environment,
                    request=lambda *_args, value=reply: value,
                    execute=lambda *args: executed.append(args),
                    probe_activation_capability=matching_activation_probe,
                )
                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")
                    ),
                    probe_activation_capability=matching_activation_probe,
                ),
                1,
            )

    def test_registration_exceptions_fail_closed(self) -> None:
        failures = [ValueError("bad"), OSError("down"), json.JSONDecodeError("bad", "x", 0)]
        for failure in failures:
            with self.subTest(failure=type(failure).__name__), redirect_stderr(io.StringIO()):
                def request(*_args, error=failure):
                    raise error

                self.assertEqual(
                    LAUNCHER.main(
                        ["--runtime", "claude", "--", "claude"],
                        environ={"MOSAIC_LEASE_BROKER_SOCKET": "/x"},
                        request=request,
                        execute=lambda *_args: self.fail("must not execute"),
                        probe_activation_capability=matching_activation_probe,
                    ),
                    1,
                )

    def test_exec_failure_is_fail_closed(self) -> None:
        with redirect_stderr(io.StringIO()):
            self.assertEqual(
                LAUNCHER.main(
                    ["--runtime", "pi", "--", "pi"],
                    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,
                    probe_activation_capability=matching_activation_probe,
                ),
                1,
            )

    def test_broker_reply_framing_and_shape_validation(self) -> None:
        replies = [
            (b'{"ok":true}\n', {"ok": True}),
            (b'{"ok":true}', ValueError),
            (b'[]\n', ValueError),
            (b"x" * (LAUNCHER.MAX_FRAME + 1), ValueError),
        ]
        for wire_reply, expected in replies:
            with self.subTest(size=len(wire_reply)):
                fake = FakeSocket(wire_reply)
                with patch.object(LAUNCHER.socket, "socket", return_value=fake):
                    if isinstance(expected, type) and issubclass(expected, Exception):
                        with self.assertRaises(expected):
                            LAUNCHER.broker_request(Path("/broker"), {"action": "register_anchor"})
                    else:
                        self.assertEqual(
                            LAUNCHER.broker_request(Path("/broker"), {"action": "register_anchor"}),
                            expected,
                        )
                        self.assertEqual(fake.timeout, LAUNCHER.BROKER_TIMEOUT_SECONDS)
                        self.assertEqual(fake.connected, "/broker")
                        self.assertEqual(fake.shutdown_how, socket.SHUT_WR)


class ExecutableEntrypointTest(unittest.TestCase):
    def test_real_claude_and_pi_gates_fail_closed_on_empty_or_truncated_reply(self) -> None:
        for runtime in ("claude", "pi"):
            for wire_reply in (b"", b'{"ok":true'):
                with self.subTest(runtime=runtime, wire_reply=wire_reply), tempfile.TemporaryDirectory() as root:
                    socket_path = Path(root) / "broker.sock"
                    server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
                    server.bind(str(socket_path))
                    server.listen(1)

                    def serve_reply() -> None:
                        with server:
                            connection, _ = server.accept()
                            with connection:
                                while connection.recv(4096):
                                    pass
                                if wire_reply:
                                    connection.sendall(wire_reply)

                    thread = threading.Thread(target=serve_reply, daemon=True)
                    thread.start()
                    environment = {
                        **os.environ,
                        "MOSAIC_LEASE_BROKER_SOCKET": str(socket_path),
                        "MOSAIC_LEASE_SESSION_ID": "d" * 64,
                        "MOSAIC_RUNTIME_GENERATION": "1",
                    }
                    try:
                        result = subprocess.run(
                            [sys.executable, str(TOOLS_DIR / "mutator-gate.py"), "--runtime", runtime],
                            input=b'{"tool_name":"Read"}\n',
                            capture_output=True,
                            env=environment,
                            check=False,
                            timeout=5,
                        )
                    except subprocess.TimeoutExpired as exc:
                        server.close()
                        thread.join(timeout=2)
                        self.fail(
                            f"{runtime} gate hung on wire reply {wire_reply!r}: {exc}"
                        )
                    thread.join(timeout=2)

                    self.assertFalse(thread.is_alive())
                    self.assertEqual(result.returncode, 2)
                    self.assertIn(b"GATE_UNAVAILABLE", result.stderr)

    def test_launcher_entrypoint_returns_usage_without_a_command(self) -> None:
        with patch.object(
            sys,
            "argv",
            [str(TOOLS_DIR / "launch-runtime.py"), "--runtime", "claude"],
        ), redirect_stderr(io.StringIO()), self.assertRaises(SystemExit) as raised:
            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"}')

        with patch.object(
            sys,
            "argv",
            [str(TOOLS_DIR / "mutator-gate.py"), "--runtime", "claude"],
        ), patch.object(sys, "stdin", Stdin()), patch.dict(
            os.environ, {}, clear=True
        ), redirect_stderr(io.StringIO()), self.assertRaises(SystemExit) as raised:
            runpy.run_path(str(TOOLS_DIR / "mutator-gate.py"), run_name="__main__")
        self.assertEqual(raised.exception.code, 2)


class MutatorGateTest(unittest.TestCase):
    @staticmethod
    def environment() -> dict[str, str]:
        return {
            "MOSAIC_LEASE_BROKER_SOCKET": "/run/test/broker.sock",
            "MOSAIC_LEASE_SESSION_ID": "d" * 64,
            "MOSAIC_RUNTIME_GENERATION": "2",
        }

    def run_main(self, *, tool: object = "Bash", reply: dict[str, object] | None = None):
        calls: list[tuple[Path, dict[str, object]]] = []

        def request(path: Path, payload: dict[str, object]) -> dict[str, object]:
            calls.append((path, payload))
            return reply if reply is not None else {"ok": True, "decision": "allow"}

        stderr = io.StringIO()
        with redirect_stderr(stderr):
            result = GATE.main(
                ["--runtime", "claude"],
                environ=self.environment(),
                stream=io.BytesIO(json.dumps({"tool_name": tool}).encode()),
                request=request,
            )
        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)
        self.assertEqual(allowed_stderr, "")
        self.assertEqual(calls[0][0], Path("/run/test/broker.sock"))
        self.assertEqual(
            calls[0][1],
            {
                "action": "authorize_tool",
                "session_id": "d" * 64,
                "runtime_generation": 2,
                "runtime": "claude",
                "tool_name": "Bash",
            },
        )

        denied, denied_stderr, _ = self.run_main(reply={"ok": False, "code": "LEASE_EXPIRED"})
        self.assertEqual(denied, 2)
        self.assertIn("LEASE_EXPIRED", denied_stderr)

        defaulted, defaulted_stderr, _ = self.run_main(reply={"ok": False, "code": 4})
        self.assertEqual(defaulted, 2)
        self.assertIn("MUTATOR_UNVERIFIED", defaulted_stderr)

    def test_input_validation_fails_closed(self) -> None:
        payloads = [
            b"x" * (GATE.MAX_FRAME + 1),
            b"[]",
            b"{}",
            json.dumps({"tool_name": ""}).encode(),
            json.dumps({"tool_name": 4}).encode(),
            json.dumps({"tool_name": "x" * 257}).encode(),
            b"not-json",
        ]
        for payload in payloads:
            with self.subTest(size=len(payload)), redirect_stderr(io.StringIO()):
                self.assertEqual(
                    GATE.main(
                        ["--runtime", "pi"],
                        environ=self.environment(),
                        stream=io.BytesIO(payload),
                        request=lambda *_args: self.fail("invalid input reached broker"),
                    ),
                    2,
                )

    def test_environment_generation_and_request_failures_deny(self) -> None:
        environments = [
            {},
            {**self.environment(), "MOSAIC_RUNTIME_GENERATION": "bad"},
            {**self.environment(), "MOSAIC_RUNTIME_GENERATION": "-1"},
        ]
        for environment in environments:
            with self.subTest(environment=environment), redirect_stderr(io.StringIO()):
                self.assertEqual(
                    GATE.main(
                        ["--runtime", "claude"],
                        environ=environment,
                        stream=io.BytesIO(b'{"tool_name":"Read"}'),
                        request=lambda *_args: {},
                    ),
                    2,
                )

        failures = [ValueError("bad"), OSError("down"), json.JSONDecodeError("bad", "x", 0)]
        for failure in failures:
            with self.subTest(failure=type(failure).__name__), redirect_stderr(io.StringIO()):
                def request(*_args, error=failure):
                    raise error

                self.assertEqual(
                    GATE.main(
                        ["--runtime", "claude"],
                        environ=self.environment(),
                        stream=io.BytesIO(b'{"tool_name":"Read"}'),
                        request=request,
                    ),
                    2,
                )

    def test_broker_request_framing_payload_and_shape_validation(self) -> None:
        with self.assertRaises(ValueError):
            GATE.broker_request(Path("/broker"), {"session_id": "x" * GATE.MAX_FRAME})

        replies = [
            (b'{"ok":true,"decision":"allow"}\n', {"ok": True, "decision": "allow"}),
            (b'{"ok":true}', ValueError),
            (b'[]\n', ValueError),
            (b"x" * (GATE.MAX_FRAME + 1), ValueError),
        ]
        for wire_reply, expected in replies:
            with self.subTest(size=len(wire_reply)):
                fake = FakeSocket(wire_reply)
                with patch.object(GATE.socket, "socket", return_value=fake):
                    if isinstance(expected, type) and issubclass(expected, Exception):
                        with self.assertRaises(expected):
                            GATE.broker_request(Path("/broker"), {"action": "authorize_tool"})
                    else:
                        self.assertEqual(
                            GATE.broker_request(Path("/broker"), {"action": "authorize_tool"}),
                            expected,
                        )
                        self.assertEqual(fake.timeout, GATE.BROKER_TIMEOUT_SECONDS)
                        self.assertEqual(fake.connected, "/broker")
                        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()
