feat(cloud): enforce host governance budgets

apps/cloud-api/cloud_api/app.py

@@ -301,6 +301,12 @@ def create_app(
             auth_provider=auth_provider,
             lease_duration_seconds=control_config.lease_duration_seconds,
             scheduler=scheduler,
+            planner_token_reservation_ceiling=(
+                control_config.planner_token_reservation_ceiling
+            ),
+            planner_token_reservation_ttl_seconds=(
+                control_config.planner_token_reservation_ttl_seconds
+            ),
         )
     )
 
@@ -397,6 +403,10 @@ async def _run_lease_reaper_loop(
                 now=utc_now(),
                 max_attempts=max_attempts,
             )
+            services.repository.cleanup_expired_token_reservations(
+                now=utc_now(),
+                limit=100,
+            )
         except Exception:
             logger.exception(
                 "cloud lifecycle iteration failed",

apps/device-host-agent/host_agent/app.py

@@ -19,6 +19,7 @@ from host_agent.history import ConsoleHistoryStore
 from host_agent.identity import HostIdentityStore
 from host_agent.lease import ActiveAssignmentRunner
 from host_agent.local_account import LocalAccountStore
+from host_agent.policy_cache import HostPolicyCacheStore
 from host_agent.processor import AssignmentProcessingResult, AssignmentProcessor
 from host_agent.status import AgentStatusTracker
 from host_agent.web.app import create_console_app
@@ -195,6 +196,14 @@ def create_application(
             if history_store is not None
             else None
         ),
+        policy_cache=HostPolicyCacheStore(
+            resolved_config.identity_path.parent / "host_governance_policy.json"
+        ),
+        on_policy_sync=(
+            (lambda revision: history_store.record_policy_sync(revision=revision))
+            if history_store is not None
+            else None
+        ),
     )
     executor = AssignmentExecutor(
         create_execution_factories(

apps/device-host-agent/host_agent/heartbeat.py

@@ -8,6 +8,7 @@ from core.errors import DeviceRuntimeError
 from device.manager import DeviceManager
 from host_agent.client import HostAgentClient
 from host_agent.config import HostAgentConfig
+from host_agent.policy_cache import HostPolicyCacheStore
 from host_agent.status import AgentStatusTracker
 
 if TYPE_CHECKING:
@@ -37,6 +38,8 @@ class HeartbeatSynchronizer:
         sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
         status_tracker: AgentStatusTracker | None = None,
         on_sync: Callable[[int], None] | None = None,
+        policy_cache: HostPolicyCacheStore | None = None,
+        on_policy_sync: Callable[[int], None] | None = None,
     ) -> None:
         self.manager = manager
         self.client = client
@@ -45,8 +48,12 @@ class HeartbeatSynchronizer:
         self._sleep = sleep
         self.status_tracker = status_tracker
         self.on_sync = on_sync
-        self.policy_revision = 0
-        self.policy = None
+        self.policy_cache = policy_cache
+        self.on_policy_sync = on_policy_sync
+        self.policy = policy_cache.load() if policy_cache is not None else None
+        self.policy_revision = self.policy.revision if self.policy is not None else 0
+        if self.status_tracker is not None:
+            self.status_tracker.mark_host_policy(self.policy)
 
     async def sync_once(self) -> HeartbeatResponse:
         snapshot = build_device_snapshot(self.manager)
@@ -58,6 +65,16 @@ class HeartbeatSynchronizer:
         self.policy_revision = response.policy_revision
         if response.policy is not None:
             self.policy = response.policy
+            if self.policy_cache is not None:
+                self.policy_cache.save(response.policy)
+            if self.on_policy_sync is not None:
+                self.on_policy_sync(response.policy.revision)
+        elif response.policy_revision == 0:
+            self.policy = None
+            if self.policy_cache is not None:
+                self.policy_cache.clear()
+        if self.status_tracker is not None:
+            self.status_tracker.mark_host_policy(self.policy)
         if self.status_tracker is not None:
             self.status_tracker.mark_heartbeat(ok=True, device_count=len(snapshot))
         if self.on_sync is not None:

apps/device-host-agent/host_agent/history.py

@@ -48,6 +48,13 @@ class ConsoleHistoryStore:
         detail = {"device_count": device_count}
         self._insert("heartbeat", summary, detail)
 
+    def record_policy_sync(self, *, revision: int) -> None:
+        self._insert(
+            "host_policy",
+            f"host policy synchronized: revision {revision}",
+            {"revision": revision},
+        )
+
     def list_recent(self, limit: int | None = None) -> list[dict[str, Any]]:
         effective_limit = limit if limit is not None else self.limit
         with self._connect() as connection:

apps/device-host-agent/host_agent/policy_cache.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+from uuid import uuid4
+
+from cloud.internal_api.models import HostGovernancePolicyModel
+
+
+class HostPolicyCacheError(RuntimeError):
+    """Raised when the locally cached non-secret Host policy is invalid."""
+
+
+class HostPolicyCacheStore:
+    """Atomically persists only the Cloud-supplied, non-secret policy cache."""
+
+    def __init__(self, path: str | Path) -> None:
+        self.path = Path(path)
+
+    def load(self) -> HostGovernancePolicyModel | None:
+        if not self.path.exists():
+            return None
+        try:
+            payload = json.loads(self.path.read_text(encoding="utf-8"))
+            return HostGovernancePolicyModel.model_validate(payload)
+        except (OSError, ValueError, json.JSONDecodeError) as exc:
+            raise HostPolicyCacheError("Host policy cache is invalid") from exc
+
+    def save(self, policy: HostGovernancePolicyModel) -> None:
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        temporary = self.path.with_name(f".{self.path.name}.{uuid4().hex}.tmp")
+        try:
+            temporary.write_text(
+                policy.model_dump_json(indent=2) + "\n",
+                encoding="utf-8",
+            )
+            os.replace(temporary, self.path)
+        finally:
+            if temporary.exists():
+                temporary.unlink()
+
+    def clear(self) -> None:
+        try:
+            self.path.unlink()
+        except FileNotFoundError:
+            return

apps/device-host-agent/host_agent/status.py

@@ -33,6 +33,7 @@ class AgentStatusTracker:
         self._lock = threading.Lock()
         self._current_assignment: _CurrentAssignment | None = None
         self._last_heartbeat: _LastHeartbeat | None = None
+        self._host_policy: dict[str, Any] | None = None
 
     def mark_assignment_started(self, assignment: AssignmentModel) -> None:
         with self._lock:
@@ -56,6 +57,19 @@ class AgentStatusTracker:
                 at=self._now(),
             )
 
