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

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

@@ -11,6 +11,7 @@ from concurrent.futures import ThreadPoolExecutor
 import json
 import os
 import secrets
+import select
 import signal
 import socket
 import stat
@@ -28,7 +29,11 @@ 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
+from receipt_observer import (
+    FileTestReceiptObserver,
+    ReceiptObserver,
+    RuntimeReceiptObserver,
+)
 
 MAX_FRAME: Final = 64 * 1024
 MAX_STATE: Final = 4 * 1024 * 1024
@@ -302,7 +307,9 @@ class StateStore:
 class Broker:
     def __init__(self, store: StateStore, observer: ReceiptObserver | None = None) -> None:
         self.store = store
-        self.observer: ReceiptObserver = observer if observer is not None else UnavailableReceiptObserver()
+        # 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.
@@ -390,6 +397,40 @@ class Broker:
             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()
@@ -400,7 +441,7 @@ class Broker:
             response = self._handle(peer, request)
             if self.store.value != previous:
                 self.store.commit()
-            if request.get("action") == "promote_lease":
+            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")
@@ -469,6 +510,71 @@ class Broker:
                 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")
@@ -718,14 +824,63 @@ def handle_connection(
             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()
@@ -740,16 +895,28 @@ def serve(
     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) -> None:
+    def process_connection(connection: socket.socket, is_observer: bool) -> None:
         try:
-            handle_connection(connection, broker, broker_lock)
+            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:
@@ -764,57 +931,73 @@ def serve(
     signal.signal(signal.SIGTERM, stop)
     signal.signal(signal.SIGINT, stop)
     server.listen(MAX_IN_FLIGHT_CONNECTIONS)
-    server.settimeout(0.1)
+    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
-            if not slots.acquire(timeout=0.1):
-                continue
+            listeners = [server, *([observer_server] if observer_server is not None else [])]
             try:
-                connection, _ = server.accept()
-            except socket.timeout:
-                slots.release()
-                continue
-            except OSError:
-                slots.release()
-                failure = fatal_error()
-                if failure is not None:
-                    raise failure
+                ready, _, _ = select.select(listeners, [], [], 0.1)
+            except (OSError, ValueError):
                 if stopping:
                     break
                 raise
-            try:
-                executor.submit(process_connection, connection)
-            except Exception:
-                slots.release()
-                connection.close()
-                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)
-        try:
-            current = socket_path.stat()
-            if (current.st_dev, current.st_ino) == owned:
-                socket_path.unlink()
-        except FileNotFoundError:
-            pass
+        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)
+    serve(arguments.socket, arguments.state, observer, arguments.observer_socket)
 
 
 if __name__ == "__main__":