thoxmesh_node/agent.py

"""The node agent: serve a model, join the mesh, hold the lease, leave cleanly.

Lifecycle
---------
1. Resolve config and (optionally) a device identity.
2. Load the backend and start the OpenAI-compatible server on a background thread.
3. Determine the public URL — a Space origin, a tunnel, or an override.
4. Register with the mesh, retrying transport failures with backoff.
5. Heartbeat on a cadence strictly shorter than the TTL, refreshing telemetry.
6. On shutdown, deregister so the endpoint leaves routing immediately.

Ephemerality is the design constraint that shapes all of it. A Colab session can
vanish without warning — the VM is reclaimed, the browser tab closes, the
12-hour cap hits — and there is no shutdown hook that reliably fires. Clean
deregistration is therefore treated as an optimisation, not a guarantee, and
correctness rests on the lease in ``capabilities.thox_lease``: the node keeps
pushing an ``expires_at`` a TTL into the future, and any consumer treats a lapsed
lease as ineligible. That way a node that dies mid-sentence stops receiving
traffic within one TTL whether or not it ever got to say goodbye.
"""

from __future__ import annotations

import logging
import os
import signal
import threading
import time
from typing import Any

import uvicorn

from .backends import Backend, build_backend
from .config import NodeConfig, detect_node_kind, load_config, space_public_url
from .errors import ConfigError, MeshRejectedError, MeshTransportError, PublicUrlRejectedError
from .mesh_client import ControllerClient, MeshClient, Registration, build_client
from .server import create_app
from .signing import DeviceIdentity, load_identity_from_env
from .telemetry import TelemetryRecorder, build_capabilities

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


class NodeAgent:
    """Owns the server, the mesh registration and the heartbeat loop."""

    def __init__(
        self,
        config: NodeConfig,
        *,
        backend: Backend,
        identity: DeviceIdentity | None = None,
        client: MeshClient | None = None,
    ) -> None:
        self.config = config
        self.backend = backend
        self.recorder = TelemetryRecorder()
        self.state: dict[str, Any] = {
            "started_at": time.time(),
            "node_kind": config.node_kind,
            "transport": None,
            "endpoint_id": None,
            "public_url": config.public_url or None,
            "last_heartbeat_at": None,
            "last_error": None,
        }
        self._identity = identity
        self._client = client or build_client(config, identity)
        self._registration: Registration | None = None
        self._stop = threading.Event()
        self._heartbeat_thread: threading.Thread | None = None
        self._server: uvicorn.Server | None = None
        self._server_thread: threading.Thread | None = None

        self.app = create_app(
            backend=backend,
            recorder=self.recorder,
            model_id=config.model_id,
            context_window=config.context_window,
            max_tokens_cap=config.max_tokens_cap,
            node_state=self.state,
        )

    # ── Capabilities ─────────────────────────────────────────────────────

    def _capabilities(self) -> dict[str, Any]:
        """Current capability + telemetry payload for register/heartbeat."""
        return build_capabilities(
            base=self.config.capabilities,
            recorder=self.recorder,
            node_kind=self.config.node_kind,
            ephemeral=self.config.ephemeral,
            heartbeat_seconds=self.config.heartbeat_seconds,
            ttl_seconds=self.config.ttl_seconds,
        )

    # ── Server ───────────────────────────────────────────────────────────

    def start_server(self) -> None:
        """Run uvicorn on a daemon thread and wait until it accepts requests."""
        uvicorn_config = uvicorn.Config(
            self.app,
            host=self.config.host,
            port=self.config.port,
            log_level="info",
            access_log=False,
        )
        self._server = uvicorn.Server(uvicorn_config)
        self._server_thread = threading.Thread(target=self._server.run, daemon=True, name="thox-node-server")
        self._server_thread.start()

        deadline = time.time() + 60
        while time.time() < deadline:
            if getattr(self._server, "started", False):
                logger.info("serving on %s:%d", self.config.host, self.config.port)
                return
            if not self._server_thread.is_alive():
                raise ConfigError("the HTTP server thread exited during startup")
            time.sleep(0.25)
        raise ConfigError("the HTTP server did not start within 60s")

    # ── Registration ─────────────────────────────────────────────────────

    def register(self, public_url: str, *, max_attempts: int = 5) -> Registration:
        """Publish this node, retrying transport errors with linear backoff.

        A ``PublicUrlRejectedError`` on the canonical transport is not retried:
        it means the mesh's private-address policy refused this URL, and no
        number of retries changes that. When a controller URL is configured the
        agent degrades to it, which is what keeps an ephemeral public node usable
        while the policy question is resolved with the Mesh backend lane.
        """
        self.state["public_url"] = public_url
        last_error: Exception | None = None

        for attempt in range(1, max_attempts + 1):
            try:
                registration = self._client.register(
                    self.config, public_url, self._capabilities(), self.recorder.status()
                )
                self._registration = registration
                self.state.update(
                    transport=registration.transport,
                    endpoint_id=registration.endpoint_id,
                    last_error=None,
                )
                return registration
            except PublicUrlRejectedError as exc:
                logger.error("mesh refused the public URL: %s", exc)
                self.state["last_error"] = str(exc)
                if self.config.controller_url and not isinstance(self._client, ControllerClient):
                    logger.warning(
                        "falling back to controller transport at %s because the v2 control "
                        "plane only accepts private addresses",
                        self.config.controller_url,
                    )
                    self._client = ControllerClient(self.config.controller_url)
                    continue
                raise
            except MeshTransportError as exc:
                last_error = exc
                self.state["last_error"] = str(exc)
                wait = self.config.register_retry_seconds * attempt
                logger.warning("register attempt %d/%d failed: %s (retrying in %ds)", attempt, max_attempts, exc, wait)
                if attempt < max_attempts:
                    time.sleep(wait)
            except MeshRejectedError as exc:
                self.state["last_error"] = str(exc)
                logger.error("mesh rejected registration: %s", exc)
                raise

        raise MeshTransportError(f"could not register after {max_attempts} attempts: {last_error}")

    # ── Heartbeat ────────────────────────────────────────────────────────

    def _heartbeat_loop(self) -> None:
        """Renew the lease until stopped, tolerating transient mesh outages."""
        consecutive_failures = 0
        while not self._stop.wait(self.config.heartbeat_seconds):
            registration = self._registration
            if registration is None:
                continue
            try:
                self._client.heartbeat(
                    self.config, registration, self._capabilities(), self.recorder.status()
                )
                self.state["last_heartbeat_at"] = time.time()
                self.state["last_error"] = None
                consecutive_failures = 0
            except PublicUrlRejectedError as exc:
                # Policy, not an outage; stop hammering and let the lease lapse.
                logger.error("heartbeat refused by URL policy, stopping renewal: %s", exc)
                self.state["last_error"] = str(exc)
                return
            except (MeshTransportError, MeshRejectedError) as exc:
                consecutive_failures += 1
                self.state["last_error"] = str(exc)
                logger.warning("heartbeat failed (%d in a row): %s", consecutive_failures, exc)
                if consecutive_failures >= 3:
                    logger.error(
                        "3 consecutive heartbeat failures; the lease will lapse in ~%ds "
                        "and consumers will stop routing here",
                        self.config.ttl_seconds,
                    )

    def start_heartbeat(self) -> None:
        if self._heartbeat_thread is not None:
            return
        self._heartbeat_thread = threading.Thread(target=self._heartbeat_loop, daemon=True, name="thox-node-heartbeat")
        self._heartbeat_thread.start()
        logger.info(
            "heartbeat every %ds, lease TTL %ds", self.config.heartbeat_seconds, self.config.ttl_seconds
        )

    # ── Shutdown ─────────────────────────────────────────────────────────

    def shutdown(self) -> None:
        """Leave the mesh and stop serving. Safe to call more than once."""
        if self._stop.is_set():
            return
        self._stop.set()
        logger.info("shutting down")
        if self._registration is not None:
            self._client.deregister(self._registration)
            self._registration = None
            self.state["endpoint_id"] = None
        if self._server is not None:
            self._server.should_exit = True
        try:
            self.backend.close()
        except Exception:  # noqa: BLE001
            logger.debug("backend close failed", exc_info=True)

    def install_signal_handlers(self) -> None:
        """Deregister on SIGINT/SIGTERM where the platform supports it.

        Colab interrupts land as SIGINT and Spaces stop containers with SIGTERM,
        so both give us a chance to leave cleanly. A hard VM reclaim gives no
        signal at all, which is exactly why the lease exists.
        """
        def _handler(signum, _frame):  # type: ignore[no-untyped-def]
            logger.info("signal %s received", signum)
            self.shutdown()

        for sig in (signal.SIGINT, signal.SIGTERM):
            try:
                signal.signal(sig, _handler)
            except (ValueError, OSError):
                # Not the main thread (common in notebooks); the lease still covers us.
                logger.debug("could not install handler for %s", sig)

    def wait(self) -> None:
        """Block until shutdown. Used by the CLI entry point."""
        try:
            while not self._stop.is_set():
                time.sleep(1)
        except KeyboardInterrupt:
            self.shutdown()


def resolve_public_url(config: NodeConfig, *, tunnel_provider: str = "cloudflared") -> tuple[str, Any]:
    """Determine how the mesh will reach this node.

    Order: explicit override, then a Space's own origin, then a tunnel. Returns
    the URL and the tunnel handle (``None`` when no tunnel was started) so the
    caller can close it.
    """
    if config.public_url:
        logger.info("using configured public URL %s", config.public_url)
        return config.public_url, None

    if config.node_kind == "hf-space":
        url = space_public_url()
        if not url:
            raise ConfigError(
                "running on a Space but neither SPACE_HOST nor SPACE_ID is set; "
                "set THOX_PUBLIC_URL explicitly rather than registering an unreachable address"
            )
        return url, None

    from .tunnel import start_tunnel

    tunnel = start_tunnel(tunnel_provider, config.port)
    return tunnel.public_url, tunnel


def run(
    *,
    env: dict[str, str] | None = None,
    echo: bool = False,
    tunnel_provider: str = "cloudflared",
    block: bool = True,
) -> NodeAgent:
    """Start a fully wired node. The single entry point notebooks and Spaces call."""
    logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(name)s] %(levelname)s %(message)s")

    config = load_config(env)
    if config.node_kind == "unknown":
        config = type(config)(**{**config.__dict__, "node_kind": detect_node_kind(env)})

    identity = load_identity_from_env(env)
    backend = build_backend(
        model_id=config.model_id,
        repo_id=config.model_repo,
        filename=config.model_file,
        n_ctx=config.context_window,
        n_threads=config.n_threads,
        n_gpu_layers=config.n_gpu_layers,
        token=(env or os.environ).get("HF_TOKEN"),
        echo=echo,
    )

    agent = NodeAgent(config, backend=backend, identity=identity)
    agent.install_signal_handlers()
    agent.start_server()

    public_url, tunnel = resolve_public_url(config, tunnel_provider=tunnel_provider)
    try:
        agent.register(public_url)
    except Exception:
        if tunnel is not None:
            tunnel.close()
        raise
    agent.start_heartbeat()

    logger.info("node live: %s -> %s", config.model_id, public_url)
    if block:
        try:
            agent.wait()
        finally:
            agent.shutdown()
            if tunnel is not None:
                tunnel.close()
    return agent