feat: image generation service — ComfyUI cluster integration

- POST /api/image/generate proxy endpoint (admin required) - GET /api/image/status lists available image gen nodes - Cluster AMQP protocol: cmd.image_generate, image_generated, image_failed - Node agent auto-detects ComfyUI, registers image_gen capability - ComfyUI workflow builder: CheckpointLoader → KSampler → VAEDecode → SaveImage - Hardware probe checks ComfyUI reachability + checkpoint model list - 27 tests covering cluster handlers, router, node agent, hardware, capabilities - Config: CAIC_COMFYUI_BASE, CAIC_COMFYUI_TIMEOUT, comfyui_port in agent.ini - Version bump to v1.1.0 - Documentation: ai.md, wiki/Developer-Architecture.md, current-wip.md, README.md, .env.example
2026-07-27 08:06:03 -07:00
parent 576d9333b3
commit aecd3330fd
15 changed files with 2115 additions and 1005 deletions
@@ -0,0 +1,689 @@
+"""Tests for image generation — cluster handlers, router, node agent, hardware probe."""
+import asyncio
+import base64
+import json
+import os
+from contextlib import asynccontextmanager
+from pathlib import Path
+from unittest.mock import AsyncMock, patch
+
+import httpx
+import psutil
+from fastapi.testclient import TestClient
+
+import app as app_module
+import cluster
+import config
+import db
+import hardware
+import node_agent.agent as agent
+from config import AMQP_EXCHANGE_ADMIN, AMQP_EXCHANGE_SYSTEM
+from security import SESSIONS, PIN_ATTEMPTS, RATE_EVENTS
+
+
+# ── helpers ──────────────────────────────────────────────────────────────
+
+def _reset():
+    cluster.CLUSTER_NODES.clear()
+    cluster.CLUSTER_EVENTS.clear()
+    cluster.CLUSTER_COORDINATOR = None
+    cluster._pending_pings.clear()
+    cluster._pending_image.clear()
+
+
+_published = []
+
+
+async def _fake_publish(exchange, routing_key, payload):
+    _published.append((exchange, routing_key, payload))
+
+
+def make_client(tmp_path: Path) -> TestClient:
+    os.environ["CAIC_ADMIN_PIN"] = "1234"
+    db.DB_PATH = tmp_path / "caic-image.db"
+    hardware.HARDWARE_STATE_PATH = tmp_path / "hardware_state.json"
+    SESSIONS.clear()
+    PIN_ATTEMPTS.clear()
+    RATE_EVENTS.clear()
+    db.init_db()
+    return TestClient(app_module.app, raise_server_exceptions=False)
+
+
+def _guest_headers(client: TestClient) -> dict:
+    sid = client.post("/api/auth/guest", headers={"Origin": "http://testserver"}).json()["session_id"]
+    return {"X-Session-ID": sid, "Origin": "http://testserver"}
+
+
+def _admin_headers(client: TestClient) -> dict:
+    resp = client.post("/api/auth/login", json={"pin": "1234"}, headers={"Origin": "http://testserver"})
+    sid = resp.json()["session_id"]
+    return {"X-Session-ID": sid, "Origin": "http://testserver"}
+
+
+class FakeMsg:
+    def __init__(self, body_dict: dict):
+        self.body = json.dumps(body_dict).encode()
+
+    @asynccontextmanager
+    async def process(self):
+        yield
+
+
+class FakeExchange:
+    def __init__(self, name=""):
+        self.name = name
+        self.published = []
+
+    async def publish(self, msg, routing_key):
+        self.published.append((msg, routing_key))
+
+
+class FakeChannel:
+    def __init__(self):
+        self.exchanges = {}
+        self.is_closed = False
+
+    async def declare_exchange(self, name, typ, durable=True):
+        self.exchanges[name] = FakeExchange(name)
+        return self.exchanges[name]
+
+    async def declare_queue(self, name="", exclusive=True):
+        return self
+
+    async def bind(self, exchange, routing_key):
+        pass
+
+
+# ── 1. cluster.handle_image_generated resolves pending request ───────────
+
+
+def test_handle_image_generated_resolves_pending(monkeypatch):
+    _reset()
+    monkeypatch.setattr(cluster, "publish", _fake_publish)
+
+    cluster.CLUSTER_NODES["corsair"] = {
+        "name": "corsair", "type": "worker", "status": "active",
+        "capabilities": ["image_gen"],
+    }
+
+    event = asyncio.Event()
+    cluster._pending_image["req-123"] = ("", event)
+
+    asyncio.run(cluster.handle_image_generated(
+        AMQP_EXCHANGE_SYSTEM, "node.corsair.image_generated",
+        {"node_name": "corsair", "request_id": "req-123", "image_base64": "aW1hZ2U="},
+    ))
+
+    assert event.is_set()
+    result = cluster._pending_image.get("req-123")
+    assert result is not None
+    assert result[0] == "aW1hZ2U="
+    assert cluster.CLUSTER_NODES["corsair"]["last_seen"] is not None
+
+
+# ── 2. cluster.handle_image_failed resolves pending request ─────────────
+
+
+def test_handle_image_failed_resolves_pending(monkeypatch):
+    _reset()
+    monkeypatch.setattr(cluster, "publish", _fake_publish)
+
+    cluster.CLUSTER_NODES["corsair"] = {
+        "name": "corsair", "type": "worker", "status": "active",
+    }
+
+    event = asyncio.Event()
+    cluster._pending_image["req-456"] = ("", event)
+
+    asyncio.run(cluster.handle_image_failed(
+        AMQP_EXCHANGE_SYSTEM, "node.corsair.image_failed",
+        {"node_name": "corsair", "request_id": "req-456", "error": "timeout"},
+    ))
+
+    assert event.is_set()
+    result = cluster._pending_image.get("req-456")
+    assert result is not None
+    assert result[0] == ""
+
+
+# ── 3. cluster.handle_image_failed unknown node ─────────────────────────
+
+
+def test_handle_image_failed_unknown_node(caplog, monkeypatch):
+    _reset()
+    caplog.set_level("WARNING")
+    monkeypatch.setattr(cluster, "publish", _fake_publish)
+
+    asyncio.run(cluster.handle_image_failed(
+        AMQP_EXCHANGE_SYSTEM, "node.ghost.image_failed",
+        {"node_name": "ghost", "request_id": "x", "error": "boom"},
+    ))
+
+    assert not any("unknown node" in rec.message for rec in caplog.records)
+
+
+# ── 4. cluster.request_image_generate publishes command ──────────────────
+
+
+def test_request_image_generate_publishes_command(monkeypatch):
+    _reset()
+    _published.clear()
+    monkeypatch.setattr(cluster, "publish", _fake_publish)
+
+    cluster.CLUSTER_NODES["corsair"] = {
+        "name": "corsair", "type": "worker", "status": "active",
+        "capabilities": ["image_gen"],
+    }
+
+    # Simulate immediate completion
+    async def fake_wait():
+        cluster._pending_image.clear()
+
+    original_wait_for = asyncio.wait_for
+
+    async def patched_wait_for(coro, timeout):
+        cluster._pending_image["fake-id"] = ("aW1hZ2U=", asyncio.Event())
+        cluster._pending_image["fake-id"][1].set()
+        return None
+
+    monkeypatch.setattr(asyncio, "wait_for", patched_wait_for)
+
+    result = asyncio.run(cluster.request_image_generate(
+        "corsair", "a red dragon", width=512, height=512, steps=10,
+    ))
+
+    assert len(_published) == 1
+    exchange, rk, payload = _published[0]
+    assert exchange == AMQP_EXCHANGE_ADMIN
+    assert rk == "node.corsair.cmd.image_generate"
+    assert payload["prompt"] == "a red dragon"
+    assert payload["width"] == 512
+    assert payload["height"] == 512
+    assert payload["steps"] == 10
+    assert "request_id" in payload
+
+
+# ── 5. cluster.request_image_generate unknown node ──────────────────────
+
+
+def test_request_image_generate_unknown_node(monkeypatch):
+    _reset()
+    _published.clear()
+    monkeypatch.setattr(cluster, "publish", _fake_publish)
+
+    result = asyncio.run(cluster.request_image_generate("ghost", "prompt"))
+    assert result is None
+    assert len(_published) == 0
+
+
+# ── 6. cluster.request_image_generate node lacks capability ─────────────
+
+
+def test_request_image_generate_no_capability(monkeypatch):
+    _reset()
+    _published.clear()
+    monkeypatch.setattr(cluster, "publish", _fake_publish)
+
+    cluster.CLUSTER_NODES["jarvis"] = {
+        "name": "jarvis", "type": "worker", "status": "active",
+        "capabilities": ["llm"],
+    }
+
+    result = asyncio.run(cluster.request_image_generate("jarvis", "prompt"))
+    assert result is None
+    assert len(_published) == 0
+
+
+# ── 7. cluster.request_image_generate timeout ───────────────────────────
+
+
+def test_request_image_generate_timeout(monkeypatch):
+    _reset()
+    _published.clear()
+    monkeypatch.setattr(cluster, "publish", _fake_publish)
+
+    cluster.CLUSTER_NODES["corsair"] = {
+        "name": "corsair", "type": "worker", "status": "active",
+        "capabilities": ["image_gen"],
+    }
+
+    async def timeout_wait(coro, timeout):
+        raise asyncio.TimeoutError()
+
+    monkeypatch.setattr(asyncio, "wait_for", timeout_wait)
+
+    result = asyncio.run(cluster.request_image_generate("corsair", "prompt", timeout=1))
+    assert result is None
+    assert len(cluster._pending_image) == 0
+
+
+# ── 8. _find_image_node selects active image_gen node ───────────────────
+
+
+def test_find_image_node_selects_active():
+    from routers.image import _find_image_node
+
+    _reset()
+    cluster.CLUSTER_NODES["corsair"] = {
+        "name": "corsair", "type": "worker", "status": "active",
+        "capabilities": ["image_gen"],
+    }
+    cluster.CLUSTER_NODES["jarvis"] = {
+        "name": "jarvis", "type": "worker", "status": "active",
+        "capabilities": ["llm"],
+    }
+
+    assert _find_image_node() == "corsair"
+
+
+def test_find_image_node_skips_inactive():
+    from routers.image import _find_image_node
+
+    _reset()
+    cluster.CLUSTER_NODES["corsair"] = {
+        "name": "corsair", "type": "worker", "status": "error",
+        "capabilities": ["image_gen"],
+    }
+
+    assert _find_image_node() is None
+
+
+def test_find_image_node_no_image_gen():
+    from routers.image import _find_image_node
+
+    _reset()
+    cluster.CLUSTER_NODES["jarvis"] = {
+        "name": "jarvis", "type": "worker", "status": "active",
+        "capabilities": ["llm"],
+    }
+
+    assert _find_image_node() is None
+
+
+# ── 9. POST /api/image/generate — no node available ─────────────────────
+
+
+def test_image_generate_no_node_503(tmp_path):
+    _reset()
+    with make_client(tmp_path) as client:
+        headers = _admin_headers(client)
+        resp = client.post("/api/image/generate", json={"prompt": "test"}, headers=headers)
+        assert resp.status_code == 503
+        assert "No image generation service" in resp.json()["detail"]
+
+
+# ── 10. POST /api/image/generate — empty prompt ─────────────────────────
+
+
+def test_image_generate_empty_prompt_400(tmp_path):
+    _reset()
+    with make_client(tmp_path) as client:
+        headers = _admin_headers(client)
+        resp = client.post("/api/image/generate", json={"prompt": ""}, headers=headers)
+        assert resp.status_code == 400
+        assert "Prompt is required" in resp.json()["detail"]
+
+
+# ── 11. POST /api/image/generate — happy path ──────────────────────────
+
+
+def test_image_generate_happy_path(tmp_path, monkeypatch):
+    _reset()
+    cluster.CLUSTER_NODES["corsair"] = {
+        "name": "corsair", "type": "worker", "status": "active",
+        "capabilities": ["image_gen"],
+    }
+
+    fake_png = b"\x89PNG\r\n\x1a\n" + b"\x00" * 100
+    fake_b64 = base64.b64encode(fake_png).decode()
+
+    async def fake_request_image_generate(**kwargs):
+        return fake_b64
+
+    monkeypatch.setattr("routers.image.request_image_generate", fake_request_image_generate)
+
+    with make_client(tmp_path) as client:
+        headers = _admin_headers(client)
+        resp = client.post("/api/image/generate", json={
+            "prompt": "a red dragon",
+            "width": 512,
+            "height": 512,
+        }, headers=headers)
+        assert resp.status_code == 200
+        assert resp.headers["content-type"] == "image/png"
+        assert resp.content == fake_png
+
+
+# ── 12. POST /api/image/generate — generation failed ───────────────────
+
+
+def test_image_generate_timeout_504(tmp_path, monkeypatch):
+    _reset()
+    cluster.CLUSTER_NODES["corsair"] = {
+        "name": "corsair", "type": "worker", "status": "active",
+        "capabilities": ["image_gen"],
+    }
+
+    async def fake_request_image_generate(**kwargs):
+        return None
+
+    monkeypatch.setattr("routers.image.request_image_generate", fake_request_image_generate)
+
+    with make_client(tmp_path) as client:
+        headers = _admin_headers(client)
+        resp = client.post("/api/image/generate", json={"prompt": "test"}, headers=headers)
+        assert resp.status_code == 504
+
+
+# ── 13. GET /api/image/status — available ──────────────────────────────
+
+
+def test_image_status_available(tmp_path):
+    _reset()
+    cluster.CLUSTER_NODES["corsair"] = {
+        "name": "corsair", "type": "worker", "status": "active",
+        "capabilities": ["image_gen"],
+        "load": {"gpu_pct": 30},
+        "last_seen": "2026-07-27T00:00:00Z",
+    }
+
+    with make_client(tmp_path) as client:
+        headers = _guest_headers(client)
+        resp = client.get("/api/image/status", headers=headers)
+        assert resp.status_code == 200
+        data = resp.json()
+        assert data["available"] is True
+        assert len(data["nodes"]) == 1
+        assert data["nodes"][0]["name"] == "corsair"
+
+
+# ── 14. GET /api/image/status — no nodes ───────────────────────────────
+
+
+def test_image_status_unavailable(tmp_path):
+    _reset()
+    with make_client(tmp_path) as client:
+        headers = _guest_headers(client)
+        resp = client.get("/api/image/status", headers=headers)
+        assert resp.status_code == 200
+        data = resp.json()
+        assert data["available"] is False
+        assert len(data["nodes"]) == 0
+
+
+# ── 15. node_agent.detect_capabilities — ComfyUI present ───────────────
+
+
+def test_detect_capabilities_comfyui_present(monkeypatch):
+    cfg = agent.AgentConfig()
+    cfg.comfyui_port = 8188
+
+    monkeypatch.setattr(agent, "HAS_HTTPX", True)
+
+    def fake_get(url, timeout=3):
+        if "system_stats" in url:
+            class R:
+                status_code = 200
+            return R()
+        raise httpx.ConnectError("refused")
+
+    monkeypatch.setattr(httpx, "get", fake_get)
+
+    caps = agent.detect_capabilities(cfg)
+    assert "image_gen" in caps
+    assert "llm" in caps
+
+
+# ── 16. node_agent.detect_capabilities — ComfyUI absent ────────────────
+
+
+def test_detect_capabilities_comfyui_absent(monkeypatch):
+    cfg = agent.AgentConfig()
+    cfg.comfyui_port = 8188
+
+    monkeypatch.setattr(agent, "HAS_HTTPX", True)
+    monkeypatch.setattr(httpx, "get", lambda url, timeout=3: (_ for _ in ()).throw(httpx.ConnectError("refused")))
+
+    caps = agent.detect_capabilities(cfg)
+    assert "image_gen" not in caps
+    assert "llm" in caps
+
+
+# ── 17. node_agent.detect_capabilities — httpx not installed ────────────
+
+
+def test_detect_capabilities_no_httpx(monkeypatch):
+    cfg = agent.AgentConfig()
+    monkeypatch.setattr(agent, "HAS_HTTPX", False)
+
+    caps = agent.detect_capabilities(cfg)
+    assert "image_gen" not in caps
+
+
+# ── 18. node_agent.handle_image_generate — success ─────────────────────
+
+
+def test_node_agent_handle_image_generate_success(monkeypatch):
+    monkeypatch.setattr(agent, "HAS_AIO_PIKA", True)
+    monkeypatch.setattr(agent, "HAS_HTTPX", True)
+
+    cfg = agent.AgentConfig()
+    cfg.node_name = "corsair"
+    cfg.comfyui_port = 8188
+
+    fake_png = b"\x89PNG" + b"\x00" * 50
+    fake_b64 = base64.b64encode(fake_png).decode()
+
+    async def fake_comfyui_generate(*a, **kw):
+        return fake_b64
+
+    monkeypatch.setattr(agent, "_comfyui_generate", fake_comfyui_generate)
+
+    channel = FakeChannel()
+    system_ex = FakeExchange("jc.system")
+    channel.exchanges["jc.system"] = system_ex
+
+    asyncio.run(agent.handle_image_generate(
+        cfg, channel, (FakeExchange(), system_ex),
+        FakeMsg({
+            "request_id": "req-789",
+            "prompt": "a castle",
+            "negative_prompt": "",
+            "width": 1024,
+            "height": 1024,
+            "steps": 20,
+            "seed": 42,
+            "model": "",
+        }),
+    ))
+
+    assert len(system_ex.published) == 1
+    msg, rk = system_ex.published[0]
+    assert rk == "node.corsair.image_generated"
+    payload = json.loads(msg.body)
+    assert payload["type"] == "image_generated"
+    assert payload["request_id"] == "req-789"
+    assert payload["image_base64"] == fake_b64
+
+
+# ── 19. node_agent.handle_image_generate — failure ─────────────────────
+
+
+def test_node_agent_handle_image_generate_failure(monkeypatch):
+    monkeypatch.setattr(agent, "HAS_AIO_PIKA", True)
+    monkeypatch.setattr(agent, "HAS_HTTPX", True)
+
+    cfg = agent.AgentConfig()
+    cfg.node_name = "corsair"
+
+    async def fake_comfyui_generate(*a, **kw):
+        raise RuntimeError("ComfyUI crashed")
+
+    monkeypatch.setattr(agent, "_comfyui_generate", fake_comfyui_generate)
+
+    channel = FakeChannel()
+    system_ex = FakeExchange("jc.system")
+    channel.exchanges["jc.system"] = system_ex
+
+    asyncio.run(agent.handle_image_generate(
+        cfg, channel, (FakeExchange(), system_ex),
+        FakeMsg({"request_id": "req-fail", "prompt": "test"}),
+    ))
+
+    assert len(system_ex.published) == 1
+    msg, rk = system_ex.published[0]
+    assert rk == "node.corsair.image_failed"
+    payload = json.loads(msg.body)
+    assert payload["type"] == "image_failed"
+    assert "ComfyUI crashed" in payload["error"]
+
+
+# ── 20. node_agent.handle_image_generate — empty prompt ────────────────
+
+
+def test_node_agent_handle_image_generate_empty_prompt(monkeypatch):
+    monkeypatch.setattr(agent, "HAS_AIO_PIKA", True)
+
+    cfg = agent.AgentConfig()
+    cfg.node_name = "corsair"
+
+    channel = FakeChannel()
+    system_ex = FakeExchange("jc.system")
+
+    asyncio.run(agent.handle_image_generate(
+        cfg, channel, (FakeExchange(), system_ex),
+        FakeMsg({"request_id": "req-x", "prompt": ""}),
+    ))
+
+    assert len(system_ex.published) == 0
+
+
+# ── 21. hardware.py — ComfyUI reachable ────────────────────────────────
+
+
+def test_assess_hardware_comfyui_reachable(tmp_path, monkeypatch):
+    hardware.HARDWARE_STATE_PATH = tmp_path / "hw.json"
+    monkeypatch.setattr(psutil, "virtual_memory", lambda: type("M", (), {"total": 16 * 1024 ** 3, "available": 8 * 1024 ** 3})())
+    monkeypatch.setattr(psutil, "cpu_count", lambda: 8)
+
+    class MockProc:
+        returncode = 1
+        stdout = ""
+
+    monkeypatch.setattr(hardware.subprocess, "run", lambda cmd, **kw: MockProc())
+
+    async def mock_get(self, url, *args, **kwargs):
+        class R:
+            status_code = 200
+            def json(self):
+                if "CheckpointLoaderSimple" in url:
+                    return {"CheckpointLoaderSimple": {"input": {"required": {"ckpt_name": [["model.safetensors", "other.ckpt"]]}}}}
+                return {}
+            def raise_for_status(self):
+                pass
+        if "8188" in url:
+            return R()
+        if "v1/models" in url:
+            class R2:
+                status_code = 200
+                def json(self):
+                    return {"data": []}
+            return R2()
+        if "6333" in url:
+            class R3:
+                status_code = 200
+                def json(self):
+                    return {"result": {"collections": []}}
+            return R3()
+        if "8888" in url:
+            class R4:
+                status_code = 200
+            return R4()
+        class R5:
+            status_code = 200
+            def json(self):
+                return {}
+        return R5()
+
+    monkeypatch.setattr(httpx.AsyncClient, "get", mock_get)
+
+    state = asyncio.run(hardware.assess_hardware())
+    assert state["comfyui_reachable"] is True
+    assert "model.safetensors" in state["comfyui_models"]
+
+
+# ── 22. hardware.py — ComfyUI unreachable ──────────────────────────────
+
+
+def test_assess_hardware_comfyui_unreachable(tmp_path, monkeypatch):
+    hardware.HARDWARE_STATE_PATH = tmp_path / "hw.json"
+    monkeypatch.setattr(psutil, "virtual_memory", lambda: type("M", (), {"total": 16 * 1024 ** 3, "available": 8 * 1024 ** 3})())
+    monkeypatch.setattr(psutil, "cpu_count", lambda: 8)
+
+    class MockProc:
+        returncode = 1
+        stdout = ""
+
+    monkeypatch.setattr(hardware.subprocess, "run", lambda cmd, **kw: MockProc())
+
+    async def mock_get(self, url, *args, **kwargs):
+        if "8188" in url:
+            raise httpx.ConnectError("refused")
+        if "v1/models" in url:
+            class R2:
+                status_code = 200
+                def json(self):
+                    return {"data": []}
+            return R2()
+        if "6333" in url:
+            class R3:
+                status_code = 200
+                def json(self):
+                    return {"result": {"collections": []}}
+            return R3()
+        if "8888" in url:
+            class R4:
+                status_code = 200
+            return R4()
+        class R5:
+            status_code = 200
+            def json(self):
+                return {}
+        return R5()
+
+    monkeypatch.setattr(httpx.AsyncClient, "get", mock_get)
+
+    state = asyncio.run(hardware.assess_hardware())
+    assert state["comfyui_reachable"] is False
+    assert state["comfyui_models"] == []
+
+
+# ── 23. node_agent config reads comfyui_port ───────────────────────────
+
+
+def test_config_from_ini_comfyui_port(tmp_path):
+    ini = tmp_path / "caic-node-agent.conf"
+    ini.write_text(
+        "[agent]\n"
+        "node_name = corsair\n"
+        "capabilities = llm,image_gen\n"
+        "comfyui_port = 8188\n"
+    )
+    cfg = agent.AgentConfig.from_ini(str(ini))
+    assert cfg.comfyui_port == 8188
+    assert "image_gen" in cfg.capabilities
+
+
+def test_config_from_ini_comfyui_port_default():
+    cfg = agent.AgentConfig()
+    assert cfg.comfyui_port == 8188
+
+
+# ── 24. SUBSCRIBE_TABLE includes image gen handlers ────────────────────
+
+
+def test_subscribe_table_includes_image_handlers():
+    routing_keys = [rks for _, rks, _ in cluster.SUBSCRIBE_TABLE]
+    all_keys = [rk for rks in routing_keys for rk in rks]
+    assert "node.*.image_generated" in all_keys
+    assert "node.*.image_failed" in all_keys