From f6bc8f6b98f4fdde898eb4459903f1b3cb2ef870 Mon Sep 17 00:00:00 2001 From: Jerry Yan <792602257@qq.com> Date: Sun, 12 Jul 2026 18:44:18 +0800 Subject: [PATCH] feat(host-agent): add control plane client --- apps/device-host-agent/host_agent/client.py | 170 ++++++++++++++++++ apps/device-host-agent/host_agent/config.py | 35 +++- apps/device-host-agent/tests/test_client.py | 133 ++++++++++++++ .../cloud-control-plane-integration/tasks.md | 2 +- 4 files changed, 335 insertions(+), 5 deletions(-) create mode 100644 apps/device-host-agent/host_agent/client.py create mode 100644 apps/device-host-agent/tests/test_client.py diff --git a/apps/device-host-agent/host_agent/client.py b/apps/device-host-agent/host_agent/client.py new file mode 100644 index 0000000..221baa1 --- /dev/null +++ b/apps/device-host-agent/host_agent/client.py @@ -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)) diff --git a/apps/device-host-agent/host_agent/config.py b/apps/device-host-agent/host_agent/config.py index 05410ba..d758f35 100644 --- a/apps/device-host-agent/host_agent/config.py +++ b/apps/device-host-agent/host_agent/config.py @@ -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 diff --git a/apps/device-host-agent/tests/test_client.py b/apps/device-host-agent/tests/test_client.py new file mode 100644 index 0000000..c7e4f99 --- /dev/null +++ b/apps/device-host-agent/tests/test_client.py @@ -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 diff --git a/openspec/changes/cloud-control-plane-integration/tasks.md b/openspec/changes/cloud-control-plane-integration/tasks.md index 7bd991b..bd5ce4e 100644 --- a/openspec/changes/cloud-control-plane-integration/tasks.md +++ b/openspec/changes/cloud-control-plane-integration/tasks.md @@ -52,7 +52,7 @@ ## 7. Device Host Agent Execution Loop -- [ ] 7.1 Implement an authenticated Host Agent client for heartbeat, long-poll claim, renewal, and result operations with bounded retry/backoff. +- [x] 7.1 Implement an authenticated Host Agent client for heartbeat, long-poll claim, renewal, and result operations with bounded retry/backoff. - [ ] 7.2 Build complete device snapshots from the local `DeviceManager` and synchronize them at the configured interval. - [ ] 7.3 Compose local `TaskRunner` and `WorkflowRunner` factories without importing cloud concerns into Runtime-owned packages. - [ ] 7.4 Execute goal assignments through the configured Runtime Planner/Executor and workflow assignments through the existing workflow runner.