agentic-mobile-control/apps/device-host-agent/host_agent/client.py

from __future__ import annotations

import asyncio
import time
from collections.abc import Awaitable, Callable
from typing import Any

import httpx

from cloud.internal_api.models import (
    AssignmentModel,
    ClaimResponse,
    DeviceEnrollmentResponse,
    DeviceSnapshotModel,
    HeartbeatResponse,
    HostEnrollmentResponse,
    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 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:
        response = self._request(
            "POST",
            "/internal/v1/enrollments",
            token=self.config.enrollment_token or None,
            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 | None,
        json: dict[str, Any],
    ) -> httpx.Response:
        headers = {"Authorization": f"Bearer {token}"} if token else {}
        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=headers,
                )
            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,
        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
                    _raise_api_error(response, stale_lease=True)
                if attempt == self.config.max_retry_attempts:
                    _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")


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))