#!/usr/bin/env python3
"""Mosaic external lease broker for Linux SO_PEERCRED authenticated clients."""

from __future__ import annotations

import argparse
import copy
import errno
import hmac
from concurrent.futures import ThreadPoolExecutor
import json
import os
import secrets
import select
import signal
import socket
import stat
import struct
import sys
import threading
import time
from pathlib import Path
from typing import Final

# This script is loaded both as an executable and through isolated stdlib tests.
# Keep its co-located receipt implementation importable in both modes.
_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
from receipt_challenge import is_verbatim_receipt, latest_assistant_digest, receipt_for
from receipt_observer import (
    FileTestReceiptObserver,
    ReceiptObserver,
    RuntimeReceiptObserver,
)

MAX_FRAME: Final = 64 * 1024
MAX_STATE: Final = 4 * 1024 * 1024
MAX_PENDING_TOKENS: Final = 256
MAX_IN_FLIGHT_CONNECTIONS: Final = 16
MAX_LEASE_TTL_SECONDS: Final = 300
STATE_VERSION: Final = 1
READ_DEADLINE_SECONDS: Final = 1.0
HANDLE_QUEUE_TIMEOUT_SECONDS: Final = 1.0
SEND_TIMEOUT_SECONDS: Final = 1.0
HEX_256_LENGTH: Final = 64
LEASE_UNVERIFIED: Final = "UNVERIFIED"
LEASE_PENDING: Final = "PENDING_VERIFICATION"
LEASE_PENDING_PROMOTION: Final = "PENDING_PROMOTION"
LEASE_VERIFIED: Final = "VERIFIED"
READ_ONLY_TOOLS: Final = {
    "claude": frozenset({"Read", "Grep", "Glob", "Ls", "Find"}),
    "pi": frozenset({"read", "grep", "find", "ls"}),
}
RECOVERY_TOOL: Final = "mosaic_context_recover"


class BrokerFailure(Exception):
    def __init__(self, code: str) -> None:
        super().__init__(code)
        self.code = code


class StateCommitUncertain(RuntimeError):
    def __init__(self) -> None:
        super().__init__("STATE_COMMIT_UNCERTAIN")


def is_non_negative_integer(value: object) -> bool:
    return type(value) is int and value >= 0


def is_hex_256(value: object) -> bool:
    return (
        isinstance(value, str)
        and len(value) == HEX_256_LENGTH
        and all(character in "0123456789abcdef" for character in value)
    )


def is_positive_decimal(value: object) -> bool:
    return (
        isinstance(value, str)
        and len(value) > 0
        and value[0] in "123456789"
        and all(character in "0123456789" for character in value)
    )


def valid_binding(binding: object) -> bool:
    if not isinstance(binding, dict):
        return False
    required = {"compaction_epoch", "request_epoch", "h_source", "h_payload", "schema_version"}
    if set(binding) != required:
        return False
    if not all(
        is_non_negative_integer(binding[field])
        for field in ("compaction_epoch", "request_epoch", "schema_version")
    ):
        return False
    return all(
        is_hex_256(binding[field])
        for field in ("h_source", "h_payload")
    )


def valid_receipt_evidence(evidence: object) -> bool:
    return (
        isinstance(evidence, dict)
        and set(evidence) == {"h_latest_assistant"}
        and is_hex_256(evidence["h_latest_assistant"])
    )


def validate_state(value: object) -> dict[str, object]:
    if not isinstance(value, dict) or set(value) != {"version", "sessions", "tokens"}:
        raise BrokerFailure("STATE_INTEGRITY")
    if type(value["version"]) is not int or value["version"] != STATE_VERSION:
        raise BrokerFailure("STATE_INTEGRITY")
    sessions = value["sessions"]
    tokens = value["tokens"]
    if not isinstance(sessions, dict) or not isinstance(tokens, dict):
        raise BrokerFailure("STATE_INTEGRITY")
    if len(tokens) > MAX_PENDING_TOKENS:
        raise BrokerFailure("STATE_INTEGRITY")

    anchors: set[tuple[int, str]] = set()
    for session_id, session in sessions.items():
        if not is_hex_256(session_id) or not isinstance(session, dict):
            raise BrokerFailure("STATE_INTEGRITY")
        if set(session) != {"anchor_pid", "anchor_starttime", "runtime_generation"}:
            raise BrokerFailure("STATE_INTEGRITY")
        anchor_pid = session["anchor_pid"]
        anchor_starttime = session["anchor_starttime"]
        if type(anchor_pid) is not int or anchor_pid <= 0:
            raise BrokerFailure("STATE_INTEGRITY")
        if not is_positive_decimal(anchor_starttime):
            raise BrokerFailure("STATE_INTEGRITY")
        if not is_non_negative_integer(session["runtime_generation"]):
            raise BrokerFailure("STATE_INTEGRITY")
        anchor = (anchor_pid, anchor_starttime)
        if anchor in anchors:
            raise BrokerFailure("STATE_INTEGRITY")
        anchors.add(anchor)

    for token_value, token in tokens.items():
        if not is_hex_256(token_value) or not isinstance(token, dict):
            raise BrokerFailure("STATE_INTEGRITY")
        token_fields = set(token)
        if token_fields not in (
            {"session_id", "runtime_generation", "binding", "consumed"},
            {"session_id", "runtime_generation", "binding", "consumed", "evidence"},
        ):
            raise BrokerFailure("STATE_INTEGRITY")
        session_id = token["session_id"]
        generation = token["runtime_generation"]
        session = sessions.get(session_id) if isinstance(session_id, str) else None
        if not is_hex_256(session_id) or not isinstance(session, dict):
            raise BrokerFailure("STATE_INTEGRITY")
        if not is_non_negative_integer(generation):
            raise BrokerFailure("STATE_INTEGRITY")
        if generation != session["runtime_generation"]:
            raise BrokerFailure("STATE_INTEGRITY")
        if not valid_binding(token["binding"]) or token["consumed"] is not False:
            raise BrokerFailure("STATE_INTEGRITY")
        evidence = token.get("evidence")
        if evidence is not None and not valid_receipt_evidence(evidence):
            raise BrokerFailure("STATE_INTEGRITY")
    return value


