diff --git a/apps/device-host-agent/host_agent/assignment.py b/apps/device-host-agent/host_agent/assignment.py index 3a3d687..df5156e 100644 --- a/apps/device-host-agent/host_agent/assignment.py +++ b/apps/device-host-agent/host_agent/assignment.py @@ -7,6 +7,7 @@ from typing import Any from cloud.internal_api.models import AssignmentModel from core.models import Task from host_agent.execution import ExecutionFactories +from host_agent.planner_context import bind_planner_execution_context @dataclass(frozen=True) @@ -26,19 +27,20 @@ class AssignmentExecutor: *, should_stop: Callable[[], bool] | None = None, ) -> AssignmentExecutionResult: - if should_stop is not None and should_stop(): + with bind_planner_execution_context(assignment): + if should_stop is not None and should_stop(): + return AssignmentExecutionResult( + status="failed", + failure_reason="execution interrupted", + ) + if assignment.workflow_definition_id is not None: + return self._execute_workflow(assignment, should_stop=should_stop) + if assignment.goal is not None: + return self._execute_goal(assignment, should_stop=should_stop) return AssignmentExecutionResult( status="failed", - failure_reason="execution interrupted", + failure_reason="assignment has neither goal nor workflow definition", ) - if assignment.workflow_definition_id is not None: - return self._execute_workflow(assignment, should_stop=should_stop) - if assignment.goal is not None: - return self._execute_goal(assignment, should_stop=should_stop) - return AssignmentExecutionResult( - status="failed", - failure_reason="assignment has neither goal nor workflow definition", - ) def _execute_goal( self, diff --git a/apps/device-host-agent/host_agent/cloud_planner_client.py b/apps/device-host-agent/host_agent/cloud_planner_client.py index 04e449a..1bd63e5 100644 --- a/apps/device-host-agent/host_agent/cloud_planner_client.py +++ b/apps/device-host-agent/host_agent/cloud_planner_client.py @@ -27,6 +27,7 @@ import httpx from cloud.internal_api.models import PlannerDecisionError, PlannerDecisionResponse from host_agent.config import HostAgentConfig +from host_agent.planner_context import current_planner_execution_context from runtime.tool_calling_client import ToolCallDecision, ToolCallUnavailable, ToolCallUsage from runtime.tool_specs import ToolSpec @@ -70,6 +71,15 @@ class CloudProxyToolCallingClient: ], "timeout_seconds": timeout, } + context = current_planner_execution_context() + if context is not None: + payload.update( + { + "task_id": context.task_id, + "attempt": context.attempt, + "lease_id": context.lease_id, + } + ) try: response = self._client.post( f"/internal/v1/hosts/{self.config.host_id}/planner/decide", diff --git a/apps/device-host-agent/host_agent/planner_context.py b/apps/device-host-agent/host_agent/planner_context.py new file mode 100644 index 0000000..3eb5324 --- /dev/null +++ b/apps/device-host-agent/host_agent/planner_context.py @@ -0,0 +1,47 @@ +"""Host-local execution metadata for Cloud planner accounting.""" + +from __future__ import annotations + +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Iterator + + from cloud.internal_api.models import AssignmentModel + + +@dataclass(frozen=True) +class PlannerExecutionContext: + task_id: str + attempt: int + lease_id: str + + +_context: ContextVar[PlannerExecutionContext | None] = ContextVar( + "host_agent_planner_execution_context", + default=None, +) + + +def current_planner_execution_context() -> PlannerExecutionContext | None: + return _context.get() + + +@contextmanager +def bind_planner_execution_context( + assignment: "AssignmentModel", +) -> "Iterator[None]": + token = _context.set( + PlannerExecutionContext( + task_id=assignment.task_id, + attempt=assignment.attempt, + lease_id=assignment.lease_id, + ) + ) + try: + yield + finally: + _context.reset(token) diff --git a/apps/device-host-agent/tests/test_cloud_planner_client.py b/apps/device-host-agent/tests/test_cloud_planner_client.py index de47a5c..90bc19c 100644 --- a/apps/device-host-agent/tests/test_cloud_planner_client.py +++ b/apps/device-host-agent/tests/test_cloud_planner_client.py @@ -7,6 +7,7 @@ import pytest from host_agent.cloud_planner_client import CloudProxyToolCallingClient from host_agent.config import HostAgentConfig +from host_agent.planner_context import PlannerExecutionContext, _context from runtime.tool_calling_client import ToolCallDecision, ToolCallUnavailable from runtime.tool_specs import ToolSpec @@ -78,6 +79,34 @@ def test_decide_base64_encodes_screenshot() -> None: assert body["screenshot_base64"] == "aGVsbG8=" +def test_decide_includes_bound_assignment_context() -> None: + seen_requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + seen_requests.append(request) + return httpx.Response(200, json={"tool_name": "tap", "arguments": {}}) + + client = _client(handler) + token = _context.set( + PlannerExecutionContext(task_id="task-a", attempt=2, lease_id="lease-a") + ) + try: + client.decide( + system_prompt="sp", + user_prompt="up", + screenshot=None, + tools=_TOOLS, + timeout=10.0, + ) + finally: + _context.reset(token) + + body = json.loads(seen_requests[0].content) + assert body["task_id"] == "task-a" + assert body["attempt"] == 2 + assert body["lease_id"] == "lease-a" + + def test_decide_raises_tool_call_unavailable_on_network_error() -> None: def handler(request: httpx.Request) -> httpx.Response: raise httpx.ConnectError("connection refused", request=request) diff --git a/cloud-console/src/api.ts b/cloud-console/src/api.ts index 293d136..bafb6d2 100644 --- a/cloud-console/src/api.ts +++ b/cloud-console/src/api.ts @@ -10,6 +10,7 @@ import type { TaskListResponse, TaskSubmissionPayload, TaskStatus, + TokenUsageEvent, UserListResponse, UserSubmissionPolicy, } from "./types"; @@ -204,6 +205,12 @@ export function getHostTokenUsage(hostId: string): Promise { + return request( + `/v1/hosts/${encodeURIComponent(hostId)}/token-usage-events?limit=20&offset=0`, + ); +} + export function listTasks(options?: { status?: TaskStatus; limit?: number; diff --git a/cloud-console/src/types.ts b/cloud-console/src/types.ts index 7648f32..5eea4aa 100644 --- a/cloud-console/src/types.ts +++ b/cloud-console/src/types.ts @@ -134,3 +134,17 @@ export interface HostTokenUsageSummary { reserved_tokens: number; remaining_tokens: number | null; } + +export interface TokenUsageEvent { + id: string; + host_id: string; + usage_day: string; + task_id: string | null; + attempt: number | null; + provider: string; + model: string; + input_tokens: number | null; + output_tokens: number | null; + total_tokens: number; + occurred_at: string; +} diff --git a/cloud-console/src/views/UsersView.vue b/cloud-console/src/views/UsersView.vue index c8e4618..82a8be1 100644 --- a/cloud-console/src/views/UsersView.vue +++ b/cloud-console/src/views/UsersView.vue @@ -5,6 +5,7 @@ import { createUser, getHostGovernancePolicy, getHostTokenUsage, + listHostTokenUsageEvents, getUserSubmissionPolicy, listDevices, listHosts, @@ -20,6 +21,7 @@ import type { DeviceRecord, HostRecord, HostTokenUsageSummary, + TokenUsageEvent, UserRole, } from "../types"; @@ -54,6 +56,7 @@ const hostSelfSubmissionEnabled = ref(true); const hostMaxActiveTasks = ref(""); const hostDailyTokenBudget = ref(""); const hostUsage = ref(null); +const hostUsageEvents = ref([]); const selectedUser = computed( () => users.value.find((user) => user.id === selectedUserId.value) ?? null, @@ -103,11 +106,13 @@ async function loadHostPolicy() { if (!props.canAdminGovernance || !selectedHostId.value) return; errorMessage.value = ""; try { - const [policy, usage] = await Promise.all([ + const [policy, usage, events] = await Promise.all([ getHostGovernancePolicy(selectedHostId.value), getHostTokenUsage(selectedHostId.value), + listHostTokenUsageEvents(selectedHostId.value), ]); hostUsage.value = usage; + hostUsageEvents.value = events; hostPolicyRevision.value = policy.revision; hostSelfSubmissionEnabled.value = policy.self_submission_enabled; hostMaxActiveTasks.value = policy.max_active_tasks?.toString() ?? ""; @@ -118,7 +123,12 @@ async function loadHostPolicy() { hostSelfSubmissionEnabled.value = true; hostMaxActiveTasks.value = ""; hostDailyTokenBudget.value = ""; - hostUsage.value = await getHostTokenUsage(selectedHostId.value); + const [usage, events] = await Promise.all([ + getHostTokenUsage(selectedHostId.value), + listHostTokenUsageEvents(selectedHostId.value), + ]); + hostUsage.value = usage; + hostUsageEvents.value = events; return; } showError(error, "failed to load Host policy"); @@ -264,7 +274,12 @@ async function saveHostPolicy() { expected_revision: hostPolicyRevision.value ?? 0, }); hostPolicyRevision.value = policy.revision; - hostUsage.value = await getHostTokenUsage(selectedHostId.value); + const [usage, events] = await Promise.all([ + getHostTokenUsage(selectedHostId.value), + listHostTokenUsageEvents(selectedHostId.value), + ]); + hostUsage.value = usage; + hostUsageEvents.value = events; successMessage.value = `Host policy saved (revision ${policy.revision})`; } catch (error) { showError(error, "failed to save Host policy"); @@ -375,6 +390,18 @@ onMounted(() => void refresh()); {{ hostUsage.usage_day }}: used {{ hostUsage.used_tokens }}, reserved {{ hostUsage.reserved_tokens }}, remaining {{ hostUsage.remaining_tokens ?? "unmetered" }} tokens.

+ + + + + + + + + + +
TimeProvider / modelTokensTask
{{ new Date(event.occurred_at).toLocaleString() }}{{ event.provider }} {{ event.model }}{{ event.total_tokens }}{{ event.task_id ?? "local" }} / #{{ event.attempt }}
+

No Cloud-proxy usage events recorded for this Host.

diff --git a/openspec/changes/cloud-console-governance/tasks.md b/openspec/changes/cloud-console-governance/tasks.md index fb3f47f..350abdc 100644 --- a/openspec/changes/cloud-console-governance/tasks.md +++ b/openspec/changes/cloud-console-governance/tasks.md @@ -68,7 +68,7 @@ - [x] 5.1 Extend the dual-provider tool-calling result with optional non-secret provider usage fields while preserving `AIPlanner`'s existing single-decision behavior and direct transport compatibility. -- [ ] 5.2 Extend planner-proxy request context and Host-Agent-local context +- [x] 5.2 Extend planner-proxy request context and Host-Agent-local context binding so Cloud-proxied calls carry known task/attempt metadata without importing Host or Cloud concerns into `runtime`. - [x] 5.3 Add Cloud proxy preflight reservation, configured conservative diff --git a/packages/cloud-platform/cloud/internal_api/api.py b/packages/cloud-platform/cloud/internal_api/api.py index ef0cfbe..9375414 100644 --- a/packages/cloud-platform/cloud/internal_api/api.py +++ b/packages/cloud-platform/cloud/internal_api/api.py @@ -373,6 +373,7 @@ def create_internal_router( status_code=status.HTTP_422_UNPROCESSABLE_CONTENT, detail="planner-decision host_id must match the request path", ) + _validate_planner_context(pool, host_id=host_id, payload=payload) screenshot: bytes | None = None if payload.screenshot_base64 is not None: @@ -401,8 +402,8 @@ def create_internal_router( host_id=host_id, usage_day=now.date().isoformat(), reserved_tokens=planner_token_reservation_ceiling, - task_id=None, - attempt=None, + task_id=payload.task_id, + attempt=payload.attempt, created_at=now, expires_at=now + timedelta(seconds=planner_token_reservation_ttl_seconds), ) @@ -485,6 +486,28 @@ def _validate_assignment_identity( ) +def _validate_planner_context(pool, *, host_id: str, payload: PlannerDecisionRequest) -> None: + context_values = (payload.task_id, payload.attempt, payload.lease_id) + if not any(value is not None for value in context_values): + return + if any(value is None for value in context_values): + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_CONTENT, + detail="planner context requires task_id, attempt, and lease_id together", + ) + attempts = pool.store.list_task_attempts(payload.task_id or "") + if not any( + attempt.attempt == payload.attempt + and attempt.host_id == host_id + and attempt.lease_id == payload.lease_id + for attempt in attempts + ): + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_CONTENT, + detail="planner context does not match a Host assignment", + ) + + def _stale_lease_conflict(detail: str) -> JSONResponse: return JSONResponse( status_code=status.HTTP_409_CONFLICT, diff --git a/packages/cloud-platform/cloud/internal_api/models.py b/packages/cloud-platform/cloud/internal_api/models.py index 0ab1059..48dfa3a 100644 --- a/packages/cloud-platform/cloud/internal_api/models.py +++ b/packages/cloud-platform/cloud/internal_api/models.py @@ -131,6 +131,9 @@ class PlannerDecisionRequest(BaseModel): screenshot_base64: str | None = None tools: list[PlannerToolSpecModel] = Field(default_factory=list) timeout_seconds: float = Field(default=30.0, gt=0, le=120) + task_id: str | None = Field(default=None, min_length=1) + attempt: int | None = Field(default=None, ge=1) + lease_id: str | None = Field(default=None, min_length=1) class PlannerDecisionResponse(BaseModel): diff --git a/packages/cloud-platform/cloud/repository.py b/packages/cloud-platform/cloud/repository.py index 08d3641..d62aab7 100644 --- a/packages/cloud-platform/cloud/repository.py +++ b/packages/cloud-platform/cloud/repository.py @@ -336,6 +336,10 @@ class CloudRepository(Protocol): self, *, host_id: str, usage_day: str, now: datetime, ) -> TokenUsageSummary: ... + def list_host_token_usage_events( + self, *, host_id: str, limit: int, offset: int, + ) -> list[TokenUsageEvent]: ... + def list_reserved_device_ids(self, *, now: datetime) -> set[str]: ... def assign_task( diff --git a/packages/cloud-platform/cloud/sdk/client.py b/packages/cloud-platform/cloud/sdk/client.py index d939db6..5dac689 100644 --- a/packages/cloud-platform/cloud/sdk/client.py +++ b/packages/cloud-platform/cloud/sdk/client.py @@ -240,6 +240,18 @@ class CloudClient: "PUT", f"/hosts/{host_id}/governance-policy", json=policy ).json() + def get_host_token_usage(self, host_id: str) -> dict[str, Any]: + return self._request("GET", f"/hosts/{host_id}/token-usage").json() + + def list_host_token_usage_events( + self, host_id: str, *, limit: int = 50, offset: int = 0 + ) -> list[dict[str, Any]]: + return self._request( + "GET", + f"/hosts/{host_id}/token-usage-events", + params={"limit": limit, "offset": offset}, + ).json() + # ------------------------------------------------------------------ helpers def _url(self, path: str) -> str: diff --git a/packages/cloud-platform/cloud/sdk/governance_api.py b/packages/cloud-platform/cloud/sdk/governance_api.py index 8b7387b..cb71961 100644 --- a/packages/cloud-platform/cloud/sdk/governance_api.py +++ b/packages/cloud-platform/cloud/sdk/governance_api.py @@ -2,7 +2,7 @@ from __future__ import annotations from uuid import uuid4 -from fastapi import APIRouter, HTTPException, Request, status +from fastapi import APIRouter, HTTPException, Query, Request, status from cloud.auth import AuthProvider, GOVERNANCE_ADMIN_SCOPE, GOVERNANCE_READ_SCOPE from cloud.governance import GovernancePolicyConflictError @@ -11,6 +11,7 @@ from cloud.sdk.models import ( HostGovernancePolicyRequest, HostGovernancePolicyResponse, HostTokenUsageSummaryResponse, + TokenUsageEventResponse, UserSubmissionPolicyRequest, UserSubmissionPolicyResponse, ) @@ -137,6 +138,36 @@ def create_governance_router(*, repository, auth_provider: AuthProvider) -> APIR remaining_tokens=summary.remaining_tokens, ) + @router.get( + "/hosts/{host_id}/token-usage-events", + response_model=list[TokenUsageEventResponse], + ) + def list_host_token_usage_events( + host_id: str, + request: Request, + limit: int = Query(default=50, ge=1, le=100), + offset: int = Query(default=0, ge=0), + ) -> list[TokenUsageEventResponse]: + authorize(request, GOVERNANCE_READ_SCOPE) + return [ + TokenUsageEventResponse( + id=event.id, + host_id=event.host_id, + usage_day=event.usage_day, + task_id=event.task_id, + attempt=event.attempt, + provider=event.provider, + model=event.model, + input_tokens=event.input_tokens, + output_tokens=event.output_tokens, + total_tokens=event.total_tokens, + occurred_at=event.occurred_at, + ) + for event in repository.list_host_token_usage_events( + host_id=host_id, limit=limit, offset=offset + ) + ] + @router.put( "/hosts/{host_id}/governance-policy", response_model=HostGovernancePolicyResponse, diff --git a/packages/cloud-platform/cloud/sdk/models.py b/packages/cloud-platform/cloud/sdk/models.py index 6ba6bf6..e98612a 100644 --- a/packages/cloud-platform/cloud/sdk/models.py +++ b/packages/cloud-platform/cloud/sdk/models.py @@ -197,3 +197,17 @@ class HostTokenUsageSummaryResponse(BaseModel): used_tokens: int reserved_tokens: int remaining_tokens: int | None = None + + +class TokenUsageEventResponse(BaseModel): + id: str + host_id: str + usage_day: str + task_id: str | None = None + attempt: int | None = None + provider: str + model: str + input_tokens: int | None = None + output_tokens: int | None = None + total_tokens: int + occurred_at: datetime diff --git a/packages/cloud-platform/cloud/sql_repository.py b/packages/cloud-platform/cloud/sql_repository.py index 2c66418..be739cd 100644 --- a/packages/cloud-platform/cloud/sql_repository.py +++ b/packages/cloud-platform/cloud/sql_repository.py @@ -1046,6 +1046,19 @@ class SQLAlchemyCloudRepository: used_tokens=int(used or 0), reserved_tokens=int(reserved or 0), ) + def list_host_token_usage_events( + self, *, host_id: str, limit: int, offset: int, + ) -> list[Any]: + with self._sessions() as session: + rows = session.scalars( + select(TokenUsageEventRow) + .where(TokenUsageEventRow.host_id == host_id) + .order_by(TokenUsageEventRow.occurred_at.desc(), TokenUsageEventRow.id.desc()) + .limit(limit) + .offset(offset) + ).all() + return [_token_usage_event_from_row(row) for row in rows] + def list_reserved_device_ids(self, *, now: datetime) -> set[str]: with self._sessions() as session: device_ids = session.scalars( diff --git a/tests/test_host_agent_internal_api.py b/tests/test_host_agent_internal_api.py index ca93e37..6ce2ce0 100644 --- a/tests/test_host_agent_internal_api.py +++ b/tests/test_host_agent_internal_api.py @@ -290,6 +290,12 @@ def test_planner_proxy_reserves_and_enforces_host_daily_token_budget(tmp_path) - assert first.json()["total_tokens"] == 2 assert second.status_code == 429 assert planner.calls == 1 + events = pool.store.list_host_token_usage_events( + host_id="host-a", limit=10, offset=0 + ) + assert len(events) == 1 + assert events[0].total_tokens == 2 + assert events[0].task_id is None def test_long_poll_claim_returns_at_most_one_owned_assignment(tmp_path) -> None: