@@ -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",
|
||||||
|
|||||||
@@ -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(
|
||||||
|
|||||||
@@ -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")
|
||||||
|
|||||||
@@ -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>
|
||||||
|
|||||||
@@ -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;
|
||||||
|
|||||||
@@ -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);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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;
|
||||||
|
}
|
||||||
|
|||||||
@@ -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")
|
||||||
@@ -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(
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|
||||||
|
|||||||
@@ -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()
|
||||||
|
|
||||||
|
|||||||
@@ -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 = {
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|||||||
@@ -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"
|
||||||
|
|||||||
Reference in New Issue
Block a user