diff --git a/comfy_cli/command/run/__init__.py b/comfy_cli/command/run/__init__.py index ad299286..d73b5c7e 100644 --- a/comfy_cli/command/run/__init__.py +++ b/comfy_cli/command/run/__init__.py @@ -186,14 +186,7 @@ def execute( # `preloaded` short-circuits file loading: an in-memory API-format graph # (e.g. the `comfy run --prompt` injected default) is handed straight in as # (workflow_dict, display_name, is_ui). Everything downstream is unchanged. - if preloaded is not None: - raw_workflow, workflow_name, is_ui = preloaded - else: - try: - raw_workflow, workflow_name, is_ui = _load_workflow_file(workflow) - except WorkflowLoadError as e: - renderer.error(code=e.code, message=str(e), hint=e.hint) - raise typer.Exit(code=1) from e + raw_workflow, workflow_name, is_ui = _load_workflow_or_exit(renderer, workflow, preloaded) if not print_prompt and not check_comfy_server_running(port, host, timeout=timeout): renderer.error( @@ -447,52 +440,31 @@ def execute( # Async path (the default). Write the initial state file and # spawn a detached watcher to keep it updated; the foreground # caller returns immediately with the prompt_id. - state = jobs_state.new( - prompt_id=execution.prompt_id, - client_id=execution.client_id, - workflow=workflow_name, + _, state_file = _init_queued_state( + execution.prompt_id, + execution.client_id, + workflow_name, where="local", + compose_meta=compose_meta, host=host, port=port, ) - state.item_map = (compose_meta or {}).get("items") - state_file = jobs_state.write(state) watcher_spawned = _spawn_watcher(execution.prompt_id, where="local", host=host, port=port, notify=notify) - - if renderer.is_pretty(): - from comfy_cli.output.glyphs import status_glyph - - pprint( - f"{status_glyph('queued')} [dim]{execution.prompt_id}[/dim]\n" - f" [dim]workflow [/dim]{workflow_name}\n" - f" [dim]watch [/dim][cyan]comfy jobs watch {execution.prompt_id}[/cyan]\n" - f" [dim]state [/dim]{state_file}" - ) - if not watcher_spawned: - pprint( - "[yellow]⚠ Background watcher could not start; poll manually with `comfy jobs status`[/yellow]" - ) - renderer.emit( - { - "workflow": workflow_name, - "status": "queued", - "prompt_id": execution.prompt_id, - "client_id": execution.client_id, - "outputs": [], - "elapsed_seconds": None, + _emit_queued( + execution.prompt_id, + execution.client_id, + workflow_name, + where="local", + watch_command=f"comfy jobs watch {execution.prompt_id}", + state_file=state_file, + watcher_spawned=watcher_spawned, + extra_envelope_fields={ "host": host, "port": port, "state_file": str(state_file) if state_file else None, "watcher_spawned": watcher_spawned, }, - command="run", - where="local", ) - # Pretty mode: brief live tail so the user can see the job - # move through "allocated → executing → completed" without - # having to run `comfy jobs watch`. The background watcher - # keeps writing the state file after we return. - _tail_state_file(execution.prompt_id) except KeyboardInterrupt: if progress is not None: progress.stop() @@ -681,6 +653,91 @@ def _journal_run(workflow: str, prompt_id, where: str) -> None: pass +def _load_workflow_or_exit(renderer, workflow, preloaded): + """Resolve the ``(raw_workflow, workflow_name, is_ui)`` triple that opens + both :func:`execute` and :func:`execute_cloud`: honor an in-memory + ``preloaded`` graph when present, else load from disk — mapping a load + failure to a rendered error + ``typer.Exit(1)``.""" + if preloaded is not None: + return preloaded + try: + return _load_workflow_file(workflow) + except WorkflowLoadError as e: + renderer.error(code=e.code, message=str(e), hint=e.hint) + raise typer.Exit(code=1) from e + + +def _init_queued_state(prompt_id, client_id, workflow_name, *, where, compose_meta, **where_fields): + """Create + persist the initial ``queued`` job state. Shared by the local + async path, the cloud non-wait path, and the cloud ``--wait`` path. + + ``where_fields`` carries the target-specific locators (``host``/``port`` for + local, ``base_url`` for cloud). Returns the ``(state, state_file)`` pair — + ``--wait`` callers keep mutating ``state`` and re-writing it later.""" + state = jobs_state.new( + prompt_id=prompt_id, + client_id=client_id, + workflow=workflow_name, + where=where, + **where_fields, + ) + state.item_map = (compose_meta or {}).get("items") + state_file = jobs_state.write(state) + return state, state_file + + +def _emit_queued( + prompt_id, + client_id, + workflow_name, + *, + where, + watch_command, + state_file, + watcher_spawned, + extra_envelope_fields, +): + """Async-submit ``queued`` tail shared by the local async path and the + cloud non-wait path: the pretty status panel (+ watcher-failure warning), + the machine-readable ``queued`` envelope, and the brief live state tail. + + The divergent bits stay at the call sites and arrive via parameters: + ``watch_command`` (the ``comfy jobs watch …`` hint, with ``--where cloud`` + for cloud) and ``extra_envelope_fields`` (the target locator plus the + local-only ``watcher_spawned`` key). The caller owns envelope field + membership + ordering so the public agent-mode JSON contract stays verbatim + per target.""" + renderer = get_renderer() + if renderer.is_pretty(): + from comfy_cli.output.glyphs import status_glyph + + pprint( + f"{status_glyph('queued')} [dim]{prompt_id}[/dim]\n" + f" [dim]workflow [/dim]{workflow_name}\n" + f" [dim]watch [/dim][cyan]{watch_command}[/cyan]\n" + f" [dim]state [/dim]{state_file}" + ) + if not watcher_spawned: + pprint("[yellow]⚠ Background watcher could not start; poll manually with `comfy jobs status`[/yellow]") + renderer.emit( + { + "workflow": workflow_name, + "status": "queued", + "prompt_id": prompt_id, + "client_id": client_id, + "outputs": [], + "elapsed_seconds": None, + **extra_envelope_fields, + }, + command="run", + where=where, + ) + # Pretty mode: brief live tail so the user sees the job move through + # "allocated → executing → completed" before the foreground exits. The + # background watcher keeps writing the state file after we return. + _tail_state_file(prompt_id) + + def _count_output_nodes(workflow: dict, object_info: dict) -> int | None: """Count nodes in ``workflow`` whose class is an output node, per ``object_info``. Returns None when object_info is empty/unknown so callers @@ -730,14 +787,7 @@ def execute_cloud( from comfy_cli.target import resolve_target renderer = get_renderer() - if preloaded is not None: - raw_workflow, workflow_name, is_ui = preloaded - else: - try: - raw_workflow, workflow_name, is_ui = _load_workflow_file(workflow) - except WorkflowLoadError as e: - renderer.error(code=e.code, message=str(e), hint=e.hint) - raise typer.Exit(code=1) from e + raw_workflow, workflow_name, is_ui = _load_workflow_or_exit(renderer, workflow, preloaded) if is_ui: # Frontend-format workflows (the `nodes`+`links` shape from the canvas @@ -886,60 +936,42 @@ def execute_cloud( raise typer.Exit(code=1) if not wait: - state = jobs_state.new( - prompt_id=submit.prompt_id, - client_id=client_id, - workflow=workflow_name, + _, state_file = _init_queued_state( + submit.prompt_id, + client_id, + workflow_name, where="cloud", + compose_meta=compose_meta, base_url=target.base_url, ) - state.item_map = (compose_meta or {}).get("items") - state_file = jobs_state.write(state) _journal_run(workflow_name, submit.prompt_id, "cloud") watcher_spawned = _spawn_watcher(submit.prompt_id, where="cloud", notify=notify) - - if renderer.is_pretty(): - from comfy_cli.output.glyphs import status_glyph - - pprint( - f"{status_glyph('queued')} [dim]{submit.prompt_id}[/dim]\n" - f" [dim]workflow [/dim]{workflow_name}\n" - f" [dim]watch [/dim][cyan]comfy jobs watch {submit.prompt_id} --where cloud[/cyan]\n" - f" [dim]state [/dim]{state_file}" - ) - if not watcher_spawned: - pprint("[yellow]⚠ Background watcher could not start; poll manually with `comfy jobs status`[/yellow]") - renderer.emit( - { - "workflow": workflow_name, - "status": "queued", - "prompt_id": submit.prompt_id, - "client_id": client_id, - "outputs": [], - "elapsed_seconds": None, + _emit_queued( + submit.prompt_id, + client_id, + workflow_name, + where="cloud", + watch_command=f"comfy jobs watch {submit.prompt_id} --where cloud", + state_file=state_file, + watcher_spawned=watcher_spawned, + extra_envelope_fields={ "base_url": target.base_url, "state_file": str(state_file) if state_file else None, }, - command="run", - where="cloud", ) - # Pretty mode only: short live tail of the state file so the human - # sees status transitions before the foreground exits. - _tail_state_file(submit.prompt_id) return # --wait: poll the cloud API directly from the foreground process. # No watcher subprocess needed — simpler, no liveness/crash concerns, # and the state file is written exactly once when the job finishes. - state = jobs_state.new( - prompt_id=submit.prompt_id, - client_id=client_id, - workflow=workflow_name, + state, state_file = _init_queued_state( + submit.prompt_id, + client_id, + workflow_name, where="cloud", + compose_meta=compose_meta, base_url=target.base_url, ) - state.item_map = (compose_meta or {}).get("items") - state_file = jobs_state.write(state) _journal_run(workflow_name, submit.prompt_id, "cloud") try: diff --git a/tests/comfy_cli/command/test_run_json.py b/tests/comfy_cli/command/test_run_json.py index 65af669e..3065bff0 100644 --- a/tests/comfy_cli/command/test_run_json.py +++ b/tests/comfy_cli/command/test_run_json.py @@ -299,6 +299,83 @@ def test_no_wait_emits_prompt_preview_then_queued(self, workflow_file, capsys): assert env["data"]["status"] == "queued" assert env["data"]["prompt_id"] == "p123" + def test_queued_envelope_field_set_local_golden(self, workflow_file, capsys): + """Golden field set of the local async ``queued`` envelope. The shared + ``_emit_queued`` helper (BE-3265) must never silently drop/add a key — + the envelope is a public agent-mode JSON contract, and local carries + ``host``/``port`` plus the local-only ``watcher_spawned``.""" + with ( + patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), + patch("comfy_cli.command.run.request.urlopen") as mock_open, + patch("comfy_cli.command.run._spawn_watcher", return_value=True), + ): + mock_open.return_value.read.return_value = json.dumps({"prompt_id": "pL"}).encode() + lines, exit_code = _run_execute_capture(workflow_file, capsys, wait=False) + assert exit_code == 0 + data = _envelope(lines)["data"] + assert set(data) == { + "workflow", + "status", + "prompt_id", + "client_id", + "outputs", + "elapsed_seconds", + "host", + "port", + "state_file", + "watcher_spawned", + } + assert data["status"] == "queued" + assert data["host"] == "127.0.0.1" + assert data["port"] == 8188 + assert data["watcher_spawned"] is True + assert data["outputs"] == [] + assert data["elapsed_seconds"] is None + + def test_queued_envelope_field_set_cloud_golden(self, workflow_file, capsys): + """Golden field set of the cloud non-wait ``queued`` envelope. Cloud + carries ``base_url`` (not ``host``/``port``) and — unlike local — + deliberately omits ``watcher_spawned`` (BE-3265: field membership is + owned per-target by the call site, not the shared helper).""" + from comfy_cli.comfy_client import SubmitResult + from comfy_cli.command.run import execute_cloud + from comfy_cli.target import Target + + target = Target( + kind="cloud", + base_url="https://cloud.example.com", + path_prefix="/api", + history_path="history_v2", + jobs_path="jobs", + api_key="test-api-key", + ) + mock_client = MagicMock() + mock_client.submit_prompt.return_value = SubmitResult(prompt_id="pC", number=1, node_errors={}) + with ( + patch("comfy_cli.target.resolve_target", return_value=target), + patch("comfy_cli.cql.engine._load_from_target", return_value={}), + patch("comfy_cli.comfy_client.Client", return_value=mock_client), + patch("comfy_cli.command.run._spawn_watcher", return_value=True), + ): + execute_cloud(workflow_file, wait=False, timeout=5) + out, _err = capsys.readouterr() + data = _envelope(_parse_lines(out))["data"] + assert set(data) == { + "workflow", + "status", + "prompt_id", + "client_id", + "outputs", + "elapsed_seconds", + "base_url", + "state_file", + } + assert data["status"] == "queued" + assert data["base_url"] == "https://cloud.example.com" + assert "watcher_spawned" not in data + assert data["outputs"] == [] + assert data["elapsed_seconds"] is None + def test_envelope_after_success(self, workflow_file, capsys): """Mocked WS flow → queued + executing/executed/output events + ok envelope.""" with (