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: