feat(mosaic): add constrained context recovery command
This commit is contained in:
wjarvis mos-comms
@@ -400,7 +400,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 +469,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")
|
||||
|
||||
146
packages/mosaic/framework/tools/lease-broker/recover-context.py
Normal file
146
packages/mosaic/framework/tools/lease-broker/recover-context.py
Normal file
@@ -0,0 +1,146 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Constrained recovery command: the sole ungated Mosaic mutator.
|
||||
|
||||
This is deliberately a thin client of the broker's recovery entrypoint. It
|
||||
never accepts receipt text or a caller-provided challenge: the broker mints the
|
||||
fresh challenge, delivers its exact receipt envelope, and later asks the
|
||||
trusted ReceiptObserver seam to observe that same pending cycle.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import socket
|
||||
import sys
|
||||
from collections.abc import Mapping, Sequence
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
from normative_fragments import build_payload_from_wire
|
||||
|
||||
MAX_FRAME: Final = 64 * 1024
|
||||
BROKER_TIMEOUT_SECONDS: Final = 1.5
|
||||
|
||||
|
||||
def broker_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("recovery 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 broker reply")
|
||||
value = json.loads(response)
|
||||
if not isinstance(value, dict):
|
||||
raise ValueError("invalid broker reply")
|
||||
return value
|
||||
|
||||
|
||||
def load_construction(path: Path) -> dict[str, object]:
|
||||
raw = path.read_bytes()
|
||||
if len(raw) > MAX_FRAME:
|
||||
raise ValueError("construction exceeds broker frame limit")
|
||||
value = json.loads(raw)
|
||||
if not isinstance(value, dict):
|
||||
raise ValueError("construction must be an object")
|
||||
return value
|
||||
|
||||
|
||||
def identity(environ: Mapping[str, str]) -> tuple[Path, str, int, str]:
|
||||
socket_path = Path(environ["MOSAIC_LEASE_BROKER_SOCKET"])
|
||||
session_id = environ["MOSAIC_LEASE_SESSION_ID"]
|
||||
generation = int(environ["MOSAIC_RUNTIME_GENERATION"])
|
||||
runtime = environ["MOSAIC_LEASE_RUNTIME"]
|
||||
if generation < 0 or runtime not in {"claude", "pi"}:
|
||||
raise ValueError("invalid runtime identity")
|
||||
return socket_path, session_id, generation, runtime
|
||||
|
||||
|
||||
def begin(
|
||||
construction_path: Path,
|
||||
compaction_epoch: int,
|
||||
request_epoch: int,
|
||||
environ: Mapping[str, str],
|
||||
) -> dict[str, object]:
|
||||
if compaction_epoch < 0 or request_epoch < 0:
|
||||
raise ValueError("epochs must be non-negative")
|
||||
construction = load_construction(construction_path)
|
||||
# Invoke the shared WI-5 construction before asking the broker to repeat
|
||||
# its authoritative admission/build. No digest or receipt enters via CLI.
|
||||
built = build_payload_from_wire(construction)
|
||||
if (
|
||||
built.injectionDecision != "ACCEPTED"
|
||||
or not built.promotion
|
||||
or not isinstance(built.h_source, str)
|
||||
or not isinstance(built.h_payload, str)
|
||||
):
|
||||
raise ValueError("payload construction refused")
|
||||
socket_path, session_id, generation, runtime = identity(environ)
|
||||
return broker_request(socket_path, {
|
||||
"action": "begin_recovery",
|
||||
"session_id": session_id,
|
||||
"runtime_generation": generation,
|
||||
"runtime": runtime,
|
||||
"binding": {
|
||||
"compaction_epoch": compaction_epoch,
|
||||
"request_epoch": request_epoch,
|
||||
"h_source": built.h_source,
|
||||
"h_payload": built.h_payload,
|
||||
"schema_version": 1,
|
||||
},
|
||||
"construction": construction,
|
||||
})
|
||||
|
||||
|
||||
def complete(environ: Mapping[str, str]) -> dict[str, object]:
|
||||
socket_path, session_id, generation, _runtime = identity(environ)
|
||||
# No receipt or challenge argument exists: recovery completion can only use
|
||||
# the broker's current recovery cycle and its trusted observer seam.
|
||||
return broker_request(socket_path, {
|
||||
"action": "complete_recovery",
|
||||
"session_id": session_id,
|
||||
"runtime_generation": generation,
|
||||
})
|
||||
|
||||
|
||||
def main(argv: Sequence[str] | None = None, *, environ: Mapping[str, str] | None = None) -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
subcommands = parser.add_subparsers(dest="phase", required=True)
|
||||
begin_parser = subcommands.add_parser("begin", help="mint and deliver a fresh recovery receipt")
|
||||
begin_parser.add_argument("--construction", required=True, type=Path)
|
||||
begin_parser.add_argument("--compaction-epoch", required=True, type=int)
|
||||
begin_parser.add_argument("--request-epoch", required=True, type=int)
|
||||
subcommands.add_parser("complete", help="observe and promote only the current recovery receipt")
|
||||
arguments = parser.parse_args(argv)
|
||||
source_environment = os.environ if environ is None else environ
|
||||
try:
|
||||
reply = (
|
||||
begin(
|
||||
arguments.construction,
|
||||
arguments.compaction_epoch,
|
||||
arguments.request_epoch,
|
||||
source_environment,
|
||||
)
|
||||
if arguments.phase == "begin"
|
||||
else complete(source_environment)
|
||||
)
|
||||
except (KeyError, OSError, ValueError, json.JSONDecodeError) as error:
|
||||
print(f"Mosaic constrained recovery refused: {error}", file=sys.stderr)
|
||||
return 2
|
||||
print(json.dumps(reply, separators=(",", ":")))
|
||||
return 0 if reply.get("ok") is True else 2
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user