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

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