wild-cloud/wild-pc
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(mesh): NATS JetStream transport replacing MQTT

Browse Source 
Replace the Mosquitto/paho mesh transport with NATS JetStream KV:
- mesh_wire.py: transport-agnostic registry (de)serialization (secret-stripped)
- nats_client.py: async CastleNATSClient — registry in the castle-registry KV
  bucket, heartbeat liveness, delete-on-stop offline, watch-driven mesh state
- rewire main/config/models/nodes; MeshStatus -> connected/nats_url
- frontend MeshStatus + MeshPanel to the new fields
- drop paho-mqtt/mqtt_client + _mqtt._tcp mDNS browse; add nats-py
- keep the secret-stripping invariant; mDNS peer discovery retained

Single-node verified; two-node parity pending the second node.
This commit is contained in:
Paul Payne
2026-07-07 04:55:02 -07:00
parent f826d972b8
commit 6c00a9d33c
12 changed files with 436 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
castle-api/pyproject.toml
View File

@@ -9,7 +9,7 @@ dependencies = [
"pydantic-settings>=2.0.0",
"httpx>=0.27.0",
"castle-core",
"paho-mqtt>=2.0.0",
"nats-py>=2.9.0",
"zeroconf>=0.131.0",
]

5
castle-api/src/castle_api/config.py
View File

@@ -15,9 +15,8 @@ class Settings(BaseSettings):
port: int = 9020
# Mesh coordination (all off by default — single-node works without them)
mqtt_enabled: bool = False
mqtt_host: str = "localhost"
mqtt_port: int = 1883
nats_enabled: bool = False
nats_url: str = "nats://localhost:4222"
mdns_enabled: bool = False
model_config = {

23
castle-api/src/castle_api/main.py
View File

@@ -56,25 +56,24 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
poll_task = asyncio.create_task(health_poll_loop())
# --- Mesh coordination (opt-in) ---
mqtt_client = None
nats_client = None
mdns_service = None
if settings.mqtt_enabled:
if settings.nats_enabled:
try:
from castle_api.mqtt_client import CastleMQTTClient
from castle_api.nats_client import CastleNATSClient
registry = get_registry()
mqtt_client = CastleMQTTClient(
nats_client = CastleNATSClient(
local_hostname=registry.node.hostname,
local_registry=registry,
broker_host=settings.mqtt_host,
broker_port=settings.mqtt_port,
servers=settings.nats_url,
)
await mqtt_client.start()
app.state.mqtt_client = mqtt_client
await nats_client.start()
app.state.nats_client = nats_client
except Exception:
logger.exception("Failed to start MQTT client")
mqtt_client = None
logger.exception("Failed to start NATS mesh client")
nats_client = None
if settings.mdns_enabled:
try:
@@ -97,8 +96,8 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
_shutting_down = True
poll_task.cancel()
if mqtt_client:
await mqtt_client.stop()
if nats_client:
await nats_client.stop()
if mdns_service:
mdns_service.stop()

36
castle-api/src/castle_api/mdns.py
View File

@@ -1,8 +1,7 @@
"""mDNS discovery for castle mesh — zero-config LAN peer and broker discovery.
"""mDNS discovery for castle mesh — zero-config LAN peer discovery.
Advertises this node as _castle._tcp and browses for:
- _castle._tcp — peer castle nodes
- _mqtt._tcp — MQTT broker address
Advertises this node as _castle._tcp and browses for _castle._tcp peers. This is
the on-LAN convenience path; cross-network meshes use explicit NATS seed URLs.
"""
from __future__ import annotations
@@ -15,11 +14,10 @@ from zeroconf import ServiceBrowser, ServiceInfo, ServiceStateChange, Zeroconf
logger = logging.getLogger(__name__)
CASTLE_SERVICE_TYPE = "_castle._tcp.local."
MQTT_SERVICE_TYPE = "_mqtt._tcp.local."
class CastleMDNS:
"""Advertise this castle node and discover peers + MQTT broker via mDNS."""
"""Advertise this castle node and discover peers via mDNS."""
def __init__(
self,
@@ -38,7 +36,6 @@ class CastleMDNS:
self.peers: dict[
str, dict
] = {} # hostname -> {gateway_port, api_port, addresses}
self.mqtt_broker: dict | None = None # {host, port} or None
def _on_service_state_change(
self,
@@ -55,8 +52,6 @@ class CastleMDNS:
if service_type == CASTLE_SERVICE_TYPE:
self._handle_castle_peer(info)
elif service_type == MQTT_SERVICE_TYPE:
self._handle_mqtt_broker(info)
elif state_change == ServiceStateChange.Removed:
if service_type == CASTLE_SERVICE_TYPE:
@@ -89,20 +84,6 @@ class CastleMDNS:
}
logger.info("mDNS: discovered peer %s at %s", peer_hostname, addresses)
def _handle_mqtt_broker(self, info: ServiceInfo) -> None:
"""Process a discovered MQTT broker."""
addresses = [
socket.inet_ntoa(addr) for addr in info.addresses if len(addr) == 4
]
if addresses:
self.mqtt_broker = {
"host": addresses[0],
"port": info.port,
}
logger.info(
"mDNS: discovered MQTT broker at %s:%d", addresses[0], info.port
)
def start(self) -> None:
"""Start advertising and browsing."""
self._zeroconf = Zeroconf()
@@ -123,7 +104,7 @@ class CastleMDNS:
"mDNS: advertising %s on port %d", self._hostname, self._gateway_port
)
# Browse for peers and MQTT broker
# Browse for peer castle nodes
self._browsers.append(
ServiceBrowser(
self._zeroconf,
@@ -131,13 +112,6 @@ class CastleMDNS:
handlers=[self._on_service_state_change],
)
)
self._browsers.append(
ServiceBrowser(
self._zeroconf,
MQTT_SERVICE_TYPE,
handlers=[self._on_service_state_change],
)
)
def stop(self) -> None:
"""Stop advertising and close."""

99
castle-api/src/castle_api/mesh_wire.py Normal file
View File

@@ -0,0 +1,99 @@
"""Mesh wire format — (de)serialize a NodeRegistry for cross-node transport.
Transport-agnostic (no MQTT/NATS imports). Only the fields needed for mesh
routing are included; env vars, run_cmd, and castle_root are **excluded** to
avoid leaking secrets — this invariant is load-bearing and must be preserved by
any transport that carries this payload.
"""
from __future__ import annotations
import json
from castle_core.registry import (
Deployment,
NodeConfig,
NodeRegistry,
)
def registry_to_json(registry: NodeRegistry) -> str:
"""Serialize a NodeRegistry to JSON (secret-stripped)."""
data: dict = {
"node": {
"hostname": registry.node.hostname,
"gateway_port": registry.node.gateway_port,
# acme domain — lets peers build launch URLs (<subdomain>.<gateway_domain>)
# for this node's exposed apps. Omitted when the node has no domain.
"gateway_domain": registry.node.gateway_domain,
},
"deployed": {},
}
for _kind, name, comp in registry.all():
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
# Socket surface + external target — so a peer can resolve cross-node
# consumption endpoints (still no secrets: only ports/URLs).
if getattr(comp, "tcp_port", None) is not None:
entry["tcp_port"] = comp.tcp_port
if getattr(comp, "base_url", None):
entry["base_url"] = comp.base_url
# requires — deployment refs (no secrets), so peers can draw cross-node deps.
if getattr(comp, "requires", None):
entry["requires"] = comp.requires
data["deployed"][NodeRegistry.key(comp.kind, name)] = entry
return json.dumps(data)
def json_to_registry(payload: str) -> NodeRegistry:
"""Deserialize a NodeRegistry from a 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),
gateway_domain=node_data.get("gateway_domain"),
)
deployed: dict[str, Deployment] = {}
for key, comp_data in data.get("deployed", {}).items():
key_kind, name = key.split("/", 1) if "/" in key else (None, key)
kind = comp_data.get("kind") or key_kind or "service"
deployed[NodeRegistry.key(kind, 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"),
name=name,
kind=kind,
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),
tcp_port=comp_data.get("tcp_port"),
base_url=comp_data.get("base_url"),
requires=comp_data.get("requires", []),
)
return NodeRegistry(node=node, deployed=deployed)

5
castle-api/src/castle_api/models.py
View File

@@ -218,9 +218,8 @@ class MeshStatus(BaseModel):
"""Current state of the mesh coordination layer."""
enabled: bool = False
mqtt_connected: bool = False
mqtt_broker_host: str | None = None
mqtt_broker_port: int | None = None
connected: bool = False
nats_url: str | None = None
mdns_enabled: bool = False
peer_count: int = 0
peers: list[str] = []

184
castle-api/src/castle_api/nats_client.py Normal file
View File

@@ -0,0 +1,184 @@
"""NATS JetStream client for inter-node mesh coordination.
Replaces the MQTT transport. State lives in a JetStream **KV bucket** rather than
retained MQTT messages:
castle-registry — key=<hostname>, value=secret-stripped NodeRegistry JSON
Lifecycle:
* on connect: ensure the bucket, PUT our registry, seed local state from every
existing key, then watch the bucket for peer changes.
* heartbeat: re-PUT our registry every HEARTBEAT_SEC so peers refresh their
last-seen clock (crash liveness rides the existing stale-TTL).
* graceful stop: DELETE our key → peers get an immediate offline signal.
Being asyncio-native, the watch callback calls ``broadcast`` directly — no
cross-thread ``run_coroutine_threadsafe`` hop (which the paho client needed).
"""
from __future__ import annotations
import asyncio
import contextlib
import logging
import nats
from nats.js.api import KeyValueConfig
from castle_core.registry import NodeRegistry
from castle_api.mesh import mesh_state
from castle_api.mesh_wire import json_to_registry, registry_to_json
from castle_api.stream import broadcast
logger = logging.getLogger(__name__)
REGISTRY_BUCKET = "castle-registry"
HEARTBEAT_SEC = 30.0
PRUNE_SEC = 30.0
class CastleNATSClient:
"""Async NATS/JetStream mesh client."""
def __init__(
self,
local_hostname: str,
local_registry: NodeRegistry,
servers: str | list[str] = "nats://localhost:4222",
) -> None:
self._local_hostname = local_hostname
self._local_registry = local_registry
self._servers = servers
self._nc: nats.NATS | None = None
self._kv = None
self._tasks: list[asyncio.Task] = []
self._last_json: dict[str, str] = {}
self._online: set[str] = set()
@property
def connected(self) -> bool:
return self._nc is not None and self._nc.is_connected
@property
def servers(self) -> str | list[str]:
return self._servers
async def start(self) -> None:
"""Connect, publish our registry, seed state, and start watchers."""
self._nc = await nats.connect(
self._servers,
name=f"castle-{self._local_hostname}",
max_reconnect_attempts=-1, # reconnect forever — nodes come and go
)
js = self._nc.jetstream()
try:
self._kv = await js.key_value(REGISTRY_BUCKET)
except Exception:
self._kv = await js.create_key_value(
config=KeyValueConfig(bucket=REGISTRY_BUCKET, history=1)
)
await self.publish_registry(self._local_registry)
await self._seed_existing()
self._tasks = [
asyncio.create_task(self._watch_loop()),
asyncio.create_task(self._heartbeat_loop()),
asyncio.create_task(self._prune_loop()),
]
logger.info("NATS mesh client started (servers=%s)", self._servers)
async def stop(self) -> None:
"""Delete our key (immediate offline to peers) and disconnect."""
for t in self._tasks:
t.cancel()
for t in self._tasks:
with contextlib.suppress(asyncio.CancelledError, Exception):
await t
self._tasks = []
if self._kv is not None:
with contextlib.suppress(Exception):
await self._kv.delete(self._local_hostname)
if self._nc is not None:
with contextlib.suppress(Exception):
await self._nc.drain()
self._nc = None
logger.info("NATS mesh client stopped")
async def publish_registry(self, registry: NodeRegistry) -> None:
"""PUT (or refresh) our local registry into the KV bucket."""
self._local_registry = registry
if self._kv is None:
return
await self._kv.put(
self._local_hostname, registry_to_json(registry).encode()
)
async def _seed_existing(self) -> None:
"""Load every peer key already present in the bucket."""
if self._kv is None:
return
try:
keys = await self._kv.keys()
except Exception:
keys = [] # empty bucket raises NoKeysError in nats-py
for key in keys:
if key == self._local_hostname:
continue
with contextlib.suppress(Exception):
entry = await self._kv.get(key)
if entry.value:
self._apply_put(key, entry.value.decode())
async def _watch_loop(self) -> None:
assert self._kv is not None
watcher = await self._kv.watchall()
async for entry in watcher:
if entry is None: # "caught up with current values" sentinel
continue
key = entry.key
if key == self._local_hostname:
continue
try:
if entry.operation in ("DEL", "PURGE"):
await self._apply_delete(key)
elif entry.value:
changed = self._apply_put(key, entry.value.decode())
if changed:
await broadcast(
"mesh", {"event": "node_updated", "hostname": key}
)
except Exception:
logger.exception("Error handling mesh entry for %s", key)
def _apply_put(self, hostname: str, payload: str) -> bool:
"""Update mesh state from a peer PUT. Returns True if content changed."""
registry = json_to_registry(payload)
mesh_state.update_node(hostname, registry) # always refresh last-seen
self._online.add(hostname)
changed = self._last_json.get(hostname) != payload
self._last_json[hostname] = payload
return changed
async def _apply_delete(self, hostname: str) -> None:
mesh_state.set_offline(hostname)
self._online.discard(hostname)
self._last_json.pop(hostname, None)
await broadcast("mesh", {"event": "node_offline", "hostname": hostname})
async def _heartbeat_loop(self) -> None:
while True:
await asyncio.sleep(HEARTBEAT_SEC)
with contextlib.suppress(Exception):
await self.publish_registry(self._local_registry)
async def _prune_loop(self) -> None:
"""Mark crashed peers (no refresh within the stale TTL) offline."""
while True:
await asyncio.sleep(PRUNE_SEC)
all_nodes = mesh_state.all_nodes(include_stale=True)
for host in list(self._online):
node = all_nodes.get(host)
if node is None or node.is_stale:
await self._apply_delete(host)

