feat(cloud): attribute planner token usage

This commit is contained in:
2026-07-13 23:10:21 +08:00
parent b4803f90e6
commit 40efa53411
16 changed files with 259 additions and 17 deletions
+12 -10
View File
@@ -7,6 +7,7 @@ from typing import Any
from cloud.internal_api.models import AssignmentModel from cloud.internal_api.models import AssignmentModel
from core.models import Task from core.models import Task
from host_agent.execution import ExecutionFactories from host_agent.execution import ExecutionFactories
from host_agent.planner_context import bind_planner_execution_context
@dataclass(frozen=True) @dataclass(frozen=True)
@@ -26,19 +27,20 @@ class AssignmentExecutor:
*, *,
should_stop: Callable[[], bool] | None = None, should_stop: Callable[[], bool] | None = None,
) -> AssignmentExecutionResult: ) -> 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( return AssignmentExecutionResult(
status="failed", 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( def _execute_goal(
self, self,
@@ -27,6 +27,7 @@ import httpx
from cloud.internal_api.models import PlannerDecisionError, PlannerDecisionResponse from cloud.internal_api.models import PlannerDecisionError, PlannerDecisionResponse
from host_agent.config import HostAgentConfig 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_calling_client import ToolCallDecision, ToolCallUnavailable, ToolCallUsage
from runtime.tool_specs import ToolSpec from runtime.tool_specs import ToolSpec
@@ -70,6 +71,15 @@ class CloudProxyToolCallingClient:
], ],
"timeout_seconds": timeout, "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: try:
response = self._client.post( response = self._client.post(
f"/internal/v1/hosts/{self.config.host_id}/planner/decide", f"/internal/v1/hosts/{self.config.host_id}/planner/decide",
@@ -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)
@@ -7,6 +7,7 @@ import pytest
from host_agent.cloud_planner_client import CloudProxyToolCallingClient from host_agent.cloud_planner_client import CloudProxyToolCallingClient
from host_agent.config import HostAgentConfig 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_calling_client import ToolCallDecision, ToolCallUnavailable
from runtime.tool_specs import ToolSpec from runtime.tool_specs import ToolSpec
@@ -78,6 +79,34 @@ def test_decide_base64_encodes_screenshot() -> None:
assert body["screenshot_base64"] == "aGVsbG8=" 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 test_decide_raises_tool_call_unavailable_on_network_error() -> None:
def handler(request: httpx.Request) -> httpx.Response: def handler(request: httpx.Request) -> httpx.Response:
raise httpx.ConnectError("connection refused", request=request) raise httpx.ConnectError("connection refused", request=request)
+7
View File
@@ -10,6 +10,7 @@ import type {
TaskListResponse, TaskListResponse,
TaskSubmissionPayload, TaskSubmissionPayload,
TaskStatus, TaskStatus,
TokenUsageEvent,
UserListResponse, UserListResponse,
UserSubmissionPolicy, UserSubmissionPolicy,
} from "./types"; } from "./types";
@@ -204,6 +205,12 @@ export function getHostTokenUsage(hostId: string): Promise<HostTokenUsageSummary
); );
} }
export function listHostTokenUsageEvents(hostId: string): Promise<TokenUsageEvent[]> {
return request<TokenUsageEvent[]>(
`/v1/hosts/${encodeURIComponent(hostId)}/token-usage-events?limit=20&offset=0`,
);
}
export function listTasks(options?: { export function listTasks(options?: {
status?: TaskStatus; status?: TaskStatus;
limit?: number; limit?: number;
+14
View File
@@ -134,3 +134,17 @@ export interface HostTokenUsageSummary {
reserved_tokens: number; reserved_tokens: number;
remaining_tokens: number | null; 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;
}
+30 -3
View File
@@ -5,6 +5,7 @@ import {
createUser, createUser,
getHostGovernancePolicy, getHostGovernancePolicy,
getHostTokenUsage, getHostTokenUsage,
listHostTokenUsageEvents,
getUserSubmissionPolicy, getUserSubmissionPolicy,
listDevices, listDevices,
listHosts, listHosts,
@@ -20,6 +21,7 @@ import type {
DeviceRecord, DeviceRecord,
HostRecord, HostRecord,
HostTokenUsageSummary, HostTokenUsageSummary,
TokenUsageEvent,
UserRole, UserRole,
} from "../types"; } from "../types";
@@ -54,6 +56,7 @@ const hostSelfSubmissionEnabled = ref(true);
const hostMaxActiveTasks = ref(""); const hostMaxActiveTasks = ref("");
const hostDailyTokenBudget = ref(""); const hostDailyTokenBudget = ref("");
const hostUsage = ref<HostTokenUsageSummary | null>(null); const hostUsage = ref<HostTokenUsageSummary | null>(null);
const hostUsageEvents = ref<TokenUsageEvent[]>([]);
const selectedUser = computed( const selectedUser = computed(
() => users.value.find((user) => user.id === selectedUserId.value) ?? null, () => users.value.find((user) => user.id === selectedUserId.value) ?? null,
@@ -103,11 +106,13 @@ async function loadHostPolicy() {
if (!props.canAdminGovernance || !selectedHostId.value) return; if (!props.canAdminGovernance || !selectedHostId.value) return;
errorMessage.value = ""; errorMessage.value = "";
try { try {
const [policy, usage] = await Promise.all([ const [policy, usage, events] = await Promise.all([
getHostGovernancePolicy(selectedHostId.value), getHostGovernancePolicy(selectedHostId.value),
getHostTokenUsage(selectedHostId.value), getHostTokenUsage(selectedHostId.value),
listHostTokenUsageEvents(selectedHostId.value),
]); ]);
hostUsage.value = usage; hostUsage.value = usage;
hostUsageEvents.value = events;
hostPolicyRevision.value = policy.revision; hostPolicyRevision.value = policy.revision;
hostSelfSubmissionEnabled.value = policy.self_submission_enabled; hostSelfSubmissionEnabled.value = policy.self_submission_enabled;
hostMaxActiveTasks.value = policy.max_active_tasks?.toString() ?? ""; hostMaxActiveTasks.value = policy.max_active_tasks?.toString() ?? "";
@@ -118,7 +123,12 @@ async function loadHostPolicy() {
hostSelfSubmissionEnabled.value = true; hostSelfSubmissionEnabled.value = true;
hostMaxActiveTasks.value = ""; hostMaxActiveTasks.value = "";
hostDailyTokenBudget.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; return;
} }
showError(error, "failed to load Host policy"); showError(error, "failed to load Host policy");
@@ -264,7 +274,12 @@ async function saveHostPolicy() {
expected_revision: hostPolicyRevision.value ?? 0, expected_revision: hostPolicyRevision.value ?? 0,
}); });
hostPolicyRevision.value = policy.revision; 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})`; successMessage.value = `Host policy saved (revision ${policy.revision})`;
} catch (error) { } catch (error) {
showError(error, "failed to save Host policy"); 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 }}, {{ hostUsage.usage_day }}: used {{ hostUsage.used_tokens }}, reserved {{ hostUsage.reserved_tokens }},
remaining {{ hostUsage.remaining_tokens ?? "unmetered" }} tokens. remaining {{ hostUsage.remaining_tokens ?? "unmetered" }} tokens.
</p> </p>
<table v-if="hostUsageEvents.length">
<thead><tr><th>Time</th><th>Provider / model</th><th>Tokens</th><th>Task</th></tr></thead>
<tbody>
<tr v-for="event in hostUsageEvents" :key="event.id">
<td class="dim">{{ new Date(event.occurred_at).toLocaleString() }}</td>
<td>{{ event.provider }} <span class="dim">{{ event.model }}</span></td>
<td>{{ event.total_tokens }}</td>
<td class="dim">{{ event.task_id ?? "local" }}<span v-if="event.attempt"> / #{{ event.attempt }}</span></td>
</tr>
</tbody>
</table>
<p v-else class="muted">No Cloud-proxy usage events recorded for this Host.</p>
</div> </div>
</div> </div>
</template> </template>
@@ -68,7 +68,7 @@
- [x] 5.1 Extend the dual-provider tool-calling result with optional - [x] 5.1 Extend the dual-provider tool-calling result with optional
non-secret provider usage fields while preserving `AIPlanner`'s existing non-secret provider usage fields while preserving `AIPlanner`'s existing
single-decision behavior and direct transport compatibility. 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 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`.
- [x] 5.3 Add Cloud proxy preflight reservation, configured conservative - [x] 5.3 Add Cloud proxy preflight reservation, configured conservative
@@ -373,6 +373,7 @@ def create_internal_router(
status_code=status.HTTP_422_UNPROCESSABLE_CONTENT, status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
detail="planner-decision host_id must match the request path", detail="planner-decision host_id must match the request path",
) )
_validate_planner_context(pool, host_id=host_id, payload=payload)
screenshot: bytes | None = None screenshot: bytes | None = None
if payload.screenshot_base64 is not None: if payload.screenshot_base64 is not None:
@@ -401,8 +402,8 @@ def create_internal_router(
host_id=host_id, host_id=host_id,
usage_day=now.date().isoformat(), usage_day=now.date().isoformat(),
reserved_tokens=planner_token_reservation_ceiling, reserved_tokens=planner_token_reservation_ceiling,
task_id=None, task_id=payload.task_id,
attempt=None, attempt=payload.attempt,
created_at=now, created_at=now,
expires_at=now + timedelta(seconds=planner_token_reservation_ttl_seconds), 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: def _stale_lease_conflict(detail: str) -> JSONResponse:
return JSONResponse( return JSONResponse(
status_code=status.HTTP_409_CONFLICT, status_code=status.HTTP_409_CONFLICT,
@@ -131,6 +131,9 @@ class PlannerDecisionRequest(BaseModel):
screenshot_base64: str | None = None screenshot_base64: str | None = None
tools: list[PlannerToolSpecModel] = Field(default_factory=list) tools: list[PlannerToolSpecModel] = Field(default_factory=list)
timeout_seconds: float = Field(default=30.0, gt=0, le=120) 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): class PlannerDecisionResponse(BaseModel):
@@ -336,6 +336,10 @@ class CloudRepository(Protocol):
self, *, host_id: str, usage_day: str, now: datetime, self, *, host_id: str, usage_day: str, now: datetime,
) -> TokenUsageSummary: ... ) -> 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 list_reserved_device_ids(self, *, now: datetime) -> set[str]: ...
def assign_task( def assign_task(
@@ -240,6 +240,18 @@ class CloudClient:
"PUT", f"/hosts/{host_id}/governance-policy", json=policy "PUT", f"/hosts/{host_id}/governance-policy", json=policy
).json() ).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 # ------------------------------------------------------------------ helpers
def _url(self, path: str) -> str: def _url(self, path: str) -> str:
@@ -2,7 +2,7 @@ from __future__ import annotations
from uuid import uuid4 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.auth import AuthProvider, GOVERNANCE_ADMIN_SCOPE, GOVERNANCE_READ_SCOPE
from cloud.governance import GovernancePolicyConflictError from cloud.governance import GovernancePolicyConflictError
@@ -11,6 +11,7 @@ from cloud.sdk.models import (
HostGovernancePolicyRequest, HostGovernancePolicyRequest,
HostGovernancePolicyResponse, HostGovernancePolicyResponse,
HostTokenUsageSummaryResponse, HostTokenUsageSummaryResponse,
TokenUsageEventResponse,
UserSubmissionPolicyRequest, UserSubmissionPolicyRequest,
UserSubmissionPolicyResponse, UserSubmissionPolicyResponse,
) )
@@ -137,6 +138,36 @@ def create_governance_router(*, repository, auth_provider: AuthProvider) -> APIR
remaining_tokens=summary.remaining_tokens, 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( @router.put(
"/hosts/{host_id}/governance-policy", "/hosts/{host_id}/governance-policy",
response_model=HostGovernancePolicyResponse, response_model=HostGovernancePolicyResponse,
@@ -197,3 +197,17 @@ class HostTokenUsageSummaryResponse(BaseModel):
used_tokens: int used_tokens: int
reserved_tokens: int reserved_tokens: int
remaining_tokens: int | None = None 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
@@ -1046,6 +1046,19 @@ class SQLAlchemyCloudRepository:
used_tokens=int(used or 0), reserved_tokens=int(reserved or 0), 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]: 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(
+6
View File
@@ -290,6 +290,12 @@ def test_planner_proxy_reserves_and_enforces_host_daily_token_budget(tmp_path) -
assert first.json()["total_tokens"] == 2 assert first.json()["total_tokens"] == 2
assert second.status_code == 429 assert second.status_code == 429
assert planner.calls == 1 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: def test_long_poll_claim_returns_at_most_one_owned_assignment(tmp_path) -> None: