fix(run): bail out of run --wait on a wall clock so system sleep can't hang it

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,

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 <file> --timeout 300[/yellow]"
+                "[yellow]For long-running workflows, increase the timeout: comfy run --workflow <file> --timeout 300[/yellow]\n"
+                "[yellow]If the machine slept mid-run, the job may still be running — check "
+                "`comfy jobs status <id>`. 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 <id>`",
             details=details,
         )
         raise typer.Exit(code=1)

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:

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 <id>`, and use `caffeinate` for long local batches",
     ),
     ErrorCode(
         "prompt_rejected",

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