Skip to content

refactor(run): extract shared queued tail + state-init + load helpers (BE-3265) #526

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mattmillerai wants to merge 1 commit into
base: main
Choose a base branch
from
+482 −90
Open

refactor(run): extract shared queued tail + state-init + load helpers (BE-3265) #526

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
208 changes: 120 additions & 88 deletions comfy_cli/command/run/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -124,14 +124,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(
Expand Down Expand Up @@ -346,52 +339,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()
Expand Down Expand Up @@ -459,6 +431,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
Expand Down Expand Up @@ -507,14 +564,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
Expand Down Expand Up @@ -651,60 +701,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:
Expand Down
77 changes: 77 additions & 0 deletions tests/comfy_cli/command/test_run_json.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down
Loading
Loading