wild-pc/castle-api/src/castle_api/mqtt_client.py

"""MQTT client for inter-node mesh coordination.

Topics:
  castle/{hostname}/registry  — retained JSON NodeRegistry, published on connect
  castle/{hostname}/status    — "online" (retained) / "offline" (LWT)

On incoming messages from other hostnames, updates MeshStateManager.
"""

from __future__ import annotations

import asyncio
import json
import logging

import paho.mqtt.client as mqtt

from castle_core.registry import (
    Deployment,
    NodeConfig,
    NodeRegistry,
)

from castle_api.mesh import mesh_state
from castle_api.stream import broadcast

logger = logging.getLogger(__name__)


def _registry_to_json(registry: NodeRegistry) -> str:
    """Serialize a NodeRegistry to JSON for MQTT publishing.

    Only includes fields needed for mesh routing — env vars, run_cmd,
    and castle_root are excluded to avoid leaking secrets.
    """
    data: dict = {
        "node": {
            "hostname": registry.node.hostname,
            "gateway_port": registry.node.gateway_port,
        },
        "deployed": {},
    }

    for name, comp in registry.deployed.items():
        entry: dict = {
            "manager": comp.manager,
            "launcher": comp.launcher,
            "kind": comp.kind,
        }
        if comp.stack:
            entry["stack"] = comp.stack
        if comp.description:
            entry["description"] = comp.description
        if comp.port is not None:
            entry["port"] = comp.port
        if comp.health_path:
            entry["health_path"] = comp.health_path
        if comp.subdomain:
            entry["subdomain"] = comp.subdomain
        if comp.schedule:
            entry["schedule"] = comp.schedule
        if comp.managed:
            entry["managed"] = comp.managed
        data["deployed"][name] = entry

    return json.dumps(data)


def _json_to_registry(payload: str) -> NodeRegistry:
    """Deserialize a NodeRegistry from MQTT JSON payload."""
    data = json.loads(payload)
    node_data = data.get("node", {})
    node = NodeConfig(
        hostname=node_data.get("hostname", ""),
        castle_root=node_data.get("castle_root"),
        gateway_port=node_data.get("gateway_port", 9000),
    )
    deployed: dict[str, Deployment] = {}
    for name, comp_data in data.get("deployed", {}).items():
        deployed[name] = Deployment(
            manager=comp_data.get("manager", "systemd"),
            launcher=comp_data.get("launcher"),
            run_cmd=comp_data.get("run_cmd", []),
            env=comp_data.get("env", {}),
            description=comp_data.get("description"),
            kind=comp_data.get("kind", "service"),
            stack=comp_data.get("stack"),
            port=comp_data.get("port"),
            health_path=comp_data.get("health_path"),
            subdomain=comp_data.get("subdomain"),
            schedule=comp_data.get("schedule"),
            managed=comp_data.get("managed", False),
        )
    return NodeRegistry(node=node, deployed=deployed)


class CastleMQTTClient:
    """Async wrapper around paho-mqtt for castle mesh coordination."""

    def __init__(
        self,
        local_hostname: str,
        local_registry: NodeRegistry,
        broker_host: str = "localhost",
        broker_port: int = 1883,
        loop: asyncio.AbstractEventLoop | None = None,
    ) -> None:
        self._local_hostname = local_hostname
        self._local_registry = local_registry
        self._broker_host = broker_host
        self._broker_port = broker_port
        self._loop = loop or asyncio.get_event_loop()

        self._client = mqtt.Client(
            callback_api_version=mqtt.CallbackAPIVersion.VERSION2,
            client_id=f"castle-{local_hostname}",
            clean_session=True,
        )

        # LWT: if we disconnect unexpectedly, broker publishes "offline"
        self._client.will_set(
            f"castle/{local_hostname}/status",
            payload="offline",
            qos=1,
            retain=True,
        )

        self._connected = False
        self._client.on_connect = self._on_connect
        self._client.on_disconnect = self._on_disconnect
        self._client.on_message = self._on_message

    @property
    def connected(self) -> bool:
        return self._connected

    @property
    def broker_host(self) -> str:
        return self._broker_host

    @property
    def broker_port(self) -> int:
        return self._broker_port

    def _on_disconnect(
        self,
        client: mqtt.Client,
        userdata: object,
        flags: mqtt.DisconnectFlags,
        rc: mqtt.ReasonCode,
        properties: mqtt.Properties | None = None,
    ) -> None:
        self._connected = False
        logger.info("Disconnected from MQTT broker (rc=%s)", rc)

    def _on_connect(
        self,
        client: mqtt.Client,
        userdata: object,
        flags: mqtt.ConnectFlags,
        rc: mqtt.ReasonCode,
        properties: mqtt.Properties | None = None,
    ) -> None:
        """Called when connected to broker — publish our state and subscribe."""
        if rc.is_failure:
            logger.error("MQTT connect failed: %s", rc)
            self._connected = False
            return

        self._connected = True
        logger.info("Connected to MQTT broker at %s:%d", self._broker_host, self._broker_port)

        # Publish our status as online (retained)
        client.publish(
            f"castle/{self._local_hostname}/status",
            payload="online",
            qos=1,
            retain=True,
        )

        # Publish our registry (retained)
        self.publish_registry(self._local_registry)

        # Subscribe to all castle nodes
        client.subscribe("castle/+/registry", qos=1)
        client.subscribe("castle/+/status", qos=1)

    def _on_message(
        self,
        client: mqtt.Client,
        userdata: object,
        msg: mqtt.MQTTMessage,
    ) -> None:
        """Process incoming MQTT messages from other nodes."""
        try:
            parts = msg.topic.split("/")
            if len(parts) != 3 or parts[0] != "castle":
                return

            hostname = parts[1]
            msg_type = parts[2]

            # Ignore our own messages
            if hostname == self._local_hostname:
                return

            payload = msg.payload.decode()

            if msg_type == "registry":
                registry = _json_to_registry(payload)
                mesh_state.update_node(hostname, registry)
                # Notify SSE clients about mesh change
                asyncio.run_coroutine_threadsafe(
                    broadcast("mesh", {"event": "node_updated", "hostname": hostname}),
                    self._loop,
                )

            elif msg_type == "status":
                if payload == "offline":
                    mesh_state.set_offline(hostname)
                    asyncio.run_coroutine_threadsafe(
                        broadcast("mesh", {"event": "node_offline", "hostname": hostname}),
                        self._loop,
                    )

        except Exception:
            logger.exception("Error processing MQTT message on %s", msg.topic)

    def publish_registry(self, registry: NodeRegistry) -> None:
        """Publish (or re-publish) our local registry."""
        self._local_registry = registry
        self._client.publish(
            f"castle/{self._local_hostname}/registry",
            payload=_registry_to_json(registry),
            qos=1,
            retain=True,
        )

    async def start(self) -> None:
        """Connect to the broker and start the network loop."""
        self._client.connect_async(self._broker_host, self._broker_port)
        self._client.loop_start()
        logger.info(
            "MQTT client starting (broker=%s:%d)",
            self._broker_host,
            self._broker_port,
        )

    async def stop(self) -> None:
        """Disconnect and stop the network loop."""
        # Publish offline status before disconnecting
        self._client.publish(
            f"castle/{self._local_hostname}/status",
            payload="offline",
            qos=1,
            retain=True,
        )
        self._client.loop_stop()
        self._client.disconnect()
        logger.info("MQTT client stopped")