diff --git a/amqp.py b/amqp.py index bfa1d75..64a8457 100644 --- a/amqp.py +++ b/amqp.py @@ -26,6 +26,51 @@ log = logging.getLogger("jarvischat") _connection = None _channel = None _lock = asyncio.Lock() +_subscriptions = [] # list of (exchange, routing_keys, handler_fn) + + +async def subscribe(exchange: str, routing_keys: list[str], handler) -> None: + """Subscribe to routing keys on an exchange. + + Creates an exclusive anonymous queue bound to the given routing keys. + Handler receives (exchange, routing_key, payload_dict). + On reconnect: _rebind_subscriptions recreates all subscriptions. + """ + if not HAS_AIO_PIKA: + log.warning("aio-pika not installed — cannot subscribe") + return + ch = await get_channel() + if ch is None: + log.error("cannot subscribe — no AMQP channel") + return + try: + queue = await ch.declare_queue("", exclusive=True) + ex = await ch.get_exchange(exchange) + for rk in routing_keys: + await queue.bind(ex, rk) + + async def _dispatch(msg: aio_pika.IncomingMessage): + async with msg.process(): + try: + payload = json.loads(msg.body.decode()) + await handler(exchange, msg.routing_key, payload) + except Exception: + log.exception("AMQP handler error for %s %s", exchange, msg.routing_key) + + await queue.consume(_dispatch) + _subscriptions.append((exchange, routing_keys, handler)) + except Exception: + log.exception("AMQP subscribe failed for %s %s", exchange, routing_keys) + + +async def _rebind_subscriptions() -> None: + """Re-create all subscriptions after reconnect.""" + subs = list(_subscriptions) + _subscriptions.clear() + for exchange, routing_keys, handler in subs: + await subscribe(exchange, routing_keys, handler) + if subs: + log.info("AMQP subscriptions rebound") async def connect() -> None: @@ -48,6 +93,7 @@ async def connect() -> None: await ch.declare_exchange(ex, ExchangeType.TOPIC, durable=True) _connection = conn _channel = ch + await _rebind_subscriptions() log.info("AMQP connected, exchanges declared") diff --git a/app.py b/app.py index 59548e3..ae6492b 100644 --- a/app.py +++ b/app.py @@ -15,6 +15,7 @@ from fastapi.staticfiles import StaticFiles from fastapi.templating import Jinja2Templates from amqp import connect as amqp_connect, disconnect as amqp_disconnect +from cluster import start_cluster_subscriptions from config import VERSION, RATE_WINDOW_SECONDS, UPLOAD_DIR, RAG_MAX_VECTORS, RAG_EVICTION_HIGH_WATER, RAG_EVICTION_LOW_WATER, RAG_EVICTION_BATCH from db import init_db from hardware import assess_hardware @@ -39,6 +40,7 @@ import routers.upload as upload import routers.ingest as ingest import routers.hardware as hardware import routers.rag_admin as rag_admin +import routers.cluster as cluster_router # --- Logging --- log = logging.getLogger("jarvischat") @@ -59,6 +61,7 @@ async def lifespan(app: FastAPI): log.info(f"Memory system: {get_memory_count()} memories loaded") await assess_hardware() await amqp_connect() + await start_cluster_subscriptions() if RAG_MAX_VECTORS > 0: if RAG_EVICTION_HIGH_WATER <= RAG_EVICTION_LOW_WATER: @@ -162,7 +165,7 @@ for router_module in [ auth_router, conversations.router, memories.router, models.router, presets.router, profile.router, settings.router, skills.router, chat.router, search_route.router, completions.router, upload.router, ingest.router, hardware.router, - rag_admin.router, + rag_admin.router, cluster_router.router, ]: app.include_router(router_module) diff --git a/cluster.py b/cluster.py new file mode 100644 index 0000000..ee863a6 --- /dev/null +++ b/cluster.py @@ -0,0 +1,189 @@ +""" +JarvisChat — Cluster protocol implementation. +Maintains node registry, event log, coordinator state, and ping-based health checks. +""" +import asyncio +import logging +import uuid +from collections import deque +from datetime import datetime, timezone + +from amqp import publish, subscribe +from config import AMQP_EXCHANGE_ADMIN, AMQP_EXCHANGE_SYSTEM + +log = logging.getLogger("jarvischat") + +CLUSTER_NODES: dict[str, dict] = {} +CLUSTER_EVENTS: deque = deque(maxlen=1000) +CLUSTER_COORDINATOR: str | None = None +_pending_pings: dict[str, tuple[str, asyncio.Event]] = {} +NODE_NAME: str = "ultron" +PING_TIMEOUT: float = 5.0 + + +def _push_event(category: str, severity: str, node_name: str | None, message: str, details: dict | None = None) -> dict: + record = { + "category": category, + "severity": severity, + "node": node_name, + "message": message, + "details": details or {}, + "timestamp": datetime.now(timezone.utc).isoformat() + "Z", + } + CLUSTER_EVENTS.append(record) + level = {"info": logging.INFO, "warn": logging.WARNING, "error": logging.ERROR, "critical": logging.CRITICAL}.get(severity, logging.INFO) + log.log(level, "[cluster] %s: %s", node_name or "system", message) + return record + + +async def handle_registration(exchange: str, routing_key: str, payload: dict) -> None: + global CLUSTER_COORDINATOR + node_name = payload.get("node_name", routing_key.split(".")[1] if "." in routing_key else "unknown") + node_type = payload.get("node_type", "worker") + + if node_name in CLUSTER_NODES: + _push_event("cluster", "warn", node_name, "Duplicate registration rejected") + await publish(AMQP_EXCHANGE_ADMIN, f"node.{node_name}.rejected", { + "from": NODE_NAME, "node_name": node_name, "type": "rejected", + "reason": "duplicate_node_name", + "timestamp": datetime.now(timezone.utc).isoformat() + "Z", + }) + return + + if "node_type" not in payload or "capabilities" not in payload: + _push_event("cluster", "warn", node_name, "Malformed registration rejected") + await publish(AMQP_EXCHANGE_ADMIN, f"node.{node_name}.rejected", { + "from": NODE_NAME, "node_name": node_name, "type": "rejected", + "reason": "malformed_payload", + "timestamp": datetime.now(timezone.utc).isoformat() + "Z", + }) + return + + now = datetime.now(timezone.utc).isoformat() + "Z" + node = { + "name": node_name, + "type": node_type, + "status": "active", + "capabilities": payload.get("capabilities", []), + "active_model": payload.get("active_model"), + "load": payload.get("load"), + "registered_at": now, + "last_seen": now, + } + CLUSTER_NODES[node_name] = node + _push_event("cluster", "info", node_name, f"Node registered (type={node_type})") + + await publish(AMQP_EXCHANGE_ADMIN, f"node.{node_name}.admitted", { + "from": NODE_NAME, "node_name": node_name, "type": "admitted", + "coordinator": CLUSTER_COORDINATOR, + "timestamp": now, + }) + + if node_type == "coordinator" and CLUSTER_COORDINATOR is None: + CLUSTER_COORDINATOR = node_name + _push_event("cluster", "info", node_name, "Elected as coordinator") + await publish(AMQP_EXCHANGE_SYSTEM, "cluster.coordinator.response", { + "from": NODE_NAME, "type": "coord_response", + "coordinator": node_name, "nodes": list(CLUSTER_NODES.keys()), + "timestamp": now, + }) + + +async def handle_deregistration(exchange: str, routing_key: str, payload: dict) -> None: + global CLUSTER_COORDINATOR + node_name = payload.get("node_name", routing_key.split(".")[1] if "." in routing_key else "unknown") + + removed = CLUSTER_NODES.pop(node_name, None) + if not removed: + log.warning("[cluster] deregistration for unknown node %s", node_name) + return + + _push_event("cluster", "info", node_name, "Node deregistered") + + if CLUSTER_COORDINATOR == node_name: + CLUSTER_COORDINATOR = None + _push_event("cluster", "warn", node_name, "Coordinator deregistered — no coordinator active") + + +async def handle_pong(exchange: str, routing_key: str, payload: dict) -> None: + node_name = payload.get("node_name", routing_key.split(".")[1] if "." in routing_key else "unknown") + correlation_id = payload.get("correlation_id") + + if correlation_id and correlation_id in _pending_pings: + _, event = _pending_pings.pop(correlation_id) + event.set() + + if node_name in CLUSTER_NODES: + now = datetime.now(timezone.utc).isoformat() + "Z" + CLUSTER_NODES[node_name]["last_seen"] = now + if "status" in payload: + CLUSTER_NODES[node_name]["status"] = payload["status"] + if "active_model" in payload: + CLUSTER_NODES[node_name]["active_model"] = payload["active_model"] + if "load" in payload: + CLUSTER_NODES[node_name]["load"] = payload["load"] + else: + log.warning("[cluster] pong from unknown node %s", node_name) + + +async def handle_event(exchange: str, routing_key: str, payload: dict) -> None: + node_name = payload.get("node_name", routing_key.split(".")[1] if "." in routing_key else "unknown") + severity = payload.get("severity", "info") + message = payload.get("message", "") + details = payload.get("details") + _push_event("application", severity, node_name, message, details) + + +async def handle_coordinator_query(exchange: str, routing_key: str, payload: dict) -> None: + if CLUSTER_COORDINATOR is None: + return + now = datetime.now(timezone.utc).isoformat() + "Z" + await publish(AMQP_EXCHANGE_SYSTEM, "cluster.coordinator.response", { + "from": NODE_NAME, "type": "coord_response", + "coordinator": CLUSTER_COORDINATOR, + "nodes": list(CLUSTER_NODES.keys()), + "timestamp": now, + }) + + +async def ping_node(node_name: str) -> bool: + global CLUSTER_COORDINATOR + if node_name not in CLUSTER_NODES: + return False + + correlation_id = str(uuid.uuid4()) + event = asyncio.Event() + _pending_pings[correlation_id] = (node_name, event) + + now = datetime.now(timezone.utc).isoformat() + "Z" + await publish(AMQP_EXCHANGE_ADMIN, f"node.{node_name}.ping", { + "from": NODE_NAME, "node_name": node_name, "type": "ping", + "correlation_id": correlation_id, "timestamp": now, + }) + + try: + await asyncio.wait_for(event.wait(), timeout=PING_TIMEOUT) + return True + except asyncio.TimeoutError: + _pending_pings.pop(correlation_id, None) + CLUSTER_NODES.pop(node_name, None) + _push_event("cluster", "warn", node_name, "Node unresponsive — deregistered after ping timeout") + if CLUSTER_COORDINATOR == node_name: + CLUSTER_COORDINATOR = None + _push_event("cluster", "warn", node_name, "Coordinator unresponsive — no coordinator active") + return False + + +SUBSCRIBE_TABLE = [ + (AMQP_EXCHANGE_ADMIN, ["node.*.register"], handle_registration), + (AMQP_EXCHANGE_ADMIN, ["node.*.deregister"], handle_deregistration), + (AMQP_EXCHANGE_ADMIN, ["node.*.pong"], handle_pong), + (AMQP_EXCHANGE_SYSTEM, ["node.*.event"], handle_event), + (AMQP_EXCHANGE_SYSTEM, ["cluster.coordinator.query"], handle_coordinator_query), +] + + +async def start_cluster_subscriptions() -> None: + for exchange, routing_keys, handler in SUBSCRIBE_TABLE: + await subscribe(exchange, routing_keys, handler) + log.info("cluster subscriptions started") diff --git a/routers/cluster.py b/routers/cluster.py new file mode 100644 index 0000000..7c186a8 --- /dev/null +++ b/routers/cluster.py @@ -0,0 +1,23 @@ +"""JarvisChat routers - Cluster status API.""" +from fastapi import APIRouter + +import cluster + +router = APIRouter() + + +@router.get("/api/cluster") +async def cluster_status(): + return { + "nodes": {name: _strip_internal(node) for name, node in cluster.CLUSTER_NODES.items()}, + "node_count": len(cluster.CLUSTER_NODES), + "coordinator": cluster.CLUSTER_COORDINATOR, + "events": list(cluster.CLUSTER_EVENTS), + } + + +def _strip_internal(node: dict) -> dict: + return {k: v for k, v in node.items() if k in { + "name", "type", "status", "capabilities", "active_model", "load", + "registered_at", "last_seen", + }} diff --git a/tests/test_cluster.py b/tests/test_cluster.py new file mode 100644 index 0000000..5cbcba5 --- /dev/null +++ b/tests/test_cluster.py @@ -0,0 +1,313 @@ +"""Tests for cluster.py — no live AMQP, all handlers called directly.""" +import asyncio +from collections import deque + +import cluster +from config import AMQP_EXCHANGE_ADMIN, AMQP_EXCHANGE_SYSTEM + + +def _reset(): + cluster.CLUSTER_NODES.clear() + cluster.CLUSTER_EVENTS.clear() + cluster.CLUSTER_COORDINATOR = None + cluster._pending_pings.clear() + + +_published = [] + + +async def _fake_publish(exchange, routing_key, payload): + _published.append((exchange, routing_key, payload)) + + +# ---------- 1. Valid worker registration ---------- + + +def test_valid_worker_registration(monkeypatch): + _reset() + _published.clear() + monkeypatch.setattr(cluster, "publish", _fake_publish) + + asyncio.run(cluster.handle_registration( + AMQP_EXCHANGE_ADMIN, "node.jarvis.register", + {"node_name": "jarvis", "node_type": "worker", "capabilities": ["llm"]}, + )) + + assert "jarvis" in cluster.CLUSTER_NODES + assert cluster.CLUSTER_NODES["jarvis"]["type"] == "worker" + assert cluster.CLUSTER_NODES["jarvis"]["status"] == "active" + assert len(cluster.CLUSTER_EVENTS) == 1 + assert cluster.CLUSTER_EVENTS[0]["category"] == "cluster" + assert cluster.CLUSTER_EVENTS[0]["message"] == "Node registered (type=worker)" + + assert len(_published) == 1 + exchange, rk, payload = _published[0] + assert exchange == AMQP_EXCHANGE_ADMIN + assert rk == "node.jarvis.admitted" + assert payload["type"] == "admitted" + assert payload["node_name"] == "jarvis" + + +# ---------- 2. First coordinator auto-promotion ---------- + + +def test_first_coordinator_auto_promotion(monkeypatch): + _reset() + _published.clear() + monkeypatch.setattr(cluster, "publish", _fake_publish) + + asyncio.run(cluster.handle_registration( + AMQP_EXCHANGE_ADMIN, "node.ultron.register", + {"node_name": "ultron", "node_type": "coordinator", "capabilities": ["llm", "rag"]}, + )) + + assert cluster.CLUSTER_COORDINATOR == "ultron" + assert len(cluster.CLUSTER_EVENTS) == 2 + assert cluster.CLUSTER_EVENTS[1]["message"] == "Elected as coordinator" + + assert len(_published) == 2 + # First publish: admitted + assert _published[0][1] == "node.ultron.admitted" + # Second publish: coord_response + exchange, rk, payload = _published[1] + assert exchange == AMQP_EXCHANGE_SYSTEM + assert rk == "cluster.coordinator.response" + assert payload["type"] == "coord_response" + assert payload["coordinator"] == "ultron" + + +# ---------- 3. Duplicate node name rejected ---------- + + +def test_duplicate_node_rejected(monkeypatch): + _reset() + _published.clear() + monkeypatch.setattr(cluster, "publish", _fake_publish) + + asyncio.run(cluster.handle_registration( + AMQP_EXCHANGE_ADMIN, "node.jarvis.register", + {"node_name": "jarvis", "node_type": "worker", "capabilities": ["llm"]}, + )) + _published.clear() + + asyncio.run(cluster.handle_registration( + AMQP_EXCHANGE_ADMIN, "node.jarvis.register", + {"node_name": "jarvis", "node_type": "worker", "capabilities": ["llm"]}, + )) + + assert len(cluster.CLUSTER_NODES) == 1 + assert len(_published) == 1 + exchange, rk, payload = _published[0] + assert exchange == AMQP_EXCHANGE_ADMIN + assert rk == "node.jarvis.rejected" + assert payload["reason"] == "duplicate_node_name" + + +# ---------- 4. Malformed payload rejected ---------- + + +def test_malformed_payload_rejected(monkeypatch): + _reset() + _published.clear() + monkeypatch.setattr(cluster, "publish", _fake_publish) + + asyncio.run(cluster.handle_registration( + AMQP_EXCHANGE_ADMIN, "node.jarvis.register", + {"node_name": "jarvis"}, # missing node_type and capabilities + )) + + assert len(cluster.CLUSTER_NODES) == 0 + assert len(_published) == 1 + exchange, rk, payload = _published[0] + assert rk == "node.jarvis.rejected" + assert payload["reason"] == "malformed_payload" + + +# ---------- 5. Graceful deregistration ---------- + + +def test_graceful_deregistration(monkeypatch): + _reset() + _published.clear() + monkeypatch.setattr(cluster, "publish", _fake_publish) + + asyncio.run(cluster.handle_registration( + AMQP_EXCHANGE_ADMIN, "node.jarvis.register", + {"node_name": "jarvis", "node_type": "worker", "capabilities": ["llm"]}, + )) + _published.clear() + + asyncio.run(cluster.handle_deregistration( + AMQP_EXCHANGE_ADMIN, "node.jarvis.deregister", + {"node_name": "jarvis"}, + )) + + assert "jarvis" not in cluster.CLUSTER_NODES + assert len(cluster.CLUSTER_EVENTS) == 2 + assert cluster.CLUSTER_EVENTS[1]["message"] == "Node deregistered" + + +def test_deregister_coordinator_clears(monkeypatch): + _reset() + _published.clear() + monkeypatch.setattr(cluster, "publish", _fake_publish) + + asyncio.run(cluster.handle_registration( + AMQP_EXCHANGE_ADMIN, "node.ultron.register", + {"node_name": "ultron", "node_type": "coordinator", "capabilities": ["llm"]}, + )) + assert cluster.CLUSTER_COORDINATOR == "ultron" + _published.clear() + + asyncio.run(cluster.handle_deregistration( + AMQP_EXCHANGE_ADMIN, "node.ultron.deregister", + {"node_name": "ultron"}, + )) + + assert "ultron" not in cluster.CLUSTER_NODES + assert cluster.CLUSTER_COORDINATOR is None + + +# ---------- 6. Pong from known node ---------- + + +def test_pong_updates_known_node(monkeypatch): + _reset() + monkeypatch.setattr(cluster, "publish", _fake_publish) + + asyncio.run(cluster.handle_registration( + AMQP_EXCHANGE_ADMIN, "node.jarvis.register", + {"node_name": "jarvis", "node_type": "worker", "capabilities": ["llm"]}, + )) + original_seen = cluster.CLUSTER_NODES["jarvis"]["last_seen"] + + asyncio.run(cluster.handle_pong( + AMQP_EXCHANGE_ADMIN, "node.jarvis.pong", + {"node_name": "jarvis", "status": "busy", "load": {"cpu_pct": 80}}, + )) + + assert cluster.CLUSTER_NODES["jarvis"]["status"] == "busy" + assert cluster.CLUSTER_NODES["jarvis"]["load"] == {"cpu_pct": 80} + assert cluster.CLUSTER_NODES["jarvis"]["last_seen"] != original_seen + + +# ---------- 7. Pong from unknown node ---------- + + +def test_pong_from_unknown_node(caplog, monkeypatch): + _reset() + caplog.set_level("WARNING") + monkeypatch.setattr(cluster, "publish", _fake_publish) + + asyncio.run(cluster.handle_pong( + AMQP_EXCHANGE_ADMIN, "node.stranger.pong", + {"node_name": "stranger", "status": "active"}, + )) + + assert "stranger" not in cluster.CLUSTER_NODES + assert any("pong from unknown node" in rec.message for rec in caplog.records) + + +# ---------- 8. Event stored in log ---------- + + +def test_event_appended_to_log(monkeypatch): + _reset() + monkeypatch.setattr(cluster, "publish", _fake_publish) + + asyncio.run(cluster.handle_event( + AMQP_EXCHANGE_SYSTEM, "node.jarvis.event", + {"node_name": "jarvis", "severity": "error", "message": "OOM on GPU"}, + )) + + assert len(cluster.CLUSTER_EVENTS) == 1 + ev = cluster.CLUSTER_EVENTS[0] + assert ev["category"] == "application" + assert ev["severity"] == "error" + assert ev["node"] == "jarvis" + + +def test_event_log_bounded(monkeypatch): + _reset() + monkeypatch.setattr(cluster, "publish", _fake_publish) + + for i in range(1001): + cluster._push_event("application", "info", "jarvis", f"event {i}") + + assert len(cluster.CLUSTER_EVENTS) == 1000 + assert cluster.CLUSTER_EVENTS[0]["message"] == "event 1" + assert cluster.CLUSTER_EVENTS[-1]["message"] == "event 1000" + + +# ---------- 9. Coordinator query produces response ---------- + + +def test_coordinator_query_response(monkeypatch): + _reset() + _published.clear() + monkeypatch.setattr(cluster, "publish", _fake_publish) + + asyncio.run(cluster.handle_registration( + AMQP_EXCHANGE_ADMIN, "node.ultron.register", + {"node_name": "ultron", "node_type": "coordinator", "capabilities": ["llm"]}, + )) + _published.clear() + + asyncio.run(cluster.handle_coordinator_query( + AMQP_EXCHANGE_SYSTEM, "cluster.coordinator.query", + {"from": "jarvis", "type": "coord_query"}, + )) + + assert len(_published) == 1 + exchange, rk, payload = _published[0] + assert exchange == AMQP_EXCHANGE_SYSTEM + assert rk == "cluster.coordinator.response" + assert payload["coordinator"] == "ultron" + assert "nodes" in payload + + +def test_coordinator_query_no_coordinator(monkeypatch): + _reset() + _published.clear() + monkeypatch.setattr(cluster, "publish", _fake_publish) + + asyncio.run(cluster.handle_coordinator_query( + AMQP_EXCHANGE_SYSTEM, "cluster.coordinator.query", + {"from": "jarvis", "type": "coord_query"}, + )) + + assert len(_published) == 0 + + +# ---------- 10. GET /api/cluster shape ---------- + + +def test_cluster_api_shape(): + _reset() + cluster.CLUSTER_NODES["jarvis"] = { + "name": "jarvis", "type": "worker", "status": "active", + "capabilities": ["llm"], "active_model": None, "load": None, + "registered_at": "2026-01-01T00:00:00Z", "last_seen": "2026-01-01T00:00:00Z", + } + cluster.CLUSTER_COORDINATOR = "ultron" + cluster._push_event("cluster", "info", "ultron", "Elected") + + from routers.cluster import cluster_status + + resp = asyncio.run(cluster_status()) + + assert "nodes" in resp + assert "node_count" in resp + assert "coordinator" in resp + assert "events" in resp + assert resp["node_count"] == 1 + assert resp["coordinator"] == "ultron" + assert "jarvis" in resp["nodes"] + assert len(resp["events"]) == 1 + # No internal keys leaked + for node in resp["nodes"].values(): + for key in node: + assert key in { + "name", "type", "status", "capabilities", "active_model", + "load", "registered_at", "last_seen", + }