feat(#829): enforce whole mutator-class lease gate

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

@@ -24,9 +24,19 @@ MAX_FRAME: Final = 64 * 1024
 MAX_STATE: Final = 4 * 1024 * 1024
 MAX_PENDING_TOKENS: Final = 256
 MAX_IN_FLIGHT_CONNECTIONS: Final = 16
+MAX_LEASE_TTL_SECONDS: Final = 300
 STATE_VERSION: Final = 1
 CONNECTION_DEADLINE_SECONDS: Final = 1.0
 HEX_256_LENGTH: Final = 64
+LEASE_UNVERIFIED: Final = "UNVERIFIED"
+LEASE_PENDING: Final = "PENDING_VERIFICATION"
+LEASE_PENDING_PROMOTION: Final = "PENDING_PROMOTION"
+LEASE_VERIFIED: Final = "VERIFIED"
+READ_ONLY_TOOLS: Final = {
+    "claude": frozenset({"Read", "Grep", "Glob", "Ls", "Find"}),
+    "pi": frozenset({"read", "grep", "find", "ls"}),
+}
+RECOVERY_TOOL: Final = "mosaic_context_recover"
 
 
 class BrokerFailure(Exception):
@@ -265,6 +275,9 @@ class StateStore:
 class Broker:
     def __init__(self, store: StateStore) -> None:
         self.store = store
+        # 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]] = {}
 
     def authenticate(self, peer_pid: int, request: dict[str, object]) -> tuple[str, dict[str, object]]:
         session_id = request.get("session_id")
@@ -289,7 +302,7 @@ class Broker:
             raise BrokerFailure("STALE_GENERATION")
         if generation > current_generation:
             session["runtime_generation"] = generation
-            self.revoke_session_tokens(session_id)
+            self.revoke_session_authority(session_id)
         return session_id, session
 
     def session_for_anchor(
@@ -310,19 +323,59 @@ class Broker:
         ]:
             del tokens[token_value]
 
+    def revoke_session_authority(self, session_id: str) -> None:
+        self.revoke_session_tokens(session_id)
+        session = self.store.sessions().get(session_id)
+        generation = session.get("runtime_generation") if isinstance(session, dict) else None
+        self.leases[session_id] = {
+            "state": LEASE_UNVERIFIED,
+            "runtime_generation": generation,
+        }
+
+    def mint_token(
+        self,
+        session_id: str,
+        generation: int,
+        binding: dict[str, object],
+    ) -> str:
+        if len(self.store.tokens()) >= MAX_PENDING_TOKENS:
+            raise BrokerFailure("TOKEN_CAPACITY")
+        token = secrets.token_hex(32)
+        self.store.tokens()[token] = {
+            "session_id": session_id,
+            "runtime_generation": generation,
+            "binding": copy.deepcopy(binding),
+            "consumed": False,
+        }
+        return token
+
+    def finish_promotion(self, session_id: str) -> None:
+        lease = self.leases.get(session_id)
+        if not isinstance(lease, dict) or lease.get("state") != LEASE_PENDING_PROMOTION:
+            raise BrokerFailure("STATE_INTEGRITY")
+        lease["state"] = LEASE_VERIFIED
+
     def handle(self, peer: tuple[int, int, int], request: dict[str, object]) -> dict[str, object]:
         if self.store.poisoned:
             raise StateCommitUncertain()
         previous = copy.deepcopy(self.store.value)
+        previous_leases = copy.deepcopy(self.leases)
         try:
             response = self._handle(peer, request)
             if self.store.value != previous:
                 self.store.commit()
+            if request.get("action") == "promote_lease":
+                session_id = request.get("session_id")
+                if not isinstance(session_id, str):
+                    raise BrokerFailure("INVALID_IDENTITY")
+                self.finish_promotion(session_id)
+                response["state"] = LEASE_VERIFIED
             return response
         except StateCommitUncertain:
             raise
         except Exception:
             self.store.value = previous
+            self.leases = previous_leases
             raise
 
     def _handle(self, peer: tuple[int, int, int], request: dict[str, object]) -> dict[str, object]:
@@ -351,7 +404,7 @@ class Broker:
                     raise BrokerFailure("STALE_GENERATION")
                 if generation > current_generation:
                     session["runtime_generation"] = generation
-                    self.revoke_session_tokens(session_id)
+                    self.revoke_session_authority(session_id)
             return {"ok": True, "session_id": session_id, "peer": {"pid": peer_pid, "uid": peer_uid, "gid": peer_gid, "starttime": anchor["starttime"]}}
         if action == "authenticate":
             self.authenticate(peer_pid, request)
@@ -361,15 +414,7 @@ class Broker:
             binding = request.get("binding")
             if not valid_binding(binding):
                 raise BrokerFailure("INVALID_BINDING")
-            if len(self.store.tokens()) >= MAX_PENDING_TOKENS:
-                raise BrokerFailure("TOKEN_CAPACITY")
-            token = secrets.token_hex(32)
-            self.store.tokens()[token] = {
-                "session_id": session_id,
-                "runtime_generation": request["runtime_generation"],
-                "binding": binding,
-                "consumed": False,
-            }
+            token = self.mint_token(session_id, request["runtime_generation"], binding)
             return {"ok": True, "token": token}
         if action == "consume_token":
             session_id, _ = self.authenticate(peer_pid, request)
@@ -379,6 +424,112 @@ class Broker:
                 raise BrokerFailure("TOKEN_REPLAY")
             del self.store.tokens()[token_value]
             return {"ok": True}
+        if action == "begin_verification":
+            session_id, _ = self.authenticate(peer_pid, request)
+            runtime = request.get("runtime")
+            binding = request.get("binding")
+            ttl_seconds = request.get("ttl_seconds", MAX_LEASE_TTL_SECONDS)
+            if runtime not in READ_ONLY_TOOLS:
+                raise BrokerFailure("INVALID_RUNTIME")
+            if not valid_binding(binding):
+                raise BrokerFailure("INVALID_BINDING")
+            if (
+                type(ttl_seconds) is not int
+                or ttl_seconds <= 0
+                or ttl_seconds > MAX_LEASE_TTL_SECONDS
+            ):
+                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)
+            self.leases[session_id] = {
+                "state": LEASE_PENDING,
+                "runtime": runtime,
+                "runtime_generation": request["runtime_generation"],
+                "binding": copy.deepcopy(binding),
+                "promotion_token": token,
+                "ttl_seconds": ttl_seconds,
+            }
+            return {
+                "ok": True,
+                "state": LEASE_PENDING,
+                "promotion_token": token,
+            }
+        if action == "promote_lease":
+            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)
+            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]
+            lease["state"] = LEASE_PENDING_PROMOTION
+            lease["expires_at"] = time.monotonic() + int(lease["ttl_seconds"])
+            # handle() commits token consumption before finish_promotion() makes
+            # VERIFIED externally visible: promote-last by construction.
+            return {"ok": True, "state": LEASE_PENDING_PROMOTION}
+        if action == "revoke_lease":
+            session_id, _ = self.authenticate(peer_pid, request)
+            self.revoke_session_authority(session_id)
+            return {"ok": True, "state": LEASE_UNVERIFIED}
+        if action == "authorize_tool":
+            session_id, _ = self.authenticate(peer_pid, request)
+            runtime = request.get("runtime")
+            tool_name = request.get("tool_name")
+            if runtime not in READ_ONLY_TOOLS:
+                raise BrokerFailure("INVALID_RUNTIME")
+            if not isinstance(tool_name, str) or not tool_name or len(tool_name) > 256:
+                raise BrokerFailure("INVALID_TOOL")
+            lease = self.leases.get(session_id)
+            state = lease.get("state") if isinstance(lease, dict) else LEASE_UNVERIFIED
+            if tool_name in READ_ONLY_TOOLS[runtime] or tool_name == RECOVERY_TOOL:
+                return {"ok": True, "decision": "allow", "state": state}
+            if not isinstance(lease, dict) or lease.get("state") != LEASE_VERIFIED:
+                return {
+                    "ok": False,
+                    "code": "MUTATOR_UNVERIFIED",
+                    "decision": "deny",
+                    "state": LEASE_UNVERIFIED,
+                }
+            if (
+                lease.get("runtime") != runtime
+                or lease.get("runtime_generation") != request.get("runtime_generation")
+            ):
+                return {
+                    "ok": False,
+                    "code": "MUTATOR_UNVERIFIED",
+                    "decision": "deny",
+                    "state": LEASE_UNVERIFIED,
+                }
+            expires_at = lease.get("expires_at")
+            if not isinstance(expires_at, (int, float)) or time.monotonic() >= expires_at:
+                self.revoke_session_authority(session_id)
+                return {
+                    "ok": False,
+                    "code": "LEASE_EXPIRED",
+                    "decision": "deny",
+                    "state": LEASE_UNVERIFIED,
+                }
+            return {
+                "ok": True,
+                "decision": "allow",
+                "state": LEASE_VERIFIED,
+            }
         raise BrokerFailure("UNKNOWN_ACTION")