diff --git a/apps/cloud-api/cloud_api/app.py b/apps/cloud-api/cloud_api/app.py index 5eca5f2..51bd6fc 100644 --- a/apps/cloud-api/cloud_api/app.py +++ b/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", diff --git a/apps/device-host-agent/host_agent/app.py b/apps/device-host-agent/host_agent/app.py index 2a724c2..a109d49 100644 --- a/apps/device-host-agent/host_agent/app.py +++ b/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( diff --git a/apps/device-host-agent/host_agent/heartbeat.py b/apps/device-host-agent/host_agent/heartbeat.py index 25638ff..ecdf2eb 100644 --- a/apps/device-host-agent/host_agent/heartbeat.py +++ b/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: diff --git a/apps/device-host-agent/host_agent/history.py b/apps/device-host-agent/host_agent/history.py index 8ea35df..61c18d5 100644 --- a/apps/device-host-agent/host_agent/history.py +++ b/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: diff --git a/apps/device-host-agent/host_agent/policy_cache.py b/apps/device-host-agent/host_agent/policy_cache.py new file mode 100644 index 0000000..e04ebb6 --- /dev/null +++ b/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 diff --git a/apps/device-host-agent/host_agent/status.py b/apps/device-host-agent/host_agent/status.py index b0815ee..87bcc7b 100644 --- a/apps/device-host-agent/host_agent/status.py +++ b/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, } diff --git a/apps/device-host-agent/host_agent/web/app.py b/apps/device-host-agent/host_agent/web/app.py index 44a6fb9..38aa360 100644 --- a/apps/device-host-agent/host_agent/web/app.py +++ b/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"{escape(device.id)}{escape(device.name or '')}" f"{escape(device.driver_type)}{escape(device.status)}" @@ -146,6 +158,10 @@ def _dashboard_body(

Heartbeat

{escape(heartbeat_text)}

+
+

Cloud policy

+

{escape(policy_text)}

+

Current assignment

{escape(assignment_text)}

