feat: wait on ComfyUI websocket feed instead of polling for completion

Resolves the "async generation" open question from the adapter plan.
generate() now watches ComfyUI's websocket events (executing/progress/
execution_error) and reacts immediately instead of sleeping between REST
polls, with an optional on_progress callback that comfyui_video uses to
print step progress on long renders. websocket-client is an optional
import; _wait() falls back to the original poll() loop (with the
remaining time budget, not a fresh one) when it's unavailable or the
connection drops, so resume_prompt_id recovery is unaffected either way.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

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?

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)

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, {})

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: