thoxmesh_node/mesh_client.py

"""Clients that publish this node into the mesh.

Two transports exist because the mesh has two reachable control planes with
different admission rules, and exactly one of them can accept a public node today:

``DeviceV2Client`` — the canonical MeshStack control plane
    Talks to the ``meshstack-device-v2`` edge function with Ed25519 request
    signing. This is the contract owned by the Mesh backend lane and it is what
    a paired ThoxOS device uses. It requires (a) a device row on
    ``meshstack_devices_v2`` with our public key, minted through a pairing
    session by an authenticated mesh owner, and (b) a ``base_url`` that passes
    ``validateLocalBaseUrl`` — RFC1918, loopback, link-local, CGNAT or
    ``.local`` only.

    Requirement (b) is why a Colab tunnel cannot currently self-register here:
    ``https://<x>.trycloudflare.com`` and ``https://<x>.hf.space`` are public
    hostnames and the function rejects them by design (anti-SSRF: the mesh must
    not be induced to call arbitrary internet hosts). We surface that refusal as
    ``PublicUrlRejectedError`` rather than retrying it, and the agent degrades to
    the controller transport.

``ControllerClient`` — this repo's mesh controller
    The pre-existing in-memory registry (``controller/mesh.py``) that already
    accepts public Colab and Space URLs and health-checks them. It is the
    transport that makes a live demo possible today, and it is the natural place
    for a public-node broker to live even after the v2 policy question is settled.

Both implement ``MeshClient`` so the agent is written once against one interface.
"""

from __future__ import annotations

import json
import logging
import time
import urllib.error
import urllib.request
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Any, Mapping

from .config import NodeConfig
from .errors import (
    MeshRejectedError,
    MeshTransportError,
    PublicUrlRejectedError,
)
from .signing import DeviceIdentity, check_clock_skew

logger = logging.getLogger("thox.mesh")

#: Substring of the edge function's private-address refusal.
_PRIVATE_URL_REFUSAL = "private mesh, LAN, link-local, or loopback"

#: HTTP statuses worth retrying; everything else is a contract decision.
_RETRYABLE_STATUS = frozenset({408, 425, 429, 500, 502, 503, 504})


@dataclass
class Registration:
    """Result of publishing this node into the mesh."""

    endpoint_id: str
    transport: str
    base_url: str
    raw: dict[str, Any]


class MeshClient(ABC):
    """Interface the agent uses to join, hold and leave the mesh."""

    name: str = "abstract"

    @abstractmethod
    def register(self, config: NodeConfig, base_url: str, capabilities: Mapping[str, Any], status: str) -> Registration:
        """Publish (or refresh) this node's endpoint. Must be idempotent."""

    @abstractmethod
    def heartbeat(self, config: NodeConfig, registration: Registration, capabilities: Mapping[str, Any], status: str) -> None:
        """Renew liveness and refresh telemetry."""

    @abstractmethod
    def deregister(self, registration: Registration) -> None:
        """Remove the endpoint from routing. Best-effort; must never raise."""


# ── MeshStack v2 (canonical) ─────────────────────────────────────────────


