From 41793226b5d0ee5932fc39b12efac2d269bafc2d Mon Sep 17 00:00:00 2001 From: calesthio Date: Thu, 13 Aug 2026 09:37:02 -0700 Subject: [PATCH] fix: make ComfyUI history authoritative --- tests/contracts/test_comfyui_tools.py | 56 +++++++++++++++++++++- tools/_comfyui/client.py | 69 ++++++++++++++++++++------- tools/audio/comfyui_music.py | 4 +- 3 files changed, 110 insertions(+), 19 deletions(-) diff --git a/tests/contracts/test_comfyui_tools.py b/tests/contracts/test_comfyui_tools.py index ac533f30..cf678a85 100644 --- a/tests/contracts/test_comfyui_tools.py +++ b/tests/contracts/test_comfyui_tools.py @@ -305,6 +305,7 @@ class TestClientHelpers: def test_generate_resume_prompt_id_skips_resubmit(self, monkeypatch, tmp_path): from tools._comfyui.client import ComfyUIClient + import sys client = ComfyUIClient("http://comfy.test") @@ -312,6 +313,14 @@ class TestClientHelpers: raise AssertionError("submit() should not be called when resuming") monkeypatch.setattr(client, "submit", fail_submit) + _install_fake_websocket(monkeypatch, frames=[]) + monkeypatch.setattr( + sys.modules["websocket"], + "create_connection", + lambda *a, **k: (_ for _ in ()).throw( + AssertionError("resumed jobs must use history polling") + ), + ) monkeypatch.setattr(client, "poll", lambda prompt_id, **kwargs: { "outputs": {"9": {"images": [{ "filename": "resumed.png", "subfolder": "", "type": "output", @@ -502,6 +511,29 @@ def _install_fake_websocket(monkeypatch, frames): class TestWebsocketWait: + def test_wait_ws_returns_job_completed_before_connection(self, monkeypatch): + from tools._comfyui.client import ComfyUIClient + import sys + + client = ComfyUIClient("http://comfy.test") + _install_fake_websocket(monkeypatch, frames=[]) + monkeypatch.setattr( + sys.modules["websocket"], + "create_connection", + lambda *a, **k: (_ for _ in ()).throw( + AssertionError("completed history must avoid websocket connection") + ), + ) + monkeypatch.setattr( + "tools._comfyui.client.requests.get", + lambda *a, **k: type("R", (), { + "raise_for_status": lambda self: None, + "json": lambda self: {"done": {"outputs": {"9": {}}}}, + })(), + ) + + assert client.wait_ws("done", timeout=5) == {"outputs": {"9": {}}} + def test_wait_ws_completes_on_executing_none_node(self, monkeypatch, tmp_path): from tools._comfyui.client import ComfyUIClient @@ -516,11 +548,12 @@ class TestWebsocketWait: }}), ] _install_fake_websocket(monkeypatch, frames) + history_calls = iter(({}, {}, {"p1": {"outputs": {"9": {}}}})) monkeypatch.setattr( "tools._comfyui.client.requests.get", lambda *a, **k: type("R", (), { "raise_for_status": lambda self: None, - "json": lambda self: {"p1": {"outputs": {"9": {}}}}, + "json": lambda self: next(history_calls), })(), ) @@ -937,6 +970,27 @@ class TestComfyUIMusic: assert seen["workflow"]["8"]["inputs"]["seed"] == 777 assert result.data["model"] == "ace-step-v1-3.5b" + def test_bundled_generation_preserves_seed_zero(self, tmp_path): + tool = ComfyUIMusic() + tool._client.is_available = lambda: True + tool._client.check_models = lambda required: (list(required), []) + seen = {} + + def fake_generate(workflow, output_node, dest, **kwargs): + seen["seed"] = workflow["8"]["inputs"]["seed"] + return [Path(dest)] + + tool._client.generate = fake_generate + result = tool.execute({ + "prompt": "deterministic test", + "seed": 0, + "output_path": str(tmp_path / "music.mp3"), + }) + + assert result.success is True + assert result.seed == 0 + assert seen["seed"] == 0 + def test_get_status_degraded_when_model_missing(self): tool = ComfyUIMusic() tool._client.is_available = lambda: True diff --git a/tools/_comfyui/client.py b/tools/_comfyui/client.py index e280bea6..e37bc2b4 100644 --- a/tools/_comfyui/client.py +++ b/tools/_comfyui/client.py @@ -199,17 +199,8 @@ class ComfyUIClient: """Block until *prompt_id* finishes. Returns the history entry.""" deadline = time.time() + timeout while time.time() < deadline: - resp = requests.get( - f"{self.server_url}/history/{prompt_id}", timeout=10 - ) - resp.raise_for_status() - history = resp.json() - if prompt_id in history: - entry = history[prompt_id] - status = entry.get("status", {}) - if status.get("status_str") == "error": - msgs = status.get("messages", []) - raise ComfyUIError(f"Execution error: {msgs}", prompt_id=prompt_id) + entry = self._history_entry(prompt_id) + if entry is not None: return entry time.sleep(interval) raise ComfyUIError( @@ -223,6 +214,28 @@ class ComfyUIClient: prompt_id=prompt_id, ) + def _history_entry(self, prompt_id: str) -> dict | None: + """Return a completed history entry, or ``None`` while it is absent.""" + 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: + return None + status = entry.get("status", {}) + if status.get("status_str") == "error": + msgs = status.get("messages", []) + raise ComfyUIError(f"Execution error: {msgs}", prompt_id=prompt_id) + return entry + + def _history_entry_if_reachable(self, prompt_id: str) -> dict | None: + """Best-effort history probe while the websocket remains usable.""" + try: + return self._history_entry(prompt_id) + except ComfyUIError: + raise + except Exception: + return None + def wait_ws( self, prompt_id: str, @@ -250,6 +263,12 @@ class ComfyUIClient: """ import websocket # websocket-client; optional, see docstring + # History is authoritative and websocket events are not replayed. The + # job may already have finished between submit() and this wait call. + entry = self._history_entry_if_reachable(prompt_id) + if entry is not None: + return entry + ws_url = self.server_url.replace("http://", "ws://", 1).replace( "https://", "wss://", 1 ) @@ -260,10 +279,19 @@ class ComfyUIClient: conn.settimeout(interval) deadline = time.time() + timeout finished = False + # Close the remaining race between the first history probe and + # websocket connection establishment. Events after this point are + # queued on the open socket; earlier completion is in history. + entry = self._history_entry_if_reachable(prompt_id) + if entry is not None: + return entry while time.time() < deadline: try: raw = conn.recv() except websocket.WebSocketTimeoutException: + entry = self._history_entry_if_reachable(prompt_id) + if entry is not None: + return entry continue if not isinstance(raw, str): continue # binary preview-image frame, not a status message @@ -289,6 +317,9 @@ class ComfyUIClient: conn.close() if not finished: + entry = self._history_entry_if_reachable(prompt_id) + if entry is not None: + return entry raise ComfyUIError( f"Prompt {prompt_id} did not complete within {timeout}s " f"(websocket wait). The job was not cancelled — resume with " @@ -296,9 +327,7 @@ class ComfyUIClient: 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) + entry = self._history_entry(prompt_id) if entry is None: raise ComfyUIError( f"No history entry for {prompt_id} after completion", @@ -397,9 +426,15 @@ class ComfyUIClient: the connection can't be used. See :meth:`_wait`. """ prompt_id = resume_prompt_id or self.submit(workflow) - entry = self._wait( - prompt_id, timeout=timeout, interval=interval, on_progress=on_progress - ) + if resume_prompt_id: + # A prompt resumed by a new client instance was submitted with the + # original instance's client_id, so its websocket events are not + # guaranteed to reach this socket. Poll authoritative history. + entry = self.poll(prompt_id, timeout=timeout, interval=interval) + else: + 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/audio/comfyui_music.py b/tools/audio/comfyui_music.py index 9751ad46..05dfddf0 100644 --- a/tools/audio/comfyui_music.py +++ b/tools/audio/comfyui_music.py @@ -227,7 +227,9 @@ class ComfyUIMusic(BaseTool): ) start = time.time() - seed = inputs.get("seed") or ComfyUIClient.random_seed() + seed = inputs.get("seed") + if seed is None: + seed = ComfyUIClient.random_seed() output_path = Path(inputs.get("output_path", f"comfyui_music_{seed}.mp3")) try: