feat(host-agent): add control plane client
This commit is contained in:
@@ -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
|
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
|
||||||
|
|||||||
@@ -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.
|
||||||
|
|||||||
Reference in New Issue
Block a user