diff --git a/comfy_cli/command/run/__init__.py b/comfy_cli/command/run/__init__.py index c1ed383d..ce0dff94 100644 --- a/comfy_cli/command/run/__init__.py +++ b/comfy_cli/command/run/__init__.py @@ -385,6 +385,12 @@ def execute( wait_state.status = "running" _write_state(wait_state) + # Only now, with the durable record on disk, announce the queue. + # The event-order contract (docs/json-output.md) is + # `prompt_preview → queued → node events`, so this must also + # precede `watch_execution`'s WebSocket stream. + _emit_queued(renderer, execution) + # `watch_execution` reports a terminal server event by rendering # the error and raising `typer.Exit` (1 for `execution_error`, 130 # for `execution_interrupted`) — the ordinary failure path, not an @@ -460,6 +466,10 @@ def execute( state_file = jobs_state.write(state) watcher_spawned = _spawn_watcher(execution.prompt_id, where="local", host=host, port=port, notify=notify) + # Emitted after the state file and the watcher spawn attempt, so a + # consumer reading this line can already `comfy jobs status` it. + _emit_queued(renderer, execution) + if renderer.is_pretty(): from comfy_cli.output.glyphs import status_glyph @@ -527,6 +537,12 @@ def execute( ) raise typer.Exit(code=1) except (WebSocketException, ConnectionError, OSError) as e: + if isinstance(e, BrokenPipeError): + # A closed stdout is the consumer leaving, not the server dying — + # let click/__main__ turn it into the documented silent exit 0 and + # leave the just-written state record untouched. Same principle as + # `completed_payload` being emitted outside this try. + raise # If we closed the WebSocket ourselves in response to Ctrl-C, the recv # loop exits with a WebSocketException that *looks* like the server # vanished. Check the cancellation token first so we emit the right @@ -594,6 +610,24 @@ def execute( renderer.emit(completed_payload, command="run", where="local") +def _emit_queued(renderer, execution) -> None: + """Emit the contractual local `queued` event (docs/json-output.md#queued). + + Lives in the caller rather than ``WorkflowExecution.queue()`` so it can be + emitted *after* the job's state file is persisted — a consumer that sees + this line can immediately `comfy jobs status ` / see the job in + `jobs ls`, and a `BrokenPipeError` here can no longer abort setup before + the durable record exists. The field set is unchanged. + """ + renderer.event( + "queued", + prompt_id=execution.prompt_id, + client_id=execution.client_id, + validation_warnings=execution.validation_warnings, + nodes=execution.workflow_manifest(), + ) + + def _write_state(state): """Best-effort ``jobs_state.write``. Returns the path, or None if the write was skipped, the state dir was unwritable, or the prompt_id was @@ -860,7 +894,8 @@ def execute_cloud( # with `queued` (async) / `executing` (--wait) carrying {workflow, base_url} # — both wrong per docs/json-output.md: `queued` means "the server accepted # the prompt" (so it cannot precede the POST), and `executing` is a per-node - # event. The contractual `queued` is emitted after submit, below. + # event. The contractual `queued` is emitted after submit *and* after the + # job state file is written, below. try: if not wait and renderer.is_pretty(): @@ -915,19 +950,27 @@ def execute_cloud( raise typer.Exit(code=1) # The contractual `queued`: the server has the prompt. Same shape the local - # path emits from `WorkflowExecution.queue()` (docs/json-output.md#queued), - # plus `base_url` so a cloud consumer still learns the target. - # `validation_warnings` is always empty here — unlike local, the cloud - # treats any `node_errors` on an accepted submit as a hard `prompt_rejected` - # above, so a partially-valid graph never reaches this line. - renderer.event( - "queued", - prompt_id=submit.prompt_id, - client_id=client_id, - validation_warnings=[], - nodes=workflow_manifest(parsed_workflow), - base_url=target.base_url, - ) + # path emits (docs/json-output.md#queued), plus `base_url` so a cloud + # consumer still learns the target. `validation_warnings` is always empty + # here — unlike local, the cloud treats any `node_errors` on an accepted + # submit as a hard `prompt_rejected` above, so a partially-valid graph + # never reaches this line. + # + # Deliberately NOT emitted at this point: both branches below emit it only + # once the job's state file (and, on the async branch, the watcher) exist. + # A `BrokenPipeError` here — `comfy run --where cloud --json-stream … | + # head -n1`, where the NDJSON renderer flushes every line — used to abort + # setup before `jobs_state.write` ran while `__main__` still exited 0, + # leaving a billable cloud job with no journal entry, state file or watcher. + def _emit_queued_cloud() -> None: + renderer.event( + "queued", + prompt_id=submit.prompt_id, + client_id=client_id, + validation_warnings=[], + nodes=workflow_manifest(parsed_workflow), + base_url=target.base_url, + ) if not wait: state = jobs_state.new( @@ -942,6 +985,9 @@ def execute_cloud( _journal_run(workflow_name, submit.prompt_id, "cloud") watcher_spawned = _spawn_watcher(submit.prompt_id, where="cloud", notify=notify) + # Durable record + watcher exist: safe to announce. + _emit_queued_cloud() + if renderer.is_pretty(): from comfy_cli.output.glyphs import status_glyph @@ -991,6 +1037,12 @@ def execute_cloud( state_file = jobs_state.write(state) _journal_run(workflow_name, submit.prompt_id, "cloud") + # Announced once the record is durable, and *outside* the polling try below + # — none of its handlers should ever see (and rewrite state for) a + # BrokenPipeError raised by this emit. It propagates to click/`__main__`, + # which exits 0 silently, with the job already recorded. + _emit_queued_cloud() + try: def _probe(): diff --git a/comfy_cli/command/run/execution.py b/comfy_cli/command/run/execution.py index 7c043b23..7f4c62ec 100644 --- a/comfy_cli/command/run/execution.py +++ b/comfy_cli/command/run/execution.py @@ -163,6 +163,12 @@ def __init__( # compute nodes that never fire a server-side `executed` event. self.cached_node_ids: list[str] = [] self.executed_node_ids: list[str] = [] + # Per-node issues the server reported alongside a successful (HTTP 200) + # queue, stashed by ``queue()`` for the caller's `queued` event. The + # event itself is emitted by ``comfy_cli.command.run.execute`` — only + # *after* the job's state file is persisted — so a consumer sees it and + # can immediately rely on `comfy jobs status` / `jobs ls` finding it. + self.validation_warnings: list[dict] = [] # Classified verdict of a terminal server-side `execution_error`, # stashed by `on_error` before it raises `typer.Exit`. The `--wait` # caller only sees the exit, so this is how the real cause reaches @@ -289,16 +295,11 @@ def queue(self): # 200 may still carry node_errors if some output chains failed # validation but others passed — surface as warnings, not a failure. + # Stored rather than emitted here: the contractual `queued` event is + # emitted by the caller once the job state file exists (see + # ``validation_warnings`` above and ``run._emit_queued``). node_errors = body.get("node_errors") if isinstance(body, dict) else None - validation_warnings = _node_errors_to_list(node_errors) - - self.renderer.event( - "queued", - prompt_id=prompt_id, - client_id=self.client_id, - validation_warnings=validation_warnings, - nodes=self.workflow_manifest(), - ) + self.validation_warnings = _node_errors_to_list(node_errors) def watch_execution(self): if self.ws is None: diff --git a/docs/json-output.md b/docs/json-output.md index 5c8f0b7b..68f7f282 100644 --- a/docs/json-output.md +++ b/docs/json-output.md @@ -238,7 +238,8 @@ final envelope and exits 0 without queuing. ### `queued` Emitted after the submit request returns success — `POST /prompt` returning -200 locally, the equivalent cloud submit under `--where cloud`. +200 locally, the equivalent cloud submit under `--where cloud` — **and** after +the job's state file has been persisted. ```json { @@ -262,8 +263,12 @@ Emitted after the submit request returns success — `POST /prompt` returning | `nodes` | array of dict | Manifest of every node in the submitted (post-conversion) workflow: `node_id` (str), `class_type` (str), `title` (str). Lets piped consumers render a per-node UI without the workflow file. | | `base_url` | str | **Cloud only.** The cloud endpoint the prompt was submitted to. Absent on `--where local`. | -`queued` is emitted **after** the submit call returns successfully, on both -targets — it means the server has the prompt. A run whose submit fails emits +`queued` is emitted **after** the submit call returns successfully **and after +the job's state file has been persisted** — on the async (`--no-wait`) paths, +after the background watcher spawn attempt too. On both targets it therefore +means more than "the server has the prompt": the durable record exists, so a +consumer may run `comfy jobs status ` — or expect the job in +`comfy jobs ls` — the moment it reads this line. A run whose submit fails emits its error envelope with no `queued` line at all. ### `executing` diff --git a/tests/comfy_cli/command/test_run_json.py b/tests/comfy_cli/command/test_run_json.py index 3e32622d..89ae833f 100644 --- a/tests/comfy_cli/command/test_run_json.py +++ b/tests/comfy_cli/command/test_run_json.py @@ -22,6 +22,7 @@ import io import json import os +import sys import tempfile import urllib.error from unittest.mock import MagicMock, patch @@ -2035,3 +2036,240 @@ def test_non_dict_error_items_do_not_crash_the_hint_builder(self, monkeypatch, w assert exit_code == 1 assert _envelope(lines)["error"]["code"] == "prompt_rejected" assert "node 1 (X): bad" in _envelope(lines)["error"]["hint"] + + +# -------------------------------------------------------------------------- +# `queued` is emitted only after the job's durable state exists (BE-6072) +# +# The event used to precede `jobs_state.write` on every path. That made +# state-file-backed follow-ups (`jobs ls`, `jobs wait`, compose's item_map) +# raceable against a just-queued job, and — worse — let a `BrokenPipeError` at +# the emit (`comfy run … --json-stream | head -n1`) abort setup before the +# record was ever written, while `__main__`'s broken-pipe handling still exited +# 0: a billable cloud job with no journal entry, no state file and no watcher. +# These tests pin the ordering and the propagation on all four paths. +# -------------------------------------------------------------------------- + + +class _TapStream: + """Stand-in for the renderer's machine stream. + + `on_line` sees every NDJSON line before it is written and may raise to + simulate a closed stdout. Writes delegate to ``sys.stdout`` at call time so + `capsys` still captures the stream. + """ + + def __init__(self, on_line): + self.on_line = on_line + + def write(self, s: str) -> int: + self.on_line(s) + return sys.stdout.write(s) + + def flush(self) -> None: + sys.stdout.flush() + + +def _tap_queued(renderer, callback): + """Invoke `callback` just before each `queued` line reaches the stream.""" + + def on_line(s: str) -> None: + try: + payload = json.loads(s) + except (json.JSONDecodeError, ValueError): + return + if isinstance(payload, dict) and payload.get("type") == "queued": + callback() + + renderer.machine_stream = _TapStream(on_line) + + +def _raise_broken_pipe() -> None: + """What writing an NDJSON line to a closed stdout (`… | head -n1`) does.""" + raise BrokenPipeError(32, "Broken pipe") + + +def _recording_write(writes: list): + """A ``jobs_state.write`` stand-in that records each persisted prompt_id.""" + + def fake_write(state): + writes.append(state.prompt_id) + return "/tmp/s.json" + + return fake_write + + +def _recording_spawn(spawned: list): + """A ``_spawn_watcher`` stand-in that records each spawn attempt.""" + + def fake_spawn(prompt_id, **_kw): + spawned.append(prompt_id) + return True + + return fake_spawn + + +def _log_state_writes(monkeypatch, log): + """Record every ``jobs_state.write`` (which ``_write_state`` delegates to) + in `log`, keeping the real return contract (a path).""" + from comfy_cli import jobs_state as jobs_state_mod + + def fake_write(state): + log.append(("state_write",)) + return "/tmp/state.json" + + monkeypatch.setattr(jobs_state_mod, "write", fake_write) + + +def _local_ws_that_finishes(MockWs): + ws_instance = MagicMock() + MockWs.return_value = ws_instance + ws_instance.recv.side_effect = [ + json.dumps({"type": "executing", "data": {"prompt_id": "p", "node": None}}), + ] + return ws_instance + + +class TestQueuedFollowsStatePersistence: + """Ordering spies: `state_write` must precede `queued_emit` on all four + target × wait-mode combinations.""" + + def test_local_async(self, monkeypatch, workflow_file, capsys, ndjson_renderer): + log: list[tuple] = [] + _log_state_writes(monkeypatch, log) + _tap_queued(ndjson_renderer, lambda: log.append(("queued_emit",))) + with ( + patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), + patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open, + patch("comfy_cli.command.run._spawn_watcher", return_value=True), + ): + mock_open.return_value.__enter__.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() + _lines, exit_code = _run_execute_capture(workflow_file, capsys, wait=False) + assert exit_code == 0 + assert ("queued_emit",) in log + assert log.index(("state_write",)) < log.index(("queued_emit",)) + + def test_local_wait(self, monkeypatch, workflow_file, capsys, ndjson_renderer): + log: list[tuple] = [] + _log_state_writes(monkeypatch, log) + _tap_queued(ndjson_renderer, lambda: log.append(("queued_emit",))) + with ( + patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), + patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open, + patch("comfy_cli.command.run.WebSocket") as MockWs, + ): + mock_open.return_value.__enter__.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() + _local_ws_that_finishes(MockWs) + _lines, exit_code = _run_execute_capture(workflow_file, capsys, wait=True) + assert exit_code == 0 + assert ("queued_emit",) in log + assert log.index(("state_write",)) < log.index(("queued_emit",)) + + @pytest.mark.parametrize("wait", [False, True]) + def test_cloud(self, monkeypatch, workflow_file, capsys, ndjson_renderer, wait): + log: list[tuple] = [] + _install_cloud_stubs(monkeypatch, client_cls=_fake_client()) + # After the stubs, so this write spy wins over their no-op one. + _log_state_writes(monkeypatch, log) + _tap_queued(ndjson_renderer, lambda: log.append(("queued_emit",))) + _lines, exit_code = _cloud_capture(capsys, workflow_file, wait=wait, timeout=5) + assert exit_code == 0 + assert ("queued_emit",) in log + assert log.index(("state_write",)) < log.index(("queued_emit",)) + + +class TestBrokenPipeAtTheQueuedEmit: + """A closed stdout at the `queued` line is the consumer leaving. It must + propagate as itself — `__main__` turns that into the documented silent exit + 0 — and only ever after the durable record (and the watcher) exist.""" + + def test_local_async_persists_before_it_raises(self, monkeypatch, workflow_file, capsys, ndjson_renderer): + writes: list[str] = [] + spawned: list[str] = [] + from comfy_cli import jobs_state as jobs_state_mod + + monkeypatch.setattr(jobs_state_mod, "write", _recording_write(writes)) + _tap_queued(ndjson_renderer, _raise_broken_pipe) + + with ( + patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), + patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open, + patch("comfy_cli.command.run._spawn_watcher", side_effect=_recording_spawn(spawned)), + ): + mock_open.return_value.__enter__.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() + with pytest.raises(BrokenPipeError): + execute(workflow_file, host="127.0.0.1", port=8188, wait=False, verbose=False, timeout=30) + capsys.readouterr() + assert writes == ["p"] + assert spawned == ["p"] + + def test_cloud_async_persists_before_it_raises(self, monkeypatch, workflow_file, capsys, ndjson_renderer): + from comfy_cli.command.run import execute_cloud + + writes: list[str] = [] + spawned: list[str] = [] + _install_cloud_stubs(monkeypatch, client_cls=_fake_client()) + from comfy_cli import jobs_state as jobs_state_mod + from comfy_cli.command import run as run_pkg + + monkeypatch.setattr(jobs_state_mod, "write", _recording_write(writes)) + monkeypatch.setattr(run_pkg, "_spawn_watcher", _recording_spawn(spawned)) + _tap_queued(ndjson_renderer, _raise_broken_pipe) + + with pytest.raises(BrokenPipeError): + execute_cloud(workflow_file, wait=False, timeout=5) + capsys.readouterr() + assert writes == ["cloud-pid"] + assert spawned == ["cloud-pid"] + + def test_local_wait_leaves_the_running_record_intact(self, monkeypatch, workflow_file, capsys, ndjson_renderer): + """The relocated emit sits inside the `except (WebSocketException, + ConnectionError, OSError)` block's try. Without the handler's + BrokenPipeError guard a closed stdout would be misreported as + `ws_disconnected` *and* rewrite this healthy record to error/server_died.""" + from comfy_cli import jobs_state as jobs_state_mod + + _tap_queued(ndjson_renderer, _raise_broken_pipe) + with ( + patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), + patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open, + patch("comfy_cli.command.run.WebSocket") as MockWs, + ): + mock_open.return_value.__enter__.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() + _local_ws_that_finishes(MockWs) + with pytest.raises(BrokenPipeError): + execute(workflow_file, host="127.0.0.1", port=8188, wait=True, verbose=False, timeout=30) + capsys.readouterr() + record = jobs_state_mod.read("p") + assert record is not None, "the submit-time state file must exist" + assert record.status == "running" + assert record.error is None + + +class TestQueueStashesWarningsWithoutEmitting: + """`WorkflowExecution.queue()` no longer emits — it stashes the warnings so + the caller can emit the identical event after persisting state.""" + + def test_queue_emits_nothing_and_stashes_the_warnings(self, simple_workflow, capsys): + ex = _make_workflow_execution(simple_workflow) + body = json.dumps( + { + "prompt_id": "p", + "node_errors": {"3": {"errors": [{"type": "x", "message": "skipped"}], "class_type": "X"}}, + } + ).encode() + with patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open: + mock_open.return_value.__enter__.return_value.read.return_value = body + ex.queue() + out, _err = capsys.readouterr() + assert "queued" not in out, f"queue() must not emit the event itself: {out!r}" + assert ex.prompt_id == "p" + assert ex.validation_warnings[0]["node_id"] == "3" + + def test_no_node_errors_leaves_the_warnings_empty(self, simple_workflow, capsys): + ex = _make_workflow_execution(simple_workflow) + with patch("comfy_cli.http._AUTHED_OPENER.open") as mock_open: + mock_open.return_value.__enter__.return_value.read.return_value = json.dumps({"prompt_id": "p"}).encode() + ex.queue() + capsys.readouterr() + assert ex.validation_warnings == []