fix(run): close the run WebSocket on all exit paths of local --wait (BE-3404)

[mattmillerai]

ELI-5

When you run comfy run --wait against a local ComfyUI server, the CLI opens a WebSocket to watch the job. On the happy path that socket was never explicitly closed — it lingered open (holding a server-side ws session) until the Python interpreter tore down at exit. This adds a best-effort close in the function's existing finally: block so the socket is shut on every exit path — success, error, or interrupt.

What & why

On the local comfy run --wait success path, the WebSocket opened by execution.connect() was only closed via the SIGINT cancellation token (token.on_cancel(lambda: _safe_close(execution))). After watch_execution() returned and the terminal envelope was emitted, the socket stayed open until interpreter teardown — a resource-hygiene defect surfaced during BE-3354.

The fix adds _safe_close(execution) to the execute() finally: block (which already stops the Rich progress):

finally:
    if progress is not None:
        progress.stop()
    _safe_close(execution)

Why it's safe

• execution is constructed before the try:, so it is always bound when finally runs.
• _safe_close (in execution.py) already guards ws is None and swallows exceptions, so it is a no-op on the async (no---wait) path where connect() never ran, and idempotent with the SIGINT-token close (which is left untouched — it's what breaks the blocking recv() on Ctrl-C).
• Early exits (workflow-load errors, server-not-running, --print-prompt) occur before execution is constructed and before any ws is opened, so there is nothing to leak there.

Tests

• test_successful_execution now asserts mock_exec.ws.close.assert_called_once() after a successful --wait run.
• New test_websocket_closed_on_watch_failure asserts close() is also called when watch_execution raises (WebSocketTimeoutException).
• Full run-module suite green: pytest tests/ -k run → 314 passed, 2 skipped.

Verification note

The unit tests assert close() is invoked on both the success and watch-failure paths (non-vacuous: absent the fix, close() is never called on the success path since the only other wiring is the SIGINT on_cancel, which doesn't fire without a signal). The live lsof/server-log confirmation in the ticket's Verification section requires a running ComfyUI server and was not run in this unattended environment.

[coderabbitai]

Warning

Review limit reached

You’ve reached a temporary PR review limit under our Fair Usage Limits Policy.

Your recent review volume is higher than typical usage, so adaptive limits are currently applied.

Next review available in: 59 minutes

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.
You're only billed for reviews past your plan's rate limits ($0.25/file).

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

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 configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: a312e456-2cde-4fe6-a1a4-092749613eed

📥 Commits

Reviewing files that changed from the base of the PR and between a732f5c and 6b07274.

📒 Files selected for processing (2)

• comfy_cli/command/run/__init__.py
• tests/comfy_cli/command/test_run.py

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch matt/be-3404-close-run-ws

✨ Simplify code

• Create PR with simplified code
• Commit simplified code in branch matt/be-3404-close-run-ws

---

_{Comment @coderabbitai help to get the list of available commands.}

[github-actions]

🔍 Cursor Review — Consolidated panel

Triggered by @mattmillerai.

⚠️ Review failed

Judge call failed (status=parse_error): Could not parse JSON findings from output. First 500 chars:
I've verified the code structure. Let me confirm the exact placement of the `execution` assignment relative to the `try` block.

The key facts:
- Line 255: `execution = WorkflowExecution(...)` — assigned **unconditionally**, before the `try`.
- Line 269: `try:` begins.
- Lines 443–449: `finally:` block calls `_safe_close(execution)`.

The single panel finding (gemini-3.1-pro edge-case) claims `execution` could be `UnboundLocalError` in the `finally` if an exception/return happens "during early a

Re-trigger by removing and re-adding the cursor-review label.