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

packages/mosaic/framework/tools/lease-broker/revoke-lease.py

@@ -0,0 +1,100 @@
+#!/usr/bin/env python3
+"""Revoke a runtime lease from a compaction or lifecycle observer."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import socket
+import sys
+from collections.abc import Callable, Mapping, Sequence
+from pathlib import Path
+from typing import Final
+
+from lease_generation import bump_runtime_generation, read_runtime_generation
+
+MAX_FRAME: Final = 64 * 1024
+BROKER_TIMEOUT_SECONDS: Final = 1.5
+
+
+def broker_request(socket_path: Path, request: dict[str, object]) -> dict[str, object]:
+    payload = (json.dumps(request, separators=(",", ":")) + "\n").encode()
+    if len(payload) > MAX_FRAME:
+        raise ValueError("invalid revoke request")
+    response = bytearray()
+    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection:
+        connection.settimeout(BROKER_TIMEOUT_SECONDS)
+        connection.connect(str(socket_path))
+        connection.sendall(payload)
+        connection.shutdown(socket.SHUT_WR)
+        while len(response) <= MAX_FRAME:
+            chunk = connection.recv(min(4096, MAX_FRAME + 1 - len(response)))
+            if not chunk:
+                break
+            response.extend(chunk)
+    if len(response) > MAX_FRAME or not response.endswith(b"\n"):
+        raise ValueError("invalid broker reply")
+    value = json.loads(response)
+    if not isinstance(value, dict):
+        raise ValueError("invalid broker reply")
+    return value
+
+
+def main(
+    argv: Sequence[str] | None = None,
+    *,
+    environ: Mapping[str, str] | None = None,
+    request: Callable[[Path, dict[str, object]], dict[str, object]] = broker_request,
+) -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--runtime", required=True, choices=("claude", "pi"))
+    parser.add_argument("--reason", required=True)
+    parser.add_argument("--bump-generation", action="store_true")
+    arguments = parser.parse_args(argv)
+    source_environment = os.environ if environ is None else environ
+
+    try:
+        if not arguments.reason or len(arguments.reason) > 128:
+            raise ValueError("invalid revoke reason")
+        socket_path = Path(source_environment["MOSAIC_LEASE_BROKER_SOCKET"])
+        session_id = source_environment["MOSAIC_LEASE_SESSION_ID"]
+        if (
+            len(session_id) != 64
+            or any(character not in "0123456789abcdef" for character in session_id)
+        ):
+            raise ValueError("invalid broker session")
+        generation = (
+            bump_runtime_generation(source_environment)
+            if arguments.bump_generation
+            else read_runtime_generation(source_environment)
+        )
+        reply = request(
+            socket_path,
+            {
+                "action": "revoke_lease",
+                "session_id": session_id,
+                "runtime_generation": generation,
+                "reason": arguments.reason,
+                "runtime": arguments.runtime,
+            },
+        )
+        if reply.get("ok") is not True or reply.get("state") != "UNVERIFIED":
+            raise ValueError("revocation refused")
+    except (KeyError, ValueError, OSError, json.JSONDecodeError):
+        # A fired observer must remain fail-closed even if the broker transport
+        # is unavailable. Advancing the private local generation fences every
+        # later tool check; the broker revokes the old lease when it next sees
+        # that higher generation. Explicit rollover already advanced it above.
+        if not arguments.bump_generation:
+            try:
+                bump_runtime_generation(source_environment)
+            except (KeyError, ValueError, OSError):
+                pass
+        print("Mosaic lease revocation failed; lifecycle transition denied.", file=sys.stderr)
+        return 2
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())