@@ -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) {{ diff --git a/apps/device-host-agent/tests/test_heartbeat.py b/apps/device-host-agent/tests/test_heartbeat.py index 2a478b5..72e75e3 100644 --- a/apps/device-host-agent/tests/test_heartbeat.py +++ b/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") diff --git a/cloud-console/src/App.vue b/cloud-console/src/App.vue index 7846129..d9ddc0d 100644 --- a/cloud-console/src/App.vue +++ b/cloud-console/src/App.vue @@ -7,6 +7,7 @@ import { LogOut, MonitorSmartphone, Puzzle, + UsersRound, } from "@lucide/vue"; import { AUTH_INVALID_EVENT, @@ -19,8 +20,9 @@ import PasswordChangeScreen from "./views/PasswordChangeScreen.vue"; import TasksView from "./views/TasksView.vue"; import DevicesView from "./views/DevicesView.vue"; import PluginsView from "./views/PluginsView.vue"; +import UsersView from "./views/UsersView.vue"; -type ViewId = "tasks" | "devices" | "plugins"; +type ViewId = "tasks" | "devices" | "plugins" | "users"; const activeView = ref("tasks"); const currentUser = ref(null); @@ -37,6 +39,16 @@ const canSubmitTasks = computed( currentUser.value?.scopes.includes("*") || currentUser.value?.scopes.includes("tasks:submit"), ); +const canAdminUsers = computed( + () => Boolean( + currentUser.value?.scopes.includes("*") || + currentUser.value?.scopes.includes("users:admin")), +); +const canAdminGovernance = computed( + () => Boolean( + currentUser.value?.scopes.includes("*") || + currentUser.value?.scopes.includes("governance:admin")), +); const isAuthenticated = computed(() => currentUser.value !== null); const currentUserLabel = computed(() => currentUser.value ? `${currentUser.value.display_name} (${currentUser.value.role})` : "", @@ -48,6 +60,9 @@ const navItems = computed<{ id: ViewId; label: string; icon: Component }[]>(() = { id: "devices", label: "Devices", icon: MonitorSmartphone }, { id: "plugins", label: "Plugins", icon: Puzzle }, ]; + if (canAdminUsers.value || canAdminGovernance.value) { + items.push({ id: "users", label: "Users & limits", icon: UsersRound }); + } return items; }); @@ -102,6 +117,8 @@ const activeComponent = computed(() => { return DevicesView; case "plugins": return PluginsView; + case "users": + return UsersView; default: return TasksView; } @@ -133,6 +150,11 @@ const activeComponent = computed(() => {
+
diff --git a/cloud-console/src/api.ts b/cloud-console/src/api.ts index 12cf667..293d136 100644 --- a/cloud-console/src/api.ts +++ b/cloud-console/src/api.ts @@ -1,6 +1,8 @@ import type { CloudUser, DeviceRecord, + HostGovernancePolicy, + HostTokenUsageSummary, HostRecord, PluginRecord, PluginRegistrationPayload, @@ -8,6 +10,8 @@ import type { TaskListResponse, TaskSubmissionPayload, TaskStatus, + UserListResponse, + UserSubmissionPolicy, } from "./types"; const configuredBaseUrl = import.meta.env.VITE_CLOUD_API_BASE_URL as @@ -119,6 +123,87 @@ export function changePassword( }); } +export function listUsers(): Promise { + return request("/v1/users?limit=100&offset=0"); +} + +export function createUser(payload: { + username: string; + display_name: string; + role: "viewer" | "operator" | "admin"; + password: string; +}): Promise { + return request("/v1/users", { + method: "POST", + body: JSON.stringify(payload), + }); +} + +export function updateUser( + userId: string, + payload: { display_name?: string; role?: "viewer" | "operator" | "admin"; enabled?: boolean }, +): Promise { + return request(`/v1/users/${encodeURIComponent(userId)}`, { + method: "PATCH", + body: JSON.stringify(payload), + }); +} + +export function resetUserPassword(userId: string, password: string): Promise { + return request(`/v1/users/${encodeURIComponent(userId)}/password`, { + method: "POST", + body: JSON.stringify({ password }), + }); +} + +export function revokeUserSessions(userId: string): Promise { + return request(`/v1/users/${encodeURIComponent(userId)}/sessions`, { + method: "DELETE", + }); +} + +export function getUserSubmissionPolicy(userId: string): Promise { + return request( + `/v1/users/${encodeURIComponent(userId)}/submission-policy`, + ); +} + +export function updateUserSubmissionPolicy( + userId: string, + payload: Omit & { + expected_revision?: number; + }, +): Promise { + return request( + `/v1/users/${encodeURIComponent(userId)}/submission-policy`, + { method: "PUT", body: JSON.stringify(payload) }, + ); +} + +export function getHostGovernancePolicy(hostId: string): Promise { + return request( + `/v1/hosts/${encodeURIComponent(hostId)}/governance-policy`, + ); +} + +export function updateHostGovernancePolicy( + hostId: string, + payload: Omit & { + expected_revision?: number; + }, +): Promise { + return request( + `/v1/hosts/${encodeURIComponent(hostId)}/governance-policy`, + { method: "PUT", body: JSON.stringify(payload) }, + ); +} + +export function getHostTokenUsage(hostId: string): Promise { + return request( + `/v1/hosts/${encodeURIComponent(hostId)}/token-usage`, + ); +} + export function listTasks(options?: { status?: TaskStatus; limit?: number; diff --git a/cloud-console/src/style.css b/cloud-console/src/style.css index 7f6f69f..b9a15fb 100644 --- a/cloud-console/src/style.css +++ b/cloud-console/src/style.css @@ -319,6 +319,19 @@ tr.row-selected { grid-column: 1 / -1; } +.form-grid input:not([type="checkbox"]), +.form-grid select { + width: 100%; +} + +.form-grid .actions { + align-self: end; +} + +.policy-grid select { + min-height: 92px; +} + .muted { color: var(--text-muted); } diff --git a/cloud-console/src/types.ts b/cloud-console/src/types.ts index 2f9d684..7648f32 100644 --- a/cloud-console/src/types.ts +++ b/cloud-console/src/types.ts @@ -96,3 +96,41 @@ export interface CloudUser { updated_at: string; last_login_at: string | null; } + +export interface UserListResponse { + items: CloudUser[]; + limit: number; + offset: number; +} + +export interface DeviceTarget { + host_id: string; + device_id: string; +} + +export interface UserSubmissionPolicy { + user_id: string; + revision: number; + submission_enabled: boolean; + allowed_host_ids: string[] | null; + allowed_device_targets: DeviceTarget[] | null; + updated_at: string; +} + +export interface HostGovernancePolicy { + host_id: string; + revision: number; + self_submission_enabled: boolean; + max_active_tasks: number | null; + daily_token_budget: number | null; + updated_at: string; +} + +export interface HostTokenUsageSummary { + host_id: string; + usage_day: string; + daily_token_budget: number | null; + used_tokens: number; + reserved_tokens: number; + remaining_tokens: number | null; +} diff --git a/cloud-console/src/views/UsersView.vue b/cloud-console/src/views/UsersView.vue new file mode 100644 index 0000000..c8e4618 --- /dev/null +++ b/cloud-console/src/views/UsersView.vue @@ -0,0 +1,380 @@ + + + diff --git a/openspec/changes/cloud-console-governance/tasks.md b/openspec/changes/cloud-console-governance/tasks.md index 1ac1eab..fb3f47f 100644 --- a/openspec/changes/cloud-console-governance/tasks.md +++ b/openspec/changes/cloud-console-governance/tasks.md @@ -51,13 +51,13 @@ - [x] 4.1 Extend shared internal heartbeat request/response models with the Host's last policy revision and a revision-aware effective-policy reply. -- [ ] 4.2 Update Host heartbeat synchronization to persist only safe cached +- [x] 4.2 Update Host heartbeat synchronization to persist only safe cached policy state and expose it through local status/history without adding an inbound Cloud connection. - [x] 4.3 Add the Host-scoped goal-only task-submission route and `HostAgentClient` method; derive Host targeting from authenticated credentials and validate any named local Device ownership. -- [ ] 4.4 Enforce the Host policy's self-submission and active-task limits in +- [x] 4.4 Enforce the Host policy's self-submission and active-task limits in the Cloud service/scheduler, not only in Host-local code. - [ ] 4.5 Add Cloud API and Host Agent tests for revision convergence, unchanged-policy replies, self-targeted task creation, foreign target @@ -71,7 +71,7 @@ - [ ] 5.2 Extend planner-proxy request context and Host-Agent-local context binding so Cloud-proxied calls carry known task/attempt metadata without importing Host or Cloud concerns into `runtime`. -- [ ] 5.3 Add Cloud proxy preflight reservation, configured conservative +- [x] 5.3 Add Cloud proxy preflight reservation, configured conservative per-call ceiling, provider invocation, actual-usage settlement, and bounded unknown-usage reservation expiry. - [ ] 5.4 Record non-secret usage events and expose accurate @@ -87,7 +87,7 @@ - [x] 6.1 Add task composer API bindings and a scope-aware Console form for goal/workflow submission, Host/Device target selection, and target/policy validation errors. -- [ ] 6.2 Complete or reconcile the admin Users view, then add user- +- [x] 6.2 Complete or reconcile the admin Users view, then add user- submission-policy editing with safe refresh and conflict/error handling. - [ ] 6.3 Add Host policy administration and AI-usage/budget views, including revision display, unmetered direct Hosts, and no rendering of prompts, diff --git a/packages/cloud-platform/cloud/control_config.py b/packages/cloud-platform/cloud/control_config.py index 9e8f4dc..7414180 100644 --- a/packages/cloud-platform/cloud/control_config.py +++ b/packages/cloud-platform/cloud/control_config.py @@ -36,6 +36,8 @@ class CloudControlConfig: login_block_seconds: int = 900 session_cookie_secure: bool = False trust_proxy_headers: bool = False + planner_token_reservation_ceiling: int = 4096 + planner_token_reservation_ttl_seconds: int = 300 def load_control_config( @@ -115,6 +117,16 @@ def load_control_config( values.get("CLOUD_TRUST_PROXY_HEADERS"), default=False, ), + planner_token_reservation_ceiling=_positive_int( + values, + "CLOUD_PLANNER_TOKEN_RESERVATION_CEILING", + 4096, + ), + planner_token_reservation_ttl_seconds=_positive_int( + values, + "CLOUD_PLANNER_TOKEN_RESERVATION_TTL_SECONDS", + 300, + ), ) validate_control_config(config) return config diff --git a/packages/cloud-platform/cloud/db_models.py b/packages/cloud-platform/cloud/db_models.py index b5da46b..ea9703c 100644 --- a/packages/cloud-platform/cloud/db_models.py +++ b/packages/cloud-platform/cloud/db_models.py @@ -241,3 +241,44 @@ class HostGovernancePolicyRow(Base): max_active_tasks: Mapped[int | None] = mapped_column(Integer, nullable=True) daily_token_budget: Mapped[int | None] = mapped_column(Integer, nullable=True) updated_at: Mapped[str] = mapped_column(String, nullable=False) + + +class TokenReservationRow(Base): + __tablename__ = "cloud_token_reservations" + __table_args__ = ( + Index("ix_cloud_token_reservations_host_day", "host_id", "usage_day"), + Index("ix_cloud_token_reservations_expires_at", "expires_at"), + ) + + id: Mapped[str] = mapped_column(String, primary_key=True) + host_id: Mapped[str] = mapped_column( + ForeignKey("host_registrations.host_id", ondelete="CASCADE"), nullable=False + ) + usage_day: Mapped[str] = mapped_column(String, nullable=False) + reserved_tokens: Mapped[int] = mapped_column(Integer, nullable=False) + task_id: Mapped[str | None] = mapped_column(String, nullable=True) + attempt: Mapped[int | None] = mapped_column(Integer, nullable=True) + created_at: Mapped[str] = mapped_column(String, nullable=False) + expires_at: Mapped[str] = mapped_column(String, nullable=False) + + +class TokenUsageEventRow(Base): + __tablename__ = "cloud_token_usage_events" + __table_args__ = ( + Index("ix_cloud_token_usage_events_host_day", "host_id", "usage_day"), + Index("ix_cloud_token_usage_events_occurred_at", "occurred_at"), + ) + + id: Mapped[str] = mapped_column(String, primary_key=True) + host_id: Mapped[str] = mapped_column( + ForeignKey("host_registrations.host_id", ondelete="CASCADE"), nullable=False + ) + usage_day: Mapped[str] = mapped_column(String, nullable=False) + task_id: Mapped[str | None] = mapped_column(String, nullable=True) + attempt: Mapped[int | None] = mapped_column(Integer, nullable=True) + provider: Mapped[str] = mapped_column(String, nullable=False) + model: Mapped[str] = mapped_column(String, nullable=False) + input_tokens: Mapped[int | None] = mapped_column(Integer, nullable=True) + output_tokens: Mapped[int | None] = mapped_column(Integer, nullable=True) + total_tokens: Mapped[int] = mapped_column(Integer, nullable=False) + occurred_at: Mapped[str] = mapped_column(String, nullable=False) diff --git a/packages/cloud-platform/cloud/governance.py b/packages/cloud-platform/cloud/governance.py index b0a2682..e52e2f5 100644 --- a/packages/cloud-platform/cloud/governance.py +++ b/packages/cloud-platform/cloud/governance.py @@ -10,6 +10,14 @@ class TaskSubmissionPolicyError(PermissionError): """Raised when a human user's policy disallows a task submission.""" +class GovernancePolicyConflictError(RuntimeError): + """Raised when a policy write was based on an obsolete revision.""" + + +class TokenBudgetExceededError(RuntimeError): + """Raised before a Cloud-proxied call would exceed a Host's token budget.""" + + @dataclass(frozen=True) class UserSubmissionPolicy: user_id: str @@ -30,6 +38,48 @@ class HostGovernancePolicy: updated_at: datetime +@dataclass(frozen=True) +class TokenReservation: + id: str + host_id: str + usage_day: str + reserved_tokens: int + task_id: str | None + attempt: int | None + created_at: datetime + expires_at: datetime + + +@dataclass(frozen=True) +class TokenUsageEvent: + id: str + host_id: str + usage_day: str + task_id: str | None + attempt: int | None + provider: str + model: str + input_tokens: int | None + output_tokens: int | None + total_tokens: int + occurred_at: datetime + + +@dataclass(frozen=True) +class TokenUsageSummary: + host_id: str + usage_day: str + daily_token_budget: int | None + used_tokens: int + reserved_tokens: int + + @property + def remaining_tokens(self) -> int | None: + if self.daily_token_budget is None: + return None + return max(0, self.daily_token_budget - self.used_tokens - self.reserved_tokens) + + def enforce_user_submission_policy( policy: UserSubmissionPolicy | None, *, diff --git a/packages/cloud-platform/cloud/internal_api/api.py b/packages/cloud-platform/cloud/internal_api/api.py index e7fd033..ef0cfbe 100644 --- a/packages/cloud-platform/cloud/internal_api/api.py +++ b/packages/cloud-platform/cloud/internal_api/api.py @@ -44,6 +44,7 @@ from cloud.repository import ( DeviceEnrollmentConflictError, HostEnrollmentConflictError, ) +from cloud.governance import TokenBudgetExceededError from core.models import Device, utc_now from runtime.tool_calling_client import ToolCallingClient, ToolCallUnavailable from runtime.tool_specs import ToolSpec @@ -65,11 +66,17 @@ def create_internal_router( sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, planner_client_factory: Callable[[], ToolCallingClient] | None = None, scheduler: TaskScheduler | None = None, + planner_token_reservation_ceiling: int = 4096, + planner_token_reservation_ttl_seconds: float = 300.0, ) -> APIRouter: if claim_poll_interval_seconds <= 0: raise ValueError("claim_poll_interval_seconds must be greater than zero") if lease_duration_seconds <= 0: raise ValueError("lease_duration_seconds must be greater than zero") + if planner_token_reservation_ceiling <= 0: + raise ValueError("planner_token_reservation_ceiling must be greater than zero") + if planner_token_reservation_ttl_seconds <= 0: + raise ValueError("planner_token_reservation_ttl_seconds must be greater than zero") router = APIRouter(prefix=version_prefix, tags=["host-agent"]) build_planner_client = planner_client_factory or _default_planner_client_factory @@ -352,6 +359,7 @@ def create_internal_router( response_model_exclude_none=True, responses={ status.HTTP_502_BAD_GATEWAY: {"model": PlannerDecisionError}, + status.HTTP_429_TOO_MANY_REQUESTS: {"model": PlannerDecisionError}, }, ) def decide_planner_call( @@ -385,6 +393,25 @@ def create_internal_router( for tool in payload.tools ] + now = utc_now() + reservation = None + try: + reservation = pool.store.reserve_host_token_budget( + reservation_id=uuid4().hex, + host_id=host_id, + usage_day=now.date().isoformat(), + reserved_tokens=planner_token_reservation_ceiling, + task_id=None, + attempt=None, + created_at=now, + expires_at=now + timedelta(seconds=planner_token_reservation_ttl_seconds), + ) + except TokenBudgetExceededError as exc: + return JSONResponse( + status_code=status.HTTP_429_TOO_MANY_REQUESTS, + content=PlannerDecisionError(detail=str(exc)).model_dump(), + ) + started_at = monotonic() client = build_planner_client() try: @@ -408,6 +435,19 @@ def create_internal_router( status_code=status.HTTP_502_BAD_GATEWAY, content=PlannerDecisionError(detail=str(exc)).model_dump(), ) + usage = decision.usage + if reservation is not None and usage is not None and usage.total_tokens is not None: + planner_config = load_cloud_planner_config() + pool.store.settle_host_token_reservation( + reservation_id=reservation.id, + event_id=uuid4().hex, + provider=planner_config.provider, + model=planner_config.resolved_model(), + input_tokens=usage.input_tokens, + output_tokens=usage.output_tokens, + total_tokens=usage.total_tokens, + occurred_at=utc_now(), + ) logger.info( "planner-decision request resolved", extra={ diff --git a/packages/cloud-platform/cloud/migrations/versions/0005_cloud_token_usage.py b/packages/cloud-platform/cloud/migrations/versions/0005_cloud_token_usage.py new file mode 100644 index 0000000..1659e74 --- /dev/null +++ b/packages/cloud-platform/cloud/migrations/versions/0005_cloud_token_usage.py @@ -0,0 +1,53 @@ +"""Add Cloud-proxy token reservations and usage events.""" + +from __future__ import annotations + +from alembic import op +import sqlalchemy as sa + + +revision = "0005_cloud_token_usage" +down_revision = "0004_cloud_governance" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.create_table( + "cloud_token_reservations", + sa.Column("id", sa.String(), primary_key=True), + sa.Column("host_id", sa.String(), sa.ForeignKey("host_registrations.host_id", ondelete="CASCADE"), nullable=False), + sa.Column("usage_day", sa.String(), nullable=False), + sa.Column("reserved_tokens", sa.Integer(), nullable=False), + sa.Column("task_id", sa.String(), nullable=True), + sa.Column("attempt", sa.Integer(), nullable=True), + sa.Column("created_at", sa.String(), nullable=False), + sa.Column("expires_at", sa.String(), nullable=False), + ) + op.create_index("ix_cloud_token_reservations_host_day", "cloud_token_reservations", ["host_id", "usage_day"]) + op.create_index("ix_cloud_token_reservations_expires_at", "cloud_token_reservations", ["expires_at"]) + op.create_table( + "cloud_token_usage_events", + sa.Column("id", sa.String(), primary_key=True), + sa.Column("host_id", sa.String(), sa.ForeignKey("host_registrations.host_id", ondelete="CASCADE"), nullable=False), + sa.Column("usage_day", sa.String(), nullable=False), + sa.Column("task_id", sa.String(), nullable=True), + sa.Column("attempt", sa.Integer(), nullable=True), + sa.Column("provider", sa.String(), nullable=False), + sa.Column("model", sa.String(), nullable=False), + sa.Column("input_tokens", sa.Integer(), nullable=True), + sa.Column("output_tokens", sa.Integer(), nullable=True), + sa.Column("total_tokens", sa.Integer(), nullable=False), + sa.Column("occurred_at", sa.String(), nullable=False), + ) + op.create_index("ix_cloud_token_usage_events_host_day", "cloud_token_usage_events", ["host_id", "usage_day"]) + op.create_index("ix_cloud_token_usage_events_occurred_at", "cloud_token_usage_events", ["occurred_at"]) + + +def downgrade() -> None: + op.drop_index("ix_cloud_token_usage_events_occurred_at", table_name="cloud_token_usage_events") + op.drop_index("ix_cloud_token_usage_events_host_day", table_name="cloud_token_usage_events") + op.drop_table("cloud_token_usage_events") + op.drop_index("ix_cloud_token_reservations_expires_at", table_name="cloud_token_reservations") + op.drop_index("ix_cloud_token_reservations_host_day", table_name="cloud_token_reservations") + op.drop_table("cloud_token_reservations") diff --git a/packages/cloud-platform/cloud/repository.py b/packages/cloud-platform/cloud/repository.py index 2ea6ddb..08d3641 100644 --- a/packages/cloud-platform/cloud/repository.py +++ b/packages/cloud-platform/cloud/repository.py @@ -15,7 +15,13 @@ if TYPE_CHECKING: UserAccount, UserSession, ) - from cloud.governance import HostGovernancePolicy, UserSubmissionPolicy + from cloud.governance import ( + HostGovernancePolicy, + TokenReservation, + TokenUsageSummary, + TokenUsageEvent, + UserSubmissionPolicy, + ) AttemptStatus = Literal["assigned", "dispatched", "done", "failed", "expired"] @@ -293,6 +299,7 @@ class CloudRepository(Protocol): submission_enabled: bool, allowed_host_ids: tuple[str, ...] | None, allowed_device_targets: tuple[tuple[str, str], ...] | None, + expected_revision: int | None, updated_at: datetime, ) -> UserSubmissionPolicy: ... @@ -305,9 +312,30 @@ class CloudRepository(Protocol): self_submission_enabled: bool, max_active_tasks: int | None, daily_token_budget: int | None, + expected_revision: int | None, updated_at: datetime, ) -> HostGovernancePolicy: ... + def count_active_tasks_for_host(self, host_id: str) -> int: ... + + def reserve_host_token_budget( + self, *, reservation_id: str, host_id: str, usage_day: str, + reserved_tokens: int, task_id: str | None, attempt: int | None, + created_at: datetime, expires_at: datetime, + ) -> TokenReservation | None: ... + + def settle_host_token_reservation( + self, *, reservation_id: str, event_id: str, provider: str, model: str, + input_tokens: int | None, output_tokens: int | None, total_tokens: int, + occurred_at: datetime, + ) -> TokenUsageEvent | None: ... + + def cleanup_expired_token_reservations(self, *, now: datetime, limit: int) -> int: ... + + def get_host_token_usage_summary( + self, *, host_id: str, usage_day: str, now: datetime, + ) -> TokenUsageSummary: ... + def list_reserved_device_ids(self, *, now: datetime) -> set[str]: ... def assign_task( diff --git a/packages/cloud-platform/cloud/schema.py b/packages/cloud-platform/cloud/schema.py index 60d1fd6..a6f0914 100644 --- a/packages/cloud-platform/cloud/schema.py +++ b/packages/cloud-platform/cloud/schema.py @@ -9,7 +9,7 @@ from alembic.runtime.migration import MigrationContext from cloud.database import create_database_engine, normalize_database_url -HEAD_REVISION = "0004_cloud_governance" +HEAD_REVISION = "0005_cloud_token_usage" class SchemaVersionError(RuntimeError): diff --git a/packages/cloud-platform/cloud/sdk/governance_api.py b/packages/cloud-platform/cloud/sdk/governance_api.py index ee710d6..8b7387b 100644 --- a/packages/cloud-platform/cloud/sdk/governance_api.py +++ b/packages/cloud-platform/cloud/sdk/governance_api.py @@ -5,10 +5,12 @@ from uuid import uuid4 from fastapi import APIRouter, HTTPException, Request, status from cloud.auth import AuthProvider, GOVERNANCE_ADMIN_SCOPE, GOVERNANCE_READ_SCOPE +from cloud.governance import GovernancePolicyConflictError from cloud.observability import current_correlation_id from cloud.sdk.models import ( HostGovernancePolicyRequest, HostGovernancePolicyResponse, + HostTokenUsageSummaryResponse, UserSubmissionPolicyRequest, UserSubmissionPolicyResponse, ) @@ -79,10 +81,13 @@ def create_governance_router(*, repository, auth_provider: AuthProvider) -> APIR if payload.allowed_device_targets is not None else None ), + expected_revision=payload.expected_revision, updated_at=utc_now(), ) except KeyError as exc: raise HTTPException(status_code=404, detail="user not found") from exc + except GovernancePolicyConflictError as exc: + raise HTTPException(status_code=409, detail=str(exc)) from exc _audit(repository, principal.id, user_id, "user_submission_policy_update") return UserSubmissionPolicyResponse( user_id=policy.user_id, @@ -111,6 +116,27 @@ def create_governance_router(*, repository, auth_provider: AuthProvider) -> APIR raise HTTPException(status_code=404, detail="Host policy not found") return _host_response(policy) + @router.get( + "/hosts/{host_id}/token-usage", + response_model=HostTokenUsageSummaryResponse, + ) + def get_host_token_usage( + host_id: str, request: Request + ) -> HostTokenUsageSummaryResponse: + authorize(request, GOVERNANCE_READ_SCOPE) + now = utc_now() + summary = repository.get_host_token_usage_summary( + host_id=host_id, usage_day=now.date().isoformat(), now=now + ) + return HostTokenUsageSummaryResponse( + host_id=summary.host_id, + usage_day=summary.usage_day, + daily_token_budget=summary.daily_token_budget, + used_tokens=summary.used_tokens, + reserved_tokens=summary.reserved_tokens, + remaining_tokens=summary.remaining_tokens, + ) + @router.put( "/hosts/{host_id}/governance-policy", response_model=HostGovernancePolicyResponse, @@ -127,10 +153,13 @@ def create_governance_router(*, repository, auth_provider: AuthProvider) -> APIR self_submission_enabled=payload.self_submission_enabled, max_active_tasks=payload.max_active_tasks, daily_token_budget=payload.daily_token_budget, + expected_revision=payload.expected_revision, updated_at=utc_now(), ) except KeyError as exc: raise HTTPException(status_code=404, detail="Host not found") from exc + except GovernancePolicyConflictError as exc: + raise HTTPException(status_code=409, detail=str(exc)) from exc _audit(repository, principal.id, host_id, "host_governance_policy_update") return _host_response(policy) diff --git a/packages/cloud-platform/cloud/sdk/models.py b/packages/cloud-platform/cloud/sdk/models.py index 5f6a8fe..6ba6bf6 100644 --- a/packages/cloud-platform/cloud/sdk/models.py +++ b/packages/cloud-platform/cloud/sdk/models.py @@ -162,11 +162,15 @@ class UserSubmissionPolicyRequest(BaseModel): submission_enabled: bool = True allowed_host_ids: list[str] | None = None allowed_device_targets: list[DeviceTargetModel] | None = None + expected_revision: int | None = Field(default=None, ge=0) -class UserSubmissionPolicyResponse(UserSubmissionPolicyRequest): +class UserSubmissionPolicyResponse(BaseModel): user_id: str revision: int + submission_enabled: bool + allowed_host_ids: list[str] | None = None + allowed_device_targets: list[DeviceTargetModel] | None = None updated_at: datetime @@ -174,9 +178,22 @@ class HostGovernancePolicyRequest(BaseModel): self_submission_enabled: bool = True max_active_tasks: int | None = Field(default=None, ge=1) daily_token_budget: int | None = Field(default=None, ge=1) + expected_revision: int | None = Field(default=None, ge=0) -class HostGovernancePolicyResponse(HostGovernancePolicyRequest): +class HostGovernancePolicyResponse(BaseModel): host_id: str revision: int + self_submission_enabled: bool + max_active_tasks: int | None = None + daily_token_budget: int | None = None updated_at: datetime + + +class HostTokenUsageSummaryResponse(BaseModel): + host_id: str + usage_day: str + daily_token_budget: int | None = None + used_tokens: int + reserved_tokens: int + remaining_tokens: int | None = None diff --git a/packages/cloud-platform/cloud/sql_repository.py b/packages/cloud-platform/cloud/sql_repository.py index 83d92f3..2c66418 100644 --- a/packages/cloud-platform/cloud/sql_repository.py +++ b/packages/cloud-platform/cloud/sql_repository.py @@ -24,6 +24,8 @@ from cloud.db_models import ( UserRow, UserSessionRow, UserSubmissionPolicyRow, + TokenReservationRow, + TokenUsageEventRow, ) from cloud.observability import current_correlation_id from core.models import utc_now @@ -824,12 +826,21 @@ class SQLAlchemyCloudRepository: allowed_host_ids: tuple[str, ...] | None, allowed_device_targets: tuple[tuple[str, str], ...] | None, updated_at: datetime, + expected_revision: int | None = None, ) -> Any: + from cloud.governance import GovernancePolicyConflictError + with self._sessions.begin() as session: if session.get(UserRow, user_id) is None: raise KeyError(f"unknown user {user_id!r}") - row = session.get(UserSubmissionPolicyRow, user_id) + row = session.get( + UserSubmissionPolicyRow, + user_id, + with_for_update=self.engine.dialect.name == "postgresql", + ) if row is None: + if expected_revision not in {None, 0}: + raise GovernancePolicyConflictError("submission policy revision changed") row = UserSubmissionPolicyRow( user_id=user_id, revision=1, @@ -842,6 +853,11 @@ class SQLAlchemyCloudRepository: ) session.add(row) else: + if ( + expected_revision is not None + and expected_revision != row.revision + ): + raise GovernancePolicyConflictError("submission policy revision changed") row.revision += 1 row.submission_enabled = 1 if submission_enabled else 0 row.allowed_host_ids_json = _dump_optional_list(allowed_host_ids) @@ -865,12 +881,21 @@ class SQLAlchemyCloudRepository: max_active_tasks: int | None, daily_token_budget: int | None, updated_at: datetime, + expected_revision: int | None = None, ) -> Any: + from cloud.governance import GovernancePolicyConflictError + with self._sessions.begin() as session: if session.get(HostRow, host_id) is None: raise KeyError(f"unknown host {host_id!r}") - row = session.get(HostGovernancePolicyRow, host_id) + row = session.get( + HostGovernancePolicyRow, + host_id, + with_for_update=self.engine.dialect.name == "postgresql", + ) if row is None: + if expected_revision not in {None, 0}: + raise GovernancePolicyConflictError("Host policy revision changed") row = HostGovernancePolicyRow( host_id=host_id, revision=1, @@ -881,6 +906,11 @@ class SQLAlchemyCloudRepository: ) session.add(row) else: + if ( + expected_revision is not None + and expected_revision != row.revision + ): + raise GovernancePolicyConflictError("Host policy revision changed") row.revision += 1 row.self_submission_enabled = 1 if self_submission_enabled else 0 row.max_active_tasks = max_active_tasks @@ -889,6 +919,133 @@ class SQLAlchemyCloudRepository: session.flush() return _host_governance_policy_from_row(row) + def count_active_tasks_for_host(self, host_id: str) -> int: + with self._sessions() as session: + count = session.scalar( + select(func.count()) + .select_from(ScheduledTaskRow) + .where( + ScheduledTaskRow.assigned_host_id == host_id, + ScheduledTaskRow.status.in_(("assigned", "dispatched")), + ) + ) + return int(count or 0) + + def reserve_host_token_budget( + self, + *, + reservation_id: str, + host_id: str, + usage_day: str, + reserved_tokens: int, + task_id: str | None, + attempt: int | None, + created_at: datetime, + expires_at: datetime, + ) -> Any | None: + from cloud.governance import TokenBudgetExceededError + + with self._sessions.begin() as session: + statement = select(HostGovernancePolicyRow).where( + HostGovernancePolicyRow.host_id == host_id + ) + if self.engine.dialect.name == "postgresql": + statement = statement.with_for_update() + policy = session.scalars(statement).first() + if policy is None or policy.daily_token_budget is None: + return None + used = session.scalar( + select(func.coalesce(func.sum(TokenUsageEventRow.total_tokens), 0)).where( + TokenUsageEventRow.host_id == host_id, + TokenUsageEventRow.usage_day == usage_day, + ) + ) + reserved = session.scalar( + select(func.coalesce(func.sum(TokenReservationRow.reserved_tokens), 0)).where( + TokenReservationRow.host_id == host_id, + TokenReservationRow.usage_day == usage_day, + TokenReservationRow.expires_at > _iso(created_at), + ) + ) + if int(used or 0) + int(reserved or 0) + reserved_tokens > policy.daily_token_budget: + raise TokenBudgetExceededError("Host daily token budget is exhausted") + row = TokenReservationRow( + id=reservation_id, host_id=host_id, usage_day=usage_day, + reserved_tokens=reserved_tokens, task_id=task_id, attempt=attempt, + created_at=_iso(created_at), expires_at=_iso(expires_at), + ) + session.add(row) + session.flush() + return _token_reservation_from_row(row) + + def settle_host_token_reservation( + self, + *, + reservation_id: str, + event_id: str, + provider: str, + model: str, + input_tokens: int | None, + output_tokens: int | None, + total_tokens: int, + occurred_at: datetime, + ) -> Any | None: + with self._sessions.begin() as session: + row = session.get( + TokenReservationRow, reservation_id, + with_for_update=self.engine.dialect.name == "postgresql", + ) + if row is None: + return None + event = TokenUsageEventRow( + id=event_id, host_id=row.host_id, usage_day=row.usage_day, + task_id=row.task_id, attempt=row.attempt, provider=provider, model=model, + input_tokens=input_tokens, output_tokens=output_tokens, + total_tokens=total_tokens, occurred_at=_iso(occurred_at), + ) + session.add(event) + session.delete(row) + session.flush() + return _token_usage_event_from_row(event) + + def cleanup_expired_token_reservations(self, *, now: datetime, limit: int) -> int: + with self._sessions.begin() as session: + rows = session.scalars( + select(TokenReservationRow) + .where(TokenReservationRow.expires_at <= _iso(now)) + .order_by(TokenReservationRow.expires_at) + .limit(limit) + ).all() + for row in rows: + session.delete(row) + return len(rows) + + def get_host_token_usage_summary( + self, *, host_id: str, usage_day: str, now: datetime, + ) -> Any: + from cloud.governance import TokenUsageSummary + + with self._sessions() as session: + policy = session.get(HostGovernancePolicyRow, host_id) + used = session.scalar( + select(func.coalesce(func.sum(TokenUsageEventRow.total_tokens), 0)).where( + TokenUsageEventRow.host_id == host_id, + TokenUsageEventRow.usage_day == usage_day, + ) + ) + reserved = session.scalar( + select(func.coalesce(func.sum(TokenReservationRow.reserved_tokens), 0)).where( + TokenReservationRow.host_id == host_id, + TokenReservationRow.usage_day == usage_day, + TokenReservationRow.expires_at > _iso(now), + ) + ) + return TokenUsageSummary( + host_id=host_id, usage_day=usage_day, + daily_token_budget=(policy.daily_token_budget if policy else None), + used_tokens=int(used or 0), reserved_tokens=int(reserved or 0), + ) + def list_reserved_device_ids(self, *, now: datetime) -> set[str]: with self._sessions() as session: device_ids = session.scalars( @@ -930,6 +1087,24 @@ class SQLAlchemyCloudRepository: if device is None or device.status != "idle": return None + policy_statement = select(HostGovernancePolicyRow).where( + HostGovernancePolicyRow.host_id == host_id + ) + if self.engine.dialect.name == "postgresql": + policy_statement = policy_statement.with_for_update() + policy = session.scalars(policy_statement).first() + if policy is not None and policy.max_active_tasks is not None: + active_count = session.scalar( + select(func.count()) + .select_from(ScheduledTaskRow) + .where( + ScheduledTaskRow.assigned_host_id == host_id, + ScheduledTaskRow.status.in_(("assigned", "dispatched")), + ) + ) + if int(active_count or 0) >= policy.max_active_tasks: + return None + active_reservation = session.scalar( select(ScheduledTaskRow.id) .where( @@ -1376,6 +1551,28 @@ def _host_governance_policy_from_row(row: HostGovernancePolicyRow) -> Any: ) +def _token_reservation_from_row(row: TokenReservationRow) -> Any: + from cloud.governance import TokenReservation + + return TokenReservation( + id=row.id, host_id=row.host_id, usage_day=row.usage_day, + reserved_tokens=row.reserved_tokens, task_id=row.task_id, attempt=row.attempt, + created_at=_parse_dt(row.created_at) or utc_now(), + expires_at=_parse_dt(row.expires_at) or utc_now(), + ) + + +def _token_usage_event_from_row(row: TokenUsageEventRow) -> Any: + from cloud.governance import TokenUsageEvent + + return TokenUsageEvent( + id=row.id, host_id=row.host_id, usage_day=row.usage_day, + task_id=row.task_id, attempt=row.attempt, provider=row.provider, model=row.model, + input_tokens=row.input_tokens, output_tokens=row.output_tokens, + total_tokens=row.total_tokens, occurred_at=_parse_dt(row.occurred_at) or utc_now(), + ) + + def _task_attempt_from_row(row: TaskAttemptRow) -> Any: from cloud.repository import TaskAttemptRecord diff --git a/tests/test_cloud_governance.py b/tests/test_cloud_governance.py index 7695c62..edf7607 100644 --- a/tests/test_cloud_governance.py +++ b/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() diff --git a/tests/test_cloud_migrations.py b/tests/test_cloud_migrations.py index 43cb9a9..b09e870 100644 --- a/tests/test_cloud_migrations.py +++ b/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 = { diff --git a/tests/test_host_agent_internal_api.py b/tests/test_host_agent_internal_api.py index fc5b21c..ca93e37 100644 --- a/tests/test_host_agent_internal_api.py +++ b/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) diff --git a/tests/test_task_scheduler.py b/tests/test_task_scheduler.py index 20ba799..f41df96 100644 --- a/tests/test_task_scheduler.py +++ b/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"