storage-labos/flashops/services/host-agent/flashops_agent/client.py

"""Small HTTP client for the control-plane Agent contract."""
from __future__ import annotations

from typing import Any, Dict, List, Optional, Tuple

import httpx


class ControlPlaneError(RuntimeError):
    pass


class AgentUnauthorized(ControlPlaneError):
    pass


class StepLeaseConflict(ControlPlaneError):
    pass


class ControlPlaneClient:
    def __init__(
        self,
        server: str,
        *,
        enrollment_token: Optional[str] = None,
        basic_auth: Optional[Tuple[str, str]] = None,
        timeout_s: float = 15.0,
        transport: Optional[httpx.BaseTransport] = None,
    ) -> None:
        self.server = server.rstrip("/")
        self.enrollment_token = enrollment_token
        self._client = httpx.Client(
            base_url=self.server,
            auth=httpx.BasicAuth(*basic_auth) if basic_auth else None,
            timeout=timeout_s,
            transport=transport,
        )

    def close(self) -> None:
        self._client.close()

    def __enter__(self) -> "ControlPlaneClient":
        return self

    def __exit__(self, *_: object) -> None:
        self.close()

    @staticmethod
    def _payload(response: httpx.Response) -> Dict[str, Any]:
        try:
            value = response.json()
        except ValueError as exc:
            raise ControlPlaneError(
                "控制平面返回了非 JSON 响应: HTTP {0}".format(response.status_code)
            ) from exc
        if not isinstance(value, dict):
            raise ControlPlaneError("控制平面响应不是 JSON object")
        return value

    def register(self, fingerprint: Dict[str, Any]) -> Dict[str, Any]:
        headers = {}
        if self.enrollment_token:
            headers["X-FlashOps-Enrollment-Token"] = self.enrollment_token
        response = self._client.post(
            "/api/v1/agent/register",
            json=fingerprint,
            headers=headers,
        )
        if response.status_code == 401:
            raise AgentUnauthorized("Agent 注册口令无效")
        if response.status_code >= 400:
            raise ControlPlaneError(
                "Agent 注册失败: HTTP {0} {1}".format(
                    response.status_code, response.text[:200]
                )
            )
        payload = self._payload(response)
        for field in ("agent_id", "host_id", "token"):
            if not payload.get(field):
                raise ControlPlaneError("注册响应缺少字段: {0}".format(field))
        return payload

    def heartbeat(
        self,
        agent_id: str,
        token: str,
        heartbeat: Dict[str, Any],
    ) -> Dict[str, Any]:
        response = self._client.post(
            "/api/v1/agent/{0}/heartbeat".format(agent_id),
            json=heartbeat,
            headers={"Authorization": "Bearer " + token},
        )
        if response.status_code == 401:
            raise AgentUnauthorized("Agent token 已失效")
        if response.status_code >= 400:
            raise ControlPlaneError(
                "Agent 心跳失败: HTTP {0} {1}".format(
                    response.status_code, response.text[:200]
                )
            )
        payload = self._payload(response)
        if payload.get("ok") is not True:
            raise ControlPlaneError("控制平面未确认心跳")
        return payload

    @staticmethod
    def _agent_headers(token: str, idempotency_key: Optional[str] = None) -> Dict[str, str]:
        headers = {"Authorization": "Bearer " + token}
        if idempotency_key:
            headers["Idempotency-Key"] = idempotency_key
        return headers

    def lease(
        self,
        agent_id: str,
        token: str,
        *,
        wait_timeout_s: float = 0.0,
    ) -> Optional[Dict[str, Any]]:
        response = self._client.post(
            "/api/v1/agent/{0}/lease".format(agent_id),
            json={"wait_timeout_s": wait_timeout_s},
            headers=self._agent_headers(token),
        )
        if response.status_code == 204:
            return None
        return self._step_response(response, "领取步骤")

    def _step_response(self, response: httpx.Response, action: str) -> Dict[str, Any]:
        if response.status_code == 401:
            raise AgentUnauthorized("Agent token 已失效")
        if response.status_code == 409:
            detail = response.text[:300]
            try:
                detail = str(response.json().get("detail", detail))
            except (ValueError, AttributeError):
                pass
            raise StepLeaseConflict("{0}冲突: {1}".format(action, detail))
        if response.status_code >= 400:
            raise ControlPlaneError(
                "{0}失败: HTTP {1} {2}".format(
                    action, response.status_code, response.text[:200]
                )
            )
        return self._payload(response)

    def _step_post(
        self,
        agent_id: str,
        token: str,
        step_id: str,
        action: str,
        payload: Dict[str, Any],
        idempotency_key: str,
    ) -> Dict[str, Any]:
        response = self._client.post(
            "/api/v1/agent/{0}/steps/{1}/{2}".format(
                agent_id, step_id, action
            ),
            json=payload,
            headers=self._agent_headers(token, idempotency_key),
        )
        return self._step_response(response, action)

    def ack(
        self,
        agent_id: str,
        token: str,
        step_id: str,
        attempt: int,
        idempotency_key: str,
    ) -> Dict[str, Any]:
        return self._step_post(
            agent_id,
            token,
            step_id,
            "ack",
            {"attempt": attempt},
            idempotency_key,
        )

    def renew(
        self,
        agent_id: str,
        token: str,
        step_id: str,
        attempt: int,
        idempotency_key: str,
    ) -> Dict[str, Any]:
        return self._step_post(
            agent_id,
            token,
            step_id,
            "renew",
            {"attempt": attempt},
            idempotency_key,
        )

    def report_events(
        self,
        agent_id: str,
        token: str,
        step_id: str,
        attempt: int,
        events: List[Dict[str, Any]],
        idempotency_key: str,
    ) -> Dict[str, Any]:
        return self._step_post(
            agent_id,
            token,
            step_id,
            "events",
            {"attempt": attempt, "events": events},
            idempotency_key,
        )

    def complete(
        self,
        agent_id: str,
        token: str,
        step_id: str,
        attempt: int,
        *,
        status: str,
        idempotency_key: str,
        exit_code: Optional[int] = None,
        error_class: Optional[str] = None,
        result: Optional[Dict[str, Any]] = None,
        log_uri: Optional[str] = None,
    ) -> Dict[str, Any]:
        return self._step_post(
            agent_id,
            token,
            step_id,
            "complete",
            {
                "attempt": attempt,
                "status": status,
                "exit_code": exit_code,
                "error_class": error_class,
                "result": result or {},
                "log_uri": log_uri,
            },
            idempotency_key,
        )