@@ -12,9 +12,15 @@ from cloud.internal_api.api import create_internal_router
|
||||
from cloud.pool import DevicePool, PooledDevice
|
||||
from cloud.scheduler import ScheduledTask, TaskConstraints, TaskScheduler
|
||||
from cloud.store import CloudStore
|
||||
from runtime.tool_calling_client import ToolCallDecision, ToolCallUsage
|
||||
|
||||
|
||||
def _build_client(tmp_path) -> tuple[TestClient, DevicePool]:
|
||||
def _build_client(
|
||||
tmp_path,
|
||||
*,
|
||||
planner_client_factory=None,
|
||||
planner_token_reservation_ceiling: int = 4096,
|
||||
) -> tuple[TestClient, DevicePool]:
|
||||
pool = DevicePool(
|
||||
CloudStore(tmp_path / "internal.sqlite3"),
|
||||
CloudConfig(stale_after_seconds=60),
|
||||
@@ -39,6 +45,8 @@ def _build_client(tmp_path) -> tuple[TestClient, DevicePool]:
|
||||
pool=pool,
|
||||
auth_provider=auth_provider,
|
||||
scheduler=TaskScheduler(pool, pool.store, CloudConfig(stale_after_seconds=60)),
|
||||
planner_client_factory=planner_client_factory,
|
||||
planner_token_reservation_ceiling=planner_token_reservation_ceiling,
|
||||
)
|
||||
)
|
||||
return TestClient(app), pool
|
||||
@@ -232,6 +240,58 @@ def test_heartbeat_and_self_submission_preserve_host_isolation(tmp_path) -> None
|
||||
assert foreign.status_code == 403
|
||||
|
||||
|
||||
def test_planner_proxy_reserves_and_enforces_host_daily_token_budget(tmp_path) -> None:
|
||||
class FakePlannerClient:
|
||||
calls = 0
|
||||
|
||||
def decide(self, **_kwargs):
|
||||
self.calls += 1
|
||||
return ToolCallDecision(
|
||||
tool_name="tap",
|
||||
arguments={"x": 1, "y": 2},
|
||||
usage=ToolCallUsage(input_tokens=1, output_tokens=1, total_tokens=2),
|
||||
)
|
||||
|
||||
planner = FakePlannerClient()
|
||||
client, pool = _build_client(
|
||||
tmp_path,
|
||||
planner_client_factory=lambda: planner,
|
||||
planner_token_reservation_ceiling=5,
|
||||
)
|
||||
headers = {"Authorization": "Bearer token-a"}
|
||||
client.put(
|
||||
"/internal/v1/hosts/host-a/heartbeat",
|
||||
headers=headers,
|
||||
json=_heartbeat_payload("host-a", "device-a"),
|
||||
)
|
||||
pool.store.upsert_host_governance_policy(
|
||||
host_id="host-a",
|
||||
self_submission_enabled=True,
|
||||
max_active_tasks=None,
|
||||
daily_token_budget=5,
|
||||
updated_at=datetime.now(UTC),
|
||||
)
|
||||
payload = {
|
||||
"host_id": "host-a",
|
||||
"system_prompt": "system",
|
||||
"user_prompt": "user",
|
||||
"tools": [{"name": "tap", "description": "tap", "parameters": {}}],
|
||||
"timeout_seconds": 1,
|
||||
}
|
||||
|
||||
first = client.post(
|
||||
"/internal/v1/hosts/host-a/planner/decide", headers=headers, json=payload
|
||||
)
|
||||
second = client.post(
|
||||
"/internal/v1/hosts/host-a/planner/decide", headers=headers, json=payload
|
||||
)
|
||||
|
||||
assert first.status_code == 200, first.text
|
||||
assert first.json()["total_tokens"] == 2
|
||||
assert second.status_code == 429
|
||||
assert planner.calls == 1
|
||||
|
||||
|
||||
def test_long_poll_claim_returns_at_most_one_owned_assignment(tmp_path) -> None:
|
||||
client, pool = _build_client(tmp_path)
|
||||
now = datetime.now(UTC)
|
||||
|
||||
Reference in New Issue
Block a user