9
castle-api/src/castle_api/nodes.py
View File

@@ -66,15 +66,14 @@ def _deployed_to_summaries(registry: object, hostname: str) -> list[DeploymentSu
@router.get("/mesh/status", response_model=MeshStatus)
def get_mesh_status(request: Request) -> MeshStatus:
"""Get the current state of the mesh coordination layer."""
mqtt_client = getattr(request.app.state, "mqtt_client", None)
nats_client = getattr(request.app.state, "nats_client", None)
peers = list(mesh_state.all_nodes(include_stale=True).keys())
return MeshStatus(
enabled=settings.mqtt_enabled,
mqtt_connected=mqtt_client.connected if mqtt_client else False,
mqtt_broker_host=mqtt_client.broker_host if mqtt_client else None,
mqtt_broker_port=mqtt_client.broker_port if mqtt_client else None,
enabled=settings.nats_enabled,
connected=nats_client.connected if nats_client else False,
nats_url=str(nats_client.servers) if nats_client else None,
mdns_enabled=settings.mdns_enabled,
peer_count=len(peers),
peers=peers,

110
castle-api/tests/test_mesh_wire.py Normal file
View File

@@ -0,0 +1,110 @@
"""Tests for the transport-agnostic mesh wire format (registry (de)serialization)."""
import json
from castle_core.registry import Deployment, NodeConfig, NodeRegistry
from castle_api.mesh_wire import json_to_registry, registry_to_json
def _make_registry() -> NodeRegistry:
return NodeRegistry(
node=NodeConfig(
hostname="tower", castle_root="/data/repos/castle", gateway_port=9000
),
deployed={
"my-svc": Deployment(
manager="systemd",
launcher="python",
run_cmd=["uv", "run", "my-svc"],
env={"PORT": "9001", "SECRET_KEY": "super-secret"},
description="My service",
name="my-svc",
kind="service",
stack="python-fastapi",
port=9001,
health_path="/health",
subdomain="my-svc",
managed=True,
),
"my-job": Deployment(
manager="systemd",
launcher="command",
run_cmd=["my-job"],
name="my-job",
kind="job",
stack="python-cli",
schedule="0 2 * * *",
),
},
)
class TestRegistrySerialization:
"""Round-trip serialization of NodeRegistry to/from JSON."""
def test_round_trip(self) -> None:
original = _make_registry()
json_str = registry_to_json(original)
restored = json_to_registry(json_str)
assert restored.node.hostname == "tower"
assert restored.node.gateway_port == 9000
def test_deployed_components_preserved(self) -> None:
original = _make_registry()
restored = json_to_registry(registry_to_json(original))
svc = restored.get("service", "my-svc")
assert svc is not None
assert svc.manager == "systemd"
assert svc.launcher == "python"
assert svc.port == 9001
assert svc.health_path == "/health"
assert svc.subdomain == "my-svc"
assert svc.managed is True
assert svc.kind == "service"
assert svc.stack == "python-fastapi"
def test_job_fields_preserved(self) -> None:
original = _make_registry()
restored = json_to_registry(registry_to_json(original))
job = restored.get("job", "my-job")
assert job is not None
assert job.launcher == "command"
assert job.schedule == "0 2 * * *"
assert job.kind == "job"
assert job.stack == "python-cli"
def test_optional_fields_omitted(self) -> None:
"""Fields like port, health_path are None when not set."""
reg = NodeRegistry(
node=NodeConfig(hostname="minimal"),
deployed={
"bare": Deployment(
manager="systemd", launcher="command", run_cmd=["bare"], name="bare"
),
},
)
restored = json_to_registry(registry_to_json(reg))
bare = restored.get("service", "bare")
assert bare.port is None
assert bare.health_path is None
assert bare.subdomain is None
assert bare.schedule is None
assert bare.managed is False
def test_no_secrets_in_payload(self) -> None:
"""env vars, run_cmd, and castle_root must never appear on the wire."""
original = _make_registry()
json_str = registry_to_json(original)
data = json.loads(json_str)
# No castle_root in node
assert "castle_root" not in data["node"]
# No env or run_cmd in any component
for name, comp in data["deployed"].items():
assert "env" not in comp, f"{name} has env in payload"
assert "run_cmd" not in comp, f"{name} has run_cmd in payload"