feat(cloud): enforce host governance budgets

tests/test_cloud_governance.py

@@ -132,6 +132,24 @@ def test_governance_routes_persist_revisioned_policies(tmp_path) -> None:
         assert host_policy.json()["revision"] == 1
         reread = client.get("/v1/hosts/host-a/governance-policy")
         assert reread.json()["daily_token_budget"] == 1000
+        conflict = client.put(
+            "/v1/hosts/host-a/governance-policy",
+            json={
+                "max_active_tasks": 3,
+                "daily_token_budget": 1000,
+                "expected_revision": 1,
+            },
+        )
+        assert conflict.status_code == 200
+        stale = client.put(
+            "/v1/hosts/host-a/governance-policy",
+            json={"max_active_tasks": 4, "expected_revision": 1},
+        )
+        assert stale.status_code == 409
+        usage = client.get("/v1/hosts/host-a/token-usage")
+        assert usage.status_code == 200
+        assert usage.json()["used_tokens"] == 0
+        assert usage.json()["daily_token_budget"] == 1000
     finally:
         database.close()
 

tests/test_cloud_migrations.py

@@ -38,6 +38,8 @@ def test_forward_and_downgrade_migrations_on_empty_database(tmp_path) -> None:
             "cloud_user_sessions",
             "cloud_login_throttles",
             "cloud_auth_audit_events",
+            "cloud_token_reservations",
+            "cloud_token_usage_events",
         } <= table_names
         assert current_revision(database_url) == HEAD_REVISION
         host_columns = {

tests/test_host_agent_internal_api.py

@@ -12,9 +12,15 @@ from cloud.internal_api.api import create_internal_router
 from cloud.pool import DevicePool, PooledDevice
 from cloud.scheduler import ScheduledTask, TaskConstraints, TaskScheduler
 from cloud.store import CloudStore
+from runtime.tool_calling_client import ToolCallDecision, ToolCallUsage
 
 
-def _build_client(tmp_path) -> tuple[TestClient, DevicePool]:
+def _build_client(
+    tmp_path,
+    *,
+    planner_client_factory=None,
+    planner_token_reservation_ceiling: int = 4096,
+) -> tuple[TestClient, DevicePool]:
     pool = DevicePool(
         CloudStore(tmp_path / "internal.sqlite3"),
         CloudConfig(stale_after_seconds=60),
@@ -39,6 +45,8 @@ def _build_client(tmp_path) -> tuple[TestClient, DevicePool]:
             pool=pool,
             auth_provider=auth_provider,
             scheduler=TaskScheduler(pool, pool.store, CloudConfig(stale_after_seconds=60)),
+            planner_client_factory=planner_client_factory,
+            planner_token_reservation_ceiling=planner_token_reservation_ceiling,
         )
     )
     return TestClient(app), pool
@@ -232,6 +240,58 @@ def test_heartbeat_and_self_submission_preserve_host_isolation(tmp_path) -> None
     assert foreign.status_code == 403
 
 
+def test_planner_proxy_reserves_and_enforces_host_daily_token_budget(tmp_path) -> None:
+    class FakePlannerClient:
+        calls = 0
+
+        def decide(self, **_kwargs):
+            self.calls += 1
+            return ToolCallDecision(
+                tool_name="tap",
+                arguments={"x": 1, "y": 2},
+                usage=ToolCallUsage(input_tokens=1, output_tokens=1, total_tokens=2),
+            )
+
+    planner = FakePlannerClient()
+    client, pool = _build_client(
+        tmp_path,
+        planner_client_factory=lambda: planner,
+        planner_token_reservation_ceiling=5,
+    )
+    headers = {"Authorization": "Bearer token-a"}
+    client.put(
+        "/internal/v1/hosts/host-a/heartbeat",
+        headers=headers,
+        json=_heartbeat_payload("host-a", "device-a"),
+    )
+    pool.store.upsert_host_governance_policy(
+        host_id="host-a",
+        self_submission_enabled=True,
+        max_active_tasks=None,
+        daily_token_budget=5,
+        updated_at=datetime.now(UTC),
+    )
+    payload = {
+        "host_id": "host-a",
+        "system_prompt": "system",
+        "user_prompt": "user",
+        "tools": [{"name": "tap", "description": "tap", "parameters": {}}],
+        "timeout_seconds": 1,
+    }
+
+    first = client.post(
+        "/internal/v1/hosts/host-a/planner/decide", headers=headers, json=payload
+    )
+    second = client.post(
+        "/internal/v1/hosts/host-a/planner/decide", headers=headers, json=payload
+    )
+
+    assert first.status_code == 200, first.text
+    assert first.json()["total_tokens"] == 2
+    assert second.status_code == 429
+    assert planner.calls == 1
+
+
 def test_long_poll_claim_returns_at_most_one_owned_assignment(tmp_path) -> None:
     client, pool = _build_client(tmp_path)
     now = datetime.now(UTC)

tests/test_task_scheduler.py

@@ -264,3 +264,32 @@ def test_unavailable_explicit_target_is_never_rerouted(tmp_path) -> None:
     task = pool.store.get_task(task_id)
     assert task is not None
     assert task.status == "queued"
+
+
+def test_host_active_task_policy_keeps_excess_tasks_queued(tmp_path) -> None:
+    from datetime import UTC, datetime
+
+    pool = _pool_with_devices(
+        tmp_path,
+        _device("device-a"),
+        _device("device-b"),
+        host_id="host-a",
+    )
+    pool.store.upsert_host_governance_policy(
+        host_id="host-a",
+        self_submission_enabled=True,
+        max_active_tasks=1,
+        daily_token_budget=None,
+        updated_at=datetime.now(UTC),
+    )
+    scheduler = TaskScheduler(pool, pool.store, _config())
+    first_id = scheduler.submit(goal="first")
+    second_id = scheduler.submit(goal="second")
+
+    assignments = scheduler.assign()
+
+    assert [assignment.task_id for assignment in assignments] == [first_id]
+    assert pool.store.count_active_tasks_for_host("host-a") == 1
+    second = pool.store.get_task(second_id)
+    assert second is not None
+    assert second.status == "queued"