def proc_node(pid: int) -> dict[str, int | str]:
    try:
        text = Path(f"/proc/{pid}/stat").read_text(encoding="utf-8")
    except (FileNotFoundError, PermissionError, ProcessLookupError) as exc:
        raise BrokerFailure("PID_UNAVAILABLE") from exc
    close = text.rfind(")")
    fields = text[close + 2 :].split()
    if close < 0 or len(fields) < 20:
        raise BrokerFailure("PROC_STAT_INVALID")
    return {"pid": pid, "ppid": int(fields[1]), "starttime": fields[19]}


def verified_ancestry(peer_pid: int, anchor_pid: int, anchor_starttime: str) -> bool:
    chain: list[dict[str, int | str]] = []
    seen: set[int] = set()
    current = peer_pid
    while current > 0 and current not in seen:
        seen.add(current)
        node = proc_node(current)
        chain.append(node)
        if current == anchor_pid:
            if node["starttime"] != anchor_starttime:
                return False
            break
        current = int(node["ppid"])
    else:
        return False
    if int(chain[-1]["pid"]) != anchor_pid:
        return False
    for original in chain:
        repeated = proc_node(int(original["pid"]))
        if repeated["starttime"] != original["starttime"]:
            raise BrokerFailure("PID_STARTTIME_RACE")
    return True


def secure_parent(path: Path) -> None:
    parent = path.parent
    if not parent.is_dir() or stat.S_IMODE(parent.stat().st_mode) != 0o700:
        raise BrokerFailure("INSECURE_PARENT_MODE")


class StateStore:
    def __init__(self, path: Path) -> None:
        self.path = path
        secure_parent(path)
        self.poisoned = False
        self.value: dict[str, object] = {"version": STATE_VERSION, "sessions": {}, "tokens": {}}
        flags = os.O_RDONLY | getattr(os, "O_CLOEXEC", 0) | getattr(os, "O_NOFOLLOW", 0)
        try:
            descriptor = os.open(path, flags)
        except FileNotFoundError:
            return
        except OSError as exc:
            raise BrokerFailure("STATE_INTEGRITY") from exc
        try:
            metadata = os.fstat(descriptor)
            if not stat.S_ISREG(metadata.st_mode):
                raise BrokerFailure("STATE_INTEGRITY")
            if stat.S_IMODE(metadata.st_mode) != 0o600:
                raise BrokerFailure("INSECURE_STATE_MODE")
            if metadata.st_size > MAX_STATE:
                raise BrokerFailure("STATE_INTEGRITY")
            chunks = bytearray()
            while len(chunks) <= MAX_STATE:
                chunk = os.read(descriptor, min(64 * 1024, MAX_STATE + 1 - len(chunks)))
                if not chunk:
                    break
                chunks.extend(chunk)
            if len(chunks) > MAX_STATE:
                raise BrokerFailure("STATE_INTEGRITY")
            try:
                loaded = json.loads(chunks)
            except (json.JSONDecodeError, UnicodeDecodeError) as exc:
                raise BrokerFailure("STATE_INTEGRITY") from exc
            self.value = validate_state(loaded)
        except OSError as exc:
            raise BrokerFailure("STATE_INTEGRITY") from exc
        finally:
            os.close(descriptor)

    def sessions(self) -> dict[str, dict[str, object]]:
        sessions = self.value.get("sessions")
        if not isinstance(sessions, dict):
            raise BrokerFailure("STATE_INTEGRITY")
        return sessions

    def tokens(self) -> dict[str, dict[str, object]]:
        tokens = self.value.get("tokens")
        if not isinstance(tokens, dict):
            raise BrokerFailure("STATE_INTEGRITY")
        return tokens

    def commit(self) -> None:
        if self.poisoned:
            raise StateCommitUncertain()
        payload = (
            json.dumps(self.value, sort_keys=True, separators=(",", ":")) + "\n"
        ).encode()
        if len(payload) > MAX_STATE:
            raise BrokerFailure("STATE_TOO_LARGE")
        temporary = self.path.with_name(f".{self.path.name}.{os.getpid()}.tmp")
        descriptor = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
        replaced = False
        try:
            try:
                remaining = memoryview(payload)
                while remaining:
                    written = os.write(descriptor, remaining)
                    if written == 0:
                        raise OSError(errno.EIO, "state write made no progress")
                    remaining = remaining[written:]
                os.fsync(descriptor)
            finally:
                os.close(descriptor)
            os.replace(temporary, self.path)
            replaced = True
            try:
                directory = os.open(self.path.parent, os.O_RDONLY | os.O_DIRECTORY)
                try:
                    os.fsync(directory)
                finally:
                    os.close(directory)
            except OSError as exc:
                self.poisoned = True
                raise StateCommitUncertain() from exc
        finally:
            if not replaced:
                try:
                    temporary.unlink()
                except FileNotFoundError:
                    pass


class Broker:
    def __init__(self, store: StateStore, observer: ReceiptObserver | None = None) -> None:
        self.store = store
        # Production construction always has a transport-capable observer. Test
        # fixtures may inject their controlled observer explicitly.
        self.observer: ReceiptObserver = observer if observer is not None else RuntimeReceiptObserver()
        # Set only by begin_verification after its mandatory revoke-first fence.
        # It is preserved if later cycle admission is refused; all other broker
        # actions retain the normal snapshot rollback behavior.
        self._rejected_cycle_fence: tuple[dict[str, object], dict[str, dict[str, object]]] | None = None
        # VERIFIED authority is deliberately volatile: broker restart revokes all
        # leases while preserving WI-1 identity and pending-token integrity.
        self.leases: dict[str, dict[str, object]] = {}

    def authenticate(self, peer_pid: int, request: dict[str, object]) -> tuple[str, dict[str, object]]:
        session_id = request.get("session_id")
        generation = request.get("runtime_generation")
        if not isinstance(session_id, str) or not is_non_negative_integer(generation):
            raise BrokerFailure("INVALID_IDENTITY")
        session = self.store.sessions().get(session_id)
        if not isinstance(session, dict):
            raise BrokerFailure("UNKNOWN_SESSION")
        anchor_pid = session.get("anchor_pid")
        anchor_starttime = session.get("anchor_starttime")
        current_generation = session.get("runtime_generation")
        if (
            type(anchor_pid) is not int
            or not isinstance(anchor_starttime, str)
            or not is_non_negative_integer(current_generation)
        ):
            raise BrokerFailure("STATE_INTEGRITY")
        if not verified_ancestry(peer_pid, anchor_pid, anchor_starttime):
            raise BrokerFailure("ANCESTRY_MISMATCH")
        if generation < current_generation:
            raise BrokerFailure("STALE_GENERATION")
        if generation > current_generation:
            session["runtime_generation"] = generation
            self.revoke_session_authority(session_id)
        return session_id, session

    def session_for_anchor(
        self, anchor_pid: int, anchor_starttime: str
    ) -> tuple[str, dict[str, object]] | None:
        for session_id, session in self.store.sessions().items():
            if (
                session["anchor_pid"] == anchor_pid
                and session["anchor_starttime"] == anchor_starttime
            ):
                return session_id, session
        return None

    def revoke_session_tokens(self, session_id: str) -> None:
        tokens = self.store.tokens()
        for token_value in [
            value for value, token in tokens.items() if token["session_id"] == session_id
        ]:
            del tokens[token_value]

    def revoke_session_authority(self, session_id: str) -> None:
        self.revoke_session_tokens(session_id)
        session = self.store.sessions().get(session_id)
        generation = session.get("runtime_generation") if isinstance(session, dict) else None
        self.leases[session_id] = {
            "state": LEASE_UNVERIFIED,
            "runtime_generation": generation,
        }

    def mint_token(
        self,
        session_id: str,
        generation: int,
        binding: dict[str, object],
    ) -> str:
        if len(self.store.tokens()) >= MAX_PENDING_TOKENS:
            raise BrokerFailure("TOKEN_CAPACITY")
        token = secrets.token_hex(32)
        self.store.tokens()[token] = {
            "session_id": session_id,
            "runtime_generation": generation,
            "binding": copy.deepcopy(binding),
            "consumed": False,
            # Receipt evidence is durably committed before this challenge may
            # be consumed and promotion made externally visible.
            "evidence": None,
        }
        return token

    def finish_promotion(self, session_id: str) -> None:
        lease = self.leases.get(session_id)
        if not isinstance(lease, dict) or lease.get("state") != LEASE_PENDING_PROMOTION:
            raise BrokerFailure("STATE_INTEGRITY")
        lease["state"] = LEASE_VERIFIED

    def record_runtime_observation(
        self, peer_pid: int, request: dict[str, object]
    ) -> dict[str, object]:
        """Record only an authenticated adapter's finalized assistant message.

        This method is deliberately unreachable through ``Broker.handle`` and
        its public broker socket. The production observer socket calls it after
        SO_PEERCRED/ancestry authentication, preserving the S1 rule that a
        broker request can never carry ``latest_assistant_message``.
        """

        required = {
            "action", "session_id", "runtime_generation", "runtime", "latest_assistant_message"
        }
        if set(request) != required or request.get("action") != "record_runtime_observation":
            raise BrokerFailure("INVALID_OBSERVATION")
        runtime = request.get("runtime")
        message = request.get("latest_assistant_message")
        if runtime not in READ_ONLY_TOOLS or not isinstance(message, str) or len(message.encode("utf-8")) > MAX_FRAME:
            raise BrokerFailure("INVALID_OBSERVATION")
        session_id, _ = self.authenticate(peer_pid, request)
        generation = request["runtime_generation"]
        lease = self.leases.get(session_id)
        if (
            not isinstance(lease, dict)
            or lease.get("state") != LEASE_PENDING
            or lease.get("runtime") != runtime
            or lease.get("runtime_generation") != generation
            or not isinstance(self.observer, RuntimeReceiptObserver)
        ):
            raise BrokerFailure("OBSERVATION_UNAVAILABLE")
        self.observer.record_latest_assistant_message(session_id, runtime, generation, message)
        return {"ok": True}

    def handle(self, peer: tuple[int, int, int], request: dict[str, object]) -> dict[str, object]:
        if self.store.poisoned:
            raise StateCommitUncertain()
        previous = copy.deepcopy(self.store.value)
        previous_leases = copy.deepcopy(self.leases)
        self._rejected_cycle_fence = None
        try:
            response = self._handle(peer, request)
            if self.store.value != previous:
                self.store.commit()
            if request.get("action") in {"promote_lease", "complete_recovery"}:
                session_id = request.get("session_id")
                if not isinstance(session_id, str):
                    raise BrokerFailure("INVALID_IDENTITY")
                self.finish_promotion(session_id)
                response["state"] = LEASE_VERIFIED
            return response
        except StateCommitUncertain:
            raise
        except Exception:
            fence = self._rejected_cycle_fence
            if fence is None:
                self.store.value = previous
                self.leases = previous_leases
            else:
                # A refused re-verification must never resurrect the preceding
                # VERIFIED authority. Preserve only this post-revoke fence;
                # every unrelated partial-write failure still rolls back.
                self.store.value, self.leases = fence
            raise
        finally:
            self._rejected_cycle_fence = None

    def _handle(self, peer: tuple[int, int, int], request: dict[str, object]) -> dict[str, object]:
        peer_pid, peer_uid, peer_gid = peer
        action = request.get("action")
        if action == "register_anchor":
            if "session_id" in request:
                raise BrokerFailure("CALLER_SESSION_ID_REFUSED")
            generation = request.get("runtime_generation")
            if not is_non_negative_integer(generation):
                raise BrokerFailure("INVALID_GENERATION")
            anchor = proc_node(peer_pid)
            anchor_starttime = str(anchor["starttime"])
            existing = self.session_for_anchor(peer_pid, anchor_starttime)
            if existing is None:
                session_id = secrets.token_hex(32)
                self.store.sessions()[session_id] = {
                    "anchor_pid": peer_pid,
                    "anchor_starttime": anchor_starttime,
                    "runtime_generation": generation,
                }
            else:
                session_id, session = existing
                current_generation = session["runtime_generation"]
                if generation < current_generation:
                    raise BrokerFailure("STALE_GENERATION")
                if generation > current_generation:
                    session["runtime_generation"] = generation
                    self.revoke_session_authority(session_id)
            return {"ok": True, "session_id": session_id, "peer": {"pid": peer_pid, "uid": peer_uid, "gid": peer_gid, "starttime": anchor["starttime"]}}
        if action == "authenticate":
            self.authenticate(peer_pid, request)
            return {"ok": True}
        if action == "mint_token":
            session_id, _ = self.authenticate(peer_pid, request)
            binding = request.get("binding")
            if not valid_binding(binding):
                raise BrokerFailure("INVALID_BINDING")
            token = self.mint_token(session_id, request["runtime_generation"], binding)
            return {"ok": True, "token": token}
        if action == "consume_token":
            session_id, _ = self.authenticate(peer_pid, request)
            token_value = request.get("token")
            token = self.store.tokens().get(token_value) if isinstance(token_value, str) else None
            if not isinstance(token, dict) or token.get("session_id") != session_id or token.get("runtime_generation") != request.get("runtime_generation") or token.get("consumed") is not False:
                raise BrokerFailure("TOKEN_REPLAY")
            del self.store.tokens()[token_value]
            return {"ok": True}
        if action == "begin_recovery":
            # Recovery is the one ungated mutator, but it is not a second
            # receipt protocol. It delegates to this exact normal-path
            # transition, then marks its volatile pending cycle so completion
            # can obtain the broker-minted challenge internally. Caller-supplied
            # receipt text is never an input to recovery.
            if any(field in request for field in ("receipt", "latest_assistant_message", "receipt_challenge")):
                raise BrokerFailure("INVALID_RECOVERY_REQUEST")
            normal_request = dict(request)
            normal_request["action"] = "begin_verification"
            response = self._handle(peer, normal_request)
            session_id = response.get("session_id")
            if not isinstance(session_id, str):
                # begin_verification deliberately does not return identity;
                # recover it only after its authenticated shared transition.
                candidate = request.get("session_id")
                if not isinstance(candidate, str):
                    raise BrokerFailure("INVALID_IDENTITY")
                session_id = candidate
            lease = self.leases.get(session_id)
            if not isinstance(lease, dict) or lease.get("state") != LEASE_PENDING:
                raise BrokerFailure("STATE_INTEGRITY")
            lease["cycle_kind"] = "recovery"
            response["state"] = "PENDING_DELIVERY"
            return response
        if action == "complete_recovery":
            if any(field in request for field in ("receipt", "latest_assistant_message", "receipt_challenge")):
                raise BrokerFailure("INVALID_RECOVERY_REQUEST")
            session_id, _ = self.authenticate(peer_pid, request)
            lease = self.leases.get(session_id)
            challenge = lease.get("receipt_challenge") if isinstance(lease, dict) else None
            if (
                not isinstance(lease, dict)
                or lease.get("cycle_kind") != "recovery"
                or lease.get("state") != LEASE_PENDING
                or not isinstance(challenge, str)
            ):
                raise BrokerFailure("INVALID_LEASE_TRANSITION")
            try:
                # Reuse the shipped observe -> evidence commit -> consume ->
                # promote transition. The recovery caller supplies neither a
                # normal-path receipt nor a challenge; the trusted observer
                # and current broker cycle remain the sole evidence authority.
                observed_request = {
                    "action": "observe_receipt",
                    "session_id": session_id,
                    "runtime_generation": request["runtime_generation"],
                    "receipt_challenge": challenge,
                }
                self._handle(peer, observed_request)
                return self._handle(peer, {
                    "action": "promote_lease",
                    "session_id": session_id,
                    "runtime_generation": request["runtime_generation"],
                    "receipt_challenge": challenge,
                })
            except Exception:
                # A malformed, absent, or stale observed receipt must leave
                # no pending recovery capability or live lease. A retry mints
                # a new challenge through the shared begin transition.
                self.revoke_session_authority(session_id)
                self._rejected_cycle_fence = (
                    copy.deepcopy(self.store.value), copy.deepcopy(self.leases)
                )
                raise
        if action == "begin_verification":
            session_id, _ = self.authenticate(peer_pid, request)
            runtime = request.get("runtime")
            binding = request.get("binding")
            construction = request.get("construction")
            ttl_seconds = request.get("ttl_seconds", MAX_LEASE_TTL_SECONDS)
            if runtime not in READ_ONLY_TOOLS:
                raise BrokerFailure("INVALID_RUNTIME")
            if not valid_binding(binding):
                raise BrokerFailure("INVALID_BINDING")
            if (
                type(ttl_seconds) is not int
                or ttl_seconds <= 0
                or ttl_seconds > MAX_LEASE_TTL_SECONDS
            ):
                raise BrokerFailure("INVALID_LEASE_TTL")
            # Revoke-first is a broker operation, not advisory adapter order.
            self.revoke_session_authority(session_id)
            # If subsequent construction admission rejects, handle() restores
            # this fence rather than the pre-cycle VERIFIED snapshot.
            self._rejected_cycle_fence = (
                copy.deepcopy(self.store.value), copy.deepcopy(self.leases)
            )
            try:
                constructed = build_payload_from_wire(construction)
            except ValueError as exc:
                raise BrokerFailure("INVALID_CONSTRUCTION") from exc
            if (
                constructed.injectionDecision != "ACCEPTED"
                or not constructed.promotion
                or not isinstance(constructed.h_source, str)
                or not isinstance(constructed.h_payload, str)
            ):
                raise BrokerFailure("PAYLOAD_CONSTRUCTION_REFUSED")
            if (
                not hmac.compare_digest(binding["h_source"], constructed.h_source)
                or not hmac.compare_digest(binding["h_payload"], constructed.h_payload)
            ):
                raise BrokerFailure("PAYLOAD_BINDING_MISMATCH")
            cycle_binding = copy.deepcopy(binding)
            cycle_binding["runtime_generation"] = request["runtime_generation"]
            challenge = self.mint_token(session_id, request["runtime_generation"], binding)
            self.leases[session_id] = {
                "state": LEASE_PENDING,
                "runtime": runtime,
                "runtime_generation": request["runtime_generation"],
                "binding": copy.deepcopy(binding),
                "receipt_challenge": challenge,
                "ttl_seconds": ttl_seconds,
            }
            # The broker constructs both the bound challenge and the delivery
            # text. The model's role is an exact copy, never hashing its output.
            return {
                "ok": True,
                "state": LEASE_PENDING,
                "receipt_challenge": challenge,
                "receipt": receipt_for(challenge, cycle_binding),
                "binding": cycle_binding,
            }
        if action == "observe_receipt":
            session_id, _ = self.authenticate(peer_pid, request)
            challenge = request.get("receipt_challenge")
            if "latest_assistant_message" in request:
                raise BrokerFailure("INVALID_RECEIPT")
            if not isinstance(challenge, str):
                raise BrokerFailure("INVALID_RECEIPT")
            token = self.store.tokens().get(challenge)
            if (
                not isinstance(token, dict)
                or token.get("session_id") != session_id
                or token.get("runtime_generation") != request.get("runtime_generation")
                or token.get("consumed") is not False
            ):
                raise BrokerFailure("RECEIPT_REPLAY")
            lease = self.leases.get(session_id)
            if not isinstance(lease, dict) or lease.get("state") != LEASE_PENDING:
                raise BrokerFailure("RECEIPT_REPLAY")
            expected_challenge = lease.get("receipt_challenge")
            binding = lease.get("binding")
            if (
                not isinstance(expected_challenge, str)
                or not secrets.compare_digest(challenge, expected_challenge)
                or not isinstance(binding, dict)
                or token.get("binding") != binding
            ):
                raise BrokerFailure("RECEIPT_REPLAY")
            receipt_binding = copy.deepcopy(binding)
            receipt_binding["runtime_generation"] = request["runtime_generation"]
            message = self.observer.observe_latest_assistant_message(
                session_id, str(lease["runtime"]), request["runtime_generation"], receipt_binding
            )
            if not isinstance(message, str):
                raise BrokerFailure("RECEIPT_OBSERVATION_UNAVAILABLE")
            # This accepts one exact current-cycle assistant entry only. It is
            # deliberately not a transcript search and rejects quoted/extra text.
            if not is_verbatim_receipt(message, challenge, receipt_binding):
                raise BrokerFailure("RECEIPT_MISMATCH")
            evidence = {"h_latest_assistant": latest_assistant_digest(message)}
            token["evidence"] = evidence
            lease["state"] = LEASE_PENDING_PROMOTION
            lease["evidence"] = copy.deepcopy(evidence)
            return {"ok": True, "state": LEASE_PENDING_PROMOTION}
        if action == "promote_lease":
            session_id, _ = self.authenticate(peer_pid, request)
            challenge = request.get("receipt_challenge")
            if not isinstance(challenge, str):
                raise BrokerFailure("INVALID_RECEIPT")
            lease = self.leases.get(session_id)
            if not isinstance(lease, dict) or lease.get("state") == LEASE_UNVERIFIED:
                raise BrokerFailure("INVALID_LEASE_TRANSITION")
            token = self.store.tokens().get(challenge)
            if (
                not isinstance(token, dict)
                or token.get("session_id") != session_id
                or token.get("runtime_generation") != request.get("runtime_generation")
                or token.get("consumed") is not False
            ):
                raise BrokerFailure("RECEIPT_REPLAY")
            if lease.get("state") != LEASE_PENDING_PROMOTION:
                raise BrokerFailure("INVALID_LEASE_TRANSITION")
            expected_challenge = lease.get("receipt_challenge")
            if not isinstance(expected_challenge, str) or not secrets.compare_digest(challenge, expected_challenge):
                raise BrokerFailure("RECEIPT_REPLAY")
            if token.get("binding") != lease.get("binding") or not valid_receipt_evidence(token.get("evidence")):
                raise BrokerFailure("PROMOTION_TOKEN_INVALID")
            del self.store.tokens()[challenge]
            lease["expires_at"] = time.monotonic() + int(lease["ttl_seconds"])
            # handle() commits the evidence-backed consumption before
            # finish_promotion() makes VERIFIED externally visible: promote-last.
            return {"ok": True, "state": LEASE_PENDING_PROMOTION}
        if action == "revoke_lease":
            session_id, _ = self.authenticate(peer_pid, request)
            self.revoke_session_authority(session_id)
            return {"ok": True, "state": LEASE_UNVERIFIED}
        if action == "authorize_tool":
            session_id, _ = self.authenticate(peer_pid, request)
            runtime = request.get("runtime")
            tool_name = request.get("tool_name")
            if runtime not in READ_ONLY_TOOLS:
                raise BrokerFailure("INVALID_RUNTIME")
            if not isinstance(tool_name, str) or not tool_name or len(tool_name) > 256:
                raise BrokerFailure("INVALID_TOOL")
            lease = self.leases.get(session_id)
            state = lease.get("state") if isinstance(lease, dict) else LEASE_UNVERIFIED
            if tool_name in READ_ONLY_TOOLS[runtime] or tool_name == RECOVERY_TOOL:
                return {"ok": True, "decision": "allow", "state": state}
            if not isinstance(lease, dict) or lease.get("state") != LEASE_VERIFIED:
                return {
                    "ok": False,
                    "code": "MUTATOR_UNVERIFIED",
                    "decision": "deny",
                    "state": LEASE_UNVERIFIED,
                }
            if (
                lease.get("runtime") != runtime
                or lease.get("runtime_generation") != request.get("runtime_generation")
            ):
                return {
                    "ok": False,
                    "code": "MUTATOR_UNVERIFIED",
                    "decision": "deny",
                    "state": LEASE_UNVERIFIED,
                }
            expires_at = lease.get("expires_at")
            if not isinstance(expires_at, (int, float)) or time.monotonic() >= expires_at:
                self.revoke_session_authority(session_id)
                return {
                    "ok": False,
                    "code": "LEASE_EXPIRED",
                    "decision": "deny",
                    "state": LEASE_UNVERIFIED,
                }
            return {
                "ok": True,
                "decision": "allow",
                "state": LEASE_VERIFIED,
            }
        raise BrokerFailure("UNKNOWN_ACTION")


def read_frame(connection: socket.socket, deadline: float) -> dict[str, object]:
    data = bytearray()
    while len(data) <= MAX_FRAME:
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            raise BrokerFailure("MALFORMED_REQUEST")
        connection.settimeout(remaining)
        chunk = connection.recv(min(4096, MAX_FRAME + 1 - len(data)))
        if not chunk:
            break
        data.extend(chunk)
    if len(data) > MAX_FRAME:
        while True:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise BrokerFailure("MALFORMED_REQUEST")
            connection.settimeout(remaining)
            if not connection.recv(4096):
                break
        raise BrokerFailure("MALFORMED_REQUEST")
    if not data.endswith(b"\n") or data.count(b"\n") != 1:
        raise BrokerFailure("MALFORMED_REQUEST")
    try:
        value = json.loads(data)
    except (json.JSONDecodeError, UnicodeDecodeError) as exc:
        raise BrokerFailure("MALFORMED_REQUEST") from exc
    if not isinstance(value, dict):
        raise BrokerFailure("MALFORMED_REQUEST")
    return value


def handle_connection(
    connection: socket.socket,
    broker: Broker,
    broker_lock: threading.Lock,
) -> None:
    with connection:
        read_deadline = time.monotonic() + READ_DEADLINE_SECONDS
        try:
            raw = connection.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12)
            peer = struct.unpack("3i", raw)
            request = read_frame(connection, read_deadline)
        except BrokerFailure as exc:
            reply = {"ok": False, "code": exc.code}
        except OSError:
            return
        else:
            # Queueing is bounded and fails closed before broker.handle mutates
            # state. Once handling starts it must finish atomically; its reply
            # then receives an independent send budget so a slow fsync cannot
            # consume the write opportunity and create client/broker ambiguity.
            acquired = broker_lock.acquire(timeout=HANDLE_QUEUE_TIMEOUT_SECONDS)
            if not acquired:
                reply = {"ok": False, "code": "BROKER_BUSY"}
            else:
                try:
                    try:
                        reply = broker.handle(peer, request)
                    except BrokerFailure as exc:
                        reply = {"ok": False, "code": exc.code}
                finally:
                    broker_lock.release()
        try:
            connection.settimeout(SEND_TIMEOUT_SECONDS)
            connection.sendall((json.dumps(reply, separators=(",", ":")) + "\n").encode())
        except OSError:
            return


