feat(host-protocol): accept heartbeat snapshots

2026-07-12 18:10:42 +08:00
parent b4cbd83dcd
commit efb754fbe7
3 changed files with 246 additions and 1 deletions
@@ -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))}",
+        )