"""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")