From 6d154289054cba2ed0beca0de3e1781c03ccbdba Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Wed, 16 Sep 2026 04:50:37 +0000 Subject: [PATCH] fix(run): bail out of `run --wait` on a wall clock so system sleep can't hang it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The local `--wait` watch loop bounded `ws.recv()` with a socket timeout enforced against a monotonic clock. That clock stops advancing while the machine sleeps (a laptop lid closed mid-run), so on wake the connection is often dead yet the timeout has under-counted the sleep and never fires — the CLI hangs indefinitely instead of reporting a timeout. Enforce the `--timeout` silence budget against the wall clock (`time.time()`), which does advance across sleep, and cap each `recv` at a short poll interval so the loop regains control regularly regardless of how large `--timeout` is. This mirrors the resilient waiters already in the codebase (the cloud `wait_for_completion` and `jobs watch`). On timeout the server-side job is still resumable, so the `ws_timeout` guidance now points users at `comfy jobs status ` and recommends `caffeinate` for long local batches. --- comfy_cli/cmdline.py | 10 ++++--- comfy_cli/command/run/__init__.py | 6 ++-- comfy_cli/command/run/execution.py | 39 ++++++++++++++++++++++++-- comfy_cli/error_codes.py | 5 ++-- tests/comfy_cli/command/test_run.py | 43 +++++++++++++++++++++++++++++ 5 files changed, 93 insertions(+), 10 deletions(-) diff --git a/comfy_cli/cmdline.py b/comfy_cli/cmdline.py index ef34b0df1..0ff399acd 100644 --- a/comfy_cli/cmdline.py +++ b/comfy_cli/cmdline.py @@ -1009,10 +1009,12 @@ def run( typer.Option( help=( "Per-event timeout in seconds: bails out if the server is silent " - "for this long. Also caps HTTP connect, /prompt POST, and websocket " - "handshake. NOT a wall-clock execution deadline — a workflow that " - "streams progress events faster than the timeout can run " - "indefinitely." + "for this long (measured against the wall clock, so it still fires " + "after the machine wakes from sleep). Also caps HTTP connect, " + "/prompt POST, and websocket handshake. NOT a wall-clock execution " + "deadline — a workflow that streams progress events faster than the " + "timeout can run indefinitely. For long local batches, prevent sleep " + "with `caffeinate` (macOS)." ), ), ] = 120, diff --git a/comfy_cli/command/run/__init__.py b/comfy_cli/command/run/__init__.py index 0c73a4817..ec70dd82d 100644 --- a/comfy_cli/command/run/__init__.py +++ b/comfy_cli/command/run/__init__.py @@ -668,7 +668,9 @@ def execute( if renderer.is_pretty(): pprint( 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 --timeout 300[/yellow]" + "[yellow]For long-running workflows, increase the timeout: comfy run --workflow --timeout 300[/yellow]\n" + "[yellow]If the machine slept mid-run, the job may still be running — check " + "`comfy jobs status `. Keep long local batches awake with `caffeinate`.[/yellow]" ) details = {"timeout": timeout} prompt_id = _submitted_prompt_id(execution) @@ -677,7 +679,7 @@ def execute( 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)", + hint="re-run with a larger --timeout (e.g. --timeout 300); if the machine slept, check `comfy jobs status `", details=details, ) raise typer.Exit(code=1) diff --git a/comfy_cli/command/run/execution.py b/comfy_cli/command/run/execution.py index 3b0273ece..386ab4843 100644 --- a/comfy_cli/command/run/execution.py +++ b/comfy_cli/command/run/execution.py @@ -27,6 +27,7 @@ import json import os +import time import urllib.error import urllib.parse import uuid @@ -35,6 +36,7 @@ import typer from rich.progress import BarColumn, Progress, TimeElapsedColumn from rich.table import Column, Table +from websocket import WebSocketTimeoutException from comfy_cli import execution_errors from comfy_cli.caller import usage_source @@ -47,6 +49,13 @@ workspace_manager = WorkspaceManager() +# Upper bound on how long a single ``recv`` blocks before ``watch_execution`` +# regains control and re-checks the wall clock. Deliberately independent of the +# user's ``--timeout`` (which can be large for long jobs) so that even a big +# silence budget can never leave us blocked inside one ``recv`` across a system +# sleep — see the wall-clock backstop in ``watch_execution``. +_RECV_POLL_SECONDS = 30 + def _safe_close(execution: WorkflowExecution) -> None: """Best-effort WebSocket close on cancellation.""" @@ -307,9 +316,35 @@ def queue(self): def watch_execution(self): if self.ws is None: raise RuntimeError("watch_execution called before the websocket was connected") - self.ws.settimeout(self.timeout) + # ``recv`` blocks on a socket timeout enforced against a MONOTONIC clock, + # which stops advancing while the machine is asleep (a laptop lid closed + # mid-run). On wake the connection is frequently dead, yet that timeout + # has under-counted the sleep, so it never fires and the loop hangs + # indefinitely instead of bailing out. Mirror the cloud waiter + # (``ComfyClient.wait_for_completion``): bound each ``recv`` to a short + # poll so the loop regains control regularly, and enforce the + # ``--timeout`` silence budget against the WALL clock, which DID advance + # across the sleep. A wake with a dead connection then aborts promptly + # (the caller reports ``ws_timeout`` — the server job is resumable via + # ``comfy jobs status``) rather than stalling for as long as the machine + # slept. + poll = min(self.timeout, _RECV_POLL_SECONDS) if self.timeout else self.timeout + self.ws.settimeout(poll) + last_activity = time.time() while True: - message = self.ws.recv() + try: + message = self.ws.recv() + except WebSocketTimeoutException: + # No frame this poll interval. Give up only once the silence + # budget has elapsed in REAL time — a monotonic timer frozen by + # a sleep can no longer keep us waiting past it. + if time.time() - last_activity >= self.timeout: + raise + continue + # Any frame — even a non-text control/binary frame — proves the + # connection is live, so it resets the silence budget exactly as the + # per-``recv`` socket timeout used to. + last_activity = time.time() if not isinstance(message, str): continue try: diff --git a/comfy_cli/error_codes.py b/comfy_cli/error_codes.py index 0756b69be..26db52194 100644 --- a/comfy_cli/error_codes.py +++ b/comfy_cli/error_codes.py @@ -165,8 +165,9 @@ class ErrorCode: ), ErrorCode( "ws_timeout", - "WebSocket idle past `--timeout` while waiting for the server.", - "re-run with a larger `--timeout` (e.g. `--timeout 300`)", + "WebSocket idle past `--timeout` (wall-clock) while waiting for the server.", + "re-run with a larger `--timeout` (e.g. `--timeout 300`); if the machine slept mid-run, " + "the job may still be running — check `comfy jobs status `, and use `caffeinate` for long local batches", ), ErrorCode( "prompt_rejected", diff --git a/tests/comfy_cli/command/test_run.py b/tests/comfy_cli/command/test_run.py index 2006a9ec9..12e250324 100644 --- a/tests/comfy_cli/command/test_run.py +++ b/tests/comfy_cli/command/test_run.py @@ -316,6 +316,49 @@ def test_successful_execution(self, mock_execution): mock_execution.watch_execution() assert len(mock_execution.remaining_nodes) == 0 + def test_wall_clock_backstop_aborts_after_sleep(self, mock_execution): + """A dead connection whose socket timeout under-counted a system sleep + must still abort: once the WALL clock is past the silence budget, a + recv timeout re-raises instead of hanging.""" + mock_execution.prompt_id = "p" + mock_ws = MagicMock() + # recv only ever times out (server silent / connection dead on wake). + mock_ws.recv.side_effect = WebSocketTimeoutException("timed out") + mock_execution.ws = mock_ws + + # last_activity=1000, then a wake with the wall clock jumped ~999s. + times = iter([1000.0, 1999.0]) + with patch("comfy_cli.command.run.execution.time.time", lambda: next(times)): + with pytest.raises(WebSocketTimeoutException): + mock_execution.watch_execution() + + def test_recv_timeout_within_budget_keeps_waiting(self, mock_execution): + """A poll-interval timeout BEFORE the wall-clock budget elapses must not + abort — the loop keeps waiting for the server.""" + mock_execution.prompt_id = "p" + mock_ws = MagicMock() + mock_ws.recv.side_effect = [ + WebSocketTimeoutException("poll tick"), # 5s in — under the 30s budget + _make_msg("executing", "p", node=None), # then the server finishes + ] + mock_execution.ws = mock_ws + + times = iter([1000.0, 1005.0, 1005.0]) + with patch("comfy_cli.command.run.execution.time.time", lambda: next(times)): + mock_execution.watch_execution() # returns normally, no raise + + def test_recv_poll_interval_is_capped(self, mock_execution): + """The per-recv socket timeout is capped so the loop regains control + regularly even when --timeout is large.""" + mock_execution.timeout = 3600 + mock_execution.prompt_id = "p" + mock_ws = MagicMock() + mock_ws.recv.side_effect = [_make_msg("executing", "p", node=None)] + mock_execution.ws = mock_ws + + mock_execution.watch_execution() + mock_ws.settimeout.assert_called_once_with(30) + def test_skips_other_prompt_messages(self, mock_execution): prompt_id = "my-prompt" mock_execution.prompt_id = prompt_id