Compare commits
3
Commits
b4803f90e6
...
b613a315ff
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
b613a315ff | ||
|
|
5b64efab53 | ||
|
|
40efa53411 |
@@ -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,
|
||||||
|
|||||||
@@ -157,6 +157,7 @@ class HostAgentClient:
|
|||||||
"address": address,
|
"address": address,
|
||||||
"devices": [device.model_dump(mode="json") for device in devices],
|
"devices": [device.model_dump(mode="json") for device in devices],
|
||||||
"policy_revision": policy_revision,
|
"policy_revision": policy_revision,
|
||||||
|
"planner_transport": self.config.ai_planner_transport,
|
||||||
},
|
},
|
||||||
)
|
)
|
||||||
return HeartbeatResponse.model_validate(response.json())
|
return HeartbeatResponse.model_validate(response.json())
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|||||||
@@ -20,6 +20,13 @@ Accounts have fixed roles:
|
|||||||
- `operator`: viewer access plus task submission APIs
|
- `operator`: viewer access plus task submission APIs
|
||||||
- `admin`: all API scopes
|
- `admin`: all API scopes
|
||||||
|
|
||||||
|
Administrators can use **Users & limits** to manage accounts, restrict task
|
||||||
|
submission to explicit Host/Device targets, configure Host self-submission and
|
||||||
|
active-task limits, and inspect non-secret Cloud-proxy usage. Daily token
|
||||||
|
budgets are enforced only for Hosts reporting `AI_PLANNER_TRANSPORT=cloud`;
|
||||||
|
direct-provider Hosts are labelled **unmetered** rather than budget compliant.
|
||||||
|
The configured proxy reservation ceiling must fit within any daily budget.
|
||||||
|
|
||||||
## Local development
|
## Local development
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
|
|||||||
@@ -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;
|
||||||
|
|||||||
@@ -63,6 +63,7 @@ export interface HostRecord {
|
|||||||
host_id: string;
|
host_id: string;
|
||||||
address: string | null;
|
address: string | null;
|
||||||
last_seen_at: string;
|
last_seen_at: string;
|
||||||
|
planner_transport: "direct" | "cloud";
|
||||||
}
|
}
|
||||||
|
|
||||||
export type PluginEntryPointKind = "driver" | "tool" | "skill";
|
export type PluginEntryPointKind = "driver" | "tool" | "skill";
|
||||||
@@ -134,3 +135,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;
|
||||||
|
}
|
||||||
|
|||||||
@@ -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,10 +56,14 @@ 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,
|
||||||
);
|
);
|
||||||
|
const selectedHost = computed(
|
||||||
|
() => hosts.value.find((host) => host.host_id === selectedHostId.value) ?? null,
|
||||||
|
);
|
||||||
const deviceKey = (device: Pick<DeviceRecord, "host_id" | "device_id">) =>
|
const deviceKey = (device: Pick<DeviceRecord, "host_id" | "device_id">) =>
|
||||||
`${device.host_id}\u0000${device.device_id}`;
|
`${device.host_id}\u0000${device.device_id}`;
|
||||||
|
|
||||||
@@ -103,11 +109,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 +126,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 +277,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");
|
||||||
@@ -371,10 +389,25 @@ onMounted(() => void refresh());
|
|||||||
<label>Daily Cloud-proxy token budget <input v-model="hostDailyTokenBudget" inputmode="numeric" placeholder="Unlimited" /></label>
|
<label>Daily Cloud-proxy token budget <input v-model="hostDailyTokenBudget" inputmode="numeric" placeholder="Unlimited" /></label>
|
||||||
</div>
|
</div>
|
||||||
<div class="actions"><button class="primary" :disabled="saving || !selectedHostId" @click="saveHostPolicy">Save Host policy {{ hostPolicyRevision === null ? "" : `(revision ${hostPolicyRevision})` }}</button></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">
|
<p v-if="selectedHost?.planner_transport === 'direct'" class="muted">
|
||||||
|
Direct provider transport: unmetered. Cloud cannot enforce or verify this Host's token budget.
|
||||||
|
</p>
|
||||||
|
<p v-else-if="hostUsage" class="muted">
|
||||||
{{ 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>
|
||||||
|
|||||||
@@ -229,7 +229,9 @@ The repository ships an independent Vue 3 + Vite SPA at `cloud-console/` that
|
|||||||
renders task history, devices, hosts, and plugins. Human operators sign in with
|
renders task history, devices, hosts, and plugins. Human operators sign in with
|
||||||
a username and password; the Cloud API creates an expiring, revocable
|
a username and password; the Cloud API creates an expiring, revocable
|
||||||
`HttpOnly` session cookie and uses a separate CSRF cookie/header for writes.
|
`HttpOnly` session cookie and uses a separate CSRF cookie/header for writes.
|
||||||
The Console has no bearer-token fallback or user-directory view.
|
An administrator can manage accounts, submission policies, Host operational
|
||||||
|
limits, and non-secret Cloud-proxy token usage from **Users & limits**. Bearer
|
||||||
|
tokens remain a compatibility path for API/SDK automation.
|
||||||
|
|
||||||
### HTTPS and session configuration
|
### HTTPS and session configuration
|
||||||
|
|
||||||
@@ -396,6 +398,34 @@ mode:
|
|||||||
`AI_PLANNER_TRANSPORT` unset or `direct` preserves the existing
|
`AI_PLANNER_TRANSPORT` unset or `direct` preserves the existing
|
||||||
direct-to-provider behavior with no change.
|
direct-to-provider behavior with no change.
|
||||||
|
|
||||||
|
### Host governance and Cloud-proxy token budgets
|
||||||
|
|
||||||
|
After deploying the API migration, an administrator can use **Users & limits**
|
||||||
|
in `/console/` to restrict a user's task submission to explicit Host/Device
|
||||||
|
targets, enable or disable Host self-submission, and set a Host active-task
|
||||||
|
limit. Policies are enforced by the API and scheduler; hiding Console controls
|
||||||
|
is not an authorization boundary.
|
||||||
|
|
||||||
|
Cloud-enforced token budgets require `AI_PLANNER_TRANSPORT=cloud`. Before each
|
||||||
|
proxy decision, the API atomically reserves
|
||||||
|
`CLOUD_PLANNER_TOKEN_RESERVATION_CEILING` tokens (default `4096`) for the
|
||||||
|
current UTC day. Set `CLOUD_PLANNER_TOKEN_RESERVATION_TTL_SECONDS` (default
|
||||||
|
`300`) to bound an unknown-usage reservation after provider/transport failure.
|
||||||
|
The daily Host budget must accommodate the reservation ceiling; otherwise the
|
||||||
|
proxy rejects before calling the provider. On a provider response, the
|
||||||
|
reservation is settled to reported usage and the Console retains only timestamp,
|
||||||
|
provider/model, token counts, and optional task/attempt identifiers.
|
||||||
|
|
||||||
|
Hosts reporting `AI_PLANNER_TRANSPORT=direct` are explicitly shown as
|
||||||
|
**unmetered**. Cloud cannot enforce or verify their provider token use. Do not
|
||||||
|
interpret an unmetered Host's absence of usage events as budget compliance.
|
||||||
|
|
||||||
|
Roll out in this order: migrate the Cloud database, deploy the Cloud API,
|
||||||
|
switch a pilot Host to Cloud transport, configure a budget above the reservation
|
||||||
|
ceiling, then review its usage events before enabling budgets fleet-wide. A
|
||||||
|
rollback to direct transport requires valid provider credentials on that Host;
|
||||||
|
preserve usage and policy rows rather than deleting accounting history.
|
||||||
|
|
||||||
## Operational Limitations
|
## Operational Limitations
|
||||||
|
|
||||||
Run exactly one scheduler-enabled Cloud API process. SQLite supports only the
|
Run exactly one scheduler-enabled Cloud API process. SQLite supports only the
|
||||||
|
|||||||
@@ -12,16 +12,16 @@
|
|||||||
|
|
||||||
## 2. Governance domain, persistence, and migration
|
## 2. Governance domain, persistence, and migration
|
||||||
|
|
||||||
- [ ] 2.1 Define target-selector, user-submission-policy, Host-governance-
|
- [x] 2.1 Define target-selector, user-submission-policy, Host-governance-
|
||||||
policy, token-usage-event, and token-reservation domain models with
|
policy, token-usage-event, and token-reservation domain models with
|
||||||
non-secret representations and revision semantics.
|
non-secret representations and revision semantics.
|
||||||
- [ ] 2.2 Extend the Cloud repository port with transactional CRUD/query
|
- [x] 2.2 Extend the Cloud repository port with transactional CRUD/query
|
||||||
operations for policies, active Host capacity, budget reservations,
|
operations for policies, active Host capacity, budget reservations,
|
||||||
settlement, expiry cleanup, and bounded usage summaries/events.
|
settlement, expiry cleanup, and bounded usage summaries/events.
|
||||||
- [ ] 2.3 Add SQLAlchemy rows, indexes, conversion helpers, and concurrency-
|
- [x] 2.3 Add SQLAlchemy rows, indexes, conversion helpers, and concurrency-
|
||||||
safe PostgreSQL/SQLite implementations for governance policies, usage
|
safe PostgreSQL/SQLite implementations for governance policies, usage
|
||||||
events, reservations, and safe policy audits.
|
events, reservations, and safe policy audits.
|
||||||
- [ ] 2.4 Add an Alembic forward/downgrade revision that preserves existing
|
- [x] 2.4 Add an Alembic forward/downgrade revision that preserves existing
|
||||||
users, Hosts, devices, tasks, and attempts; advance schema readiness
|
users, Hosts, devices, tasks, and attempts; advance schema readiness
|
||||||
checks to the new head.
|
checks to the new head.
|
||||||
- [ ] 2.5 Add repository and migration tests for policy revisions, null versus
|
- [ ] 2.5 Add repository and migration tests for policy revisions, null versus
|
||||||
@@ -40,7 +40,7 @@
|
|||||||
- [x] 3.3 Add `governance:read` and `governance:admin` scopes and policy-aware
|
- [x] 3.3 Add `governance:read` and `governance:admin` scopes and policy-aware
|
||||||
public task authorization that combines `tasks:submit` with the
|
public task authorization that combines `tasks:submit` with the
|
||||||
authenticated human user's effective submission policy.
|
authenticated human user's effective submission policy.
|
||||||
- [ ] 3.4 Add bounded, non-secret public governance routes and `CloudClient`
|
- [x] 3.4 Add bounded, non-secret public governance routes and `CloudClient`
|
||||||
methods for user policy, Host policy, Host AI-budget summaries, and
|
methods for user policy, Host policy, Host AI-budget summaries, and
|
||||||
paginated usage events; audit every policy mutation.
|
paginated usage events; audit every policy mutation.
|
||||||
- [ ] 3.5 Add public API/SDK tests for targeted submission, target-policy
|
- [ ] 3.5 Add public API/SDK tests for targeted submission, target-policy
|
||||||
@@ -59,7 +59,7 @@
|
|||||||
credentials and validate any named local Device ownership.
|
credentials and validate any named local Device ownership.
|
||||||
- [x] 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,
|
- [x] 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
|
||||||
rejection, disabled self-submission, and outbound-only compatibility.
|
rejection, disabled self-submission, and outbound-only compatibility.
|
||||||
|
|
||||||
@@ -68,13 +68,13 @@
|
|||||||
- [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
|
||||||
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
|
- [x] 5.4 Record non-secret usage events and expose accurate
|
||||||
used/reserved/remaining UTC-day budget summaries; explicitly report
|
used/reserved/remaining UTC-day budget summaries; explicitly report
|
||||||
direct transport as unmetered.
|
direct transport as unmetered.
|
||||||
- [ ] 5.5 Add provider-fake, repository concurrency, Cloud API, and Host
|
- [ ] 5.5 Add provider-fake, repository concurrency, Cloud API, and Host
|
||||||
@@ -89,7 +89,7 @@
|
|||||||
validation errors.
|
validation errors.
|
||||||
- [x] 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
|
- [x] 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,
|
||||||
screenshots, provider credentials, cookies, or lease secrets.
|
screenshots, provider credentials, cookies, or lease secrets.
|
||||||
- [ ] 6.4 Add frontend tests for task composer scope/policy failures,
|
- [ ] 6.4 Add frontend tests for task composer scope/policy failures,
|
||||||
@@ -98,7 +98,7 @@
|
|||||||
|
|
||||||
## 7. Documentation, verification, and rollout
|
## 7. Documentation, verification, and rollout
|
||||||
|
|
||||||
- [ ] 7.1 Update deployment and Console documentation with the dependency
|
- [x] 7.1 Update deployment and Console documentation with the dependency
|
||||||
order, migration/rollback sequence, policy semantics, Cloud transport
|
order, migration/rollback sequence, policy semantics, Cloud transport
|
||||||
prerequisite, budget reservation behavior, direct-host limitation, and
|
prerequisite, budget reservation behavior, direct-host limitation, and
|
||||||
safe administrator operations.
|
safe administrator operations.
|
||||||
|
|||||||
@@ -51,6 +51,9 @@ class HostRow(Base):
|
|||||||
display_name: Mapped[str | None] = mapped_column(String, nullable=True)
|
display_name: Mapped[str | None] = mapped_column(String, nullable=True)
|
||||||
enrolled_at: Mapped[str | None] = mapped_column(String, nullable=True)
|
enrolled_at: Mapped[str | None] = mapped_column(String, nullable=True)
|
||||||
revoked_at: Mapped[str | None] = mapped_column(String, nullable=True)
|
revoked_at: Mapped[str | None] = mapped_column(String, nullable=True)
|
||||||
|
planner_transport: Mapped[str] = mapped_column(
|
||||||
|
String, nullable=False, default="direct", server_default=text("'direct'")
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class DeviceEnrollmentRow(Base):
|
class DeviceEnrollmentRow(Base):
|
||||||
|
|||||||
@@ -177,6 +177,7 @@ def create_internal_router(
|
|||||||
devices,
|
devices,
|
||||||
address=payload.address,
|
address=payload.address,
|
||||||
allow_device_takeover=allow_device_takeover,
|
allow_device_takeover=allow_device_takeover,
|
||||||
|
planner_transport=payload.planner_transport,
|
||||||
)
|
)
|
||||||
policy = pool.store.get_host_governance_policy(host_id)
|
policy = pool.store.get_host_governance_policy(host_id)
|
||||||
policy_revision = policy.revision if policy is not None else 0
|
policy_revision = policy.revision if policy is not None else 0
|
||||||
@@ -373,6 +374,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 +403,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 +487,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,
|
||||||
|
|||||||
@@ -39,6 +39,7 @@ class HeartbeatRequest(BaseModel):
|
|||||||
address: str | None = None
|
address: str | None = None
|
||||||
devices: list[DeviceSnapshotModel] = Field(default_factory=list)
|
devices: list[DeviceSnapshotModel] = Field(default_factory=list)
|
||||||
policy_revision: int = Field(default=0, ge=0)
|
policy_revision: int = Field(default=0, ge=0)
|
||||||
|
planner_transport: Literal["direct", "cloud"] = "direct"
|
||||||
|
|
||||||
|
|
||||||
class HostGovernancePolicyModel(BaseModel):
|
class HostGovernancePolicyModel(BaseModel):
|
||||||
@@ -131,6 +132,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):
|
||||||
|
|||||||
@@ -0,0 +1,23 @@
|
|||||||
|
"""Record each Host's planner transport for governance visibility."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from alembic import op
|
||||||
|
import sqlalchemy as sa
|
||||||
|
|
||||||
|
|
||||||
|
revision = "0006_host_planner_transport"
|
||||||
|
down_revision = "0005_cloud_token_usage"
|
||||||
|
branch_labels = None
|
||||||
|
depends_on = None
|
||||||
|
|
||||||
|
|
||||||
|
def upgrade() -> None:
|
||||||
|
op.add_column(
|
||||||
|
"host_registrations",
|
||||||
|
sa.Column("planner_transport", sa.String(), nullable=False, server_default="direct"),
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def downgrade() -> None:
|
||||||
|
op.drop_column("host_registrations", "planner_transport")
|
||||||
@@ -34,6 +34,7 @@ class HostRegistration:
|
|||||||
host_id: str
|
host_id: str
|
||||||
address: str | None
|
address: str | None
|
||||||
last_seen_at: datetime
|
last_seen_at: datetime
|
||||||
|
planner_transport: Literal["direct", "cloud"] = "direct"
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True)
|
@dataclass(frozen=True)
|
||||||
@@ -61,6 +62,7 @@ class DevicePool:
|
|||||||
snapshot: list[Device],
|
snapshot: list[Device],
|
||||||
*,
|
*,
|
||||||
address: str | None = None,
|
address: str | None = None,
|
||||||
|
planner_transport: Literal["direct", "cloud"] = "direct",
|
||||||
allow_device_takeover: bool = False,
|
allow_device_takeover: bool = False,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Push a host's current device snapshot into the pool.
|
"""Push a host's current device snapshot into the pool.
|
||||||
@@ -70,7 +72,12 @@ class DevicePool:
|
|||||||
other hosts are untouched.
|
other hosts are untouched.
|
||||||
"""
|
"""
|
||||||
now = utc_now()
|
now = utc_now()
|
||||||
self.store.upsert_host(host_id, address=address, last_seen_at=now)
|
self.store.upsert_host(
|
||||||
|
host_id,
|
||||||
|
address=address,
|
||||||
|
last_seen_at=now,
|
||||||
|
planner_transport=planner_transport,
|
||||||
|
)
|
||||||
devices = [self._to_pooled(device, host_id, now) for device in snapshot]
|
devices = [self._to_pooled(device, host_id, now) for device in snapshot]
|
||||||
if allow_device_takeover:
|
if allow_device_takeover:
|
||||||
self.store.replace_host_devices(
|
self.store.replace_host_devices(
|
||||||
|
|||||||
@@ -154,6 +154,7 @@ class CloudRepository(Protocol):
|
|||||||
*,
|
*,
|
||||||
address: str | None,
|
address: str | None,
|
||||||
last_seen_at: datetime,
|
last_seen_at: datetime,
|
||||||
|
planner_transport: Literal["direct", "cloud"] = "direct",
|
||||||
) -> None: ...
|
) -> None: ...
|
||||||
|
|
||||||
def replace_host_devices(
|
def replace_host_devices(
|
||||||
@@ -336,6 +337,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(
|
||||||
|
|||||||
@@ -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 = "0005_cloud_token_usage"
|
HEAD_REVISION = "0006_host_planner_transport"
|
||||||
|
|
||||||
|
|
||||||
class SchemaVersionError(RuntimeError):
|
class SchemaVersionError(RuntimeError):
|
||||||
|
|||||||
@@ -240,6 +240,7 @@ def create_cloud_router(
|
|||||||
host_id=h.host_id,
|
host_id=h.host_id,
|
||||||
address=h.address,
|
address=h.address,
|
||||||
last_seen_at=h.last_seen_at.isoformat() if h.last_seen_at else "",
|
last_seen_at=h.last_seen_at.isoformat() if h.last_seen_at else "",
|
||||||
|
planner_transport=h.planner_transport,
|
||||||
)
|
)
|
||||||
for h in pool.list_hosts()
|
for h in pool.list_hosts()
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -86,6 +86,7 @@ class HostResponse(BaseModel):
|
|||||||
host_id: str
|
host_id: str
|
||||||
address: str | None = None
|
address: str | None = None
|
||||||
last_seen_at: str
|
last_seen_at: str
|
||||||
|
planner_transport: Literal["direct", "cloud"] = "direct"
|
||||||
|
|
||||||
|
|
||||||
class PluginRegistrationRequest(BaseModel):
|
class PluginRegistrationRequest(BaseModel):
|
||||||
@@ -197,3 +198,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
|
||||||
|
|||||||
@@ -266,6 +266,7 @@ class SQLAlchemyCloudRepository:
|
|||||||
*,
|
*,
|
||||||
address: str | None,
|
address: str | None,
|
||||||
last_seen_at: datetime,
|
last_seen_at: datetime,
|
||||||
|
planner_transport: str = "direct",
|
||||||
) -> None:
|
) -> None:
|
||||||
with self._sessions.begin() as session:
|
with self._sessions.begin() as session:
|
||||||
row = session.get(HostRow, host_id)
|
row = session.get(HostRow, host_id)
|
||||||
@@ -275,12 +276,14 @@ class SQLAlchemyCloudRepository:
|
|||||||
host_id=host_id,
|
host_id=host_id,
|
||||||
address=address,
|
address=address,
|
||||||
last_seen_at=_iso(last_seen_at),
|
last_seen_at=_iso(last_seen_at),
|
||||||
|
planner_transport=planner_transport,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
return
|
return
|
||||||
if address is not None:
|
if address is not None:
|
||||||
row.address = address
|
row.address = address
|
||||||
row.last_seen_at = _iso(last_seen_at)
|
row.last_seen_at = _iso(last_seen_at)
|
||||||
|
row.planner_transport = planner_transport
|
||||||
|
|
||||||
def replace_host_devices(
|
def replace_host_devices(
|
||||||
self,
|
self,
|
||||||
@@ -1046,6 +1049,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(
|
||||||
@@ -1406,6 +1422,9 @@ def _host_from_row(row: HostRow) -> Any:
|
|||||||
host_id=row.host_id,
|
host_id=row.host_id,
|
||||||
address=row.address,
|
address=row.address,
|
||||||
last_seen_at=_parse_dt(row.last_seen_at) or utc_now(),
|
last_seen_at=_parse_dt(row.last_seen_at) or utc_now(),
|
||||||
|
planner_transport=(
|
||||||
|
row.planner_transport if row.planner_transport in {"direct", "cloud"} else "direct"
|
||||||
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -91,6 +91,19 @@ def test_authenticated_heartbeat_replaces_complete_snapshot(tmp_path) -> None:
|
|||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def test_heartbeat_records_host_planner_transport(tmp_path) -> None:
|
||||||
|
client, pool = _build_client(tmp_path)
|
||||||
|
|
||||||
|
response = client.put(
|
||||||
|
"/internal/v1/hosts/host-a/heartbeat",
|
||||||
|
headers={"Authorization": "Bearer token-a"},
|
||||||
|
json={**_heartbeat_payload("host-a", "device-a"), "planner_transport": "cloud"},
|
||||||
|
)
|
||||||
|
|
||||||
|
assert response.status_code == 200
|
||||||
|
assert pool.store.get_host("host-a").planner_transport == "cloud" # type: ignore[union-attr]
|
||||||
|
|
||||||
|
|
||||||
def test_empty_heartbeat_removes_only_reporting_hosts_devices(tmp_path) -> None:
|
def test_empty_heartbeat_removes_only_reporting_hosts_devices(tmp_path) -> None:
|
||||||
client, pool = _build_client(tmp_path)
|
client, pool = _build_client(tmp_path)
|
||||||
client.put(
|
client.put(
|
||||||
@@ -240,6 +253,49 @@ def test_heartbeat_and_self_submission_preserve_host_isolation(tmp_path) -> None
|
|||||||
assert foreign.status_code == 403
|
assert foreign.status_code == 403
|
||||||
|
|
||||||
|
|
||||||
|
def test_host_policy_converges_and_disables_self_submission(tmp_path) -> None:
|
||||||
|
client, pool = _build_client(tmp_path)
|
||||||
|
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=False,
|
||||||
|
max_active_tasks=2,
|
||||||
|
daily_token_budget=1000,
|
||||||
|
updated_at=datetime.now(UTC),
|
||||||
|
)
|
||||||
|
|
||||||
|
stale = client.put(
|
||||||
|
"/internal/v1/hosts/host-a/heartbeat",
|
||||||
|
headers=headers,
|
||||||
|
json={**_heartbeat_payload("host-a", "device-a"), "policy_revision": 0},
|
||||||
|
)
|
||||||
|
assert stale.status_code == 200
|
||||||
|
assert stale.json()["policy"] == {
|
||||||
|
"revision": 1,
|
||||||
|
"self_submission_enabled": False,
|
||||||
|
"max_active_tasks": 2,
|
||||||
|
"daily_token_budget": 1000,
|
||||||
|
}
|
||||||
|
current = client.put(
|
||||||
|
"/internal/v1/hosts/host-a/heartbeat",
|
||||||
|
headers=headers,
|
||||||
|
json={**_heartbeat_payload("host-a", "device-a"), "policy_revision": 1},
|
||||||
|
)
|
||||||
|
assert current.status_code == 200
|
||||||
|
assert current.json()["policy"] is None
|
||||||
|
disabled = client.post(
|
||||||
|
"/internal/v1/hosts/host-a/tasks",
|
||||||
|
headers=headers,
|
||||||
|
json={"host_id": "host-a", "goal": "should be rejected"},
|
||||||
|
)
|
||||||
|
assert disabled.status_code == 403
|
||||||
|
|
||||||
|
|
||||||
def test_planner_proxy_reserves_and_enforces_host_daily_token_budget(tmp_path) -> None:
|
def test_planner_proxy_reserves_and_enforces_host_daily_token_budget(tmp_path) -> None:
|
||||||
class FakePlannerClient:
|
class FakePlannerClient:
|
||||||
calls = 0
|
calls = 0
|
||||||
@@ -290,6 +346,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:
|
||||||
|
|||||||
Reference in New Issue
Block a user