WI-5 #832: Receipt-challenge protocol (compaction-refresh, milestone 188) (#845)

packages/mosaic/framework/tools/lease-broker/daemon.py

@@ -6,6 +6,7 @@ from __future__ import annotations
 import argparse
 import copy
 import errno
+import hmac
 from concurrent.futures import ThreadPoolExecutor
 import json
 import os
@@ -20,6 +21,15 @@ 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, UnavailableReceiptObserver
+
 MAX_FRAME: Final = 64 * 1024
 MAX_STATE: Final = 4 * 1024 * 1024
 MAX_PENDING_TOKENS: Final = 256
@@ -90,6 +100,14 @@ def valid_binding(binding: object) -> bool:
     )
 
 
+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")
@@ -124,7 +142,11 @@ def validate_state(value: object) -> dict[str, object]:
     for token_value, token in tokens.items():
         if not is_hex_256(token_value) or not isinstance(token, dict):
             raise BrokerFailure("STATE_INTEGRITY")
-        if set(token) != {"session_id", "runtime_generation", "binding", "consumed"}:
+        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"]
@@ -137,6 +159,9 @@ def validate_state(value: object) -> dict[str, object]:
             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
 
 
@@ -275,8 +300,13 @@ class StateStore:
 
 
 class Broker:
-    def __init__(self, store: StateStore) -> None:
+    def __init__(self, store: StateStore, observer: ReceiptObserver | None = None) -> None:
         self.store = store
+        self.observer: ReceiptObserver = observer if observer is not None else UnavailableReceiptObserver()
+        # 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]] = {}
@@ -348,6 +378,9 @@ class Broker:
             "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
 
@@ -362,6 +395,7 @@ class Broker:
             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:
@@ -376,9 +410,18 @@ class Broker:
         except StateCommitUncertain:
             raise
         except Exception:
-            self.store.value = previous
-            self.leases = previous_leases
+            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
@@ -430,6 +473,7 @@ class Broker:
             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")
@@ -443,47 +487,117 @@ class Broker:
                 raise BrokerFailure("INVALID_LEASE_TTL")
             # Revoke-first is a broker operation, not advisory adapter order.
             self.revoke_session_authority(session_id)
-            token = self.mint_token(session_id, request["runtime_generation"], binding)
+            # 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),
-                "promotion_token": token,
+                "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,
-                "promotion_token": token,
+                "receipt_challenge": challenge,
+                "receipt": receipt_for(challenge, cycle_binding),
+                "binding": cycle_binding,
             }
-        if action == "promote_lease":
+        if action == "observe_receipt":
             session_id, _ = self.authenticate(peer_pid, request)
-            promotion_token = request.get("promotion_token")
-            lease = self.leases.get(session_id)
-            if not isinstance(lease, dict) or lease.get("state") != LEASE_PENDING:
-                raise BrokerFailure("INVALID_LEASE_TRANSITION")
-            expected_token = lease.get("promotion_token")
-            if (
-                not isinstance(promotion_token, str)
-                or not isinstance(expected_token, str)
-                or not secrets.compare_digest(promotion_token, expected_token)
-            ):
-                raise BrokerFailure("PROMOTION_TOKEN_MISMATCH")
-            token = self.store.tokens().get(promotion_token)
+            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("binding") != lease.get("binding")
                 or token.get("consumed") is not False
             ):
-                raise BrokerFailure("PROMOTION_TOKEN_INVALID")
-            del self.store.tokens()[promotion_token]
+                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 token consumption before finish_promotion() makes
-            # VERIFIED externally visible: promote-last by construction.
+            # 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)
@@ -604,12 +718,16 @@ def handle_connection(
             return
 
 
-def serve(socket_path: Path, state_path: Path) -> None:
+def serve(
+    socket_path: Path,
+    state_path: Path,
+    observer: ReceiptObserver | None = None,
+) -> None:
     secure_parent(socket_path)
     if socket_path.exists() or socket_path.is_symlink():
         raise BrokerFailure("SOCKET_ALREADY_EXISTS")
     store = StateStore(state_path)
-    broker = Broker(store)
+    broker = Broker(store, observer)
     broker_lock = threading.Lock()
     slots = threading.BoundedSemaphore(MAX_IN_FLIGHT_CONNECTIONS)
     fatal_lock = threading.Lock()
@@ -689,8 +807,14 @@ 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("--test-observer-file", type=Path)
     arguments = parser.parse_args()
-    serve(arguments.socket, arguments.state)
+    observer = (
+        FileTestReceiptObserver(arguments.test_observer_file)
+        if arguments.test_observer_file is not None
+        else None
+    )
+    serve(arguments.socket, arguments.state, observer)
 
 
 if __name__ == "__main__":