fix: comfy run --wait writes job state at submit, names prompt_id on disconnect (BE-4750)

comfy_cli/command/run/__init__.py

@@ -268,6 +268,19 @@ def execute(
     token = cancellation.get_token()
     token.on_cancel(lambda: _safe_close(execution))
 
+    # --wait only: the state written right after a successful submit, kept in
+    # scope so the exception handlers below can record what happened to a
+    # prompt that was already in flight (e.g. the server dying mid-run). Stays
+    # None on the async path — that branch writes and owns its own state — and
+    # before ``queue()`` has returned a prompt_id.
+    wait_state: jobs_state.JobState | None = None
+    # --wait success payload, rendered/emitted AFTER the try below rather than
+    # inside it. Writing to a closed stdout (`comfy run --wait … | head`)
+    # raises BrokenPipeError — a ConnectionError/OSError subclass the
+    # disconnect handler would otherwise catch, rewriting the just-persisted
+    # `completed` record to `error` and flipping a successful run's exit to 1.
+    completed_payload: dict | None = None
+
     try:
         if wait:
             execution.connect()
@@ -280,43 +293,59 @@ def execute(
             execution.queue()
         _journal_run(workflow_name, execution.prompt_id, "local")
         if wait:
-            execution.watch_execution()
+            # Write the state file at SUBMIT time, exactly like the async
+            # branch below — if the server dies mid-run the on-disk record is
+            # the only place the in-flight prompt_id survives. Status is
+            # "running" rather than "queued": this foreground process is
+            # actively watching it, not leaving it detached in the queue.
+            wait_state = jobs_state.new(
+                prompt_id=execution.prompt_id,
+                client_id=execution.client_id,
+                workflow=workflow_name,
+                where="local",
+                host=host,
+                port=port,
+            )
+            wait_state.item_map = (compose_meta or {}).get("items")
+            wait_state.status = "running"
+            _write_state(wait_state)
+
+            # `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
+            # exception the handlers below see. Finalize the submit-time
+            # record here, or a failed job is stranded as a phantom `running`
+            # forever: `jobs ls` only reaps non-terminal records whose
+            # `watcher_pid` is dead, and `--wait` never sets one.
+            try:
+                execution.watch_execution()
+            except typer.Exit as exit_exc:
+                _mark_watch_exit(wait_state, exit_exc.exit_code, execution)
+                raise
             end = time.time()
             if progress is not None:
                 progress.stop()
                 progress = None
 
             if token.is_set():
+                _mark_cancelled(wait_state)
                 renderer.error(
                     code="cancelled",
                     message="Cancelled by user",
                     exit_code=130,
                 )
                 raise typer.Exit(code=130)
 
-            # Foreground (--wait) completion path also writes the state
-            # file so the on-disk record is consistent regardless of which
-            # mode the user ran in.
-            state = jobs_state.new(
-                prompt_id=execution.prompt_id,
-                client_id=execution.client_id,
-                workflow=workflow_name,
-                where="local",
-                host=host,
-                port=port,
-            )
+            # Completion updates the record written at submit — same file,
+            # same final shape (``jobs_state.write`` stamps the timestamps).
+            state = wait_state
             state.status = "completed"
             state.outputs = list(execution.outputs)
-            state.item_map = (compose_meta or {}).get("items")
-            state_file = jobs_state.write(state)
-
-            if renderer.is_pretty():
-                if len(execution.outputs) > 0:
-                    pprint("[bold green]\nOutputs:[/bold green]")
-                    for f in execution.outputs:
-                        pprint(f)
-                elapsed = timedelta(seconds=end - start)
-                pprint(f"[bold green]\nWorkflow execution completed ({elapsed})[/bold green]")
+            # No fallback to the submit-time path: if the terminal write
+            # failed, that file still says `running`, so handing it back as
+            # the `completed` record's `state_file` would point the caller at
+            # contents contradicting what we just reported.
+            state_file = _write_state(state)
 
             # Grouped views of the same artifacts — local parity with the
             # cloud --wait envelope: by producing node always, and by
@@ -325,25 +354,21 @@ def execute(
 
             outputs_by_node, outputs_by_item = _group_outputs(list(execution.output_entries), state.item_map)
 
-            renderer.emit(
-                {
-                    "workflow": workflow_name,
-                    "status": "completed",
-                    "prompt_id": execution.prompt_id,
-                    "client_id": execution.client_id,
-                    "outputs": list(execution.outputs),
-                    "outputs_by_node": outputs_by_node,
-                    "outputs_by_item": outputs_by_item,
-                    "cached_node_ids": list(execution.cached_node_ids),
-                    "executed_node_ids": list(execution.executed_node_ids),
-                    "elapsed_seconds": end - start,
-                    "host": host,
-                    "port": port,
-                    "state_file": str(state_file) if state_file else None,
-                },
-                command="run",
-                where="local",
-            )
+            completed_payload = {
+                "workflow": workflow_name,
+                "status": "completed",
+                "prompt_id": execution.prompt_id,
+                "client_id": execution.client_id,
+                "outputs": list(execution.outputs),
+                "outputs_by_node": outputs_by_node,
+                "outputs_by_item": outputs_by_item,
+                "cached_node_ids": list(execution.cached_node_ids),
+                "executed_node_ids": list(execution.executed_node_ids),
+                "elapsed_seconds": end - start,
+                "host": host,
+                "port": port,
+                "state_file": str(state_file) if state_file else None,
+            }
         else:
             # Async path (the default). Write the initial state file and
             # spawn a detached watcher to keep it updated; the foreground
@@ -398,6 +423,7 @@ def execute(
         if progress is not None:
             progress.stop()
             progress = None
+        _mark_cancelled(wait_state)
         if renderer.is_pretty():
             pprint("[yellow]Workflow execution was interrupted[/yellow]")
         renderer.error(
@@ -412,11 +438,17 @@ def execute(
                 f"[bold red]Error: WebSocket timed out after {timeout}s waiting for server response.[/bold red]\n"
                 "[yellow]For long-running workflows, increase the timeout: comfy run --workflow <file> --timeout 300[/yellow]"
             )
+        # The job may genuinely still be running server-side, so the
+        # submit-time "running" record is left as-is — not marked terminal.
+        details = {"timeout": timeout}
+        prompt_id = _submitted_prompt_id(execution)
+        if prompt_id is not None:
+            details["prompt_id"] = prompt_id
         renderer.error(
             code="ws_timeout",
             message=f"WebSocket timed out after {timeout}s waiting for server response.",
             hint="re-run with a larger --timeout (e.g. --timeout 300)",
-            details={"timeout": timeout},
+            details=details,
         )
         raise typer.Exit(code=1)
     except (WebSocketException, ConnectionError, OSError) as e:
@@ -428,6 +460,7 @@ def execute(
         if token.is_set():
             if progress is not None:
                 progress.stop()
+            _mark_cancelled(wait_state)
             renderer.error(
                 code="cancelled",
                 message="Cancelled by user",
@@ -436,6 +469,27 @@ def execute(
             raise typer.Exit(code=130) from e
         if renderer.is_pretty():
             pprint(f"[bold red]Error: Lost connection to ComfyUI server: {e}[/bold red]")
+        # The server died with a prompt of ours in flight: record that on the
+        # submit-time state file and name the prompt in the emitted error, so
+        # `comfy jobs status <id>` has something to consult afterwards.
+        prompt_id = _submitted_prompt_id(execution)
+        if prompt_id is not None and wait_state is not None:
+            wait_state.status = "error"
+            wait_state.error = {
+                "code": "server_died",
+                "message": f"Lost connection to ComfyUI while job {prompt_id} was running: {e}",
+                "details": {},
+            }
+            # Same rule as the success path: report a state_file only when
+            # this terminal write actually landed.
+            path = _write_state(wait_state)
+            renderer.error(
+                code="ws_disconnected",
+                message=f"Lost connection to ComfyUI server while job {prompt_id} was running: {e}",
+                hint="check the server is still running; re-run the command",
+                details={"prompt_id": prompt_id, "state_file": str(path) if path else None},
+            )
+            raise typer.Exit(code=1)
         renderer.error(
             code="ws_disconnected",
             message=f"Lost connection to ComfyUI server: {e}",
@@ -450,6 +504,93 @@ def execute(
         # a no-op; it is idempotent with the SIGINT-token close wired above.
         _safe_close(execution)
 
+    # Deliberately outside the try: the job is done and persisted, so a
+    # BrokenPipeError from a closed stdout must surface as itself rather than
+    # be mistaken for the server disconnecting (see `completed_payload`).
+    if completed_payload is not None:
+        if renderer.is_pretty():
+            if completed_payload["outputs"]:
+                pprint("[bold green]\nOutputs:[/bold green]")
+                for f in completed_payload["outputs"]:
+                    pprint(f)
+            elapsed = timedelta(seconds=completed_payload["elapsed_seconds"])
+            pprint(f"[bold green]\nWorkflow execution completed ({elapsed})[/bold green]")
+        renderer.emit(completed_payload, command="run", where="local")
+
+
+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
+    unusable as a filename.
+
+    A state-file failure must never fail an otherwise-successful run — same
+    tolerance the async path gives a watcher that won't spawn. Swallowing
+    ``ValueError`` matters now that ``--wait`` writes at SUBMIT: a server
+    returning an id ``state_path()`` rejects must not stop us from watching
+    the run it just accepted.
+    """
+    try:
+        return jobs_state.write(state)
+    except (OSError, ValueError):
+        return None
+
+
+def _mark_cancelled(state):
+    """Record a user cancellation on the job state file (best effort).
+
+    Returns the path written, or None when there is nothing to write — a
+    Ctrl-C before ``queue()`` returned leaves ``state`` None, so there is no
+    prompt_id to file it under, and a record that is ALREADY terminal is left
+    alone. That last guard matters because the Ctrl-C handler is shared by the
+    whole run: a Ctrl-C after the completion write has landed must not walk a
+    ``completed`` job backwards to ``cancelled``.
+    """
+    if state is None or state.is_terminal:
+        return None
+    state.status = "cancelled"
+    state.error = {"code": "cancelled", "message": "Cancelled by user", "details": {}}
+    return _write_state(state)
+
+
+def _mark_watch_exit(state, exit_code, execution):
+    """Finalize the submit-time record when ``watch_execution`` ends the run by
+    raising ``typer.Exit``: 130 is a server-side interrupt (``cancelled``),
+    anything else is a node failure (``error``).
+
+    The error envelope has already been rendered by the watcher, so this only
+    moves the on-disk record off ``running``. Prefers the classified verdict
+    ``on_error`` stashed on the execution; falls back to a generic
+    ``execution_error`` when there isn't one (or it isn't a real dict — mocked
+    executions hand back attribute stubs).
+    """
+    if state is None:
+        return None
+    if exit_code == 130:
+        return _mark_cancelled(state)
+    verdict = getattr(execution, "last_error", None)
+    if not isinstance(verdict, dict):
+        verdict = {
+            "code": "execution_error",
+            "message": "Workflow execution failed on the server",
+            "details": {},
+        }
+    state.status = "error"
+    state.error = verdict
+    return _write_state(state)
+
+
+def _submitted_prompt_id(execution) -> str | None:
+    """The prompt id if the server really returned one, else None.
+
+    Mirrors ``jobs_state.write``'s defensiveness about mocked executions: a
+    non-string id means no usable submit happened, so callers fall back to
+    their pre-prompt_id behavior.
+    """
+    prompt_id = getattr(execution, "prompt_id", None)
+    if isinstance(prompt_id, str) and prompt_id.strip():
+        return prompt_id
+    return None
+
 
 def _journal_run(workflow: str, prompt_id, where: str) -> None:
     """Append the run-submit event to the governing project's run journal

comfy_cli/command/run/execution.py

@@ -120,6 +120,11 @@ def __init__(
         # compute nodes that never fire a server-side `executed` event.
         self.cached_node_ids: list[str] = []
         self.executed_node_ids: list[str] = []
+        # 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
+        # the job state file.
+        self.last_error: dict | None = None
 
     def connect(self):
         # Resolve via the package namespace so tests can patch
@@ -541,6 +546,11 @@ def on_error(self, data):
         # the error envelope carries the classified one-line verdict.
         self.renderer.event("execution_error", prompt_id=self.prompt_id, details=data)
         verdict = execution_errors.classify(data)
+        self.last_error = {
+            "code": verdict["code"],
+            "message": verdict["message"],
+            "details": dict(verdict["details"]),
+        }
         self.renderer.error(
             code=verdict["code"],
             message=verdict["message"],

comfy_cli/error_codes.py

@@ -329,6 +329,11 @@ class ErrorCode:
         "Background watcher gave up after max runtime without a terminal status.",
         "the job may still be running — check `comfy jobs status <id>`, or re-watch with a longer `--timeout`",
     ),
+    ErrorCode(
+        "server_died",
+        "Server connection dropped while a foreground (`--wait`) job was in flight; recorded on the job state file.",
+        "check the server (it may have been OOM-killed); the prompt_id is in `comfy jobs status <id>`",
+    ),
     ErrorCode(
         "watcher_poll_error",
         "Background watcher encountered a transient error polling the server.",