diff --git a/openspec/changes/cloud-control-plane-integration/tasks.md b/openspec/changes/cloud-control-plane-integration/tasks.md index ed8057f..e1e4918 100644 --- a/openspec/changes/cloud-control-plane-integration/tasks.md +++ b/openspec/changes/cloud-control-plane-integration/tasks.md @@ -36,7 +36,7 @@ ## 5. Host Agent Internal API - [x] 5.1 Add versioned `/internal/v1` request/response models for heartbeat snapshots, long-poll claim, lease renewal, and terminal result reporting. -- [ ] 5.2 Add atomic heartbeat/snapshot validation and device ownership-conflict handling before delegating to `DevicePool`. +- [x] 5.2 Add atomic heartbeat/snapshot validation and device ownership-conflict handling before delegating to `DevicePool`. - [ ] 5.3 Add long-poll assignment delivery that returns at most one claimed task and produces a normal empty timeout response. - [ ] 5.4 Add lease-renewal and idempotent terminal-result endpoints with typed stale-lease conflicts. - [ ] 5.5 Add internal API integration tests for multi-host isolation, duplicate device ids, timeout behavior, stale attempts, and repeated result reports. diff --git a/packages/cloud-platform/cloud/internal_api/api.py b/packages/cloud-platform/cloud/internal_api/api.py new file mode 100644 index 0000000..7f45eb2 --- /dev/null +++ b/packages/cloud-platform/cloud/internal_api/api.py @@ -0,0 +1,107 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from fastapi import APIRouter, HTTPException, Request, status + +from cloud.auth import ( + AuthProvider, + HostAuthorizationError, +) +from cloud.internal_api.models import HeartbeatRequest, HeartbeatResponse +from core.models import Device, utc_now + +if TYPE_CHECKING: + from cloud.pool import DevicePool + + +def create_internal_router( + *, + pool: DevicePool, + auth_provider: AuthProvider, + version_prefix: str = "/internal/v1", +) -> APIRouter: + router = APIRouter(prefix=version_prefix, tags=["host-agent"]) + + def authorize_host(request: Request, host_id: str) -> None: + principal = auth_provider.authenticate(request) + if principal is None: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="unauthorized", + headers={"WWW-Authenticate": "Bearer"}, + ) + try: + principal.require_host(host_id) + except HostAuthorizationError as exc: + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail=str(exc), + ) from exc + + @router.put( + "/hosts/{host_id}/heartbeat", + response_model=HeartbeatResponse, + ) + def heartbeat( + host_id: str, + payload: HeartbeatRequest, + request: Request, + ) -> HeartbeatResponse: + authorize_host(request, host_id) + _validate_snapshot(pool, host_id=host_id, payload=payload) + devices = [ + Device( + id=device.device_id, + driver_type=device.driver_type, + status=device.status, + capability_tags=list(device.capability_tags), + ) + for device in payload.devices + ] + pool.sync_host_devices(host_id, devices, address=payload.address) + return HeartbeatResponse( + host_id=host_id, + accepted_devices=len(devices), + received_at=utc_now(), + ) + + return router + + +def _validate_snapshot( + pool: DevicePool, + *, + host_id: str, + payload: HeartbeatRequest, +) -> None: + if payload.host_id != host_id: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_CONTENT, + detail="heartbeat host_id must match the request path", + ) + device_ids = [device.device_id for device in payload.devices] + if len(device_ids) != len(set(device_ids)): + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_CONTENT, + detail="heartbeat snapshot contains duplicate device ids", + ) + + now = utc_now() + hosts = {host.host_id: host for host in pool.store.list_hosts()} + conflicts: list[str] = [] + requested_ids = set(device_ids) + for device in pool.store.list_devices(): + if device.device_id not in requested_ids or device.host_id == host_id: + continue + owner = hosts.get(device.host_id) + if owner is None: + continue + age_seconds = (now - owner.last_seen_at).total_seconds() + if age_seconds <= pool.config.stale_after_seconds: + conflicts.append(device.device_id) + if conflicts: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"device ownership conflict: {sorted(set(conflicts))}", + ) diff --git a/tests/test_host_agent_internal_api.py b/tests/test_host_agent_internal_api.py new file mode 100644 index 0000000..23766eb --- /dev/null +++ b/tests/test_host_agent_internal_api.py @@ -0,0 +1,138 @@ +from __future__ import annotations + +from datetime import UTC, datetime + +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.store import CloudStore + + +def _build_client(tmp_path) -> 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)) + 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_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_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