feat(cloud): add edge host enrollment

apps/device-host-agent/host_agent/client.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import time
 from collections.abc import Awaitable, Callable
 from typing import Any
 
@@ -9,8 +10,10 @@ import httpx
 from cloud.internal_api.models import (
     AssignmentModel,
     ClaimResponse,
+    DeviceEnrollmentResponse,
     DeviceSnapshotModel,
     HeartbeatResponse,
+    HostEnrollmentResponse,
     LeaseRenewalResponse,
     TerminalResultResponse,
 )
@@ -30,6 +33,99 @@ class StaleLeaseError(HostAgentAPIError):
     pass
 
 
+class HostAgentEnrollmentClient:
+    def __init__(
+        self,
+        config: HostAgentConfig,
+        *,
+        http_client: httpx.Client | None = None,
+        sleep: Callable[[float], None] = time.sleep,
+    ) -> None:
+        self.config = config
+        self._sleep = sleep
+        self._owns_client = http_client is None
+        self._client = http_client or httpx.Client(base_url=config.control_plane_url)
+
+    def enroll_host(
+        self,
+        *,
+        agent_instance_id: str,
+        host_token: str,
+        display_name: str | None,
+    ) -> HostEnrollmentResponse:
+        if not self.config.enrollment_token:
+            raise HostAgentAPIError(0, "Host enrollment token is unavailable")
+        response = self._request(
+            "POST",
+            "/internal/v1/enrollments",
+            token=self.config.enrollment_token,
+            json={
+                "agent_instance_id": agent_instance_id,
+                "host_token": host_token,
+                "display_name": display_name,
+            },
+        )
+        return HostEnrollmentResponse.model_validate(response.json())
+
+    def enroll_device(
+        self,
+        *,
+        local_device_id: str,
+        driver_type: str,
+        name: str | None,
+        capability_tags: list[str],
+    ) -> DeviceEnrollmentResponse:
+        if not self.config.host_id or not self.config.token:
+            raise HostAgentAPIError(0, "Host identity is unresolved")
+        response = self._request(
+            "POST",
+            f"/internal/v1/hosts/{self.config.host_id}/devices/enroll",
+            token=self.config.token,
+            json={
+                "local_device_id": local_device_id,
+                "driver_type": driver_type,
+                "name": name,
+                "capability_tags": list(capability_tags),
+            },
+        )
+        return DeviceEnrollmentResponse.model_validate(response.json())
+
+    def close(self) -> None:
+        if self._owns_client:
+            self._client.close()
+
+    def _request(
+        self,
+        method: str,
+        path: str,
+        *,
+        token: str,
+        json: dict[str, Any],
+    ) -> httpx.Response:
+        backoff = self.config.retry_backoff_seconds
+        for attempt in range(1, self.config.max_retry_attempts + 1):
+            try:
+                response = self._client.request(
+                    method,
+                    path,
+                    json=json,
+                    headers={"Authorization": f"Bearer {token}"},
+                )
+            except httpx.TransportError:
+                if attempt == self.config.max_retry_attempts:
+                    raise
+            else:
+                if response.status_code < 500:
+                    if response.is_success:
+                        return response
+                    _raise_api_error(response)
+                if attempt == self.config.max_retry_attempts:
+                    _raise_api_error(response)
+            self._sleep(backoff)
+            backoff = min(backoff * 2, self.config.max_retry_backoff_seconds)
+        raise AssertionError("retry loop exited unexpectedly")
+
+
 class HostAgentClient:
     def __init__(
         self,
@@ -149,22 +245,24 @@ class HostAgentClient:
                 if response.status_code < 500:
                     if response.is_success:
                         return response
-                    self._raise_api_error(response)
+                    _raise_api_error(response, stale_lease=True)
                 if attempt == self.config.max_retry_attempts:
-                    self._raise_api_error(response)
+                    _raise_api_error(response, stale_lease=True)
 
             await self._sleep(backoff)
             backoff = min(backoff * 2, self.config.max_retry_backoff_seconds)
         raise AssertionError("retry loop exited unexpectedly")
 
-    @staticmethod
-    def _raise_api_error(response: httpx.Response) -> None:
-        try:
-            payload = response.json()
-        except ValueError:
-            payload = {}
-        detail = payload.get("detail") or payload.get("code") or "request rejected"
-        error_type = (
-            StaleLeaseError if response.status_code == 409 else HostAgentAPIError
-        )
-        raise error_type(response.status_code, str(detail))
+
+def _raise_api_error(response: httpx.Response, *, stale_lease: bool = False) -> None:
+    try:
+        payload = response.json()
+    except ValueError:
+        payload = {}
+    detail = payload.get("detail") or payload.get("code") or "request rejected"
+    error_type = (
+        StaleLeaseError
+        if stale_lease and response.status_code == 409
+        else HostAgentAPIError
+    )
+    raise error_type(response.status_code, str(detail))