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

packages/mosaic/framework/tools/lease-broker/launch-runtime.py

@@ -12,6 +12,8 @@ from collections.abc import Callable, Mapping, Sequence
 from pathlib import Path
 from typing import Final
 
+from lease_generation import initialize_runtime_generation
+
 MAX_FRAME: Final = 64 * 1024
 BROKER_TIMEOUT_SECONDS: Final = 1.5
 CLAUDE_DANGEROUS_FLAG: Final = "--dangerously-skip-permissions"
@@ -44,6 +46,7 @@ def main(
     environ: Mapping[str, str] | None = None,
     request: Callable[[Path, dict[str, object]], dict[str, object]] = broker_request,
     execute: Callable[[str, list[str], dict[str, str]], object] = os.execvpe,
+    initialize_generation: Callable[[Path, int], None] = initialize_runtime_generation,
 ) -> int:
     parser = argparse.ArgumentParser()
     parser.add_argument("--runtime", required=True, choices=("claude", "pi"))
@@ -83,6 +86,8 @@ def main(
             or any(character not in "0123456789abcdef" for character in session_id)
         ):
             raise ValueError("registration refused")
+        generation_file = socket_path.parent / f"generation-{session_id}.state"
+        initialize_generation(generation_file, generation)
     except (KeyError, ValueError, OSError, json.JSONDecodeError):
         print("Mosaic lease broker registration failed; runtime launch denied.", file=sys.stderr)
         return 1
@@ -90,6 +95,7 @@ def main(
     environment = dict(source_environment)
     environment["MOSAIC_LEASE_SESSION_ID"] = session_id
     environment["MOSAIC_RUNTIME_GENERATION"] = str(generation)
+    environment["MOSAIC_LEASE_GENERATION_FILE"] = str(generation_file)
     environment["MOSAIC_LEASE_RUNTIME"] = arguments.runtime
     try:
         execute(command[0], command, environment)

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

@@ -0,0 +1,105 @@
+#!/usr/bin/env python3
+"""Private monotonic runtime-generation state shared by runtime hook processes."""
+
+from __future__ import annotations
+
+import fcntl
+import os
+import stat
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Final
+
+MAX_GENERATION: Final = (1 << 63) - 1
+MAX_GENERATION_BYTES: Final = 32
+
+
+def parse_generation(value: object) -> int:
+    if not isinstance(value, str) or not value.isascii() or not value.isdigit():
+        raise ValueError("invalid runtime generation")
+    generation = int(value)
+    if generation < 0 or generation > MAX_GENERATION:
+        raise ValueError("invalid runtime generation")
+    return generation
+
+
+def _validate_descriptor(descriptor: int) -> None:
+    metadata = os.fstat(descriptor)
+    if not stat.S_ISREG(metadata.st_mode):
+        raise ValueError("runtime generation state is not a regular file")
+    if metadata.st_uid != os.geteuid():
+        raise ValueError("runtime generation state has the wrong owner")
+    if stat.S_IMODE(metadata.st_mode) & 0o077:
+        raise ValueError("runtime generation state permissions are not private")
+    if metadata.st_size > MAX_GENERATION_BYTES:
+        raise ValueError("runtime generation state is oversized")
+
+
+def _read_descriptor(descriptor: int) -> int:
+    os.lseek(descriptor, 0, os.SEEK_SET)
+    raw = os.read(descriptor, MAX_GENERATION_BYTES + 1)
+    if len(raw) > MAX_GENERATION_BYTES:
+        raise ValueError("runtime generation state is oversized")
+    try:
+        return parse_generation(raw.decode("ascii").strip())
+    except UnicodeDecodeError as exc:
+        raise ValueError("invalid runtime generation state") from exc
+
+
+def _write_descriptor(descriptor: int, generation: int) -> None:
+    payload = f"{generation}\n".encode("ascii")
+    os.lseek(descriptor, 0, os.SEEK_SET)
+    os.ftruncate(descriptor, 0)
+    remaining = memoryview(payload)
+    while remaining:
+        written = os.write(descriptor, remaining)
+        if written <= 0:
+            raise OSError("runtime generation write made no progress")
+        remaining = remaining[written:]
+    os.fsync(descriptor)
+
+
+def initialize_runtime_generation(path: Path, generation: int) -> None:
+    if generation < 0 or generation > MAX_GENERATION:
+        raise ValueError("invalid runtime generation")
+    descriptor = os.open(
+        path,
+        os.O_RDWR | os.O_CREAT | os.O_TRUNC | os.O_CLOEXEC | os.O_NOFOLLOW,
+        0o600,
+    )
+    try:
+        os.fchmod(descriptor, 0o600)
+        _validate_descriptor(descriptor)
+        _write_descriptor(descriptor, generation)
+    finally:
+        os.close(descriptor)
+
+
+def read_runtime_generation(environ: Mapping[str, str]) -> int:
+    state_path = environ.get("MOSAIC_LEASE_GENERATION_FILE")
+    if not state_path:
+        return parse_generation(environ["MOSAIC_RUNTIME_GENERATION"])
+    descriptor = os.open(state_path, os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW)
+    try:
+        _validate_descriptor(descriptor)
+        return _read_descriptor(descriptor)
+    finally:
+        os.close(descriptor)
+
+
+def bump_runtime_generation(environ: Mapping[str, str]) -> int:
+    state_path = environ.get("MOSAIC_LEASE_GENERATION_FILE")
+    if not state_path:
+        raise ValueError("runtime generation file is required for same-PID rollover")
+    descriptor = os.open(state_path, os.O_RDWR | os.O_CLOEXEC | os.O_NOFOLLOW)
+    try:
+        fcntl.flock(descriptor, fcntl.LOCK_EX)
+        _validate_descriptor(descriptor)
+        current = _read_descriptor(descriptor)
+        if current >= MAX_GENERATION:
+            raise ValueError("runtime generation exhausted")
+        generation = current + 1
+        _write_descriptor(descriptor, generation)
+        return generation
+    finally:
+        os.close(descriptor)

packages/mosaic/framework/tools/lease-broker/mutator-gate.py

@@ -12,6 +12,8 @@ from collections.abc import Callable, Mapping, Sequence
 from pathlib import Path
 from typing import BinaryIO, Final
 
+from lease_generation import read_runtime_generation
+
 MAX_FRAME: Final = 64 * 1024
 BROKER_TIMEOUT_SECONDS: Final = 1.5
 
@@ -64,6 +66,7 @@ def main(
     environ: Mapping[str, str] | None = None,
     stream: BinaryIO | None = None,
     request: Callable[[Path, dict[str, object]], dict[str, object]] = broker_request,
+    resolve_generation: Callable[[Mapping[str, str]], int] = read_runtime_generation,
 ) -> int:
     parser = argparse.ArgumentParser()
     parser.add_argument("--runtime", required=True, choices=("claude", "pi"))
@@ -74,9 +77,7 @@ def main(
         tool_name = read_tool_name(stream)
         socket_value = source_environment["MOSAIC_LEASE_BROKER_SOCKET"]
         session_id = source_environment["MOSAIC_LEASE_SESSION_ID"]
-        generation = int(source_environment["MOSAIC_RUNTIME_GENERATION"])
-        if generation < 0:
-            raise ValueError("INVALID_GENERATION")
+        generation = resolve_generation(source_environment)
         reply = request(
             Path(socket_value),
             {

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())