+    def mark_host_policy(self, policy: Any | None) -> None:
+        with self._lock:
+            self._host_policy = (
+                {
+                    "revision": policy.revision,
+                    "self_submission_enabled": policy.self_submission_enabled,
+                    "max_active_tasks": policy.max_active_tasks,
+                    "daily_token_budget": policy.daily_token_budget,
+                }
+                if policy is not None
+                else None
+            )
+
     def snapshot(self) -> dict[str, Any]:
         with self._lock:
             current_assignment = self._current_assignment
@@ -81,4 +95,5 @@ class AgentStatusTracker:
                 if last_heartbeat is not None
                 else None
             ),
+            "host_policy": self._host_policy.copy() if self._host_policy else None,
         }

apps/device-host-agent/host_agent/web/app.py

@@ -117,6 +117,7 @@ def _dashboard_body(
 ) -> str:
     heartbeat = snapshot.get("last_heartbeat")
     assignment = snapshot.get("current_assignment")
+    policy = snapshot.get("host_policy")
     heartbeat_text = (
         f"{'ok' if heartbeat['ok'] else 'failed'} at {heartbeat['at']} "
         f"({heartbeat['device_count']} devices)"
@@ -129,6 +130,17 @@ def _dashboard_body(
         if assignment
         else "none"
     )
+    policy_text = (
+        "revision {revision}; self-submission {self_submission}; "
+        "max active tasks {max_active}; daily token budget {daily_budget}".format(
+            revision=policy["revision"],
+            self_submission="enabled" if policy["self_submission_enabled"] else "disabled",
+            max_active=policy["max_active_tasks"] or "unlimited",
+            daily_budget=policy["daily_token_budget"] or "unmetered",
+        )
+        if policy
+        else "no Cloud policy cached"
+    )
     device_rows = "".join(
         f"<tr><td>{escape(device.id)}</td><td>{escape(device.name or '')}</td>"
         f"<td>{escape(device.driver_type)}</td><td>{escape(device.status)}</td></tr>"
@@ -146,6 +158,10 @@ def _dashboard_body(
       <h2>Heartbeat</h2>
       <p id="last-heartbeat">{escape(heartbeat_text)}</p>
     </section>
+    <section>
+      <h2>Cloud policy</h2>
+      <p id="host-policy">{escape(policy_text)}</p>
+    </section>
     <section>
       <h2>Current assignment</h2>
       <p id="current-assignment">{escape(assignment_text)}</p>
@@ -168,6 +184,13 @@ def _dashboard_body(
         document.getElementById("current-assignment").textContent = current
           ? current.task_id + " on " + current.device_id + " (started " + current.started_at + ")"
           : "none";
+        var policy = data.status.host_policy;
+        document.getElementById("host-policy").textContent = policy
+          ? "revision " + policy.revision + "; self-submission "
+            + (policy.self_submission_enabled ? "enabled" : "disabled")
+            + "; max active tasks " + (policy.max_active_tasks || "unlimited")
+            + "; daily token budget " + (policy.daily_token_budget || "unmetered")
+          : "no Cloud policy cached";
         var body = document.getElementById("device-status-body");
         body.innerHTML = "";
         data.devices.forEach(function (device) {{

apps/device-host-agent/tests/test_heartbeat.py

@@ -4,9 +4,11 @@ import asyncio
 from datetime import UTC, datetime
 
 from cloud.internal_api.models import HeartbeatResponse
+from cloud.internal_api.models import HostGovernancePolicyModel
 from device.manager import DeviceManager
 from host_agent.config import HostAgentConfig
 from host_agent.heartbeat import HeartbeatSynchronizer, build_device_snapshot
+from host_agent.policy_cache import HostPolicyCacheStore
 from host_agent.status import AgentStatusTracker
 
 
@@ -123,3 +125,56 @@ def test_sync_once_notifies_status_tracker_and_on_sync_with_device_count() -> No
         assert last_heartbeat["device_count"] == 2
 
     asyncio.run(scenario())
+
+
+def test_heartbeat_caches_safe_host_policy_and_reuses_its_revision(tmp_path) -> None:
+    manager = DeviceManager()
+    cache = HostPolicyCacheStore(tmp_path / "host_policy.json")
+    revisions: list[int] = []
+
+    class UpdatingClient:
+        async def heartbeat(self, devices, *, address=None, policy_revision=0):
+            revisions.append(policy_revision)
+            return HeartbeatResponse(
+                host_id="host-a",
+                accepted_devices=len(devices),
+                received_at=datetime.now(UTC),
+                policy_revision=4,
+                policy=HostGovernancePolicyModel(
+                    revision=4,
+                    self_submission_enabled=False,
+                    max_active_tasks=2,
+                    daily_token_budget=900,
+                ),
+            )
+
+    async def scenario() -> None:
+        tracker = AgentStatusTracker()
+        synchronizer = HeartbeatSynchronizer(
+            manager,
+            UpdatingClient(),  # type: ignore[arg-type]
+            _config(),
+            policy_cache=cache,
+            status_tracker=tracker,
+        )
+        await synchronizer.sync_once()
+        assert tracker.snapshot()["host_policy"] == {
+            "revision": 4,
+            "self_submission_enabled": False,
+            "max_active_tasks": 2,
+            "daily_token_budget": 900,
+        }
+
+        restarted = HeartbeatSynchronizer(
+            manager,
+            UpdatingClient(),  # type: ignore[arg-type]
+            _config(),
+            policy_cache=cache,
+        )
+        assert restarted.policy_revision == 4
+
+    asyncio.run(scenario())
+    assert revisions == [0]
+    assert '"token":' not in (
+        tmp_path / "host_policy.json"
+    ).read_text(encoding="utf-8")