stack/packages/mosaic/framework/tools/lease-broker/receipt-observer-client.py

#!/usr/bin/env python3
"""Authenticated adapter-to-daemon transport for finalized assistant receipts.

This is not a broker request client. It writes only to the daemon-owned observer
socket, which authenticates SO_PEERCRED/ancestry before retaining a message for
the broker's ReceiptObserver seam.
"""

from __future__ import annotations

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


MAX_FRAME: Final = 64 * 1024
BROKER_TIMEOUT_SECONDS: Final = 1.5
MAX_TRANSCRIPT_BYTES: Final = 4 * 1024 * 1024


def read_json(stream: object) -> dict[str, object]:
    raw = getattr(stream, "buffer", stream).read(MAX_FRAME + 1)
    if not isinstance(raw, bytes) or len(raw) > MAX_FRAME:
        raise ValueError("invalid observer input")
    value = json.loads(raw)
    if not isinstance(value, dict):
        raise ValueError("invalid observer input")
    return value


def assistant_text(entry: object) -> str | None:
    if not isinstance(entry, dict):
        return None
    message = entry.get("message", entry)
    if not isinstance(message, dict) or message.get("role") != "assistant":
        return None
    content = message.get("content")
    if isinstance(content, str):
        return content
    if not isinstance(content, list):
        return None
    parts: list[str] = []
    for item in content:
        if not isinstance(item, dict) or item.get("type") != "text" or not isinstance(item.get("text"), str):
            return None
        parts.append(item["text"])
    return "".join(parts)


def claude_latest_entry(value: dict[str, object]) -> str:
    transcript_path = value.get("transcript_path")
    if not isinstance(transcript_path, str) or not transcript_path:
        raise ValueError("invalid Claude observer input")
    path = Path(transcript_path)
    flags = os.O_RDONLY | getattr(os, "O_CLOEXEC", 0) | getattr(os, "O_NOFOLLOW", 0)
    descriptor = os.open(path, flags)
    try:
        metadata = os.fstat(descriptor)
        if not stat.S_ISREG(metadata.st_mode) or metadata.st_size > MAX_TRANSCRIPT_BYTES:
            raise ValueError("unsafe Claude transcript")
        raw = os.read(descriptor, MAX_TRANSCRIPT_BYTES + 1)
    finally:
        os.close(descriptor)
    if len(raw) > MAX_TRANSCRIPT_BYTES:
        raise ValueError("oversized Claude transcript")
    for line in reversed(raw.decode("utf-8").splitlines()):
        try:
            text = assistant_text(json.loads(line))
        except json.JSONDecodeError as exc:
            raise ValueError("invalid Claude transcript") from exc
        if text is not None:
            return text
    raise ValueError("Claude transcript has no assistant entry")


def pi_message_end(value: dict[str, object]) -> str:
    if set(value) != {"latest_assistant_message"} or not isinstance(value["latest_assistant_message"], str):
        raise ValueError("invalid Pi observer input")
    return value["latest_assistant_message"]


def observer_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("observer 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 observer reply")
    value = json.loads(response)
    if not isinstance(value, dict):
        raise ValueError("invalid observer reply")
    return value


def main(argv: Sequence[str] | None = None, *, environ: Mapping[str, str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--runtime", required=True, choices=("claude", "pi"))
    parser.add_argument("--latest-entry", action="store_true")
    arguments = parser.parse_args(argv)
    source_environment = os.environ if environ is None else environ
    try:
        source = read_json(sys.stdin)
        if arguments.runtime == "claude":
            if not arguments.latest_entry:
                raise ValueError("Claude observer requires --latest-entry")
            message = claude_latest_entry(source)
        else:
            if arguments.latest_entry:
                raise ValueError("Pi observer is message_end only")
            message = pi_message_end(source)
        if len(message.encode("utf-8")) > MAX_FRAME:
            raise ValueError("assistant message too large")
        reply = observer_request(Path(source_environment["MOSAIC_RECEIPT_OBSERVER_SOCKET"]), {
            "action": "record_runtime_observation",
            "session_id": source_environment["MOSAIC_LEASE_SESSION_ID"],
            "runtime_generation": int(source_environment["MOSAIC_RUNTIME_GENERATION"]),
            "runtime": arguments.runtime,
            "latest_assistant_message": message,
        })
    except (KeyError, OSError, ValueError, json.JSONDecodeError) as error:
        print(f"Mosaic receipt observer refused: {error}", file=sys.stderr)
        return 2
    return 0 if reply == {"ok": True} else 2


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