feat(cloud): attribute planner token usage
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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.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)
|
||||
|
||||
@@ -10,6 +10,7 @@ import type {
|
||||
TaskListResponse,
|
||||
TaskSubmissionPayload,
|
||||
TaskStatus,
|
||||
TokenUsageEvent,
|
||||
UserListResponse,
|
||||
UserSubmissionPolicy,
|
||||
} 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?: {
|
||||
status?: TaskStatus;
|
||||
limit?: number;
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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<HostTokenUsageSummary | null>(null);
|
||||
const hostUsageEvents = ref<TokenUsageEvent[]>([]);
|
||||
|
||||
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.
|
||||
</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>
|
||||
</template>
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user