def handle_runtime_observation_connection(
    connection: socket.socket,
    broker: Broker,
    broker_lock: threading.Lock,
) -> None:
    """Serve the authenticated production observer transport, never the broker API."""

    with connection:
        read_deadline = time.monotonic() + READ_DEADLINE_SECONDS
        try:
            raw = connection.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12)
            peer = struct.unpack("3i", raw)
            request = read_frame(connection, read_deadline)
        except BrokerFailure as exc:
            reply = {"ok": False, "code": exc.code}
        except OSError:
            return
        else:
            acquired = broker_lock.acquire(timeout=HANDLE_QUEUE_TIMEOUT_SECONDS)
            if not acquired:
                reply = {"ok": False, "code": "BROKER_BUSY"}
            else:
                try:
                    try:
                        reply = broker.record_runtime_observation(peer[0], request)
                    except BrokerFailure as exc:
                        reply = {"ok": False, "code": exc.code}
                finally:
                    broker_lock.release()
        try:
            connection.settimeout(SEND_TIMEOUT_SECONDS)
            connection.sendall((json.dumps(reply, separators=(",", ":")) + "\n").encode())
        except OSError:
            return


def serve(
    socket_path: Path,
    state_path: Path,
    observer: ReceiptObserver | None = None,
    observer_socket_path: Path | None = None,
) -> None:
    secure_parent(socket_path)
    if socket_path.exists() or socket_path.is_symlink():
        raise BrokerFailure("SOCKET_ALREADY_EXISTS")
    runtime_observer = observer is None
    if runtime_observer:
        observer = RuntimeReceiptObserver()
        observer_socket_path = observer_socket_path or socket_path.with_name("receipt-observer.sock")
        if observer_socket_path == socket_path:
            raise BrokerFailure("OBSERVER_SOCKET_CONFLICT")
        secure_parent(observer_socket_path)
        if observer_socket_path.exists() or observer_socket_path.is_symlink():
            raise BrokerFailure("OBSERVER_SOCKET_ALREADY_EXISTS")
    elif observer_socket_path is not None:
        raise BrokerFailure("TEST_OBSERVER_SOCKET_CONFLICT")

    store = StateStore(state_path)
    broker = Broker(store, observer)
    broker_lock = threading.Lock()
    slots = threading.BoundedSemaphore(MAX_IN_FLIGHT_CONNECTIONS)
    fatal_lock = threading.Lock()
    fatal_errors: list[Exception] = []
    executor = ThreadPoolExecutor(
        max_workers=MAX_IN_FLIGHT_CONNECTIONS,
        thread_name_prefix="mosaic-lease-broker",
    )
    server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    server.bind(str(socket_path))
    os.chmod(socket_path, 0o600)
    owned = (socket_path.stat().st_dev, socket_path.stat().st_ino)
    observer_server: socket.socket | None = None
    observer_owned: tuple[int, int] | None = None
    if runtime_observer and observer_socket_path is not None:
        observer_server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        observer_server.bind(str(observer_socket_path))
        os.chmod(observer_socket_path, 0o600)
        observer_owned = (observer_socket_path.stat().st_dev, observer_socket_path.stat().st_ino)
    stopping = False

    def stop(_signum: int, _frame: object) -> None:
        nonlocal stopping
        stopping = True
        server.close()
        if observer_server is not None:
            observer_server.close()

    def process_connection(connection: socket.socket, is_observer: bool) -> None:
        try:
            if is_observer:
                handle_runtime_observation_connection(connection, broker, broker_lock)
            else:
                handle_connection(connection, broker, broker_lock)
        except Exception as exc:
            with fatal_lock:
                if not fatal_errors:
                    fatal_errors.append(exc)
        finally:
            slots.release()

    def fatal_error() -> Exception | None:
        with fatal_lock:
            return fatal_errors[0] if fatal_errors else None

    signal.signal(signal.SIGTERM, stop)
    signal.signal(signal.SIGINT, stop)
    server.listen(MAX_IN_FLIGHT_CONNECTIONS)
    server.setblocking(False)
    if observer_server is not None:
        observer_server.listen(MAX_IN_FLIGHT_CONNECTIONS)
        observer_server.setblocking(False)
    print("READY", flush=True)
    try:
        while not stopping:
            failure = fatal_error()
            if failure is not None:
                raise failure
            listeners = [server, *([observer_server] if observer_server is not None else [])]
            try:
                ready, _, _ = select.select(listeners, [], [], 0.1)
            except (OSError, ValueError):
                if stopping:
                    break
                raise
            for listener in ready:
                if not slots.acquire(blocking=False):
                    continue
                try:
                    connection, _ = listener.accept()
                except BlockingIOError:
                    slots.release()
                    continue
                except OSError:
                    slots.release()
                    if stopping:
                        break
                    raise
                try:
                    executor.submit(process_connection, connection, listener is observer_server)
                except Exception:
                    slots.release()
                    connection.close()
                    raise
    finally:
        server.close()
        if observer_server is not None:
            observer_server.close()
        executor.shutdown(wait=True)
        for path, inode in ((socket_path, owned), (observer_socket_path, observer_owned)):
            if path is None or inode is None:
                continue
            try:
                current = path.stat()
                if (current.st_dev, current.st_ino) == inode:
                    path.unlink()
            except FileNotFoundError:
                pass


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--socket", required=True, type=Path)
    parser.add_argument("--state", required=True, type=Path)
    parser.add_argument("--observer-socket", type=Path)
    parser.add_argument("--test-observer-file", type=Path)
    arguments = parser.parse_args()
    if arguments.observer_socket is not None and arguments.test_observer_file is not None:
        raise BrokerFailure("TEST_OBSERVER_SOCKET_CONFLICT")
    observer = (
        FileTestReceiptObserver(arguments.test_observer_file)
        if arguments.test_observer_file is not None
        else None
    )
    serve(arguments.socket, arguments.state, observer, arguments.observer_socket)


if __name__ == "__main__":
    try:
        main()
    except BrokerFailure as failure:
        print(failure.code, file=sys.stderr)
        raise SystemExit(1)
    except StateCommitUncertain as failure:
        print(str(failure), file=sys.stderr)
        raise SystemExit(1)