feat(host-protocol): accept heartbeat snapshots

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.

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))}",
+        )

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