From 6b072746a9201096fb6e2f1c32fb46c484ebdd0e Mon Sep 17 00:00:00 2001 From: Matt Miller Date: Fri, 17 Jul 2026 17:55:27 -0700 Subject: [PATCH] fix(run): close the run WebSocket on all exit paths of local --wait (BE-3404) --- comfy_cli/command/run/__init__.py | 4 ++++ tests/comfy_cli/command/test_run.py | 19 +++++++++++++++++++ 2 files changed, 23 insertions(+) diff --git a/comfy_cli/command/run/__init__.py b/comfy_cli/command/run/__init__.py index fc230d00..032c070e 100644 --- a/comfy_cli/command/run/__init__.py +++ b/comfy_cli/command/run/__init__.py @@ -443,6 +443,10 @@ def execute( finally: if progress is not None: progress.stop() + # Best-effort close of the run WebSocket on every exit path. On the + # async (no --wait) path connect() never ran so ws is None and this is + # a no-op; it is idempotent with the SIGINT-token close wired above. + _safe_close(execution) def _journal_run(workflow: str, prompt_id, where: str) -> None: diff --git a/tests/comfy_cli/command/test_run.py b/tests/comfy_cli/command/test_run.py index d1729aa1..5bc5e59f 100644 --- a/tests/comfy_cli/command/test_run.py +++ b/tests/comfy_cli/command/test_run.py @@ -413,6 +413,25 @@ def test_successful_execution(self, workflow_file): mock_exec.connect.assert_called_once() mock_exec.queue.assert_called_once() mock_exec.watch_execution.assert_called_once() + # The run WebSocket must be closed on the success path (BE-3404) — + # the finally-block _safe_close, not left open until teardown. + mock_exec.ws.close.assert_called_once() + + def test_websocket_closed_on_watch_failure(self, workflow_file): + # BE-3404: the finally-block close also fires when watch_execution + # raises, so a mid-run error doesn't linger the server-side session. + with ( + patch("comfy_cli.command.run.check_comfy_server_running", return_value=True), + patch("comfy_cli.command.run.ExecutionProgress"), + patch("comfy_cli.command.run.WorkflowExecution") as MockExec, + ): + mock_exec = MagicMock() + MockExec.return_value = mock_exec + mock_exec.watch_execution.side_effect = WebSocketTimeoutException("timed out") + + with pytest.raises(typer.Exit): + execute(workflow_file, host="127.0.0.1", port=8188, wait=True, timeout=30) + mock_exec.ws.close.assert_called_once() def test_file_not_found_exits(self): with pytest.raises(typer.Exit) as exc_info: