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

#!/usr/bin/env python3
"""Runtime-neutral whole mutator-class gate backed by the Mosaic lease broker."""

from __future__ import annotations

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

# Isolated (`python -I`) adapter invocations must still import co-located
# framework modules; never depend on the caller's PYTHONPATH.
_MODULE_DIRECTORY = str(Path(__file__).resolve().parent)
if _MODULE_DIRECTORY not in sys.path:
    sys.path.insert(0, _MODULE_DIRECTORY)
from lease_generation import read_runtime_generation

MAX_FRAME: Final = 64 * 1024
BROKER_TIMEOUT_SECONDS: Final = 1.5
RECOVERY_TOOL: Final = "mosaic_context_recover"


def deny(code: str) -> int:
    print(f"BLOCKED: Mosaic mutator gate denied this tool ({code}).", file=sys.stderr)
    return 2


def read_tool_request(stream: BinaryIO | None = None) -> dict[str, object]:
    source = sys.stdin.buffer if stream is None else stream
    raw = source.read(MAX_FRAME + 1)
    if len(raw) > MAX_FRAME:
        raise ValueError("INVALID_GATE_INPUT")
    value = json.loads(raw)
    if not isinstance(value, dict):
        raise ValueError("INVALID_GATE_INPUT")
    tool_name = value.get("tool_name")
    if not isinstance(tool_name, str) or not tool_name or len(tool_name) > 256:
        raise ValueError("INVALID_GATE_INPUT")
    return value


def read_tool_name(stream: BinaryIO | None = None) -> str:
    """Backward-compatible strict extraction for callers that need only the name."""

    return str(read_tool_request(stream)["tool_name"])


def recovery_invocation_name(request: dict[str, object], recovery_command: Path | None) -> str:
    """Map only a direct Claude Bash invocation of the recovery executable.

    The mapping is adapter-configured and remains broker-authenticated through
    the ordinary authorization request. It never blesses Bash generally, shell
    composition, a different executable, or unsupported recovery arguments.
    """

    tool_name = request["tool_name"]
    if tool_name != "Bash" or recovery_command is None:
        return str(tool_name)
    tool_input = request.get("tool_input")
    if not isinstance(tool_input, dict) or set(tool_input) != {"command"}:
        return str(tool_name)
    command = tool_input.get("command")
    if not isinstance(command, str) or not command:
        return str(tool_name)
    try:
        argv = shlex.split(command, posix=True)
    except ValueError:
        return str(tool_name)
    if len(argv) < 3 or argv[0] != "python3":
        return str(tool_name)
    try:
        if Path(argv[1]).resolve(strict=False) != recovery_command.resolve(strict=False):
            return str(tool_name)
    except OSError:
        return str(tool_name)
    phase = argv[2]
    if phase == "complete" and len(argv) == 3:
        return RECOVERY_TOOL
    if phase != "begin" or len(argv) != 9:
        return str(tool_name)
    expected_flags = {"--construction", "--compaction-epoch", "--request-epoch"}
    supplied_flags = set(argv[3::2])
    if supplied_flags != expected_flags or any(not value for value in argv[4::2]):
        return str(tool_name)
    return RECOVERY_TOOL


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_GATE_INPUT")
    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,
    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"))
    parser.add_argument("--recovery-command", type=Path)
    arguments = parser.parse_args(argv)
    source_environment = os.environ if environ is None else environ

    try:
        request_input = read_tool_request(stream)
        tool_name = recovery_invocation_name(request_input, arguments.recovery_command)
        socket_value = source_environment["MOSAIC_LEASE_BROKER_SOCKET"]
        session_id = source_environment["MOSAIC_LEASE_SESSION_ID"]
        generation = resolve_generation(source_environment)
        reply = request(
            Path(socket_value),
            {
                "action": "authorize_tool",
                "session_id": session_id,
                "runtime_generation": generation,
                "runtime": arguments.runtime,
                "tool_name": tool_name,
            },
        )
    except (KeyError, ValueError, OSError, json.JSONDecodeError):
        return deny("GATE_UNAVAILABLE")

    if reply.get("ok") is True and reply.get("decision") == "allow":
        return 0
    code = reply.get("code")
    return deny(code if isinstance(code, str) else "MUTATOR_UNVERIFIED")


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