diff --git a/docs/comfyui-adapter-plan.md b/docs/comfyui-adapter-plan.md index c0a941e6..c8b7f79e 100644 --- a/docs/comfyui-adapter-plan.md +++ b/docs/comfyui-adapter-plan.md @@ -474,9 +474,16 @@ pipeline definition, or any schema. user-provided via a config directory? Bundling gives reproducibility; external gives flexibility. -2. **Async generation:** ComfyUI supports websocket connections for real-time - progress. Worth implementing for long video generations, or is polling - sufficient? +2. ~~**Async generation:**~~ **Resolved.** `ComfyUIClient.generate()` now + waits via ComfyUI's websocket feed (`wait_ws()`) by default, reacting to + `executing`/`execution_error` events immediately instead of sleeping + between REST polls — completion and errors are caught without the + `interval`-seconds lag, and an optional `on_progress` callback gets live + `progress` events (`comfyui_video` uses this to print step progress on + long renders). No new hard dependency: `websocket-client` is an optional + import, and `_wait()` transparently falls back to the original + `poll()` REST loop when it isn't installed or the connection fails — + `resume_prompt_id` recovery behaves identically either way. 3. **Multi-server:** Should the adapter support multiple ComfyUI instances (e.g., one for images, one for video) via per-capability URLs? diff --git a/tests/contracts/test_comfyui_tools.py b/tests/contracts/test_comfyui_tools.py index ff3e279b..21dc5166 100644 --- a/tests/contracts/test_comfyui_tools.py +++ b/tests/contracts/test_comfyui_tools.py @@ -351,6 +351,187 @@ class TestClientHelpers: assert "myhost:9999" in msg assert "COMFYUI_SERVER_URL" not in msg + def test_submit_includes_client_id_for_websocket_targeting(self, monkeypatch): + from tools._comfyui.client import ComfyUIClient + + client = ComfyUIClient("http://comfy.test") + seen = {} + + def fake_post(url, json=None, timeout=None): + seen.update(json) + return type("R", (), { + "raise_for_status": lambda self: None, + "json": lambda self: {"prompt_id": "abc"}, + })() + + monkeypatch.setattr("tools._comfyui.client.requests.post", fake_post) + client.submit({"1": {"inputs": {}}}) + assert seen["client_id"] == client.client_id + + +class _FakeWSTimeout(Exception): + pass + + +class _FakeWSConn: + def __init__(self, frames): + self._frames = list(frames) + + def settimeout(self, value): + pass + + def recv(self): + if not self._frames: + raise _FakeWSTimeout() + return self._frames.pop(0) + + def close(self): + pass + + +def _install_fake_websocket(monkeypatch, frames): + """Inject a fake `websocket` module so wait_ws() runs without the real + optional websocket-client dependency installed.""" + import sys + import types + + fake_module = types.SimpleNamespace( + WebSocketTimeoutException=_FakeWSTimeout, + create_connection=lambda url, timeout=10: _FakeWSConn(frames), + ) + monkeypatch.setitem(sys.modules, "websocket", fake_module) + + +class TestWebsocketWait: + + def test_wait_ws_completes_on_executing_none_node(self, monkeypatch, tmp_path): + from tools._comfyui.client import ComfyUIClient + + client = ComfyUIClient("http://comfy.test") + progress_events = [] + frames = [ + json.dumps({"type": "progress", "data": { + "value": 2, "max": 20, "prompt_id": "p1", + }}), + json.dumps({"type": "executing", "data": { + "node": None, "prompt_id": "p1", + }}), + ] + _install_fake_websocket(monkeypatch, frames) + monkeypatch.setattr( + "tools._comfyui.client.requests.get", + lambda *a, **k: type("R", (), { + "raise_for_status": lambda self: None, + "json": lambda self: {"p1": {"outputs": {"9": {}}}}, + })(), + ) + + entry = client.wait_ws("p1", timeout=5, on_progress=progress_events.append) + + assert entry == {"outputs": {"9": {}}} + assert progress_events == [{"value": 2, "max": 20, "prompt_id": "p1"}] + + def test_wait_ws_execution_error_raises_with_prompt_id(self, monkeypatch): + from tools._comfyui.client import ComfyUIClient, ComfyUIError + + client = ComfyUIClient("http://comfy.test") + frames = [ + json.dumps({"type": "execution_error", "data": { + "prompt_id": "p2", "exception_message": "boom", + }}), + ] + _install_fake_websocket(monkeypatch, frames) + + with pytest.raises(ComfyUIError) as excinfo: + client.wait_ws("p2", timeout=5) + + assert excinfo.value.prompt_id == "p2" + + def test_wait_ws_ignores_other_prompts_on_shared_connection(self, monkeypatch): + from tools._comfyui.client import ComfyUIClient + + client = ComfyUIClient("http://comfy.test") + frames = [ + # Another job's event on the same client_id -- must not trigger completion. + json.dumps({"type": "executing", "data": { + "node": None, "prompt_id": "someone-elses-job", + }}), + json.dumps({"type": "executing", "data": { + "node": None, "prompt_id": "p3", + }}), + ] + _install_fake_websocket(monkeypatch, frames) + monkeypatch.setattr( + "tools._comfyui.client.requests.get", + lambda *a, **k: type("R", (), { + "raise_for_status": lambda self: None, + "json": lambda self: {"p3": {"outputs": {}}}, + })(), + ) + + entry = client.wait_ws("p3", timeout=5) + assert entry == {"outputs": {}} + + def test_wait_ws_timeout_raises_comfyuierror_with_prompt_id(self, monkeypatch): + from tools._comfyui.client import ComfyUIClient, ComfyUIError + + client = ComfyUIClient("http://comfy.test") + _install_fake_websocket(monkeypatch, frames=[]) # recv() always times out + + with pytest.raises(ComfyUIError) as excinfo: + client.wait_ws("p4", timeout=0) + + assert excinfo.value.prompt_id == "p4" + + def test_wait_falls_back_to_poll_when_websocket_unavailable(self, monkeypatch): + """No websocket-client installed (or any transport failure) must + silently fall back to REST polling, not blow up the whole call.""" + from tools._comfyui.client import ComfyUIClient + import sys + + client = ComfyUIClient("http://comfy.test") + monkeypatch.delitem(sys.modules, "websocket", raising=False) + monkeypatch.setattr( + "builtins.__import__", + _raise_on_websocket_import(__import__), + ) + monkeypatch.setattr( + client, "poll", lambda prompt_id, **kwargs: {"outputs": {"used": "poll"}} + ) + + entry = client._wait("p5", timeout=5, interval=5) + assert entry == {"outputs": {"used": "poll"}} + + def test_wait_does_not_swallow_genuine_comfyuierror_from_websocket(self, monkeypatch): + """A real execution error detected over the websocket must propagate, + not be masked by a fallback-to-poll retry.""" + from tools._comfyui.client import ComfyUIClient, ComfyUIError + + client = ComfyUIClient("http://comfy.test") + frames = [ + json.dumps({"type": "execution_error", "data": { + "prompt_id": "p6", "exception_message": "bad node", + }}), + ] + _install_fake_websocket(monkeypatch, frames) + + def fail_poll(prompt_id, **kwargs): + raise AssertionError("poll() should not be called after a real ws error") + + monkeypatch.setattr(client, "poll", fail_poll) + + with pytest.raises(ComfyUIError) as excinfo: + client._wait("p6", timeout=5, interval=5) + assert excinfo.value.prompt_id == "p6" + + +def _raise_on_websocket_import(real_import): + def _import(name, *args, **kwargs): + if name == "websocket": + raise ImportError("no module named websocket") + return real_import(name, *args, **kwargs) + return _import + # ------------------------------------------------------------------ # Model discovery (offline, no server needed) diff --git a/tools/_comfyui/client.py b/tools/_comfyui/client.py index 41fee64b..85b69421 100644 --- a/tools/_comfyui/client.py +++ b/tools/_comfyui/client.py @@ -11,8 +11,9 @@ import json import os import random import time +import uuid from pathlib import Path -from typing import Any +from typing import Any, Callable import requests @@ -46,6 +47,9 @@ class ComfyUIClient: server_url or os.environ.get("COMFYUI_SERVER_URL", "http://localhost:8188") ).rstrip("/") + # Scopes websocket execution events to this client (see wait_ws) and + # is echoed back on /prompt so the server targets messages to us. + self.client_id = str(uuid.uuid4()) # ------------------------------------------------------------------ # Health @@ -147,7 +151,7 @@ class ComfyUIClient: """Queue a workflow for execution. Returns the ``prompt_id``.""" resp = requests.post( f"{self.server_url}/prompt", - json={"prompt": workflow}, + json={"prompt": workflow, "client_id": self.client_id}, timeout=30, ) try: @@ -198,6 +202,117 @@ class ComfyUIClient: prompt_id=prompt_id, ) + def wait_ws( + self, + prompt_id: str, + *, + timeout: int = 600, + interval: int = 5, + on_progress: Callable[[dict], None] | None = None, + ) -> dict: + """Block until *prompt_id* finishes, watching ComfyUI's websocket feed. + + Reacts to server-pushed ``executing``/``progress``/``execution_error`` + events instead of sleeping between REST polls, so completion and + errors are detected immediately rather than up to *interval* seconds + late. *on_progress*, if given, is called with each ``progress`` + message's ``data`` dict (``value``, ``max``, ``node``, ``prompt_id``). + + Requires the optional ``websocket-client`` package. Any transport + failure (missing dependency, connection refused, dropped socket, + malformed frame) propagates as a plain exception — callers should + catch it and fall back to :meth:`poll`, which is what :meth:`generate` + does. A genuine ComfyUI-side execution error or an unmet deadline is + raised as :class:`ComfyUIError` with ``prompt_id`` set, exactly like + :meth:`poll`, so ``resume_prompt_id`` recovery works the same way + regardless of which wait strategy was used. + """ + import websocket # websocket-client; optional, see docstring + + ws_url = self.server_url.replace("http://", "ws://", 1).replace( + "https://", "wss://", 1 + ) + conn = websocket.create_connection( + f"{ws_url}/ws?clientId={self.client_id}", timeout=10 + ) + try: + conn.settimeout(interval) + deadline = time.time() + timeout + finished = False + while time.time() < deadline: + try: + raw = conn.recv() + except websocket.WebSocketTimeoutException: + continue + if not isinstance(raw, str): + continue # binary preview-image frame, not a status message + try: + message = json.loads(raw) + except json.JSONDecodeError: + continue + data = message.get("data", {}) + if data.get("prompt_id") not in (None, prompt_id): + continue # another job sharing this connection + msg_type = message.get("type") + if msg_type == "progress": + if on_progress: + on_progress(data) + elif msg_type == "execution_error": + raise ComfyUIError( + f"Execution error: {data}", prompt_id=prompt_id + ) + elif msg_type == "executing" and data.get("node") is None: + finished = True + break + finally: + conn.close() + + if not finished: + raise ComfyUIError( + f"Prompt {prompt_id} did not complete within {timeout}s " + f"(websocket wait). The job was not cancelled — resume with " + f"resume_prompt_id={prompt_id!r} and a longer timeout.", + prompt_id=prompt_id, + ) + + resp = requests.get(f"{self.server_url}/history/{prompt_id}", timeout=10) + resp.raise_for_status() + entry = resp.json().get(prompt_id) + if entry is None: + raise ComfyUIError( + f"No history entry for {prompt_id} after completion", + prompt_id=prompt_id, + ) + return entry + + def _wait( + self, + prompt_id: str, + *, + timeout: int, + interval: int, + on_progress: Callable[[dict], None] | None = None, + ) -> dict: + """Wait for *prompt_id*, preferring the websocket feed over polling. + + Falls back to :meth:`poll` when ``websocket-client`` isn't installed + or the websocket can't be established/maintained. A genuine + :class:`ComfyUIError` (execution error or deadline reached) is never + swallowed by the fallback — only transport-level failures are. The + fallback gets whatever's left of *timeout*, not a fresh budget, so a + mid-wait websocket drop can't double the caller's worst-case wait. + """ + started = time.time() + try: + return self.wait_ws( + prompt_id, timeout=timeout, interval=interval, on_progress=on_progress + ) + except ComfyUIError: + raise + except Exception: + remaining = max(timeout - (time.time() - started), 0) + return self.poll(prompt_id, timeout=remaining, interval=interval) + def download( self, filename: str, @@ -247,15 +362,23 @@ class ComfyUIClient: timeout: int = 600, interval: int = 5, resume_prompt_id: str | None = None, + on_progress: Callable[[dict], None] | None = None, ) -> list[Path]: - """Submit → poll → download. Returns list of artifact paths. + """Submit → wait → download. Returns list of artifact paths. Pass ``resume_prompt_id`` (from a previous ``ComfyUIError.prompt_id``) to skip re-submitting an already-queued/running job and just resume waiting on it — the common recovery path after a timeout. + + Waiting prefers ComfyUI's websocket feed (immediate completion/error + detection, optional live ``on_progress`` callback) and transparently + falls back to REST polling if ``websocket-client`` isn't installed or + the connection can't be used. See :meth:`_wait`. """ prompt_id = resume_prompt_id or self.submit(workflow) - entry = self.poll(prompt_id, timeout=timeout, interval=interval) + entry = self._wait( + prompt_id, timeout=timeout, interval=interval, on_progress=on_progress + ) outputs = entry.get("outputs", {}) node_output = outputs.get(output_node, {}) diff --git a/tools/video/comfyui_video.py b/tools/video/comfyui_video.py index 5bab999d..84270c27 100644 --- a/tools/video/comfyui_video.py +++ b/tools/video/comfyui_video.py @@ -217,6 +217,23 @@ class ComfyUIVideo(BaseTool): def __init__(self) -> None: self._client = ComfyUIClient() + self._last_progress_log = 0.0 + + def _log_progress(self, data: dict) -> None: + """Print a throttled progress line for long video renders. + + Video jobs can run for tens of minutes; without this the process + looks hung. Throttled to once per 10s since ComfyUI pushes a + ``progress`` event per sampling step, which would otherwise flood + stdout on fast GPUs. + """ + now = time.monotonic() + if now - self._last_progress_log < 10: + return + self._last_progress_log = now + value, max_value = data.get("value"), data.get("max") + if value is not None and max_value: + print(f"[comfyui_video] step {value}/{max_value}") def get_status(self) -> ToolStatus: if not self._client.is_available(): @@ -341,6 +358,7 @@ class ComfyUIVideo(BaseTool): timeout=inputs.get("timeout_seconds", 3600), interval=10, resume_prompt_id=inputs.get("resume_prompt_id"), + on_progress=self._log_progress, ) except ComfyUIError as exc: