feat(host-agent): add control plane client

2026-07-12 18:44:18 +08:00
parent 1e10d5aa95
commit f6bc8f6b98
4 changed files with 335 additions and 5 deletions
@@ -0,0 +1,170 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Awaitable, Callable
+from typing import Any
+
+import httpx
+
+from cloud.internal_api.models import (
+    AssignmentModel,
+    ClaimResponse,
+    DeviceSnapshotModel,
+    HeartbeatResponse,
+    LeaseRenewalResponse,
+    TerminalResultResponse,
+)
+from host_agent.config import HostAgentConfig
+
+
+class HostAgentAPIError(RuntimeError):
+    def __init__(self, status_code: int, detail: str) -> None:
+        super().__init__(
+            f"control plane request failed with status {status_code}: {detail}"
+        )
+        self.status_code = status_code
+        self.detail = detail
+
+
+class StaleLeaseError(HostAgentAPIError):
+    pass
+
+
+class HostAgentClient:
+    def __init__(
+        self,
+        config: HostAgentConfig,
+        *,
+        http_client: httpx.AsyncClient | None = None,
+        sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
+    ) -> None:
+        self.config = config
+        self._sleep = sleep
+        self._owns_client = http_client is None
+        self._client = http_client or httpx.AsyncClient(
+            base_url=config.control_plane_url,
+            headers={"Authorization": f"Bearer {config.token}"},
+        )
+
+    async def heartbeat(
+        self,
+        devices: list[DeviceSnapshotModel],
+        *,
+        address: str | None = None,
+    ) -> HeartbeatResponse:
+        response = await self._request(
+            "PUT",
+            f"/internal/v1/hosts/{self.config.host_id}/heartbeat",
+            json={
+                "host_id": self.config.host_id,
+                "address": address,
+                "devices": [device.model_dump(mode="json") for device in devices],
+            },
+        )
+        return HeartbeatResponse.model_validate(response.json())
+
+    async def claim(self) -> AssignmentModel | None:
+        response = await self._request(
+            "POST",
+            f"/internal/v1/hosts/{self.config.host_id}/assignments/claim",
+            json={
+                "host_id": self.config.host_id,
+                "timeout_seconds": self.config.poll_timeout_seconds,
+            },
+            timeout=self.config.poll_timeout_seconds + 5,
+        )
+        return ClaimResponse.model_validate(response.json()).assignment
+
+    async def renew(
+        self,
+        assignment: AssignmentModel,
+    ) -> LeaseRenewalResponse:
+        response = await self._request(
+            "POST",
+            (
+                f"/internal/v1/hosts/{self.config.host_id}/assignments/"
+                f"{assignment.task_id}/renew"
+            ),
+            json={
+                "host_id": self.config.host_id,
+                "task_id": assignment.task_id,
+                "attempt": assignment.attempt,
+                "lease_id": assignment.lease_id,
+            },
+        )
+        return LeaseRenewalResponse.model_validate(response.json())
+
+    async def report_result(
+        self,
+        assignment: AssignmentModel,
+        *,
+        status: str,
+        failure_reason: str | None = None,
+        result: dict[str, Any] | None = None,
+    ) -> TerminalResultResponse:
+        response = await self._request(
+            "POST",
+            (
+                f"/internal/v1/hosts/{self.config.host_id}/assignments/"
+                f"{assignment.task_id}/result"
+            ),
+            json={
+                "host_id": self.config.host_id,
+                "task_id": assignment.task_id,
+                "attempt": assignment.attempt,
+                "lease_id": assignment.lease_id,
+                "status": status,
+                "failure_reason": failure_reason,
+                "result": result,
+            },
+        )
+        return TerminalResultResponse.model_validate(response.json())
+
+    async def aclose(self) -> None:
+        if self._owns_client:
+            await self._client.aclose()
+
+    async def _request(
+        self,
+        method: str,
+        path: str,
+        *,
+        json: dict[str, Any],
+        timeout: float | None = None,
+    ) -> httpx.Response:
+        backoff = self.config.retry_backoff_seconds
+        for attempt in range(1, self.config.max_retry_attempts + 1):
+            try:
+                response = await self._client.request(
+                    method,
+                    path,
+                    json=json,
+                    timeout=timeout,
+                    headers={"Authorization": f"Bearer {self.config.token}"},
+                )
+            except httpx.TransportError:
+                if attempt == self.config.max_retry_attempts:
+                    raise
+            else:
+                if response.status_code < 500:
+                    if response.is_success:
+                        return response
+                    self._raise_api_error(response)
+                if attempt == self.config.max_retry_attempts:
+                    self._raise_api_error(response)
+
+            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))
@@ -19,16 +19,21 @@ class HostAgentConfig:
    poll_timeout_seconds: float = 20.0
    retry_backoff_seconds: float = 1.0
    max_retry_backoff_seconds: float = 30.0
+    max_retry_attempts: int = 5


 def load_host_agent_config(
    env: Mapping[str, str] | None = None,
 ) -> HostAgentConfig:
    values = os.environ if env is None else env
-    control_plane_url = values.get(
-        "HOST_AGENT_CONTROL_PLANE_URL",
-        "http://127.0.0.1:8001",
-    ).strip().rstrip("/")
+    control_plane_url = (
+        values.get(
+            "HOST_AGENT_CONTROL_PLANE_URL",
+            "http://127.0.0.1:8001",
+        )
+        .strip()
+        .rstrip("/")
+    )
    parsed_url = urlparse(control_plane_url)
    if parsed_url.scheme not in {"http", "https"} or not parsed_url.netloc:
        raise HostAgentConfigurationError(
@@ -66,6 +71,11 @@ def load_host_agent_config(
            "HOST_AGENT_MAX_RETRY_BACKOFF_SECONDS",
            30.0,
        ),
+        max_retry_attempts=_positive_int(
+            values,
+            "HOST_AGENT_MAX_RETRY_ATTEMPTS",
+            5,
+        ),
    )
    if config.max_retry_backoff_seconds < config.retry_backoff_seconds:
        raise HostAgentConfigurationError(
@@ -89,3 +99,20 @@ def _positive_float(
    if value <= 0:
        raise HostAgentConfigurationError(f"{name} must be greater than zero")
    return value
+
+
+def _positive_int(
+    values: Mapping[str, str],
+    name: str,
+    default: int,
+) -> int:
+    raw_value = values.get(name)
+    if raw_value is None:
+        return default
+    try:
+        value = int(raw_value)
+    except ValueError as exc:
+        raise HostAgentConfigurationError(f"{name} must be an integer") from exc
+    if value <= 0:
+        raise HostAgentConfigurationError(f"{name} must be greater than zero")
+    return value
@@ -0,0 +1,133 @@
+from __future__ import annotations
+
+import asyncio
+from datetime import UTC, datetime
+
+import httpx
+import pytest
+
+from cloud.internal_api.models import AssignmentModel, DeviceSnapshotModel
+from host_agent.client import HostAgentClient, StaleLeaseError
+from host_agent.config import HostAgentConfig
+
+
+def _config(**overrides) -> HostAgentConfig:
+    values = {
+        "control_plane_url": "https://control.example",
+        "host_id": "host-a",
+        "token": "host-secret",
+        "retry_backoff_seconds": 0.01,
+        "max_retry_backoff_seconds": 0.02,
+        "max_retry_attempts": 3,
+    }
+    values.update(overrides)
+    return HostAgentConfig(**values)
+
+
+def _assignment() -> AssignmentModel:
+    return AssignmentModel(
+        task_id="task-a",
+        attempt=1,
+        lease_id="lease-a",
+        lease_expires_at=datetime(2026, 7, 12, tzinfo=UTC),
+        host_id="host-a",
+        device_id="device-a",
+        goal="open settings",
+    )
+
+
+def test_client_sends_authenticated_heartbeat_and_claim() -> None:
+    requests: list[httpx.Request] = []
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        requests.append(request)
+        if request.url.path.endswith("/heartbeat"):
+            return httpx.Response(
+                200,
+                json={
+                    "host_id": "host-a",
+                    "accepted_devices": 1,
+                    "received_at": "2026-07-12T00:00:00Z",
+                },
+            )
+        return httpx.Response(200, json={"assignment": None, "timed_out": True})
+
+    async def scenario() -> None:
+        async with httpx.AsyncClient(
+            transport=httpx.MockTransport(handler),
+            base_url="https://control.example",
+        ) as http_client:
+            client = HostAgentClient(_config(), http_client=http_client)
+            heartbeat = await client.heartbeat(
+                [
+                    DeviceSnapshotModel(
+                        device_id="device-a",
+                        driver_type="wda",
+                        status="idle",
+                    )
+                ]
+            )
+            assignment = await client.claim()
+            assert heartbeat.accepted_devices == 1
+            assert assignment is None
+
+    asyncio.run(scenario())
+    assert [request.url.path for request in requests] == [
+        "/internal/v1/hosts/host-a/heartbeat",
+        "/internal/v1/hosts/host-a/assignments/claim",
+    ]
+    assert all(
+        request.headers["authorization"] == "Bearer host-secret" for request in requests
+    )
+
+
+def test_transport_failures_retry_with_bounded_backoff() -> None:
+    attempts = 0
+    sleeps: list[float] = []
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal attempts
+        attempts += 1
+        if attempts < 3:
+            raise httpx.ConnectError("temporary", request=request)
+        return httpx.Response(200, json={"assignment": None, "timed_out": True})
+
+    async def sleep(delay: float) -> None:
+        sleeps.append(delay)
+
+    async def scenario() -> None:
+        async with httpx.AsyncClient(
+            transport=httpx.MockTransport(handler),
+            base_url="https://control.example",
+        ) as http_client:
+            client = HostAgentClient(_config(), http_client=http_client, sleep=sleep)
+            assert await client.claim() is None
+
+    asyncio.run(scenario())
+    assert attempts == 3
+    assert sleeps == [0.01, 0.02]
+
+
+def test_stale_lease_response_raises_typed_error_without_retry() -> None:
+    attempts = 0
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        nonlocal attempts
+        attempts += 1
+        return httpx.Response(
+            409,
+            json={"code": "stale_lease", "detail": "lease expired"},
+        )
+
+    async def scenario() -> None:
+        async with httpx.AsyncClient(
+            transport=httpx.MockTransport(handler),
+            base_url="https://control.example",
+        ) as http_client:
+            client = HostAgentClient(_config(), http_client=http_client)
+            with pytest.raises(StaleLeaseError) as error:
+                await client.renew(_assignment())
+            assert "host-secret" not in str(error.value)
+
+    asyncio.run(scenario())
+    assert attempts == 1