atlas-iac/dockerfiles/hermes-worker-hux/hux_hook/client.py

"""Thin urllib client for the HUX service on the pod loopback.

The client owns three things: the identity headers ``hux.identity`` expects,
the mapping of ``hux.error.v1`` bodies to ``HuxServiceError``, and the
per-process capabilities cache. It never logs or embeds request bodies in
exceptions (SO-07, SO-11): an error carries status, code and the service's
own message only.
"""

from __future__ import annotations

import http.client
import json
import threading
import urllib.error
import urllib.request
from collections.abc import Mapping
from typing import Any

HEADER_SLOT = "X-Hermes-Tenant-Identity"
HEADER_SUBJECT = "X-Hux-Subject"
HEADER_SURFACE = "X-Hux-Surface"
HEADER_TRUST = "X-Hux-Trust"
HEADER_KEY = "X-Hux-Relay-Key"
DEFAULT_BASE_URL = "http://127.0.0.1:8790"
MESSAGE_MAX = 280


class HuxServiceError(Exception):
    """The service answered with a ``hux.error.v1`` body (or a non-JSON failure)."""

    def __init__(self, status: int, code: str, message: str) -> None:
        super().__init__(f"{status} {code}: {message}")
        self.status = int(status)
        self.code = str(code)
        self.message = str(message)[:MESSAGE_MAX]


class HuxUnavailable(HuxServiceError):
    """The service could not be reached at all; callers decide open or closed."""

    def __init__(self, message: str = "hux service unreachable") -> None:
        super().__init__(0, "unavailable", message)


class HuxResponse:
    """Status, parsed JSON body and headers of one answer."""

    def __init__(self, status: int, body: Any, headers: Mapping[str, str]) -> None:
        self.status = status
        self.body = body
        self.headers = {k.lower(): v for k, v in headers.items()}

    def header(self, name: str) -> str:
        """Case-insensitive header lookup; empty when absent."""
        return self.headers.get(name.lower(), "")


def _error_from(status: int, body: Any) -> HuxServiceError:
    if isinstance(body, dict) and body.get("schema") == "hux.error.v1":
        return HuxServiceError(int(body.get("status", status)), str(body.get("code", "invalid")), str(body.get("message", "")))
    return HuxServiceError(status, "invalid", f"unexpected response {status}")


class HuxClient:
    """One tenant identity talking to one HUX base URL.

    ``identity`` carries ``tenant_slot``, ``subject``, ``surface`` and
    ``trust``; ``key`` is the relay or worker shared key when ``trust`` needs
    one. The agent hook normally runs as ``surface=worker, trust=worker``.
    """

    def __init__(self, base_url: str = DEFAULT_BASE_URL, identity: Mapping[str, str] | None = None,
                 key: str | None = None, timeout: float = 5) -> None:
        identity = dict(identity or {})
        self.base_url = base_url.rstrip("/")
        self.identity = {
            "tenant_slot": str(identity.get("tenant_slot", "")), "subject": str(identity.get("subject", "")),
            "surface": str(identity.get("surface", "worker")), "trust": str(identity.get("trust", "worker")),
        }
        self._key = key
        self.timeout = timeout
        self._capabilities: dict[str, Any] | None = None
        self._guard = threading.Lock()

    # -- transport -------------------------------------------------------------

    def headers(self, extra: Mapping[str, str] | None = None) -> dict[str, str]:
        """Identity headers exactly as ``hux.identity.resolve`` reads them, plus ``extra``."""
        out = {
            HEADER_SLOT: self.identity["tenant_slot"], HEADER_SUBJECT: self.identity["subject"],
            HEADER_SURFACE: self.identity["surface"], HEADER_TRUST: self.identity["trust"],
            "Accept": "application/json",
        }
        if self._key:
            out[HEADER_KEY] = self._key
        for name, value in (extra or {}).items():
            if value:
                out[name] = str(value)
        return out

    def request(self, method: str, path: str, body: Any = None, *, idempotency_key: str | None = None,
                if_match: int | str | None = None, query: Mapping[str, str] | None = None) -> HuxResponse:
        """Send one request; raise ``HuxServiceError`` on 4xx/5xx and ``HuxUnavailable`` on transport failure."""
        extra: dict[str, str] = {}
        if idempotency_key:
            extra["Idempotency-Key"] = idempotency_key
        if if_match is not None:
            extra["If-Match"] = str(if_match)
        data = None
        if body is not None:
            data = json.dumps(body, sort_keys=True, separators=(",", ":")).encode()
            extra["Content-Type"] = "application/json"
        url = self.base_url + path
        if query:
            url += "?" + "&".join(f"{k}={urllib.request.quote(str(v), safe='')}" for k, v in query.items())
        req = urllib.request.Request(url, data=data, method=method, headers=self.headers(extra))
        try:
            with urllib.request.urlopen(req, timeout=self.timeout) as raw:  # noqa: S310 - loopback only
                response = HuxResponse(raw.status, _decode(raw.read()), dict(raw.headers.items()))
        except urllib.error.HTTPError as error:
            payload = _decode(error.read())
            raise _error_from(error.code, payload) from None
        except (urllib.error.URLError, OSError, TimeoutError) as error:
            raise HuxUnavailable(f"hux service unreachable: {type(error).__name__}") from None
        except http.client.InvalidURL:
            # http.client refuses control characters and spaces in the path; treat it as our own bad request.
            raise HuxServiceError(400, "invalid", "malformed request path") from None
        except http.client.HTTPException as error:
            raise HuxUnavailable(f"hux service unreachable: {type(error).__name__}") from None
        if response.status >= 400:
            raise _error_from(response.status, response.body)
        return response

    def get(self, path: str, query: Mapping[str, str] | None = None) -> HuxResponse:
        """``GET path``."""
        return self.request("GET", path, query=query)

    def post(self, path: str, body: Any, idempotency_key: str | None = None) -> HuxResponse:
        """``POST path`` with an optional ``Idempotency-Key``."""
        return self.request("POST", path, body, idempotency_key=idempotency_key)

    def put(self, path: str, body: Any, if_match: int | str | None = None) -> HuxResponse:
        """``PUT path`` with an optional ``If-Match`` revision."""
        return self.request("PUT", path, body, if_match=if_match)

    # -- capabilities ------------------------------------------------------------

    def capabilities(self, refresh: bool = False) -> dict[str, Any]:
        """Which cards are on, cached per process; unreachable or off means every card reads as off."""
        with self._guard:
            if self._capabilities is not None and not refresh:
                return self._capabilities
            try:
                body = self.get("/hux/v1/capabilities").body
            except HuxServiceError as error:
                return {"reachable": error.code != "unavailable", "cards": {}, "contract_version": ""}
            cards = {c["card"]: bool(c.get("enabled")) for c in body.get("cards", []) if isinstance(c, dict) and "card" in c}
            self._capabilities = {"reachable": True, "cards": cards, "contract_version": str(body.get("contract_version", ""))}
            return self._capabilities

    def card_enabled(self, card: str) -> bool:
        """True only when the service answered and reports ``card`` on."""
        return bool(self.capabilities()["cards"].get(card))

    def forget_capabilities(self) -> None:
        """Drop the cache so the next call re-reads flags (used after a reload signal)."""
        with self._guard:
            self._capabilities = None


def _decode(raw: bytes) -> Any:
    if not raw:
        return None
    try:
        return json.loads(raw)
    except ValueError:
        return None