agentic-mobile-control/packages/cloud-platform/cloud/sdk/client.py

"""Thin Python SDK client for the ``/v1`` REST API.

Capability: ``platform-sdk``.

Uses ``httpx`` directly so the client can talk to any deployed cloud-runtime
process over HTTP, or to a FastAPI ``TestClient`` instance in tests.
"""

from __future__ import annotations

from typing import Any

import httpx


class CloudAPIError(httpx.HTTPStatusError):
    def __init__(self, response: httpx.Response, detail: str) -> None:
        super().__init__(
            f"cloud API request failed with status {response.status_code}: {detail}",
            request=response.request,
            response=response,
        )
        self.status_code = response.status_code
        self.detail = detail


class CloudAuthorizationError(CloudAPIError):
    pass


class CloudClient:
    """A minimal Python wrapper for the platform SDK's ``/v1`` routes."""

    def __init__(
        self,
        base_url: str,
        *,
        http_client: httpx.Client | Any = None,
        api_prefix: str = "/v1",
        token: str | None = None,
        auth: httpx.Auth | None = None,
    ) -> None:
        if token is not None and auth is not None:
            raise ValueError("token and auth are mutually exclusive")
        self._base_url = base_url.rstrip("/")
        self._api_prefix = api_prefix.rstrip("/")
        self._headers = {"Authorization": f"Bearer {token}"} if token else None
        self._auth = auth
        if http_client is None:
            self._http = httpx.Client(base_url=self._base_url)
            self._owns_client = True
        else:
            self._http = http_client
            self._owns_client = False

    def close(self) -> None:
        if self._owns_client:
            self._http.close()

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

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

    # ------------------------------------------------------------------- tasks

    def submit_task(
        self,
        *,
        goal: str | None = None,
        workflow_definition_id: str | None = None,
        driver_type: str | None = None,
        capability_tags: list[str] | None = None,
        target_host_id: str | None = None,
        target_device_id: str | None = None,
    ) -> dict[str, Any]:
        payload: dict[str, Any] = {
            "goal": goal,
            "workflow_definition_id": workflow_definition_id,
        }
        if (
            driver_type is not None
            or capability_tags is not None
            or target_host_id is not None
            or target_device_id is not None
        ):
            payload["constraints"] = {
                "driver_type": driver_type,
                "capability_tags": list(capability_tags or []),
                "target_host_id": target_host_id,
                "target_device_id": target_device_id,
            }
        resp = self._request("POST", "/tasks", json=payload)
        return resp.json()

    def get_task_status(self, task_id: str) -> dict[str, Any]:
        resp = self._request("GET", f"/tasks/{task_id}")
        return resp.json()

    def list_tasks(
        self,
        *,
        status: str | None = None,
        limit: int = 50,
        offset: int = 0,
    ) -> dict[str, Any]:
        params: dict[str, Any] = {"limit": limit, "offset": offset}
        if status is not None:
            params["status"] = status
        resp = self._request("GET", "/tasks", params=params)
        return resp.json()

    def get_task_attempts(self, task_id: str) -> list[dict[str, Any]]:
        resp = self._request("GET", f"/tasks/{task_id}/attempts")
        return resp.json()

    def cancel_task(self, task_id: str) -> dict[str, Any]:
        resp = self._request("POST", f"/tasks/{task_id}/cancel")
        return resp.json()

    # ----------------------------------------------------------------- devices

    def list_devices(self) -> list[dict[str, Any]]:
        resp = self._request("GET", "/devices")
        return resp.json()

    def list_hosts(self) -> list[dict[str, Any]]:
        resp = self._request("GET", "/hosts")
        return resp.json()

    # ----------------------------------------------------------------- plugins

    def list_plugins(self) -> list[dict[str, Any]]:
        resp = self._request("GET", "/plugins")
        return resp.json()

    def register_plugin(
        self,
        *,
        name: str,
        version: str,
        entry_point_kind: str,
        target: str,
    ) -> dict[str, Any]:
        payload = {
            "name": name,
            "version": version,
            "entry_point_kind": entry_point_kind,
            "target": target,
        }
        resp = self._request("POST", "/plugins", json=payload)
        return resp.json()

    # --------------------------------------------------------------- user auth

    def login(self, *, username: str, password: str) -> dict[str, Any]:
        resp = self._request(
            "POST",
            "/auth/login",
            json={"username": username, "password": password},
        )
        return resp.json()

    def current_user(self) -> dict[str, Any]:
        resp = self._request("GET", "/auth/me")
        return resp.json()

    def logout(self) -> None:
        self._request("POST", "/auth/logout")

    def change_password(self, *, current_password: str, new_password: str) -> None:
        self._request(
            "POST",
            "/auth/password",
            json={
                "current_password": current_password,
                "new_password": new_password,
            },
        )

    def list_users(self, *, limit: int = 50, offset: int = 0) -> dict[str, Any]:
        resp = self._request(
            "GET",
            "/users",
            params={"limit": limit, "offset": offset},
        )
        return resp.json()

    def create_user(
        self,
        *,
        username: str,
        display_name: str,
        role: str,
        password: str,
    ) -> dict[str, Any]:
        resp = self._request(
            "POST",
            "/users",
            json={
                "username": username,
                "display_name": display_name,
                "role": role,
                "password": password,
            },
        )
        return resp.json()

    def update_user(self, user_id: str, **changes: Any) -> dict[str, Any]:
        resp = self._request("PATCH", f"/users/{user_id}", json=changes)
        return resp.json()

    def reset_user_password(self, user_id: str, *, password: str) -> dict[str, Any]:
        resp = self._request(
            "POST",
            f"/users/{user_id}/password",
            json={"password": password},
        )
        return resp.json()

    def revoke_user_sessions(self, user_id: str) -> None:
        self._request("DELETE", f"/users/{user_id}/sessions")

    # ------------------------------------------------------------- governance

    def get_user_submission_policy(self, user_id: str) -> dict[str, Any]:
        return self._request("GET", f"/users/{user_id}/submission-policy").json()

    def update_user_submission_policy(
        self, user_id: str, **policy: Any
    ) -> dict[str, Any]:
        return self._request(
            "PUT", f"/users/{user_id}/submission-policy", json=policy
        ).json()

    def get_host_governance_policy(self, host_id: str) -> dict[str, Any]:
        return self._request("GET", f"/hosts/{host_id}/governance-policy").json()

    def update_host_governance_policy(
        self, host_id: str, **policy: Any
    ) -> dict[str, Any]:
        return self._request(
            "PUT", f"/hosts/{host_id}/governance-policy", json=policy
        ).json()

    def get_host_token_usage(self, host_id: str) -> dict[str, Any]:
        return self._request("GET", f"/hosts/{host_id}/token-usage").json()

    def list_host_token_usage_events(
        self, host_id: str, *, limit: int = 50, offset: int = 0
    ) -> list[dict[str, Any]]:
        return self._request(
            "GET",
            f"/hosts/{host_id}/token-usage-events",
            params={"limit": limit, "offset": offset},
        ).json()

    # ------------------------------------------------------------------ helpers

    def _url(self, path: str) -> str:
        return f"{self._base_url}{self._api_prefix}{path}"

    def _request(
        self,
        method: str,
        path: str,
        *,
        json: dict[str, Any] | None = None,
        params: dict[str, Any] | None = None,
    ) -> httpx.Response:
        headers = dict(self._headers or {})
        if method in {"POST", "PUT", "PATCH", "DELETE"} and not headers:
            csrf_token = _cookie_value(self._http, "amcp_csrf")
            if csrf_token:
                headers["X-CSRF-Token"] = csrf_token
        response = self._http.request(
            method,
            self._url(path),
            json=json,
            params=params,
            headers=headers or None,
            auth=self._auth,
        )
        if response.is_success:
            return response
        try:
            payload = response.json()
        except ValueError:
            payload = {}
        detail = payload.get("detail") or "request rejected"
        error_type = (
            CloudAuthorizationError
            if response.status_code in {401, 403}
            else CloudAPIError
        )
        raise error_type(response, str(detail))


def _cookie_value(client: Any, name: str) -> str | None:
    cookies = getattr(client, "cookies", None)
    if cookies is None:
        return None
    value = cookies.get(name)
    return value if isinstance(value, str) else None