diff --git a/comfy_cli/command/transfer.py b/comfy_cli/command/transfer.py index e32ef7af..d62a5feb 100644 --- a/comfy_cli/command/transfer.py +++ b/comfy_cli/command/transfer.py @@ -9,6 +9,7 @@ from __future__ import annotations +import http.client import json import mimetypes import re @@ -324,6 +325,22 @@ def execute_upload( details={"status": status, "filename": filename}, ) raise typer.Exit(code=1) + except (urllib.error.URLError, TimeoutError, ConnectionError, http.client.HTTPException) as e: + # A connection- or transfer-level failure — not HTTPError. A + # refused/DNS/timeout/TLS failure at connect raises URLError; a read + # timeout raises a bare TimeoutError; a reset raises ConnectionError; + # a truncated (e.g. chunked) response body raises + # http.client.IncompleteRead (an HTTPException). Surface it as a + # structured envelope instead of an unhandled traceback that breaks + # machine/NDJSON consumers. + reason = getattr(e, "reason", None) or e + renderer.error( + code="upload_failed", + message=f"Failed to upload {filename}: {reason}", + hint="check that the server is reachable", + details={"filename": filename, "reason": str(reason)}, + ) + raise typer.Exit(code=1) cloud_name = result.get("name", filename) subfolder = result.get("subfolder", "") @@ -630,6 +647,31 @@ def execute_download( details={"status": e.code, "url": url, "index": idx}, ) raise typer.Exit(code=1) + except (urllib.error.URLError, TimeoutError, ConnectionError, http.client.HTTPException) as e: + # A connection- or transfer-level failure — not HTTPError. A + # refused/DNS/timeout/TLS failure at connect raises URLError; a + # read timeout raises a bare TimeoutError; a reset raises + # ConnectionError; a truncated (e.g. chunked) response body raises + # http.client.IncompleteRead (an HTTPException). Surface it as a + # structured envelope instead of an unhandled traceback that breaks + # machine/NDJSON consumers. + reason = getattr(e, "reason", None) or e + # A mid-transfer failure (reset / read-timeout) may have already + # created a partial file at local_path via open(..., "wb") — remove + # it so a failed download never leaves a truncated artifact in the + # out-dir. missing_ok covers the connect-time case (no file yet); + # the OSError guard keeps cleanup from re-raising on this error path. + try: + local_path.unlink(missing_ok=True) + except OSError: + pass + renderer.error( + code="download_failed", + message=f"Failed to download output {idx}: {reason}", + hint="check that the server is reachable and the output URL is correct", + details={"url": url, "index": idx, "reason": str(reason)}, + ) + raise typer.Exit(code=1) file_size = local_path.stat().st_size entry: dict[str, Any] = { diff --git a/tests/comfy_cli/command/test_transfer_download.py b/tests/comfy_cli/command/test_transfer_download.py index f5b24a43..16668edf 100644 --- a/tests/comfy_cli/command/test_transfer_download.py +++ b/tests/comfy_cli/command/test_transfer_download.py @@ -9,11 +9,14 @@ from __future__ import annotations +import http.client import json +import urllib.error from pathlib import Path from unittest.mock import patch import pytest +import typer from comfy_cli import comfy_client, jobs_state from comfy_cli.comfy_client import extract_output_entries @@ -588,6 +591,104 @@ def test_loopback_file_uris_are_local(self, url): assert transfer._local_source_path(url) == Path("/abs/path.png") +class TestDownloadConnectionError: + """A connection-level failure — server unreachable, DNS failure, read + timeout, TLS error — raises the parent ``URLError`` (or a bare + ``TimeoutError`` on a read timeout), not ``HTTPError``. It must surface as a + structured ``download_failed`` envelope with stdout staying pure JSON, never + an unhandled traceback that breaks ``--json``/NDJSON consumers (BE-2454).""" + + def _run(self, fake_target, tmp_path, side_effect) -> typer.Exit: + with ( + patch("comfy_cli.command.transfer.resolve_target", return_value=fake_target), + patch.object(transfer._DOWNLOAD_OPENER, "open", side_effect=side_effect), + ): + with pytest.raises(typer.Exit) as excinfo: + transfer.execute_download(PROMPT_ID, out_dir=str(tmp_path / "out")) + return excinfo.value + + def test_urlerror_emits_download_failed_envelope(self, fake_target, tmp_path, capsys): + _write_state(fake_target) + set_renderer(Renderer(mode=OutputMode.JSON, command="download")) + + def _refused(req): + raise urllib.error.URLError(ConnectionRefusedError(111, "Connection refused")) + + exc = self._run(fake_target, tmp_path, _refused) + + assert exc.exit_code == 1 + out_lines = [ln for ln in capsys.readouterr().out.splitlines() if ln.strip()] + env = json.loads(out_lines[-1]) + assert env["ok"] is False + assert env["error"]["code"] == "download_failed" + assert "Connection refused" in env["error"]["message"] + assert "Connection refused" in env["error"]["details"]["reason"] + + def test_bare_timeout_is_also_caught(self, fake_target, tmp_path, capsys): + _write_state(fake_target) + set_renderer(Renderer(mode=OutputMode.JSON, command="download")) + + def _timeout(req): + raise TimeoutError("timed out") + + exc = self._run(fake_target, tmp_path, _timeout) + + assert exc.exit_code == 1 + env = json.loads([ln for ln in capsys.readouterr().out.splitlines() if ln.strip()][-1]) + assert env["error"]["code"] == "download_failed" + + def test_connection_reset_mid_stream_is_caught(self, fake_target, tmp_path, capsys): + _write_state(fake_target) + set_renderer(Renderer(mode=OutputMode.JSON, command="download")) + + class _ResetResp: + def __init__(self): + self._n = 0 + + def read(self, n): + self._n += 1 + if self._n == 1: + return b"partial-bytes" + raise ConnectionResetError(104, "Connection reset by peer") + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + exc = self._run(fake_target, tmp_path, lambda req: _ResetResp()) + + assert exc.exit_code == 1 + env = json.loads([ln for ln in capsys.readouterr().out.splitlines() if ln.strip()][-1]) + assert env["error"]["code"] == "download_failed" + assert "reset by peer" in env["error"]["message"] + assert list((tmp_path / "out").glob("*")) == [] + + def test_incomplete_read_is_caught(self, fake_target, tmp_path, capsys): + # A truncated (e.g. chunked) body makes resp.read() raise + # http.client.IncompleteRead — an HTTPException, not a URLError. + _write_state(fake_target) + set_renderer(Renderer(mode=OutputMode.JSON, command="download")) + + class _TruncResp: + def read(self, n): + raise http.client.IncompleteRead(b"partial", 100) + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + exc = self._run(fake_target, tmp_path, lambda req: _TruncResp()) + + assert exc.exit_code == 1 + env = json.loads([ln for ln in capsys.readouterr().out.splitlines() if ln.strip()][-1]) + assert env["error"]["code"] == "download_failed" + assert list((tmp_path / "out").glob("*")) == [] + + # --------------------------------------------------------------------------- # _default_out_dir — project/1 root wins over the legacy config key # --------------------------------------------------------------------------- diff --git a/tests/comfy_cli/command/test_transfer_upload.py b/tests/comfy_cli/command/test_transfer_upload.py index 3852dfa1..3ec93d2b 100644 --- a/tests/comfy_cli/command/test_transfer_upload.py +++ b/tests/comfy_cli/command/test_transfer_upload.py @@ -8,12 +8,14 @@ from __future__ import annotations +import http.client import io import json import urllib.error from pathlib import Path import pytest +import typer from comfy_cli.command import transfer from comfy_cli.target import Target @@ -145,3 +147,53 @@ def test_json_mode_stdout_is_pure_json_no_human_line(self, asset, monkeypatch, c assert parsed[-1]["data"]["uploads"][0]["cloud_name"] == "ab12.png" assert "uploaded" not in captured.out assert "uploaded" not in captured.err + + +class TestUploadConnectionError: + """A connection-level failure (``URLError``/``TimeoutError``) on upload must + surface as a structured ``upload_failed`` envelope, not an unhandled + traceback that breaks ``--json``/NDJSON consumers (BE-2454).""" + + @pytest.fixture(autouse=True) + def reset_renderer(self): + from comfy_cli.output.renderer import reset_renderer_for_testing + + reset_renderer_for_testing() + yield + reset_renderer_for_testing() + + def test_urlerror_emits_upload_failed_envelope(self, asset, monkeypatch, capsys): + from comfy_cli.output import Renderer, set_renderer + from comfy_cli.output.renderer import OutputMode + + err = urllib.error.URLError(ConnectionRefusedError(111, "Connection refused")) + monkeypatch.setattr(transfer, "_TRANSFER_OPENER", _FakeOpener(error=err)) + monkeypatch.setattr(transfer, "resolve_target", lambda where=None: _local_target()) + set_renderer(Renderer(mode=OutputMode.JSON, command="upload")) + + with pytest.raises(typer.Exit) as excinfo: + transfer.execute_upload([str(asset)], where="local") + + assert excinfo.value.exit_code == 1 + env = json.loads([ln for ln in capsys.readouterr().out.splitlines() if ln.strip()][-1]) + assert env["ok"] is False + assert env["error"]["code"] == "upload_failed" + assert "Connection refused" in env["error"]["message"] + assert "Connection refused" in env["error"]["details"]["reason"] + + def test_incomplete_read_emits_upload_failed_envelope(self, asset, monkeypatch, capsys): + from comfy_cli.output import Renderer, set_renderer + from comfy_cli.output.renderer import OutputMode + + # A truncated response body raises http.client.IncompleteRead — an + # HTTPException, not a URLError. + monkeypatch.setattr(transfer, "_TRANSFER_OPENER", _FakeOpener(error=http.client.IncompleteRead(b"x", 100))) + monkeypatch.setattr(transfer, "resolve_target", lambda where=None: _local_target()) + set_renderer(Renderer(mode=OutputMode.JSON, command="upload")) + + with pytest.raises(typer.Exit) as excinfo: + transfer.execute_upload([str(asset)], where="local") + + assert excinfo.value.exit_code == 1 + env = json.loads([ln for ln in capsys.readouterr().out.splitlines() if ln.strip()][-1]) + assert env["error"]["code"] == "upload_failed"