stack/packages/mosaic/framework/tools/lease-broker/recover-context.py

#!/usr/bin/env python3
"""Constrained recovery command: the sole ungated Mosaic mutator.

This is deliberately a thin client of the broker's recovery entrypoint.  It
never accepts receipt text or a caller-provided challenge: the broker mints the
fresh challenge, delivers its exact receipt envelope, and later asks the
trusted ReceiptObserver seam to observe that same pending cycle.
"""

from __future__ import annotations

import argparse
import json
import os
import socket
import sys
from collections.abc import Mapping, Sequence
from pathlib import Path
from typing import Final

# The out-of-process recovery command is intentionally runnable with `python
# -I`; locate its shipped construction module without caller-controlled paths.
_MODULE_DIRECTORY = str(Path(__file__).resolve().parent)
if _MODULE_DIRECTORY not in sys.path:
    sys.path.insert(0, _MODULE_DIRECTORY)
from normative_fragments import build_payload_from_wire

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("recovery request too large")
    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 load_construction(path: Path) -> dict[str, object]:
    raw = path.read_bytes()
    if len(raw) > MAX_FRAME:
        raise ValueError("construction exceeds broker frame limit")
    value = json.loads(raw)
    if not isinstance(value, dict):
        raise ValueError("construction must be an object")
    return value


def identity(environ: Mapping[str, str]) -> tuple[Path, str, int, str]:
    socket_path = Path(environ["MOSAIC_LEASE_BROKER_SOCKET"])
    session_id = environ["MOSAIC_LEASE_SESSION_ID"]
    generation = int(environ["MOSAIC_RUNTIME_GENERATION"])
    runtime = environ["MOSAIC_LEASE_RUNTIME"]
    if generation < 0 or runtime not in {"claude", "pi"}:
        raise ValueError("invalid runtime identity")
    return socket_path, session_id, generation, runtime


def begin(
    construction_path: Path,
    compaction_epoch: int,
    request_epoch: int,
    environ: Mapping[str, str],
) -> dict[str, object]:
    if compaction_epoch < 0 or request_epoch < 0:
        raise ValueError("epochs must be non-negative")
    construction = load_construction(construction_path)
    # Invoke the shared WI-5 construction before asking the broker to repeat
    # its authoritative admission/build. No digest or receipt enters via CLI.
    built = build_payload_from_wire(construction)
    if (
        built.injectionDecision != "ACCEPTED"
        or not built.promotion
        or not isinstance(built.h_source, str)
        or not isinstance(built.h_payload, str)
    ):
        raise ValueError("payload construction refused")
    socket_path, session_id, generation, runtime = identity(environ)
    return broker_request(socket_path, {
        "action": "begin_recovery",
        "session_id": session_id,
        "runtime_generation": generation,
        "runtime": runtime,
        "binding": {
            "compaction_epoch": compaction_epoch,
            "request_epoch": request_epoch,
            "h_source": built.h_source,
            "h_payload": built.h_payload,
            "schema_version": 1,
        },
        "construction": construction,
    })


def complete(environ: Mapping[str, str]) -> dict[str, object]:
    socket_path, session_id, generation, _runtime = identity(environ)
    # No receipt or challenge argument exists: recovery completion can only use
    # the broker's current recovery cycle and its trusted observer seam.
    return broker_request(socket_path, {
        "action": "complete_recovery",
        "session_id": session_id,
        "runtime_generation": generation,
    })


def main(argv: Sequence[str] | None = None, *, environ: Mapping[str, str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    subcommands = parser.add_subparsers(dest="phase", required=True)
    begin_parser = subcommands.add_parser("begin", help="mint and deliver a fresh recovery receipt")
    begin_parser.add_argument("--construction", required=True, type=Path)
    begin_parser.add_argument("--compaction-epoch", required=True, type=int)
    begin_parser.add_argument("--request-epoch", required=True, type=int)
    subcommands.add_parser("complete", help="observe and promote only the current recovery receipt")
    arguments = parser.parse_args(argv)
    source_environment = os.environ if environ is None else environ
    try:
        reply = (
            begin(
                arguments.construction,
                arguments.compaction_epoch,
                arguments.request_epoch,
                source_environment,
            )
            if arguments.phase == "begin"
            else complete(source_environment)
        )
    except (KeyError, OSError, ValueError, json.JSONDecodeError) as error:
        print(f"Mosaic constrained recovery refused: {error}", file=sys.stderr)
        return 2
    print(json.dumps(reply, separators=(",", ":")))
    return 0 if reply.get("ok") is True else 2


if __name__ == "__main__":
    raise SystemExit(main())