WI-6 (#833): constrained recovery command + mosaic-context-refresh skill wrapper (#846)

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

@@ -0,0 +1,143 @@
+#!/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())