167 lines
6.5 KiB
Python
167 lines
6.5 KiB
Python
#!/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
|
|
BENIGN_OBSERVATION_UNAVAILABLE_CODE: Final = "OBSERVATION_UNAVAILABLE"
|
|
|
|
|
|
def reject_duplicate_json_keys(pairs: list[tuple[str, object]]) -> dict[str, object]:
|
|
value: dict[str, object] = {}
|
|
for key, item in pairs:
|
|
if key in value:
|
|
raise ValueError("duplicate observer JSON key")
|
|
value[key] = item
|
|
return value
|
|
|
|
|
|
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, object_pairs_hook=reject_duplicate_json_keys)
|
|
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 response.count(b"\n") != 1 or not response.endswith(b"\n"):
|
|
raise ValueError("invalid observer reply")
|
|
value = json.loads(response[:-1], object_pairs_hook=reject_duplicate_json_keys)
|
|
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")
|
|
if "last_assistant_message" in source:
|
|
message = source["last_assistant_message"]
|
|
if not isinstance(message, str):
|
|
raise ValueError("invalid Claude observer input")
|
|
else:
|
|
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, RecursionError, ValueError, json.JSONDecodeError) as error:
|
|
print(f"Mosaic receipt observer refused: {error}", file=sys.stderr)
|
|
return 2
|
|
if set(reply) == {"ok"} and reply.get("ok") is True:
|
|
return 0
|
|
if (
|
|
set(reply) == {"ok", "code"}
|
|
and reply.get("ok") is False
|
|
and reply.get("code") == BENIGN_OBSERVATION_UNAVAILABLE_CODE
|
|
):
|
|
return 0
|
|
return 2
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|