feat(cloud): enforce host governance budgets
Tests / Test passed: 662

This commit is contained in:
2026-07-13 22:56:31 +08:00
parent 2cd314b183
commit b4803f90e6
28 changed files with 1311 additions and 14 deletions
+10
View File
@@ -301,6 +301,12 @@ def create_app(
auth_provider=auth_provider, auth_provider=auth_provider,
lease_duration_seconds=control_config.lease_duration_seconds, lease_duration_seconds=control_config.lease_duration_seconds,
scheduler=scheduler, 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(), now=utc_now(),
max_attempts=max_attempts, max_attempts=max_attempts,
) )
services.repository.cleanup_expired_token_reservations(
now=utc_now(),
limit=100,
)
except Exception: except Exception:
logger.exception( logger.exception(
"cloud lifecycle iteration failed", "cloud lifecycle iteration failed",
+9
View File
@@ -19,6 +19,7 @@ from host_agent.history import ConsoleHistoryStore
from host_agent.identity import HostIdentityStore from host_agent.identity import HostIdentityStore
from host_agent.lease import ActiveAssignmentRunner from host_agent.lease import ActiveAssignmentRunner
from host_agent.local_account import LocalAccountStore from host_agent.local_account import LocalAccountStore
from host_agent.policy_cache import HostPolicyCacheStore
from host_agent.processor import AssignmentProcessingResult, AssignmentProcessor from host_agent.processor import AssignmentProcessingResult, AssignmentProcessor
from host_agent.status import AgentStatusTracker from host_agent.status import AgentStatusTracker
from host_agent.web.app import create_console_app from host_agent.web.app import create_console_app
@@ -195,6 +196,14 @@ def create_application(
if history_store is not None if history_store is not None
else 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( executor = AssignmentExecutor(
create_execution_factories( create_execution_factories(
+19 -2
View File
@@ -8,6 +8,7 @@ from core.errors import DeviceRuntimeError
from device.manager import DeviceManager from device.manager import DeviceManager
from host_agent.client import HostAgentClient from host_agent.client import HostAgentClient
from host_agent.config import HostAgentConfig from host_agent.config import HostAgentConfig
from host_agent.policy_cache import HostPolicyCacheStore
from host_agent.status import AgentStatusTracker from host_agent.status import AgentStatusTracker
if TYPE_CHECKING: if TYPE_CHECKING:
@@ -37,6 +38,8 @@ class HeartbeatSynchronizer:
sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
status_tracker: AgentStatusTracker | None = None, status_tracker: AgentStatusTracker | None = None,
on_sync: Callable[[int], None] | None = None, on_sync: Callable[[int], None] | None = None,
policy_cache: HostPolicyCacheStore | None = None,
on_policy_sync: Callable[[int], None] | None = None,
) -> None: ) -> None:
self.manager = manager self.manager = manager
self.client = client self.client = client
@@ -45,8 +48,12 @@ class HeartbeatSynchronizer:
self._sleep = sleep self._sleep = sleep
self.status_tracker = status_tracker self.status_tracker = status_tracker
self.on_sync = on_sync self.on_sync = on_sync
self.policy_revision = 0 self.policy_cache = policy_cache
self.policy = None 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: async def sync_once(self) -> HeartbeatResponse:
snapshot = build_device_snapshot(self.manager) snapshot = build_device_snapshot(self.manager)
@@ -58,6 +65,16 @@ class HeartbeatSynchronizer:
self.policy_revision = response.policy_revision self.policy_revision = response.policy_revision
if response.policy is not None: if response.policy is not None:
self.policy = response.policy 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: if self.status_tracker is not None:
self.status_tracker.mark_heartbeat(ok=True, device_count=len(snapshot)) self.status_tracker.mark_heartbeat(ok=True, device_count=len(snapshot))
if self.on_sync is not None: if self.on_sync is not None:
@@ -48,6 +48,13 @@ class ConsoleHistoryStore:
detail = {"device_count": device_count} detail = {"device_count": device_count}
self._insert("heartbeat", summary, detail) 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]]: def list_recent(self, limit: int | None = None) -> list[dict[str, Any]]:
effective_limit = limit if limit is not None else self.limit effective_limit = limit if limit is not None else self.limit
with self._connect() as connection: with self._connect() as connection:
@@ -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
@@ -33,6 +33,7 @@ class AgentStatusTracker:
self._lock = threading.Lock() self._lock = threading.Lock()
self._current_assignment: _CurrentAssignment | None = None self._current_assignment: _CurrentAssignment | None = None
self._last_heartbeat: _LastHeartbeat | None = None self._last_heartbeat: _LastHeartbeat | None = None
self._host_policy: dict[str, Any] | None = None
def mark_assignment_started(self, assignment: AssignmentModel) -> None: def mark_assignment_started(self, assignment: AssignmentModel) -> None:
with self._lock: with self._lock:
@@ -56,6 +57,19 @@ class AgentStatusTracker:
at=self._now(), 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]: def snapshot(self) -> dict[str, Any]:
with self._lock: with self._lock:
current_assignment = self._current_assignment current_assignment = self._current_assignment
@@ -81,4 +95,5 @@ class AgentStatusTracker:
if last_heartbeat is not None if last_heartbeat is not None
else None else None
), ),
"host_policy": self._host_policy.copy() if self._host_policy else None,
} }
@@ -117,6 +117,7 @@ def _dashboard_body(
) -> str: ) -> str:
heartbeat = snapshot.get("last_heartbeat") heartbeat = snapshot.get("last_heartbeat")
assignment = snapshot.get("current_assignment") assignment = snapshot.get("current_assignment")
policy = snapshot.get("host_policy")
heartbeat_text = ( heartbeat_text = (
f"{'ok' if heartbeat['ok'] else 'failed'} at {heartbeat['at']} " f"{'ok' if heartbeat['ok'] else 'failed'} at {heartbeat['at']} "
f"({heartbeat['device_count']} devices)" f"({heartbeat['device_count']} devices)"
@@ -129,6 +130,17 @@ def _dashboard_body(
if assignment if assignment
else "none" 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( device_rows = "".join(
f"<tr><td>{escape(device.id)}</td><td>{escape(device.name or '')}</td>" 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>" f"<td>{escape(device.driver_type)}</td><td>{escape(device.status)}</td></tr>"
@@ -146,6 +158,10 @@ def _dashboard_body(
<h2>Heartbeat</h2> <h2>Heartbeat</h2>
<p id="last-heartbeat">{escape(heartbeat_text)}</p> <p id="last-heartbeat">{escape(heartbeat_text)}</p>
</section> </section>
<section>
<h2>Cloud policy</h2>
<p id="host-policy">{escape(policy_text)}</p>
</section>
<section> <section>
<h2>Current assignment</h2> <h2>Current assignment</h2>
<p id="current-assignment">{escape(assignment_text)}</p> <p id="current-assignment">{escape(assignment_text)}</p>
@@ -168,6 +184,13 @@ def _dashboard_body(
document.getElementById("current-assignment").textContent = current document.getElementById("current-assignment").textContent = current
? current.task_id + " on " + current.device_id + " (started " + current.started_at + ")" ? current.task_id + " on " + current.device_id + " (started " + current.started_at + ")"
: "none"; : "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"); var body = document.getElementById("device-status-body");
body.innerHTML = ""; body.innerHTML = "";
data.devices.forEach(function (device) {{ data.devices.forEach(function (device) {{
@@ -4,9 +4,11 @@ import asyncio
from datetime import UTC, datetime from datetime import UTC, datetime
from cloud.internal_api.models import HeartbeatResponse from cloud.internal_api.models import HeartbeatResponse
from cloud.internal_api.models import HostGovernancePolicyModel
from device.manager import DeviceManager from device.manager import DeviceManager
from host_agent.config import HostAgentConfig from host_agent.config import HostAgentConfig
from host_agent.heartbeat import HeartbeatSynchronizer, build_device_snapshot from host_agent.heartbeat import HeartbeatSynchronizer, build_device_snapshot
from host_agent.policy_cache import HostPolicyCacheStore
from host_agent.status import AgentStatusTracker 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 assert last_heartbeat["device_count"] == 2
asyncio.run(scenario()) 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")
+23 -1
View File
@@ -7,6 +7,7 @@ import {
LogOut, LogOut,
MonitorSmartphone, MonitorSmartphone,
Puzzle, Puzzle,
UsersRound,
} from "@lucide/vue"; } from "@lucide/vue";
import { import {
AUTH_INVALID_EVENT, AUTH_INVALID_EVENT,
@@ -19,8 +20,9 @@ import PasswordChangeScreen from "./views/PasswordChangeScreen.vue";
import TasksView from "./views/TasksView.vue"; import TasksView from "./views/TasksView.vue";
import DevicesView from "./views/DevicesView.vue"; import DevicesView from "./views/DevicesView.vue";
import PluginsView from "./views/PluginsView.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<ViewId>("tasks"); const activeView = ref<ViewId>("tasks");
const currentUser = ref<CloudUser | null>(null); const currentUser = ref<CloudUser | null>(null);
@@ -37,6 +39,16 @@ const canSubmitTasks = computed(
currentUser.value?.scopes.includes("*") || currentUser.value?.scopes.includes("*") ||
currentUser.value?.scopes.includes("tasks:submit"), 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 isAuthenticated = computed(() => currentUser.value !== null);
const currentUserLabel = computed(() => const currentUserLabel = computed(() =>
currentUser.value ? `${currentUser.value.display_name} (${currentUser.value.role})` : "", 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: "devices", label: "Devices", icon: MonitorSmartphone },
{ id: "plugins", label: "Plugins", icon: Puzzle }, { id: "plugins", label: "Plugins", icon: Puzzle },
]; ];
if (canAdminUsers.value || canAdminGovernance.value) {
items.push({ id: "users", label: "Users & limits", icon: UsersRound });
}
return items; return items;
}); });
@@ -102,6 +117,8 @@ const activeComponent = computed(() => {
return DevicesView; return DevicesView;
case "plugins": case "plugins":
return PluginsView; return PluginsView;
case "users":
return UsersView;
default: default:
return TasksView; return TasksView;
} }
@@ -133,6 +150,11 @@ const activeComponent = computed(() => {
</nav> </nav>
<main class="app-main"> <main class="app-main">
<PluginsView v-if="activeView === 'plugins'" :can-admin="canAdminPlugins" /> <PluginsView v-if="activeView === 'plugins'" :can-admin="canAdminPlugins" />
<UsersView
v-else-if="activeView === 'users'"
:can-admin-users="canAdminUsers"
:can-admin-governance="canAdminGovernance"
/>
<component v-else :is="activeComponent" :can-submit="canSubmitTasks" /> <component v-else :is="activeComponent" :can-submit="canSubmitTasks" />
</main> </main>
</div> </div>
+85
View File
@@ -1,6 +1,8 @@
import type { import type {
CloudUser, CloudUser,
DeviceRecord, DeviceRecord,
HostGovernancePolicy,
HostTokenUsageSummary,
HostRecord, HostRecord,
PluginRecord, PluginRecord,
PluginRegistrationPayload, PluginRegistrationPayload,
@@ -8,6 +10,8 @@ import type {
TaskListResponse, TaskListResponse,
TaskSubmissionPayload, TaskSubmissionPayload,
TaskStatus, TaskStatus,
UserListResponse,
UserSubmissionPolicy,
} from "./types"; } from "./types";
const configuredBaseUrl = import.meta.env.VITE_CLOUD_API_BASE_URL as const configuredBaseUrl = import.meta.env.VITE_CLOUD_API_BASE_URL as
@@ -119,6 +123,87 @@ export function changePassword(
}); });
} }
export function listUsers(): Promise<UserListResponse> {
return request<UserListResponse>("/v1/users?limit=100&offset=0");
}
export function createUser(payload: {
username: string;
display_name: string;
role: "viewer" | "operator" | "admin";
password: string;
}): Promise<CloudUser> {
return request<CloudUser>("/v1/users", {
method: "POST",
body: JSON.stringify(payload),
});
}
export function updateUser(
userId: string,
payload: { display_name?: string; role?: "viewer" | "operator" | "admin"; enabled?: boolean },
): Promise<CloudUser> {
return request<CloudUser>(`/v1/users/${encodeURIComponent(userId)}`, {
method: "PATCH",
body: JSON.stringify(payload),
});
}
export function resetUserPassword(userId: string, password: string): Promise<CloudUser> {
return request<CloudUser>(`/v1/users/${encodeURIComponent(userId)}/password`, {
method: "POST",
body: JSON.stringify({ password }),
});
}
export function revokeUserSessions(userId: string): Promise<void> {
return request<void>(`/v1/users/${encodeURIComponent(userId)}/sessions`, {
method: "DELETE",
});
}
export function getUserSubmissionPolicy(userId: string): Promise<UserSubmissionPolicy> {
return request<UserSubmissionPolicy>(
`/v1/users/${encodeURIComponent(userId)}/submission-policy`,
);
}
export function updateUserSubmissionPolicy(
userId: string,
payload: Omit<UserSubmissionPolicy, "user_id" | "revision" | "updated_at"> & {
expected_revision?: number;
},
): Promise<UserSubmissionPolicy> {
return request<UserSubmissionPolicy>(
`/v1/users/${encodeURIComponent(userId)}/submission-policy`,
{ method: "PUT", body: JSON.stringify(payload) },
);
}
export function getHostGovernancePolicy(hostId: string): Promise<HostGovernancePolicy> {
return request<HostGovernancePolicy>(
`/v1/hosts/${encodeURIComponent(hostId)}/governance-policy`,
);
}
export function updateHostGovernancePolicy(
hostId: string,
payload: Omit<HostGovernancePolicy, "host_id" | "revision" | "updated_at"> & {
expected_revision?: number;
},
): Promise<HostGovernancePolicy> {
return request<HostGovernancePolicy>(
`/v1/hosts/${encodeURIComponent(hostId)}/governance-policy`,
{ method: "PUT", body: JSON.stringify(payload) },
);
}
export function getHostTokenUsage(hostId: string): Promise<HostTokenUsageSummary> {
return request<HostTokenUsageSummary>(
`/v1/hosts/${encodeURIComponent(hostId)}/token-usage`,
);
}
export function listTasks(options?: { export function listTasks(options?: {
status?: TaskStatus; status?: TaskStatus;
limit?: number; limit?: number;
+13
View File
@@ -319,6 +319,19 @@ tr.row-selected {
grid-column: 1 / -1; 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 { .muted {
color: var(--text-muted); color: var(--text-muted);
} }
+38
View File
@@ -96,3 +96,41 @@ export interface CloudUser {
updated_at: string; updated_at: string;
last_login_at: string | null; 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;
}
+380
View File
@@ -0,0 +1,380 @@
<script setup lang="ts">
import { computed, onMounted, ref, watch } from "vue";
import {
CloudApiError,
createUser,
getHostGovernancePolicy,
getHostTokenUsage,
getUserSubmissionPolicy,
listDevices,
listHosts,
listUsers,
resetUserPassword,
revokeUserSessions,
updateUser,
updateUserSubmissionPolicy,
updateHostGovernancePolicy,
} from "../api";
import type {
CloudUser,
DeviceRecord,
HostRecord,
HostTokenUsageSummary,
UserRole,
} from "../types";
const props = defineProps<{ canAdminUsers: boolean; canAdminGovernance: boolean }>();
const users = ref<CloudUser[]>([]);
const hosts = ref<HostRecord[]>([]);
const devices = ref<DeviceRecord[]>([]);
const selectedUserId = ref("");
const loading = ref(false);
const saving = ref(false);
const errorMessage = ref("");
const successMessage = ref("");
const createUsername = ref("");
const createDisplayName = ref("");
const createRole = ref<UserRole>("operator");
const createPassword = ref("");
const selectedRole = ref<UserRole>("viewer");
const selectedEnabled = ref(true);
const resetPassword = ref("");
const submissionEnabled = ref(true);
const restrictHosts = ref(false);
const allowedHostIds = ref<string[]>([]);
const restrictDevices = ref(false);
const allowedDeviceKeys = ref<string[]>([]);
const policyRevision = ref<number | null>(null);
const selectedHostId = ref("");
const hostPolicyRevision = ref<number | null>(null);
const hostSelfSubmissionEnabled = ref(true);
const hostMaxActiveTasks = ref("");
const hostDailyTokenBudget = ref("");
const hostUsage = ref<HostTokenUsageSummary | null>(null);
const selectedUser = computed(
() => users.value.find((user) => user.id === selectedUserId.value) ?? null,
);
const deviceKey = (device: Pick<DeviceRecord, "host_id" | "device_id">) =>
`${device.host_id}\u0000${device.device_id}`;
function showError(error: unknown, fallback: string) {
successMessage.value = "";
errorMessage.value = error instanceof Error ? error.message : fallback;
}
function selectUser(user: CloudUser) {
selectedUserId.value = user.id;
selectedRole.value = user.role;
selectedEnabled.value = user.enabled;
}
async function refresh() {
loading.value = true;
errorMessage.value = "";
try {
const requests: Promise<unknown>[] = [];
if (props.canAdminUsers) requests.push(listUsers());
if (props.canAdminGovernance) requests.push(listHosts(), listDevices());
const results = await Promise.all(requests);
let index = 0;
if (props.canAdminUsers) {
users.value = (results[index++] as { items: CloudUser[] }).items;
if (!selectedUserId.value && users.value[0]) selectUser(users.value[0]);
}
if (props.canAdminGovernance) {
hosts.value = results[index++] as HostRecord[];
devices.value = results[index++] as DeviceRecord[];
if (!selectedHostId.value && hosts.value[0]) {
selectedHostId.value = hosts.value[0].host_id;
}
}
} catch (error) {
showError(error, "failed to load user administration data");
} finally {
loading.value = false;
}
}
async function loadHostPolicy() {
if (!props.canAdminGovernance || !selectedHostId.value) return;
errorMessage.value = "";
try {
const [policy, usage] = await Promise.all([
getHostGovernancePolicy(selectedHostId.value),
getHostTokenUsage(selectedHostId.value),
]);
hostUsage.value = usage;
hostPolicyRevision.value = policy.revision;
hostSelfSubmissionEnabled.value = policy.self_submission_enabled;
hostMaxActiveTasks.value = policy.max_active_tasks?.toString() ?? "";
hostDailyTokenBudget.value = policy.daily_token_budget?.toString() ?? "";
} catch (error) {
if (error instanceof CloudApiError && error.status === 404) {
hostPolicyRevision.value = null;
hostSelfSubmissionEnabled.value = true;
hostMaxActiveTasks.value = "";
hostDailyTokenBudget.value = "";
hostUsage.value = await getHostTokenUsage(selectedHostId.value);
return;
}
showError(error, "failed to load Host policy");
}
}
async function loadPolicy() {
if (!props.canAdminGovernance || !selectedUserId.value) return;
errorMessage.value = "";
try {
const policy = await getUserSubmissionPolicy(selectedUserId.value);
submissionEnabled.value = policy.submission_enabled;
restrictHosts.value = policy.allowed_host_ids !== null;
allowedHostIds.value = policy.allowed_host_ids ?? [];
restrictDevices.value = policy.allowed_device_targets !== null;
allowedDeviceKeys.value = (policy.allowed_device_targets ?? []).map(deviceKey);
policyRevision.value = policy.revision;
} catch (error) {
if (error instanceof CloudApiError && error.status === 404) {
submissionEnabled.value = true;
restrictHosts.value = false;
allowedHostIds.value = [];
restrictDevices.value = false;
allowedDeviceKeys.value = [];
policyRevision.value = null;
return;
}
showError(error, "failed to load submission policy");
}
}
async function createAccount() {
saving.value = true;
errorMessage.value = "";
try {
const user = await createUser({
username: createUsername.value.trim(),
display_name: createDisplayName.value.trim(),
role: createRole.value,
password: createPassword.value,
});
createUsername.value = "";
createDisplayName.value = "";
successMessage.value = `created ${user.username}`;
await refresh();
selectUser(user);
} catch (error) {
showError(error, "failed to create user");
} finally {
createPassword.value = "";
saving.value = false;
}
}
async function saveUser() {
if (!selectedUser.value) return;
saving.value = true;
try {
await updateUser(selectedUser.value.id, {
role: selectedRole.value,
enabled: selectedEnabled.value,
});
successMessage.value = "user updated";
await refresh();
} catch (error) {
showError(error, "failed to update user");
} finally {
saving.value = false;
}
}
async function resetPasswordForUser() {
if (!selectedUser.value || !resetPassword.value) return;
saving.value = true;
try {
await resetUserPassword(selectedUser.value.id, resetPassword.value);
successMessage.value = "password reset; existing sessions were revoked";
} catch (error) {
showError(error, "failed to reset password");
} finally {
resetPassword.value = "";
saving.value = false;
}
}
async function revokeSessions() {
if (!selectedUser.value) return;
saving.value = true;
try {
await revokeUserSessions(selectedUser.value.id);
successMessage.value = "sessions revoked";
} catch (error) {
showError(error, "failed to revoke sessions");
} finally {
saving.value = false;
}
}
async function savePolicy() {
if (!selectedUserId.value) {
errorMessage.value = "select or enter a user id before saving a policy";
return;
}
saving.value = true;
try {
const selectedDevices = new Set(allowedDeviceKeys.value);
const policy = await updateUserSubmissionPolicy(selectedUserId.value, {
submission_enabled: submissionEnabled.value,
expected_revision: policyRevision.value ?? 0,
allowed_host_ids: restrictHosts.value ? allowedHostIds.value : null,
allowed_device_targets: restrictDevices.value
? devices.value
.filter((device) => selectedDevices.has(deviceKey(device)))
.map((device) => ({ host_id: device.host_id, device_id: device.device_id }))
: null,
});
policyRevision.value = policy.revision;
successMessage.value = `submission policy saved (revision ${policy.revision})`;
} catch (error) {
showError(error, "failed to save submission policy");
} finally {
saving.value = false;
}
}
function parseOptionalPositive(value: string, label: string): number | null {
if (!value.trim()) return null;
const parsed = Number(value);
if (!Number.isInteger(parsed) || parsed < 1) {
throw new Error(`${label} must be a positive integer`);
}
return parsed;
}
async function saveHostPolicy() {
if (!selectedHostId.value) return;
saving.value = true;
try {
const policy = await updateHostGovernancePolicy(selectedHostId.value, {
self_submission_enabled: hostSelfSubmissionEnabled.value,
max_active_tasks: parseOptionalPositive(hostMaxActiveTasks.value, "max active tasks"),
daily_token_budget: parseOptionalPositive(hostDailyTokenBudget.value, "daily token budget"),
expected_revision: hostPolicyRevision.value ?? 0,
});
hostPolicyRevision.value = policy.revision;
hostUsage.value = await getHostTokenUsage(selectedHostId.value);
successMessage.value = `Host policy saved (revision ${policy.revision})`;
} catch (error) {
showError(error, "failed to save Host policy");
} finally {
saving.value = false;
}
}
watch(selectedUserId, () => {
if (selectedUser.value) {
selectedRole.value = selectedUser.value.role;
selectedEnabled.value = selectedUser.value.enabled;
}
void loadPolicy();
});
watch(selectedHostId, () => {
void loadHostPolicy();
});
onMounted(() => void refresh());
</script>
<template>
<div>
<div class="toolbar">
<h2>Users & task limits</h2>
<button :disabled="loading" @click="refresh">Refresh</button>
<span v-if="loading" class="muted">loading</span>
</div>
<div v-if="errorMessage" class="notice error">{{ errorMessage }}</div>
<div v-if="successMessage" class="notice success">{{ successMessage }}</div>
<div v-if="canAdminUsers" class="panel">
<h3>Create user</h3>
<form class="form-grid" @submit.prevent="createAccount">
<label>Username <input v-model="createUsername" required /></label>
<label>Display name <input v-model="createDisplayName" required /></label>
<label>Role <select v-model="createRole"><option value="viewer">Viewer</option><option value="operator">Operator</option><option value="admin">Admin</option></select></label>
<label>Password <input v-model="createPassword" type="password" required /></label>
<div class="field-full actions"><button class="primary" :disabled="saving">Create user</button></div>
</form>
</div>
<div v-if="canAdminUsers" class="panel">
<h3>Accounts</h3>
<table v-if="users.length">
<thead><tr><th>Username</th><th>Role</th><th>Status</th><th></th></tr></thead>
<tbody>
<tr v-for="user in users" :key="user.id" :class="{ 'row-selected': user.id === selectedUserId }">
<td>{{ user.display_name }} <span class="dim">({{ user.username }})</span></td>
<td>{{ user.role }}</td>
<td>{{ user.enabled ? "enabled" : "disabled" }}</td>
<td><button @click="selectUser(user)">Manage</button></td>
</tr>
</tbody>
</table>
<p v-else class="muted">No accounts found.</p>
</div>
<div v-if="canAdminGovernance && !canAdminUsers" class="panel">
<label>User ID <input v-model="selectedUserId" placeholder="Cloud user id" /></label>
<p class="muted">Use an existing user id to manage its task-submission policy.</p>
</div>
<div v-if="selectedUser && canAdminUsers" class="panel">
<h3>Account: {{ selectedUser.username }}</h3>
<div class="form-grid">
<label>Role <select v-model="selectedRole"><option value="viewer">Viewer</option><option value="operator">Operator</option><option value="admin">Admin</option></select></label>
<label><input v-model="selectedEnabled" type="checkbox" /> Enabled</label>
<div class="field-full actions"><button :disabled="saving" @click="saveUser">Save account</button></div>
<label>Password reset <input v-model="resetPassword" type="password" placeholder="New temporary password" /></label>
<div class="actions"><button :disabled="saving || !resetPassword" @click="resetPasswordForUser">Reset password</button><button :disabled="saving" @click="revokeSessions">Revoke sessions</button></div>
</div>
</div>
<div v-if="canAdminGovernance && selectedUserId" class="panel">
<h3>Task-submission policy <span class="dim">{{ policyRevision === null ? "not configured" : `revision ${policyRevision}` }}</span></h3>
<p class="muted">An unrestricted policy allows ordinary scope-authorized submissions. Enabling a restriction requires an explicit target.</p>
<label><input v-model="submissionEnabled" type="checkbox" /> Allow this user to submit tasks</label>
<div class="form-grid policy-grid">
<label><input v-model="restrictHosts" type="checkbox" /> Restrict Hosts</label>
<select v-model="allowedHostIds" multiple :disabled="!restrictHosts">
<option v-for="host in hosts" :key="host.host_id" :value="host.host_id">{{ host.host_id }}</option>
</select>
<label><input v-model="restrictDevices" type="checkbox" /> Restrict Devices</label>
<select v-model="allowedDeviceKeys" multiple :disabled="!restrictDevices">
<option v-for="device in devices" :key="deviceKey(device)" :value="deviceKey(device)">{{ device.host_id }} / {{ device.device_id }}</option>
</select>
</div>
<div class="actions"><button class="primary" :disabled="saving" @click="savePolicy">Save policy</button></div>
</div>
<div v-if="canAdminGovernance" class="panel">
<h3>Host operational limits</h3>
<p class="muted">Daily token budgets apply when this Host uses Cloud planner transport. Each decision reserves the Cloud-configured conservative ceiling before calling a provider, so a usable budget must cover that ceiling. Direct provider transport remains unmetered until it is moved behind the proxy.</p>
<div class="form-grid">
<label>Host
<select v-model="selectedHostId">
<option v-for="host in hosts" :key="host.host_id" :value="host.host_id">{{ host.host_id }}</option>
</select>
</label>
<label><input v-model="hostSelfSubmissionEnabled" type="checkbox" /> Allow edge self-submission</label>
<label>Maximum active tasks <input v-model="hostMaxActiveTasks" inputmode="numeric" placeholder="Unlimited" /></label>
<label>Daily Cloud-proxy token budget <input v-model="hostDailyTokenBudget" inputmode="numeric" placeholder="Unlimited" /></label>
</div>
<div class="actions"><button class="primary" :disabled="saving || !selectedHostId" @click="saveHostPolicy">Save Host policy {{ hostPolicyRevision === null ? "" : `(revision ${hostPolicyRevision})` }}</button></div>
<p v-if="hostUsage" class="muted">
{{ hostUsage.usage_day }}: used {{ hostUsage.used_tokens }}, reserved {{ hostUsage.reserved_tokens }},
remaining {{ hostUsage.remaining_tokens ?? "unmetered" }} tokens.
</p>
</div>
</div>
</template>
@@ -51,13 +51,13 @@
- [x] 4.1 Extend shared internal heartbeat request/response models with the - [x] 4.1 Extend shared internal heartbeat request/response models with the
Host's last policy revision and a revision-aware effective-policy reply. 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 policy state and expose it through local status/history without adding an
inbound Cloud connection. inbound Cloud connection.
- [x] 4.3 Add the Host-scoped goal-only task-submission route and - [x] 4.3 Add the Host-scoped goal-only task-submission route and
`HostAgentClient` method; derive Host targeting from authenticated `HostAgentClient` method; derive Host targeting from authenticated
credentials and validate any named local Device ownership. 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. the Cloud service/scheduler, not only in Host-local code.
- [ ] 4.5 Add Cloud API and Host Agent tests for revision convergence, - [ ] 4.5 Add Cloud API and Host Agent tests for revision convergence,
unchanged-policy replies, self-targeted task creation, foreign target 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 - [ ] 5.2 Extend planner-proxy request context and Host-Agent-local context
binding so Cloud-proxied calls carry known task/attempt metadata without binding so Cloud-proxied calls carry known task/attempt metadata without
importing Host or Cloud concerns into `runtime`. 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 per-call ceiling, provider invocation, actual-usage settlement, and
bounded unknown-usage reservation expiry. bounded unknown-usage reservation expiry.
- [ ] 5.4 Record non-secret usage events and expose accurate - [ ] 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 - [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 goal/workflow submission, Host/Device target selection, and target/policy
validation errors. 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. submission-policy editing with safe refresh and conflict/error handling.
- [ ] 6.3 Add Host policy administration and AI-usage/budget views, including - [ ] 6.3 Add Host policy administration and AI-usage/budget views, including
revision display, unmetered direct Hosts, and no rendering of prompts, revision display, unmetered direct Hosts, and no rendering of prompts,
@@ -36,6 +36,8 @@ class CloudControlConfig:
login_block_seconds: int = 900 login_block_seconds: int = 900
session_cookie_secure: bool = False session_cookie_secure: bool = False
trust_proxy_headers: bool = False trust_proxy_headers: bool = False
planner_token_reservation_ceiling: int = 4096
planner_token_reservation_ttl_seconds: int = 300
def load_control_config( def load_control_config(
@@ -115,6 +117,16 @@ def load_control_config(
values.get("CLOUD_TRUST_PROXY_HEADERS"), values.get("CLOUD_TRUST_PROXY_HEADERS"),
default=False, 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) validate_control_config(config)
return config return config
@@ -241,3 +241,44 @@ class HostGovernancePolicyRow(Base):
max_active_tasks: Mapped[int | None] = mapped_column(Integer, nullable=True) max_active_tasks: Mapped[int | None] = mapped_column(Integer, nullable=True)
daily_token_budget: 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) 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)
@@ -10,6 +10,14 @@ class TaskSubmissionPolicyError(PermissionError):
"""Raised when a human user's policy disallows a task submission.""" """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) @dataclass(frozen=True)
class UserSubmissionPolicy: class UserSubmissionPolicy:
user_id: str user_id: str
@@ -30,6 +38,48 @@ class HostGovernancePolicy:
updated_at: datetime 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( def enforce_user_submission_policy(
policy: UserSubmissionPolicy | None, policy: UserSubmissionPolicy | None,
*, *,
@@ -44,6 +44,7 @@ from cloud.repository import (
DeviceEnrollmentConflictError, DeviceEnrollmentConflictError,
HostEnrollmentConflictError, HostEnrollmentConflictError,
) )
from cloud.governance import TokenBudgetExceededError
from core.models import Device, utc_now from core.models import Device, utc_now
from runtime.tool_calling_client import ToolCallingClient, ToolCallUnavailable from runtime.tool_calling_client import ToolCallingClient, ToolCallUnavailable
from runtime.tool_specs import ToolSpec from runtime.tool_specs import ToolSpec
@@ -65,11 +66,17 @@ def create_internal_router(
sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
planner_client_factory: Callable[[], ToolCallingClient] | None = None, planner_client_factory: Callable[[], ToolCallingClient] | None = None,
scheduler: TaskScheduler | None = None, scheduler: TaskScheduler | None = None,
planner_token_reservation_ceiling: int = 4096,
planner_token_reservation_ttl_seconds: float = 300.0,
) -> APIRouter: ) -> APIRouter:
if claim_poll_interval_seconds <= 0: if claim_poll_interval_seconds <= 0:
raise ValueError("claim_poll_interval_seconds must be greater than zero") raise ValueError("claim_poll_interval_seconds must be greater than zero")
if lease_duration_seconds <= 0: if lease_duration_seconds <= 0:
raise ValueError("lease_duration_seconds must be greater than zero") 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"]) router = APIRouter(prefix=version_prefix, tags=["host-agent"])
build_planner_client = planner_client_factory or _default_planner_client_factory build_planner_client = planner_client_factory or _default_planner_client_factory
@@ -352,6 +359,7 @@ def create_internal_router(
response_model_exclude_none=True, response_model_exclude_none=True,
responses={ responses={
status.HTTP_502_BAD_GATEWAY: {"model": PlannerDecisionError}, status.HTTP_502_BAD_GATEWAY: {"model": PlannerDecisionError},
status.HTTP_429_TOO_MANY_REQUESTS: {"model": PlannerDecisionError},
}, },
) )
def decide_planner_call( def decide_planner_call(
@@ -385,6 +393,25 @@ def create_internal_router(
for tool in payload.tools 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() started_at = monotonic()
client = build_planner_client() client = build_planner_client()
try: try:
@@ -408,6 +435,19 @@ def create_internal_router(
status_code=status.HTTP_502_BAD_GATEWAY, status_code=status.HTTP_502_BAD_GATEWAY,
content=PlannerDecisionError(detail=str(exc)).model_dump(), 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( logger.info(
"planner-decision request resolved", "planner-decision request resolved",
extra={ extra={
@@ -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")
+29 -1
View File
@@ -15,7 +15,13 @@ if TYPE_CHECKING:
UserAccount, UserAccount,
UserSession, UserSession,
) )
from cloud.governance import HostGovernancePolicy, UserSubmissionPolicy from cloud.governance import (
HostGovernancePolicy,
TokenReservation,
TokenUsageSummary,
TokenUsageEvent,
UserSubmissionPolicy,
)
AttemptStatus = Literal["assigned", "dispatched", "done", "failed", "expired"] AttemptStatus = Literal["assigned", "dispatched", "done", "failed", "expired"]
@@ -293,6 +299,7 @@ class CloudRepository(Protocol):
submission_enabled: bool, submission_enabled: bool,
allowed_host_ids: tuple[str, ...] | None, allowed_host_ids: tuple[str, ...] | None,
allowed_device_targets: tuple[tuple[str, str], ...] | None, allowed_device_targets: tuple[tuple[str, str], ...] | None,
expected_revision: int | None,
updated_at: datetime, updated_at: datetime,
) -> UserSubmissionPolicy: ... ) -> UserSubmissionPolicy: ...
@@ -305,9 +312,30 @@ class CloudRepository(Protocol):
self_submission_enabled: bool, self_submission_enabled: bool,
max_active_tasks: int | None, max_active_tasks: int | None,
daily_token_budget: int | None, daily_token_budget: int | None,
expected_revision: int | None,
updated_at: datetime, updated_at: datetime,
) -> HostGovernancePolicy: ... ) -> 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 list_reserved_device_ids(self, *, now: datetime) -> set[str]: ...
def assign_task( def assign_task(
+1 -1
View File
@@ -9,7 +9,7 @@ from alembic.runtime.migration import MigrationContext
from cloud.database import create_database_engine, normalize_database_url from cloud.database import create_database_engine, normalize_database_url
HEAD_REVISION = "0004_cloud_governance" HEAD_REVISION = "0005_cloud_token_usage"
class SchemaVersionError(RuntimeError): class SchemaVersionError(RuntimeError):
@@ -5,10 +5,12 @@ from uuid import uuid4
from fastapi import APIRouter, HTTPException, Request, status from fastapi import APIRouter, HTTPException, Request, status
from cloud.auth import AuthProvider, GOVERNANCE_ADMIN_SCOPE, GOVERNANCE_READ_SCOPE 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.observability import current_correlation_id
from cloud.sdk.models import ( from cloud.sdk.models import (
HostGovernancePolicyRequest, HostGovernancePolicyRequest,
HostGovernancePolicyResponse, HostGovernancePolicyResponse,
HostTokenUsageSummaryResponse,
UserSubmissionPolicyRequest, UserSubmissionPolicyRequest,
UserSubmissionPolicyResponse, UserSubmissionPolicyResponse,
) )
@@ -79,10 +81,13 @@ def create_governance_router(*, repository, auth_provider: AuthProvider) -> APIR
if payload.allowed_device_targets is not None if payload.allowed_device_targets is not None
else None else None
), ),
expected_revision=payload.expected_revision,
updated_at=utc_now(), updated_at=utc_now(),
) )
except KeyError as exc: except KeyError as exc:
raise HTTPException(status_code=404, detail="user not found") from 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") _audit(repository, principal.id, user_id, "user_submission_policy_update")
return UserSubmissionPolicyResponse( return UserSubmissionPolicyResponse(
user_id=policy.user_id, 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") raise HTTPException(status_code=404, detail="Host policy not found")
return _host_response(policy) 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( @router.put(
"/hosts/{host_id}/governance-policy", "/hosts/{host_id}/governance-policy",
response_model=HostGovernancePolicyResponse, response_model=HostGovernancePolicyResponse,
@@ -127,10 +153,13 @@ def create_governance_router(*, repository, auth_provider: AuthProvider) -> APIR
self_submission_enabled=payload.self_submission_enabled, self_submission_enabled=payload.self_submission_enabled,
max_active_tasks=payload.max_active_tasks, max_active_tasks=payload.max_active_tasks,
daily_token_budget=payload.daily_token_budget, daily_token_budget=payload.daily_token_budget,
expected_revision=payload.expected_revision,
updated_at=utc_now(), updated_at=utc_now(),
) )
except KeyError as exc: except KeyError as exc:
raise HTTPException(status_code=404, detail="Host not found") from 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") _audit(repository, principal.id, host_id, "host_governance_policy_update")
return _host_response(policy) return _host_response(policy)
+19 -2
View File
@@ -162,11 +162,15 @@ class UserSubmissionPolicyRequest(BaseModel):
submission_enabled: bool = True submission_enabled: bool = True
allowed_host_ids: list[str] | None = None allowed_host_ids: list[str] | None = None
allowed_device_targets: list[DeviceTargetModel] | 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 user_id: str
revision: int revision: int
submission_enabled: bool
allowed_host_ids: list[str] | None = None
allowed_device_targets: list[DeviceTargetModel] | None = None
updated_at: datetime updated_at: datetime
@@ -174,9 +178,22 @@ class HostGovernancePolicyRequest(BaseModel):
self_submission_enabled: bool = True self_submission_enabled: bool = True
max_active_tasks: int | None = Field(default=None, ge=1) max_active_tasks: int | None = Field(default=None, ge=1)
daily_token_budget: 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 host_id: str
revision: int revision: int
self_submission_enabled: bool
max_active_tasks: int | None = None
daily_token_budget: int | None = None
updated_at: datetime 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
+199 -2
View File
@@ -24,6 +24,8 @@ from cloud.db_models import (
UserRow, UserRow,
UserSessionRow, UserSessionRow,
UserSubmissionPolicyRow, UserSubmissionPolicyRow,
TokenReservationRow,
TokenUsageEventRow,
) )
from cloud.observability import current_correlation_id from cloud.observability import current_correlation_id
from core.models import utc_now from core.models import utc_now
@@ -824,12 +826,21 @@ class SQLAlchemyCloudRepository:
allowed_host_ids: tuple[str, ...] | None, allowed_host_ids: tuple[str, ...] | None,
allowed_device_targets: tuple[tuple[str, str], ...] | None, allowed_device_targets: tuple[tuple[str, str], ...] | None,
updated_at: datetime, updated_at: datetime,
expected_revision: int | None = None,
) -> Any: ) -> Any:
from cloud.governance import GovernancePolicyConflictError
with self._sessions.begin() as session: with self._sessions.begin() as session:
if session.get(UserRow, user_id) is None: if session.get(UserRow, user_id) is None:
raise KeyError(f"unknown user {user_id!r}") 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 row is None:
if expected_revision not in {None, 0}:
raise GovernancePolicyConflictError("submission policy revision changed")
row = UserSubmissionPolicyRow( row = UserSubmissionPolicyRow(
user_id=user_id, user_id=user_id,
revision=1, revision=1,
@@ -842,6 +853,11 @@ class SQLAlchemyCloudRepository:
) )
session.add(row) session.add(row)
else: else:
if (
expected_revision is not None
and expected_revision != row.revision
):
raise GovernancePolicyConflictError("submission policy revision changed")
row.revision += 1 row.revision += 1
row.submission_enabled = 1 if submission_enabled else 0 row.submission_enabled = 1 if submission_enabled else 0
row.allowed_host_ids_json = _dump_optional_list(allowed_host_ids) row.allowed_host_ids_json = _dump_optional_list(allowed_host_ids)
@@ -865,12 +881,21 @@ class SQLAlchemyCloudRepository:
max_active_tasks: int | None, max_active_tasks: int | None,
daily_token_budget: int | None, daily_token_budget: int | None,
updated_at: datetime, updated_at: datetime,
expected_revision: int | None = None,
) -> Any: ) -> Any:
from cloud.governance import GovernancePolicyConflictError
with self._sessions.begin() as session: with self._sessions.begin() as session:
if session.get(HostRow, host_id) is None: if session.get(HostRow, host_id) is None:
raise KeyError(f"unknown host {host_id!r}") 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 row is None:
if expected_revision not in {None, 0}:
raise GovernancePolicyConflictError("Host policy revision changed")
row = HostGovernancePolicyRow( row = HostGovernancePolicyRow(
host_id=host_id, host_id=host_id,
revision=1, revision=1,
@@ -881,6 +906,11 @@ class SQLAlchemyCloudRepository:
) )
session.add(row) session.add(row)
else: else:
if (
expected_revision is not None
and expected_revision != row.revision
):
raise GovernancePolicyConflictError("Host policy revision changed")
row.revision += 1 row.revision += 1
row.self_submission_enabled = 1 if self_submission_enabled else 0 row.self_submission_enabled = 1 if self_submission_enabled else 0
row.max_active_tasks = max_active_tasks row.max_active_tasks = max_active_tasks
@@ -889,6 +919,133 @@ class SQLAlchemyCloudRepository:
session.flush() session.flush()
return _host_governance_policy_from_row(row) 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]: def list_reserved_device_ids(self, *, now: datetime) -> set[str]:
with self._sessions() as session: with self._sessions() as session:
device_ids = session.scalars( device_ids = session.scalars(
@@ -930,6 +1087,24 @@ class SQLAlchemyCloudRepository:
if device is None or device.status != "idle": if device is None or device.status != "idle":
return None 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( active_reservation = session.scalar(
select(ScheduledTaskRow.id) select(ScheduledTaskRow.id)
.where( .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: def _task_attempt_from_row(row: TaskAttemptRow) -> Any:
from cloud.repository import TaskAttemptRecord from cloud.repository import TaskAttemptRecord
+18
View File
@@ -132,6 +132,24 @@ def test_governance_routes_persist_revisioned_policies(tmp_path) -> None:
assert host_policy.json()["revision"] == 1 assert host_policy.json()["revision"] == 1
reread = client.get("/v1/hosts/host-a/governance-policy") reread = client.get("/v1/hosts/host-a/governance-policy")
assert reread.json()["daily_token_budget"] == 1000 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: finally:
database.close() database.close()
+2
View File
@@ -38,6 +38,8 @@ def test_forward_and_downgrade_migrations_on_empty_database(tmp_path) -> None:
"cloud_user_sessions", "cloud_user_sessions",
"cloud_login_throttles", "cloud_login_throttles",
"cloud_auth_audit_events", "cloud_auth_audit_events",
"cloud_token_reservations",
"cloud_token_usage_events",
} <= table_names } <= table_names
assert current_revision(database_url) == HEAD_REVISION assert current_revision(database_url) == HEAD_REVISION
host_columns = { host_columns = {
+61 -1
View File
@@ -12,9 +12,15 @@ from cloud.internal_api.api import create_internal_router
from cloud.pool import DevicePool, PooledDevice from cloud.pool import DevicePool, PooledDevice
from cloud.scheduler import ScheduledTask, TaskConstraints, TaskScheduler from cloud.scheduler import ScheduledTask, TaskConstraints, TaskScheduler
from cloud.store import CloudStore 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( pool = DevicePool(
CloudStore(tmp_path / "internal.sqlite3"), CloudStore(tmp_path / "internal.sqlite3"),
CloudConfig(stale_after_seconds=60), CloudConfig(stale_after_seconds=60),
@@ -39,6 +45,8 @@ def _build_client(tmp_path) -> tuple[TestClient, DevicePool]:
pool=pool, pool=pool,
auth_provider=auth_provider, auth_provider=auth_provider,
scheduler=TaskScheduler(pool, pool.store, CloudConfig(stale_after_seconds=60)), 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 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 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: def test_long_poll_claim_returns_at_most_one_owned_assignment(tmp_path) -> None:
client, pool = _build_client(tmp_path) client, pool = _build_client(tmp_path)
now = datetime.now(UTC) now = datetime.now(UTC)
+29
View File
@@ -264,3 +264,32 @@ def test_unavailable_explicit_target_is_never_rerouted(tmp_path) -> None:
task = pool.store.get_task(task_id) task = pool.store.get_task(task_id)
assert task is not None assert task is not None
assert task.status == "queued" 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"