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>
This commit is contained in:
Ntsako
2026-08-06 12:31:24 +02:00
parent ceee7c7d56
commit ca203e49b7
4 changed files with 336 additions and 7 deletions

View File

@@ -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: