stack/docs/compaction-refresh/probes/p6_constrained_recovery.py

#!/usr/bin/env python3
"""P6 constrained-recovery probe; BUILT ONLY and Mos-gated.

DO NOT self-fire.  Under Mos authorization only:
  python3 -I -S -B docs/compaction-refresh/probes/p6_constrained_recovery.py

The default three isolated runs each start a fresh shipped daemon on a private
Unix socket and drive the shipped recovery *command* out of process.  This
probe never resets broker state, replaces the observer seam, or reimplements
promotion: it exercises `recover-context.py` and the daemon's real socket.
"""

from __future__ import annotations

import argparse
import base64
import hashlib
import importlib.util
import json
import os
import shutil
import socket
import subprocess
import sys
import tempfile
import time
from pathlib import Path


HERE = Path(__file__).resolve().parent
REPOSITORY = HERE.parents[2]
TOOLS = REPOSITORY / "packages/mosaic/framework/tools/lease-broker"
DAEMON = TOOLS / "daemon.py"
RECOVERY_COMMAND = TOOLS / "recover-context.py"
FRAGMENTS = TOOLS / "normative_fragments.py"


def load_shipped_fragments():
    spec = importlib.util.spec_from_file_location("p6_shipped_fragments", FRAGMENTS)
    if spec is None or spec.loader is None:
        raise RuntimeError("shipped normative construction unavailable")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def request(socket_path: Path, value: dict[str, object]) -> dict[str, object]:
    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection:
        connection.settimeout(3.0)
        connection.connect(str(socket_path))
        connection.sendall((json.dumps(value, separators=(",", ":")) + "\n").encode())
        connection.shutdown(socket.SHUT_WR)
        response = bytearray()
        while True:
            chunk = connection.recv(4096)
            if not chunk:
                break
            response.extend(chunk)
    if not response.endswith(b"\n") or response.count(b"\n") != 1:
        raise AssertionError(f"unframed broker reply: {bytes(response)!r}")
    reply = json.loads(response[:-1])
    if not isinstance(reply, dict):
        raise AssertionError("broker reply is not an object")
    return reply


def wait_ready(process: subprocess.Popen[str], socket_path: Path) -> None:
    deadline = time.monotonic() + 5.0
    while time.monotonic() < deadline:
        if socket_path.exists():
            return
        if process.poll() is not None:
            output = process.stdout.read() if process.stdout is not None else ""
            raise RuntimeError(f"shipped daemon exited before READY: {output}")
        time.sleep(0.02)
    raise TimeoutError("shipped daemon did not create private probe socket")


def recovery_command(arguments: list[str], environment: dict[str, str]) -> dict[str, object]:
    completed = subprocess.run(
        [sys.executable, "-I", "-S", "-B", str(RECOVERY_COMMAND), *arguments],
        text=True,
        capture_output=True,
        env=environment,
        check=False,
    )
    if not completed.stdout.endswith("\n"):
        raise AssertionError(f"recovery command omitted framed result: {completed.stderr!r}")
    reply = json.loads(completed.stdout)
    if not isinstance(reply, dict):
        raise AssertionError("recovery command result is not an object")
    return reply


def run_once(index: int) -> None:
    # Parity guard: this is an out-of-process driver of the shipped recovery
    # entrypoint, not a shadow receipt/promotion implementation.
    source = RECOVERY_COMMAND.read_text(encoding="utf-8")
    if '"action": "begin_recovery"' not in source or '"action": "complete_recovery"' not in source:
        raise AssertionError("P6 parity guard: recovery command no longer drives shipped broker entrypoints")

    fragments = load_shipped_fragments()
    root = Path(tempfile.mkdtemp(prefix=f"mosaic-p6-recovery-{index}-"))
    os.chmod(root, 0o700)
    socket_path = root / "broker.sock"
    state_path = root / "state.json"
    observer_path = root / "test-observer.json"
    construction_path = root / "construction.json"
    content = b"P6 constrained recovery fixture\n"
    construction = {
        "manifest_version": 1,
        "generator_version": "p6-constrained-recovery",
        "fragments": [{
            "source_id": "authority/p6",
            "content_base64": base64.b64encode(content).decode("ascii"),
            "expected_sha256": hashlib.sha256(content).hexdigest(),
        }],
    }
    construction_path.write_text(json.dumps(construction), encoding="utf-8")
    os.chmod(construction_path, 0o600)
    process = subprocess.Popen(
        [sys.executable, "-I", "-S", "-B", str(DAEMON), "--socket", str(socket_path),
         "--state", str(state_path), "--test-observer-file", str(observer_path)],
        stdin=subprocess.DEVNULL,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
    )
    try:
        wait_ready(process, socket_path)
        registered = request(socket_path, {"action": "register_anchor", "runtime_generation": 1})
        session_id = registered.get("session_id")
        if registered.get("ok") is not True or not isinstance(session_id, str):
            raise AssertionError(f"broker anchor registration failed: {registered!r}")
        built = fragments.build_payload_from_wire(construction)
        normal = request(socket_path, {
            "action": "begin_verification", "session_id": session_id, "runtime_generation": 1,
            "runtime": "pi", "construction": construction,
            "binding": {"compaction_epoch": index, "request_epoch": index + 100,
                        "h_source": built.h_source, "h_payload": built.h_payload, "schema_version": 1},
        })
        normal_challenge = normal.get("receipt_challenge")
        normal_receipt = normal.get("receipt")
        if not isinstance(normal_challenge, str) or not isinstance(normal_receipt, str):
            raise AssertionError("normal path did not mint a receipt challenge")
        environment = {
            **os.environ,
            "MOSAIC_LEASE_BROKER_SOCKET": str(socket_path),
            "MOSAIC_LEASE_SESSION_ID": session_id,
            "MOSAIC_RUNTIME_GENERATION": "1",
            "MOSAIC_LEASE_RUNTIME": "pi",
        }
        recovery = recovery_command([
            "begin", "--construction", str(construction_path), "--compaction-epoch", str(index + 10),
            "--request-epoch", str(index + 110),
        ], environment)
        challenge = recovery.get("receipt_challenge")
        receipt = recovery.get("receipt")
        if recovery.get("state") != "PENDING_DELIVERY" or not isinstance(challenge, str) or not isinstance(receipt, str):
            raise AssertionError(f"recovery command did not drive pending delivery: {recovery!r}")
        if challenge == normal_challenge:
            raise AssertionError("recovery reused a normal-path challenge")

        # C4: a normal-path receipt is not an admissible recovery receipt.
        observer_path.write_text(json.dumps({
            "session_id": session_id, "runtime_generation": 1, "latest_assistant_message": normal_receipt,
        }), encoding="utf-8")
        os.chmod(observer_path, 0o600)
        refused = recovery_command(["complete"], environment)
        if refused.get("ok") is not False or refused.get("code") != "RECEIPT_MISMATCH":
            raise AssertionError(f"normal-path receipt replay was not refused: {refused!r}")

        # Begin a new recovery cycle, then prove shipped evidence -> consume ->
        # promote ordering and post-consumption non-repromotion.
        recovery = recovery_command([
            "begin", "--construction", str(construction_path), "--compaction-epoch", str(index + 20),
            "--request-epoch", str(index + 120),
        ], environment)
        receipt = recovery.get("receipt")
        if recovery.get("state") != "PENDING_DELIVERY" or not isinstance(receipt, str):
            raise AssertionError(f"fresh recovery retry did not pend: {recovery!r}")
        observer_path.write_text(json.dumps({
            "session_id": session_id, "runtime_generation": 1, "latest_assistant_message": receipt,
        }), encoding="utf-8")
        os.chmod(observer_path, 0o600)
        promoted = recovery_command(["complete"], environment)
        if promoted.get("ok") is not True or promoted.get("state") != "VERIFIED":
            raise AssertionError(f"recovery consume-before-promote failed: {promoted!r}")
        replay = recovery_command(["complete"], environment)
        if replay.get("ok") is not False or replay.get("code") != "INVALID_LEASE_TRANSITION":
            raise AssertionError(f"consumed recovery challenge re-promoted: {replay!r}")
    finally:
        if process.poll() is None:
            process.terminate()
            try:
                process.wait(timeout=3.0)
            except subprocess.TimeoutExpired:
                process.kill()
                process.wait()
        shutil.rmtree(root, ignore_errors=True)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--runs", type=int, default=3)
    arguments = parser.parse_args()
    if arguments.runs != 3:
        raise SystemExit("P6 requires exactly three isolated runs")
    for index in range(arguments.runs):
        run_once(index)
    print("P6 constrained recovery probe PASS: 3 isolated shipped recovery-command runs")


if __name__ == "__main__":
    main()