class DeviceV2Client(MeshClient):
    """Signed client for the ``meshstack-device-v2`` edge function.

    Note there is no Supabase API key anywhere in this class. The function runs
    with JWT verification disabled and authenticates purely on the Ed25519
    signature, so the device private key is the entire credential. That is a
    deliberately smaller secret surface than shipping an anon key into a public
    notebook, where it would be trivially extractable.
    """

    name = "device_v2"

    def __init__(self, identity: DeviceIdentity, function_url: str, *, timeout: int = 20) -> None:
        self._identity = identity
        self._url = function_url
        self._timeout = timeout
        self._clock_checked = False

    def _call(self, action: str, payload: Mapping[str, Any]) -> dict[str, Any]:
        """Sign and POST one action, translating failures into typed errors."""
        body_payload = dict(payload)
        body_payload["action"] = action
        signed = self._identity.sign_request(action, body_payload)
        request = urllib.request.Request(self._url, data=signed.body, headers=signed.headers, method="POST")
        try:
            with urllib.request.urlopen(request, timeout=self._timeout) as response:
                return json.loads(response.read().decode("utf-8") or "{}")
        except urllib.error.HTTPError as exc:
            detail = _read_error_body(exc)
            message = detail.get("message") or detail.get("error") or f"HTTP {exc.code}"
            if _PRIVATE_URL_REFUSAL in message:
                raise PublicUrlRejectedError(message, status=exc.code, action=action) from exc
            if exc.code in _RETRYABLE_STATUS:
                raise MeshTransportError(f"{action} failed with retryable HTTP {exc.code}: {message}") from exc
            raise MeshRejectedError(message, status=exc.code, action=action) from exc
        except urllib.error.URLError as exc:
            raise MeshTransportError(f"{action} could not reach the mesh: {exc.reason}") from exc
        except (TimeoutError, json.JSONDecodeError) as exc:
            raise MeshTransportError(f"{action} failed: {exc}") from exc

    def register(self, config: NodeConfig, base_url: str, capabilities: Mapping[str, Any], status: str) -> Registration:
        response = self._call(
            "register_model",
            {
                "model_id": config.model_id,
                "alias": config.alias,
                "protocol": config.protocol,
                "base_url": base_url,
                "context_window": config.context_window,
                "capabilities": dict(capabilities),
                "resource_requirements": {
                    "node_kind": config.node_kind,
                    "ephemeral": config.ephemeral,
                    "context_window": config.context_window,
                },
                "status": status,
            },
        )
        endpoint = response.get("endpoint") or {}
        endpoint_id = endpoint.get("id")
        if not endpoint_id:
            raise MeshRejectedError("register_model returned no endpoint id", action="register_model")
        logger.info("registered endpoint %s via device_v2 (%s)", endpoint_id, base_url)
        return Registration(endpoint_id=endpoint_id, transport=self.name, base_url=base_url, raw=endpoint)

    def heartbeat(self, config: NodeConfig, registration: Registration, capabilities: Mapping[str, Any], status: str) -> None:
        """Renew device liveness, then refresh the endpoint's telemetry.

        The backend's ``heartbeat`` action updates the **device** row only — it
        never touches ``meshstack_model_endpoints_v2``. Telemetry therefore has
        to ride on a ``register_model`` upsert, which is idempotent on
        ``(device_id, alias)`` and is the contract's own refresh path.
        """
        response = self._call(
            "heartbeat",
            {
                "capabilities": dict(capabilities),
                "resource_snapshot": capabilities.get("thox_telemetry", {}),
                "app_version": "thoxmesh-node/1.0",
            },
        )
        if not self._clock_checked:
            server_time = response.get("server_time")
            if isinstance(server_time, str):
                parsed = _parse_server_time(server_time)
                if parsed is not None:
                    skew = check_clock_skew(parsed)
                    logger.info("clock skew vs mesh: %.1fs", skew)
            self._clock_checked = True
        self.register(config, registration.base_url, capabilities, status)

    def deregister(self, registration: Registration) -> None:
        try:
            self._call("unregister_model", {"endpoint_id": registration.endpoint_id})
            logger.info("deregistered endpoint %s", registration.endpoint_id)
        except Exception as exc:  # noqa: BLE001 - shutdown must not raise
            logger.warning("deregister failed (endpoint will expire by lease): %s", exc)


# ── Controller (public-node broker, works today) ─────────────────────────


class ControllerClient(MeshClient):
    """Client for this repo's mesh controller, which accepts public URLs.

    The controller keeps endpoints in memory and health-checks them, so a lapsed
    Colab node is dropped by its own health loop as well as by our lease.
    """

    name = "controller"

    def __init__(self, controller_url: str, *, timeout: int = 20, token: str | None = None) -> None:
        self._url = controller_url.rstrip("/")
        self._timeout = timeout
        self._token = token

    def _post(self, path: str, payload: Mapping[str, Any]) -> dict[str, Any]:
        body = json.dumps(payload).encode("utf-8")
        headers = {"content-type": "application/json"}
        if self._token:
            headers["authorization"] = f"Bearer {self._token}"
        request = urllib.request.Request(f"{self._url}{path}", data=body, headers=headers, method="POST")
        try:
            with urllib.request.urlopen(request, timeout=self._timeout) as response:
                return json.loads(response.read().decode("utf-8") or "{}")
        except urllib.error.HTTPError as exc:
            detail = _read_error_body(exc)
            message = detail.get("message") or detail.get("error") or f"HTTP {exc.code}"
            if exc.code in _RETRYABLE_STATUS:
                raise MeshTransportError(f"{path} retryable HTTP {exc.code}: {message}") from exc
            raise MeshRejectedError(message, status=exc.code, action=path) from exc
        except urllib.error.URLError as exc:
            raise MeshTransportError(f"{path} could not reach the controller: {exc.reason}") from exc
        except (TimeoutError, json.JSONDecodeError) as exc:
            raise MeshTransportError(f"{path} failed: {exc}") from exc

    def register(self, config: NodeConfig, base_url: str, capabilities: Mapping[str, Any], status: str) -> Registration:
        response = self._post(
            "/mesh/register",
            {
                "model_id": config.model_id,
                "alias": config.alias,
                "base_url": base_url,
                "source": config.node_kind,
                "protocol": config.protocol,
                "context_window": config.context_window,
                "capabilities": dict(capabilities),
                "status": status,
            },
        )
        logger.info("registered %s via controller (%s)", config.model_id, base_url)
        return Registration(
            endpoint_id=response.get("model_id") or config.model_id,
            transport=self.name,
            base_url=base_url,
            raw=response,
        )

    def heartbeat(self, config: NodeConfig, registration: Registration, capabilities: Mapping[str, Any], status: str) -> None:
        # The controller's register is an upsert, so a re-register is the heartbeat.
        self.register(config, registration.base_url, capabilities, status)

    def deregister(self, registration: Registration) -> None:
        try:
            self._post("/mesh/deregister", {"model_id": registration.endpoint_id})
            logger.info("deregistered %s from controller", registration.endpoint_id)
        except Exception as exc:  # noqa: BLE001 - shutdown must not raise
            logger.warning("controller deregister failed: %s", exc)


class NullClient(MeshClient):
    """No-op transport for local development and tests.

    Chosen when no mesh credentials and no controller URL are configured, so the
    node still serves completions and can be exercised end-to-end without any
    control plane at all.
    """

    name = "none"

    def register(self, config: NodeConfig, base_url: str, capabilities: Mapping[str, Any], status: str) -> Registration:
        logger.warning("no mesh transport configured; serving locally without registration")
        return Registration(endpoint_id="local", transport=self.name, base_url=base_url, raw={})

    def heartbeat(self, config: NodeConfig, registration: Registration, capabilities: Mapping[str, Any], status: str) -> None:
        return None

    def deregister(self, registration: Registration) -> None:
        return None


def build_client(config: NodeConfig, identity: DeviceIdentity | None) -> MeshClient:
    """Pick a transport from configuration and available credentials.

    ``auto`` prefers the canonical control plane when a device identity exists,
    because that is the contract everything else in the fleet speaks; it falls
    back to the controller so an unpaired node is still useful.
    """
    if config.transport == "device_v2":
        if identity is None:
            raise MeshRejectedError(
                "transport device_v2 requires THOX_MESH_DEVICE_ID and a device key; "
                "pair the device first or use THOX_MESH_TRANSPORT=controller"
            )
        return DeviceV2Client(identity, config.device_function_url)
    if config.transport == "controller":
        return ControllerClient(config.controller_url)
    # auto
    if identity is not None:
        return DeviceV2Client(identity, config.device_function_url)
    if config.controller_url:
        return ControllerClient(config.controller_url)
    return NullClient()


def _read_error_body(exc: urllib.error.HTTPError) -> dict[str, Any]:
    """Decode a JSON error body, tolerating non-JSON responses."""
    try:
        raw = exc.read().decode("utf-8")
    except Exception:  # noqa: BLE001
        return {}
    try:
        parsed = json.loads(raw)
        return parsed if isinstance(parsed, dict) else {"message": raw[:400]}
    except json.JSONDecodeError:
        return {"message": raw[:400]}


def _parse_server_time(value: str) -> int | None:
    """Parse the backend's ISO timestamp into epoch milliseconds."""
    text = value.strip().replace("Z", "+0000")
    for fmt in ("%Y-%m-%dT%H:%M:%S.%f%z", "%Y-%m-%dT%H:%M:%S%z"):
        try:
            import datetime as _dt

            return int(_dt.datetime.strptime(text, fmt).timestamp() * 1000)
        except ValueError:
            continue
    return None