618 lines
19 KiB
Python
618 lines
19 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import UTC, datetime
|
|
from datetime import timedelta
|
|
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
|
|
from cloud.auth import BearerCredential, ConfiguredBearerAuthProvider
|
|
from cloud.config import CloudConfig
|
|
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,
|
|
*,
|
|
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),
|
|
)
|
|
auth_provider = ConfiguredBearerAuthProvider(
|
|
[
|
|
BearerCredential(
|
|
principal_id="agent-a",
|
|
token="token-a",
|
|
host_id="host-a",
|
|
),
|
|
BearerCredential(
|
|
principal_id="agent-b",
|
|
token="token-b",
|
|
host_id="host-b",
|
|
),
|
|
]
|
|
)
|
|
app = FastAPI()
|
|
app.include_router(
|
|
create_internal_router(
|
|
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
|
|
|
|
|
|
def _heartbeat_payload(host_id: str, *device_ids: str) -> dict[str, object]:
|
|
return {
|
|
"host_id": host_id,
|
|
"address": f"{host_id}.internal",
|
|
"devices": [
|
|
{
|
|
"device_id": device_id,
|
|
"driver_type": "wda",
|
|
"status": "idle",
|
|
"capability_tags": ["ios"],
|
|
}
|
|
for device_id in device_ids
|
|
],
|
|
}
|
|
|
|
|
|
def test_authenticated_heartbeat_replaces_complete_snapshot(tmp_path) -> None:
|
|
client, pool = _build_client(tmp_path)
|
|
|
|
first = client.put(
|
|
"/internal/v1/hosts/host-a/heartbeat",
|
|
headers={"Authorization": "Bearer token-a"},
|
|
json=_heartbeat_payload("host-a", "device-1", "device-2"),
|
|
)
|
|
second = client.put(
|
|
"/internal/v1/hosts/host-a/heartbeat",
|
|
headers={"Authorization": "Bearer token-a"},
|
|
json=_heartbeat_payload("host-a", "device-2", "device-3"),
|
|
)
|
|
|
|
assert first.status_code == 200
|
|
assert second.status_code == 200
|
|
assert second.json()["accepted_devices"] == 2
|
|
assert {device.device_id for device in pool.store.list_devices()} == {
|
|
"device-2",
|
|
"device-3",
|
|
}
|
|
|
|
|
|
def test_empty_heartbeat_removes_only_reporting_hosts_devices(tmp_path) -> None:
|
|
client, pool = _build_client(tmp_path)
|
|
client.put(
|
|
"/internal/v1/hosts/host-a/heartbeat",
|
|
headers={"Authorization": "Bearer token-a"},
|
|
json=_heartbeat_payload("host-a", "device-a"),
|
|
)
|
|
client.put(
|
|
"/internal/v1/hosts/host-b/heartbeat",
|
|
headers={"Authorization": "Bearer token-b"},
|
|
json=_heartbeat_payload("host-b", "device-b"),
|
|
)
|
|
|
|
response = client.put(
|
|
"/internal/v1/hosts/host-a/heartbeat",
|
|
headers={"Authorization": "Bearer token-a"},
|
|
json=_heartbeat_payload("host-a"),
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert response.json()["accepted_devices"] == 0
|
|
devices = pool.store.list_devices()
|
|
assert [(device.device_id, device.host_id) for device in devices] == [
|
|
("device-b", "host-b")
|
|
]
|
|
|
|
|
|
def test_invalid_duplicate_snapshot_preserves_previous_devices(tmp_path) -> None:
|
|
client, pool = _build_client(tmp_path)
|
|
client.put(
|
|
"/internal/v1/hosts/host-a/heartbeat",
|
|
headers={"Authorization": "Bearer token-a"},
|
|
json=_heartbeat_payload("host-a", "existing-device"),
|
|
)
|
|
|
|
response = client.put(
|
|
"/internal/v1/hosts/host-a/heartbeat",
|
|
headers={"Authorization": "Bearer token-a"},
|
|
json=_heartbeat_payload("host-a", "duplicate", "duplicate"),
|
|
)
|
|
|
|
assert response.status_code == 422
|
|
assert [device.device_id for device in pool.store.list_devices()] == [
|
|
"existing-device"
|
|
]
|
|
|
|
|
|
def test_live_device_owner_conflict_is_rejected_without_partial_sync(tmp_path) -> None:
|
|
client, pool = _build_client(tmp_path)
|
|
now = datetime.now(UTC)
|
|
pool.store.upsert_host("host-a", address=None, last_seen_at=now)
|
|
pool.store.replace_host_devices(
|
|
"host-a",
|
|
[
|
|
PooledDevice(
|
|
device_id="shared-device",
|
|
host_id="host-a",
|
|
driver_type="wda",
|
|
status="idle",
|
|
synced_at=now,
|
|
)
|
|
],
|
|
)
|
|
|
|
response = client.put(
|
|
"/internal/v1/hosts/host-b/heartbeat",
|
|
headers={"Authorization": "Bearer token-b"},
|
|
json=_heartbeat_payload("host-b", "shared-device"),
|
|
)
|
|
|
|
assert response.status_code == 409
|
|
assert pool.store.get_host("host-b") is None
|
|
devices = pool.store.list_devices()
|
|
assert len(devices) == 1
|
|
assert devices[0].host_id == "host-a"
|
|
|
|
|
|
def test_stale_device_owner_can_be_replaced_by_live_host(tmp_path) -> None:
|
|
client, pool = _build_client(tmp_path)
|
|
stale_at = datetime.now(UTC) - timedelta(seconds=61)
|
|
pool.store.upsert_host("host-a", address=None, last_seen_at=stale_at)
|
|
pool.store.replace_host_devices(
|
|
"host-a",
|
|
[
|
|
PooledDevice(
|
|
device_id="shared-device",
|
|
host_id="host-a",
|
|
driver_type="wda",
|
|
status="idle",
|
|
synced_at=stale_at,
|
|
)
|
|
],
|
|
)
|
|
|
|
response = client.put(
|
|
"/internal/v1/hosts/host-b/heartbeat",
|
|
headers={"Authorization": "Bearer token-b"},
|
|
json=_heartbeat_payload("host-b", "shared-device"),
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
device = pool.store.list_devices()[0]
|
|
assert device.device_id == "shared-device"
|
|
assert device.host_id == "host-b"
|
|
|
|
|
|
def test_host_token_cannot_submit_heartbeat_for_another_host(tmp_path) -> None:
|
|
client, pool = _build_client(tmp_path)
|
|
|
|
response = client.put(
|
|
"/internal/v1/hosts/host-b/heartbeat",
|
|
headers={"Authorization": "Bearer token-a"},
|
|
json=_heartbeat_payload("host-b", "device-b"),
|
|
)
|
|
|
|
assert response.status_code == 403
|
|
assert pool.store.get_host("host-b") is None
|
|
|
|
|
|
def test_heartbeat_and_self_submission_preserve_host_isolation(tmp_path) -> None:
|
|
client, pool = _build_client(tmp_path)
|
|
headers = {"Authorization": "Bearer token-a"}
|
|
heartbeat = client.put(
|
|
"/internal/v1/hosts/host-a/heartbeat",
|
|
headers=headers,
|
|
json=_heartbeat_payload("host-a", "device-a"),
|
|
)
|
|
assert heartbeat.status_code == 200
|
|
assert heartbeat.json()["policy_revision"] == 0
|
|
assert heartbeat.json()["policy"] is None
|
|
|
|
created = client.post(
|
|
"/internal/v1/hosts/host-a/tasks",
|
|
headers=headers,
|
|
json={"host_id": "host-a", "goal": "local work", "device_id": "device-a"},
|
|
)
|
|
assert created.status_code == 201, created.text
|
|
task = pool.store.get_task(created.json()["task_id"])
|
|
assert task is not None
|
|
assert task.constraints.target_host_id == "host-a"
|
|
assert task.constraints.target_device_id == "device-a"
|
|
foreign = client.post(
|
|
"/internal/v1/hosts/host-b/tasks",
|
|
headers=headers,
|
|
json={"host_id": "host-b", "goal": "forbidden"},
|
|
)
|
|
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
|
|
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:
|
|
client, pool = _build_client(tmp_path)
|
|
now = datetime.now(UTC)
|
|
pool.store.upsert_host("host-a", address=None, last_seen_at=now)
|
|
pool.store.replace_host_devices(
|
|
"host-a",
|
|
[
|
|
PooledDevice(
|
|
device_id="claim-device",
|
|
host_id="host-a",
|
|
driver_type="wda",
|
|
status="idle",
|
|
synced_at=now,
|
|
)
|
|
],
|
|
)
|
|
pool.store.enqueue_task(
|
|
ScheduledTask(
|
|
id="claim-task",
|
|
goal="open settings",
|
|
workflow_definition_id=None,
|
|
constraints=TaskConstraints(),
|
|
created_at=now,
|
|
)
|
|
)
|
|
pool.store.assign_task(
|
|
task_id="claim-task",
|
|
host_id="host-a",
|
|
device_id="claim-device",
|
|
lease_id="claim-lease",
|
|
lease_expires_at=now + timedelta(minutes=1),
|
|
now=now,
|
|
)
|
|
|
|
first = client.post(
|
|
"/internal/v1/hosts/host-a/assignments/claim",
|
|
headers={"Authorization": "Bearer token-a"},
|
|
json={"host_id": "host-a", "timeout_seconds": 0},
|
|
)
|
|
second = client.post(
|
|
"/internal/v1/hosts/host-a/assignments/claim",
|
|
headers={"Authorization": "Bearer token-a"},
|
|
json={"host_id": "host-a", "timeout_seconds": 0},
|
|
)
|
|
|
|
assert first.status_code == 200
|
|
assignment = first.json()["assignment"]
|
|
assert assignment["task_id"] == "claim-task"
|
|
assert assignment["lease_id"] == "claim-lease"
|
|
assert second.status_code == 200
|
|
assert second.json() == {"assignment": None, "timed_out": True}
|
|
|
|
|
|
def test_empty_long_poll_timeout_is_normal_response(tmp_path) -> None:
|
|
client, _ = _build_client(tmp_path)
|
|
|
|
response = client.post(
|
|
"/internal/v1/hosts/host-a/assignments/claim",
|
|
headers={"Authorization": "Bearer token-a"},
|
|
json={"host_id": "host-a", "timeout_seconds": 0},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert response.json() == {"assignment": None, "timed_out": True}
|
|
|
|
|
|
def _seed_active_assignment(pool: DevicePool) -> datetime:
|
|
now = datetime.now(UTC)
|
|
pool.store.upsert_host("host-a", address=None, last_seen_at=now)
|
|
pool.store.replace_host_devices(
|
|
"host-a",
|
|
[
|
|
PooledDevice(
|
|
device_id="active-device",
|
|
host_id="host-a",
|
|
driver_type="wda",
|
|
status="idle",
|
|
synced_at=now,
|
|
)
|
|
],
|
|
)
|
|
pool.store.enqueue_task(
|
|
ScheduledTask(
|
|
id="active-task",
|
|
goal="execute assignment",
|
|
workflow_definition_id=None,
|
|
constraints=TaskConstraints(),
|
|
created_at=now,
|
|
)
|
|
)
|
|
pool.store.assign_task(
|
|
task_id="active-task",
|
|
host_id="host-a",
|
|
device_id="active-device",
|
|
lease_id="active-lease",
|
|
lease_expires_at=now + timedelta(minutes=1),
|
|
now=now,
|
|
)
|
|
pool.store.claim_assignment(host_id="host-a", now=now)
|
|
return now
|
|
|
|
|
|
def test_lease_renewal_extends_active_assignment(tmp_path) -> None:
|
|
client, pool = _build_client(tmp_path)
|
|
original_time = _seed_active_assignment(pool)
|
|
|
|
response = client.post(
|
|
"/internal/v1/hosts/host-a/assignments/active-task/renew",
|
|
headers={"Authorization": "Bearer token-a"},
|
|
json={
|
|
"host_id": "host-a",
|
|
"task_id": "active-task",
|
|
"attempt": 1,
|
|
"lease_id": "active-lease",
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert response.json()["status"] == "renewed"
|
|
renewed_expiry = datetime.fromisoformat(response.json()["lease_expires_at"])
|
|
assert renewed_expiry > original_time + timedelta(seconds=30)
|
|
|
|
|
|
def test_stale_renewal_returns_typed_conflict(tmp_path) -> None:
|
|
client, pool = _build_client(tmp_path)
|
|
_seed_active_assignment(pool)
|
|
|
|
response = client.post(
|
|
"/internal/v1/hosts/host-a/assignments/active-task/renew",
|
|
headers={"Authorization": "Bearer token-a"},
|
|
json={
|
|
"host_id": "host-a",
|
|
"task_id": "active-task",
|
|
"attempt": 1,
|
|
"lease_id": "stale-lease",
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 409
|
|
assert response.json()["code"] == "stale_lease"
|
|
|
|
|
|
def test_terminal_result_is_idempotent_through_internal_api(tmp_path) -> None:
|
|
client, pool = _build_client(tmp_path)
|
|
_seed_active_assignment(pool)
|
|
payload = {
|
|
"host_id": "host-a",
|
|
"task_id": "active-task",
|
|
"attempt": 1,
|
|
"lease_id": "active-lease",
|
|
"status": "done",
|
|
"result": {"steps": 4},
|
|
}
|
|
|
|
first = client.post(
|
|
"/internal/v1/hosts/host-a/assignments/active-task/result",
|
|
headers={"Authorization": "Bearer token-a"},
|
|
json=payload,
|
|
)
|
|
repeated = client.post(
|
|
"/internal/v1/hosts/host-a/assignments/active-task/result",
|
|
headers={"Authorization": "Bearer token-a"},
|
|
json=payload,
|
|
)
|
|
|
|
assert first.status_code == 200
|
|
assert first.json()["status"] == "recorded"
|
|
assert repeated.status_code == 200
|
|
assert repeated.json()["status"] == "already_recorded"
|
|
assert pool.store.get_task("active-task").status == "done" # type: ignore[union-attr]
|
|
|
|
|
|
def test_conflicting_repeated_result_returns_stale_lease_conflict(tmp_path) -> None:
|
|
client, pool = _build_client(tmp_path)
|
|
_seed_active_assignment(pool)
|
|
base_payload = {
|
|
"host_id": "host-a",
|
|
"task_id": "active-task",
|
|
"attempt": 1,
|
|
"lease_id": "active-lease",
|
|
"status": "done",
|
|
"result": {"steps": 4},
|
|
}
|
|
client.post(
|
|
"/internal/v1/hosts/host-a/assignments/active-task/result",
|
|
headers={"Authorization": "Bearer token-a"},
|
|
json=base_payload,
|
|
)
|
|
|
|
response = client.post(
|
|
"/internal/v1/hosts/host-a/assignments/active-task/result",
|
|
headers={"Authorization": "Bearer token-a"},
|
|
json={**base_payload, "status": "failed", "failure_reason": "late failure"},
|
|
)
|
|
|
|
assert response.status_code == 409
|
|
assert response.json()["code"] == "stale_lease"
|
|
|
|
|
|
def test_multi_host_claims_are_isolated(tmp_path) -> None:
|
|
client, pool = _build_client(tmp_path)
|
|
now = datetime.now(UTC)
|
|
for host_id, device_id, task_id, lease_id in (
|
|
("host-a", "device-a", "task-a", "lease-a"),
|
|
("host-b", "device-b", "task-b", "lease-b"),
|
|
):
|
|
pool.store.upsert_host(host_id, address=None, last_seen_at=now)
|
|
pool.store.replace_host_devices(
|
|
host_id,
|
|
[
|
|
PooledDevice(
|
|
device_id=device_id,
|
|
host_id=host_id,
|
|
driver_type="wda",
|
|
status="idle",
|
|
synced_at=now,
|
|
)
|
|
],
|
|
)
|
|
pool.store.enqueue_task(
|
|
ScheduledTask(
|
|
id=task_id,
|
|
goal=f"work for {host_id}",
|
|
workflow_definition_id=None,
|
|
constraints=TaskConstraints(),
|
|
created_at=now,
|
|
)
|
|
)
|
|
pool.store.assign_task(
|
|
task_id=task_id,
|
|
host_id=host_id,
|
|
device_id=device_id,
|
|
lease_id=lease_id,
|
|
lease_expires_at=now + timedelta(minutes=1),
|
|
now=now,
|
|
)
|
|
|
|
host_a = client.post(
|
|
"/internal/v1/hosts/host-a/assignments/claim",
|
|
headers={"Authorization": "Bearer token-a"},
|
|
json={"host_id": "host-a", "timeout_seconds": 0},
|
|
)
|
|
host_b = client.post(
|
|
"/internal/v1/hosts/host-b/assignments/claim",
|
|
headers={"Authorization": "Bearer token-b"},
|
|
json={"host_id": "host-b", "timeout_seconds": 0},
|
|
)
|
|
|
|
assert host_a.json()["assignment"]["task_id"] == "task-a"
|
|
assert host_b.json()["assignment"]["task_id"] == "task-b"
|
|
|
|
|
|
def test_superseded_attempt_result_cannot_overwrite_current_lease(tmp_path) -> None:
|
|
client, pool = _build_client(tmp_path)
|
|
now = datetime.now(UTC)
|
|
pool.store.upsert_host("host-a", address=None, last_seen_at=now)
|
|
pool.store.replace_host_devices(
|
|
"host-a",
|
|
[
|
|
PooledDevice(
|
|
device_id="retry-device",
|
|
host_id="host-a",
|
|
driver_type="wda",
|
|
status="idle",
|
|
synced_at=now,
|
|
)
|
|
],
|
|
)
|
|
pool.store.enqueue_task(
|
|
ScheduledTask(
|
|
id="retry-task",
|
|
goal="retry safely",
|
|
workflow_definition_id=None,
|
|
constraints=TaskConstraints(),
|
|
created_at=now,
|
|
)
|
|
)
|
|
pool.store.assign_task(
|
|
task_id="retry-task",
|
|
host_id="host-a",
|
|
device_id="retry-device",
|
|
lease_id="old-lease",
|
|
lease_expires_at=now + timedelta(seconds=1),
|
|
now=now,
|
|
)
|
|
reaped_at = now + timedelta(seconds=2)
|
|
pool.store.reap_expired_leases(now=reaped_at, max_attempts=2)
|
|
pool.store.assign_task(
|
|
task_id="retry-task",
|
|
host_id="host-a",
|
|
device_id="retry-device",
|
|
lease_id="current-lease",
|
|
lease_expires_at=reaped_at + timedelta(minutes=1),
|
|
now=reaped_at,
|
|
)
|
|
|
|
response = client.post(
|
|
"/internal/v1/hosts/host-a/assignments/retry-task/result",
|
|
headers={"Authorization": "Bearer token-a"},
|
|
json={
|
|
"host_id": "host-a",
|
|
"task_id": "retry-task",
|
|
"attempt": 1,
|
|
"lease_id": "old-lease",
|
|
"status": "done",
|
|
"result": {"late": True},
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 409
|
|
assert response.json()["code"] == "stale_lease"
|
|
task = pool.store.get_task("retry-task")
|
|
assert task is not None
|
|
assert task.status == "assigned"
|
|
assert task.attempt_count == 2
|
|
assert task.lease_id == "current-lease"
|
|
assert task.terminal_result is None
|