fix: comfy run --wait writes job state at submit, names prompt_id on disconnect (BE-4750) - #605
Conversation
…disconnect (BE-4750) On the foreground (--wait) path the jobs state file was written only on successful completion, so a server that died mid-run (e.g. OOM-kill) left nothing on disk and emitted a ws_disconnected error with no prompt_id — no record anywhere that the prompt had been in flight. Write the state file right after a successful queue() with status "running" (mirroring the async path's submit-time write), then update that same record on completion, cancellation, and disconnect. The disconnect path now records error.code = "server_died" and names the prompt_id + state file in the emitted error; the timeout path adds the prompt_id but deliberately leaves the record non-terminal, since the job may still be running server-side. All state writes are best-effort — a state-file failure never fails an otherwise-successful run.
|
Warning Review limit reachedYou’ve reached a temporary PR review limit under our Fair Usage Limits Policy. Next review available in: 27 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (4)
✨ Finishing Touches🧪 Generate unit tests (beta)
✨ Simplify code
Comment |
There was a problem hiding this comment.
🔍 Cursor Review — Consolidated panel
Triggered by @mattmillerai.
Found 6 finding(s).
| Severity | Count |
|---|---|
| 🟠 High | 2 |
| 🟡 Medium | 1 |
| 🟢 Low | 2 |
| ⚪ Nit | 1 |
Panel: 6/8 reviewers contributed findings.
Reviewers that did not contribute: kimi-k2.5:adversarial (empty), kimi-k2.5:edge-case (empty)
…letion out of the disconnect catch (BE-4750) Cursor review panel findings on the submit-time state file: - watch_execution() ends a failed run by rendering the error and raising typer.Exit (1 for execution_error, 130 for execution_interrupted) — the ordinary failure path, not something the disconnect/KeyboardInterrupt handlers see. The submit-time `running` record was therefore never moved to a terminal status, stranding every server-side node failure as a phantom `running` that nothing reaps (`jobs ls` only reaps non-terminal records with a dead watcher_pid, which --wait never sets). Catch the exit and finalize to error/cancelled, carrying the classified verdict that on_error now stashes on the execution. - The completion rendering/emit sat inside the try that catches (WebSocketException, ConnectionError, OSError). A BrokenPipeError from a closed stdout (`comfy run --wait … | head`) is an OSError subclass, so it rewrote the just-persisted `completed` record to error/server_died, printed a false "Lost connection", and exited 1. Build the payload inside the try, render it after. - _mark_cancelled no longer walks an already-terminal record backwards: a Ctrl-C landing after the completion write must not turn `completed` into `cancelled`. - Report `state_file` only when the terminal write actually landed; the submit-time fallback pointed at a file whose contents still said `running`, contradicting the reported status. Drops the now-dead wait_state_file assignments with it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
ELI-5
When you run
comfy run --workflow x --wait, the CLI sits there and watches the job until it finishes. Until now it only wrote the little "here's what happened to this job" file at the very end, on success. So if the ComfyUI server died halfway through (say it got OOM-killed by a huge allocation), you got an error that said "lost connection" and nothing else — no job id, no file on disk. Nothing anywhere remembered that a prompt was in flight when the server went down.Now the CLI writes that file the moment the server accepts the job, and keeps it updated. If the server dies, the file says so and the error tells you which prompt it was, so
comfy jobs status <id>has something to read.What changed
All in
comfy_cli/command/run/__init__.py, localexecutepath only (execute_cloudalready wrote state at submit for both modes and is untouched):--wait. Right after a successfulqueue(), the state is created and written withstatus = "running"— mirroring the async branch's submit-time write. "running" rather than "queued" because this foreground process is actively watching it, not leaving it detached.stateobject (setsstatus = "completed"+outputs), so it is the same file with the same final shape as before.KeyboardInterrupthandler, and the token-cancelled branch inside the disconnect handler all setstatus = "cancelled"+error.code = "cancelled". Guarded so a cancel beforequeue()returned (no state yet) simply skips the write.ws_disconnectedrecordsserver_diedand names the prompt. When a prompt_id exists, the state file goes tostatus = "error"witherror.code = "server_died", and the emitted error keepscode = "ws_disconnected"but names the prompt in the message and carriesdetails = {"prompt_id": ..., "state_file": ...}. A disconnect beforequeue()returned emits today's bare error unchanged.ws_timeoutnames the prompt but stays non-terminal.detailsgainsprompt_id; the state file is deliberately left at its submit-timerunningrecord, because a timed-out watch says nothing about the job — it may genuinely still be running server-side.Every state write goes through a
_write_statehelper that swallows write failures, so a state-dir problem can never fail an otherwise-successful run (matching the async path's tolerance of a watcher that won't spawn).Tests
Seven new tests in
tests/comfy_cli/command/test_run.py(TestWaitStateFile), following the existingWorkflowExecutionmock pattern and the autouse_isolate_jobs_state_dirfixture:watch_executionside effect reads the file mid-run and assertsstatus == "running"with nocompleted_at, then the post-run file iscompletedwith outputs and the samesubmitted_at(proving it is an update, not a fresh record)watch_executionraisesConnectionError— emitted error isws_disconnectedwithdetails.prompt_idand astate_filepath; the file isstatus == "error",error.code == "server_died"queue()raises) — no state file is created, no crash, and the emitted error keeps today's bare shape (details is None)details == {"timeout": 30, "prompt_id": ...}and the file staysrunning/ non-terminalKeyboardInterrupt— both leavestatus == "cancelled"OSError,ValueError) proving a broken state file never fails the runruff check+ruff format --diffclean on the CI-pinned 0.15.15; fullpytestgreen (2777 passed, 37 skipped).Judgment calls
comfy_cli/command/run/__init__.py+ tests, but the repo's owntest_every_raised_code_is_registeredscans{"code": "..."}dict literals in source (explicitly including "the watcher and state-file paths") and fails on any code missing fromcomfy_cli/error_codes.REGISTRY. Soserver_diedneeded a five-line registry entry. That is the repo enforcing its own invariant, not scope creep, but flagging it since it is a deviation from the ticket text._write_stateswallowsValueErroras well asOSError. The ticket asked fortry/except OSError.jobs_state.state_pathraisesValueErroron a prompt_id it can't turn into a safe filename, and moving the write from completion to submit would newly let that abort the run beforewatch_execution()— i.e. we'd stop watching a job the server had already accepted. Swallowing it keeps the submit-time write strictly non-regressive. Covered by the parametrized test.state_filefalls back to the submit-time path (_write_state(state) or wait_state_file) if the completion write fails, since that file does exist on disk. It would then readrunningwhile the envelope sayscompleted— the accepted trade of best-effort state writes.runningrather than marking it terminal. That is deliberate: we don't know the job died, and it is still strictly more information than the previous behavior of no file at all.ws_disconnectederror already existed; this change only adds information to it (prompt_id, state file) and creates an on-disk record where there was none, so the user-facing outcome strictly expands what is recoverable after a mid-run server death.