feat(host-agent): add control plane client

This commit is contained in:
2026-07-12 18:44:18 +08:00
parent 1e10d5aa95
commit f6bc8f6b98
4 changed files with 335 additions and 5 deletions
+170
View File
@@ -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))
+31 -4
View File
@@ -19,16 +19,21 @@ class HostAgentConfig:
poll_timeout_seconds: float = 20.0 poll_timeout_seconds: float = 20.0
retry_backoff_seconds: float = 1.0 retry_backoff_seconds: float = 1.0
max_retry_backoff_seconds: float = 30.0 max_retry_backoff_seconds: float = 30.0
max_retry_attempts: int = 5
def load_host_agent_config( def load_host_agent_config(
env: Mapping[str, str] | None = None, env: Mapping[str, str] | None = None,
) -> HostAgentConfig: ) -> HostAgentConfig:
values = os.environ if env is None else env values = os.environ if env is None else env
control_plane_url = values.get( control_plane_url = (
"HOST_AGENT_CONTROL_PLANE_URL", values.get(
"http://127.0.0.1:8001", "HOST_AGENT_CONTROL_PLANE_URL",
).strip().rstrip("/") "http://127.0.0.1:8001",
)
.strip()
.rstrip("/")
)
parsed_url = urlparse(control_plane_url) parsed_url = urlparse(control_plane_url)
if parsed_url.scheme not in {"http", "https"} or not parsed_url.netloc: if parsed_url.scheme not in {"http", "https"} or not parsed_url.netloc:
raise HostAgentConfigurationError( raise HostAgentConfigurationError(
@@ -66,6 +71,11 @@ def load_host_agent_config(
"HOST_AGENT_MAX_RETRY_BACKOFF_SECONDS", "HOST_AGENT_MAX_RETRY_BACKOFF_SECONDS",
30.0, 30.0,
), ),
max_retry_attempts=_positive_int(
values,
"HOST_AGENT_MAX_RETRY_ATTEMPTS",
5,
),
) )
if config.max_retry_backoff_seconds < config.retry_backoff_seconds: if config.max_retry_backoff_seconds < config.retry_backoff_seconds:
raise HostAgentConfigurationError( raise HostAgentConfigurationError(
@@ -89,3 +99,20 @@ def _positive_float(
if value <= 0: if value <= 0:
raise HostAgentConfigurationError(f"{name} must be greater than zero") raise HostAgentConfigurationError(f"{name} must be greater than zero")
return value 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
+133
View File
@@ -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
@@ -52,7 +52,7 @@
## 7. Device Host Agent Execution Loop ## 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.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.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. - [ ] 7.4 Execute goal assignments through the configured Runtime Planner/Executor and workflow assignments through the existing workflow runner.