diff --git a/comfy_cli/command/models/models.py b/comfy_cli/command/models/models.py index a654a18e..f79bae7a 100644 --- a/comfy_cli/command/models/models.py +++ b/comfy_cli/command/models/models.py @@ -1,9 +1,11 @@ import contextlib +import ntpath import os import pathlib +import shutil import time from typing import Annotated -from urllib.parse import parse_qs, unquote, urlparse +from urllib.parse import parse_qs, unquote, urlparse, urlsplit, urlunsplit import requests import typer @@ -13,6 +15,7 @@ from comfy_cli.config_manager import ConfigManager from comfy_cli.constants import DEFAULT_COMFY_MODEL_PATH from comfy_cli.file_utils import DownloadException, check_unauthorized, download_file +from comfy_cli.output import get_renderer from comfy_cli.output import rprint as print # context-aware: stderr in JSON mode from comfy_cli.workspace_manager import WorkspaceManager @@ -53,6 +56,118 @@ def potentially_strip_param_url(path_name: str) -> str: return path_name.split("?")[0] +def _download_failure( + *, + code: str, + message: str, + hint: str | None = None, + details: dict | None = None, +) -> typer.Exit: + """Emit an ``envelope/1`` error for a `model download` failure and return the + ``typer.Exit`` the caller should raise. + + Every failure path in :func:`download` funnels through here so a `--json` + consumer always gets a machine-readable ``error.code`` and a non-zero exit — + never a bare exit that an envelope-synthesizing wrapper would read as a + success for a download that never happened. + + In pretty mode ``renderer.error`` renders the red error panel, so call sites + do not print the message themselves. Rich builds the panel body from + ``rich.text.Text``, which is literal, so a message containing markup + metacharacters (``[/]``) neither crashes nor gets swallowed — no ``escape()`` + needed. + + ``code`` is keyword-only on purpose: the registry test + (``tests/comfy_cli/output/test_error_code_registry.py``) AST-scans for + ``code="..."`` keywords, so a positional call site would silently drop out of + the both-ways registry enforcement. + """ + get_renderer().error( + code=code, + message=message, + hint=hint, + details=details, + command="model download", + ) + return typer.Exit(code=1) + + +def _scrub_url(url: str) -> str: + """Drop the query string, fragment and userinfo from *url*. + + CivitAI download links carry the API token as ``?token=`` — ``tracking._scrub_value`` + already strips it before telemetry, and an error envelope is a *louder* channel + than telemetry (agent/MCP wrappers capture and log stdout verbatim), so the same + scrubbing applies to anything we put in ``error.details`` or a log line. + """ + if not isinstance(url, str): + return url + try: + parts = urlsplit(url) + except ValueError: # malformed authority (e.g. an unbracketed IPv6 literal) + return url.partition("?")[0].partition("#")[0] + if not parts.scheme: + return url.partition("?")[0].partition("#")[0] + return urlunsplit((parts.scheme, parts.netloc.rpartition("@")[2], parts.path, "", "")) + + +def _reject_unsafe_component(value: str | None, *, label: str, url: str) -> None: + """Fail the download if *value* is anything but a single, inert path component. + + ``local_filename`` and ``basemodel`` are joined into the destination path, and + both can come straight off the CivitAI API response (``file["name"]``, + ``version["baseModel"]``) — which is remote input, accepted without a prompt in + non-interactive runs. ``pathlib``/``os.path.join`` do not sanitize ``..`` or an + absolute component, so an unvalidated value writes outside the workspace. + + Subdirectories are still reachable — via ``--relative-path``, which is the + option that exists for choosing the destination directory — so this rejects a + traversal without removing the capability. + """ + if not value: + return + unsafe = ( + value in (".", "..") + or os.path.isabs(value) + or bool(ntpath.splitdrive(value)[0]) + or "/" in value + or "\\" in value + ) + if unsafe: + raise _download_failure( + code="invalid_argument", + message=f"Unsafe {label}: {value!r} is not a plain name (it contains a path separator, a drive, or '..').", + hint="use `--relative-path ` to choose the destination directory and `--filename ` for a plain name", + details={"url": _scrub_url(url), label: value}, + ) + + +def _resolve_civitai_source(fetch, url: str): + """Run a CivitAI metadata lookup, converting any failure into an error envelope. + + The lookups raise on HTTP/parse errors and ``request_civitai_model_version_api`` + returns ``None`` when the version carries no primary file; both used to escape + as an unhandled traceback with no envelope. + """ + try: + resolved = fetch() + except Exception as e: + raise _download_failure( + code="download_failed", + message=f"Could not resolve a downloadable file from the CivitAI URL: {e}", + hint="check the model/version exists and is public; a private model needs --set-civitai-api-token", + details={"url": _scrub_url(url), "stage": "resolve"}, + ) from None + if resolved is None: + raise _download_failure( + code="download_failed", + message="The CivitAI model version has no primary file to download.", + hint="pick a version that has a primary file, or pass its direct download URL", + details={"url": _scrub_url(url), "stage": "resolve"}, + ) + return resolved + + def check_huggingface_url(url: str) -> tuple[bool, str | None, str | None, str | None, str | None]: """ Check if the given URL is a Hugging Face URL and extract relevant information. @@ -270,7 +385,10 @@ def download( headers["Authorization"] = f"Bearer {civitai_api_token}" if is_civitai_model_url: - local_filename, url, model_type, basemodel = request_civitai_model_api(model_id, version_id, headers) + local_filename, url, model_type, basemodel = _resolve_civitai_source( + lambda: request_civitai_model_api(model_id, version_id, headers), url + ) + _reject_unsafe_component(basemodel, label="basemodel", url=url) model_path = model_path_map.get(model_type) @@ -280,7 +398,10 @@ def download( relative_path = os.path.join(DEFAULT_COMFY_MODEL_PATH, model_path, basemodel) elif is_civitai_api_url: - local_filename, url, model_type, basemodel = request_civitai_model_version_api(version_id, headers) + local_filename, url, model_type, basemodel = _resolve_civitai_source( + lambda: request_civitai_model_version_api(version_id, headers), url + ) + _reject_unsafe_component(basemodel, label="basemodel", url=url) model_path = model_path_map.get(model_type) @@ -299,7 +420,14 @@ def download( basemodel = ui.prompt_input("Enter base model (e.g. SD1.5, SDXL, ...)", default="") relative_path = os.path.join(DEFAULT_COMFY_MODEL_PATH, model_path, basemodel) else: - print("Model source is unknown") + # Neither CivitAI nor Hugging Face — a plain file URL, which IS a supported + # source: the download proceeds via download_file() below. This is a note, + # not a failure, so it must not short-circuit. + # escape() the URL: `print` is Rich's markup-parsing print, and a URL holding + # markup metacharacters (`[/]`, or an IPv6 literal like `http://[::1]/x`) + # would otherwise raise MarkupError — an uncaught crash with no envelope, + # the exact failure mode this command is being hardened against. + print(f"Model source is unknown; treating {escape(_scrub_url(url))} as a direct file URL") if filename is None: if local_filename is None: @@ -312,25 +440,43 @@ def download( if relative_path is None: relative_path = DEFAULT_COMFY_MODEL_PATH - if local_filename is None: - raise typer.Exit(code=1) - if local_filename == "": - raise DownloadException("Filename cannot be empty") + # None (prompt cancelled / no TTY) and "" (prompting skipped for an agentic + # caller, which returns the empty default) are the same failure: nothing to + # save the file as. Both used to end without an envelope — a bare exit 1 and + # an unhandled DownloadException traceback respectively. + if not local_filename: + raise _download_failure( + code="missing_argument", + message="Could not determine a filename to save the model as.", + hint="pass `--filename `", + details={"url": _scrub_url(url)}, + ) + + _reject_unsafe_component(local_filename, label="filename", url=url) local_filepath = get_workspace() / relative_path / local_filename if local_filepath.exists(): - print(f"[bold red]File already exists: {local_filepath}[/bold red]") - return + raise _download_failure( + code="model_file_exists", + message=f"File already exists: {local_filepath}", + hint="pass `--filename` to save under a different name, or remove the existing file", + details={"path": str(local_filepath)}, + ) start_time = time.monotonic() if is_huggingface_url and check_unauthorized(url, headers): if hf_api_token is None: - print( - f"Unauthorized access to Hugging Face model. Please set the Hugging Face API token using `comfy model download --set-hf-api-token` or via the `{constants.HF_API_TOKEN_ENV_KEY}` environment variable" + raise _download_failure( + code="hf_unauthorized", + message="Unauthorized access to the Hugging Face model, and no Hugging Face API token is configured.", + hint=( + "set the token via `comfy model download --set-hf-api-token ` or the " + f"`{constants.HF_API_TOKEN_ENV_KEY}` environment variable" + ), + details={"url": _scrub_url(url), "repo_id": repo_id}, ) - return else: try: import huggingface_hub @@ -340,30 +486,96 @@ def download( from comfy_cli.resolve_python import resolve_workspace_python - python = resolve_workspace_python(str(get_workspace())) - subprocess.check_call([python, "-m", "pip", "install", "huggingface_hub"]) - import huggingface_hub - - print(f"Downloading model {model_id} from Hugging Face...") - output_path = huggingface_hub.hf_hub_download( - repo_id=repo_id, - filename=hf_filename, - subfolder=hf_folder_name, - revision=hf_branch_name, - token=hf_api_token, - local_dir=get_workspace() / relative_path, - cache_dir=get_workspace() / relative_path, - ) - print(f"Model downloaded successfully to: {output_path}") + try: + # NB: this installs into the *workspace* interpreter (pinned by + # tests/comfy_cli/test_models_python_resolution.py), which is not + # necessarily the one running this process — so the import below + # can still fail for e.g. a pipx-installed CLI. That predates this + # change and is tracked separately; what's new here is that the + # failure now ends in an envelope instead of a traceback. + python = resolve_workspace_python(str(get_workspace())) + subprocess.check_call([python, "-m", "pip", "install", "huggingface_hub"]) + import huggingface_hub + except Exception as e: + raise _download_failure( + code="download_failed", + message=f"`huggingface_hub` is required for Hugging Face downloads and could not be installed: {e}", + hint="install it manually (`pip install huggingface_hub`) and retry", + details={"url": _scrub_url(url)}, + ) from None + + print(f"Downloading model {escape(str(model_id))} from Hugging Face...") + try: + output_path = huggingface_hub.hf_hub_download( + repo_id=repo_id, + filename=hf_filename, + subfolder=hf_folder_name, + revision=hf_branch_name, + token=hf_api_token, + local_dir=get_workspace() / relative_path, + cache_dir=get_workspace() / relative_path, + ) + except Exception as e: + raise _download_failure( + code="download_failed", + message=f"Hugging Face download failed: {e}", + hint="check the repo/revision exists and the token has access to it", + details={"url": _scrub_url(url), "repo_id": repo_id}, + ) from None + + # `hf_hub_download` names the file after the *repo* path (`hf_filename`), + # so an explicit `--filename` was silently ignored on this branch only — + # the public-HF path below goes through download_file(local_filepath) and + # honours it. That split made the `local_filepath.exists()` guard above + # check a path this branch never writes, and made its `model_file_exists` + # hint ("pass `--filename`") impossible to act on. Move the result to the + # requested name so both are true on every path. + if filename is not None and pathlib.Path(output_path) != local_filepath: + try: + local_filepath.parent.mkdir(parents=True, exist_ok=True) + shutil.move(str(output_path), str(local_filepath)) + output_path = str(local_filepath) + except OSError as e: + raise _download_failure( + code="download_failed", + message=f"Downloaded the model but could not save it as {local_filepath}: {e}", + hint="check the destination directory is writable, or omit `--filename`", + details={"url": _scrub_url(url), "path": str(local_filepath)}, + ) from None + + print(f"Model downloaded successfully to: {escape(str(output_path))}") else: - print(f"Start downloading URL: {url} into {local_filepath}") + print(f"Start downloading URL: {escape(_scrub_url(url))} into {escape(str(local_filepath))}") try: download_file(url, local_filepath, headers, downloader=resolved_downloader) except DownloadException as e: - # escape() so a dynamic error message containing "[/]" or similar - # rich-markup syntax doesn't trigger MarkupError or get mis-rendered. - print(f"[bold red]{escape(str(e))}[/bold red]") - raise typer.Exit(code=1) from None + raise _download_failure( + code="download_failed", + message=str(e), + details={"url": _scrub_url(url), "path": str(local_filepath)}, + ) from None + except OSError as e: + # download_file() converts network failures to DownloadException, but a + # local filesystem failure (unwritable dir, disk full) still escapes as + # OSError — which would end the command with a traceback and no envelope. + raise _download_failure( + code="download_failed", + message=f"Could not write the downloaded file to {local_filepath}: {e}", + hint="check the destination directory exists, is writable, and has free space", + details={"url": _scrub_url(url), "path": str(local_filepath)}, + ) from None + except Exception as e: + # Backstop for the invariant this command promises: a `--json` caller must + # never get a bare traceback with no envelope. The downloader can still + # raise something neither branch above names — e.g. a malformed + # `Content-Length` from a broken proxy used to surface as ValueError out of + # `int(...)` (now guarded in file_utils, but the class of failure remains). + raise _download_failure( + code="download_failed", + message=f"Download failed unexpectedly ({type(e).__name__}): {e}", + hint="retry; if it persists, re-run without `--json` for the full traceback", + details={"url": _scrub_url(url), "path": str(local_filepath)}, + ) from None elapsed = time.monotonic() - start_time print(f"Done in {_format_elapsed(elapsed)}") diff --git a/comfy_cli/error_codes.py b/comfy_cli/error_codes.py index f0729744..b8e63036 100644 --- a/comfy_cli/error_codes.py +++ b/comfy_cli/error_codes.py @@ -228,8 +228,12 @@ class ErrorCode: # --- models / templates introspection ------------------------------------ ErrorCode( "invalid_argument", - "An argument intended for a URL path failed safe-path validation.", - "use only alphanumerics, `_`, `-`, or `.` in path-segment arguments", + "An argument intended for a URL path or for a filesystem path component failed " + "safe-path validation — e.g. a `comfy model download` filename (from `--filename` or " + "from the CivitAI API response) that carries a path separator, a drive letter or `..` " + "and would write outside the workspace.", + "use only alphanumerics, `_`, `-`, or `.` in path-segment arguments; choose the " + "destination directory with `--relative-path`", ), ErrorCode( "folder_not_found", @@ -516,8 +520,25 @@ class ErrorCode: ), ErrorCode( "download_failed", - "HTTP error while downloading an output file.", - "check that the job completed successfully and the server is reachable", + "A download failed. Either an HTTP error while fetching a job's output file, or " + "`comfy model download` failing to fetch the model (transfer error, Hugging Face " + "download error, or an unresolvable CivitAI model/version). `details.url` carries the " + "source URL; `details.stage` is `resolve` when the failure was metadata lookup, not transfer.", + "check that the source URL is reachable and the job completed successfully", + ), + ErrorCode( + "model_file_exists", + "`comfy model download` refused to overwrite an existing file at the target path " + "(`details.path`). The download was NOT performed — the command fails rather than " + "exiting 0, so a caller can't mistake the skip for a completed download.", + "pass `--filename` to save under a different name, or remove the existing file", + ), + ErrorCode( + "hf_unauthorized", + "Hugging Face returned 401 for the model URL and no Hugging Face API token is configured " + "(gated or private repo).", + "set the token via `comfy model download --set-hf-api-token ` or the `HF_API_TOKEN` " + "environment variable", ), ErrorCode( "download_no_outputs", diff --git a/comfy_cli/file_utils.py b/comfy_cli/file_utils.py index 3a84027c..92818d42 100644 --- a/comfy_cli/file_utils.py +++ b/comfy_cli/file_utils.py @@ -252,7 +252,16 @@ def _download_file_httpx( raise DownloadException(f"Failed to download file.\n{status_reason}") content_length = response.headers.get("Content-Length") - total = int(content_length) if content_length is not None else None + try: + total = int(content_length) if content_length is not None else None + except ValueError: + # A broken server/proxy can send a non-numeric Content-Length. That is not + # a reason to fail the transfer — treat it exactly like a missing header + # (indeterminate progress) instead of letting ValueError escape the whole + # download and end the command with a traceback. + total = None + if total is not None and total < 0: + total = None if total is not None: description = f"Downloading {total // 1024 // 1024} MB" else: diff --git a/tests/comfy_cli/command/models/test_model_download_json.py b/tests/comfy_cli/command/models/test_model_download_json.py new file mode 100644 index 00000000..f54c3e7e --- /dev/null +++ b/tests/comfy_cli/command/models/test_model_download_json.py @@ -0,0 +1,518 @@ +"""`comfy --json model download` failure paths must emit `envelope/1` errors. + +The invariant under test: **every** failure of `model download` exits non-zero +AND puts exactly one `envelope/1` object with `ok: false` on stdout. Two of these +paths used to exit 0 with no envelope (file-exists, HF-unauthorized), which the +local MCP's `plain_ok` synthesizer turned into a synthesized *success* for a +download that never happened. + +The success path is deliberately unchanged (prints + exit 0, no envelope) — that +same synthesizer depends on exit-0-no-envelope meaning success. +""" + +from __future__ import annotations + +import io +import json +import sys +from typing import Any +from unittest.mock import Mock, patch + +import pytest +import typer.testing + +from comfy_cli import constants +from comfy_cli.command.models.models import app +from comfy_cli.file_utils import DownloadException +from comfy_cli.output.renderer import ( + OutputMode, + Renderer, + reset_renderer_for_testing, + set_renderer, +) + +runner = typer.testing.CliRunner() + +# The renderer's machine stream (stdout in real runs) is pinned to this buffer so +# the "exactly one envelope on stdout, nothing else" assertion holds regardless of +# whether the installed Click version mixes stderr into CliRunner's captured +# output. Human log lines go to `_HUMAN` (stderr in real runs) and must never +# appear in `_MACHINE`. +_MACHINE = io.StringIO() +_HUMAN = io.StringIO() + + +@pytest.fixture(autouse=True) +def json_renderer(): + """Install a JSON-mode renderer with pinned streams, and tear it down. + + The renderer is a process-wide singleton, so the teardown matters: without it + the pretty-mode expectations of the neighbouring `test_models.py` would break + depending on test ordering. + """ + global _MACHINE, _HUMAN + reset_renderer_for_testing() + _MACHINE = io.StringIO() + _HUMAN = io.StringIO() + r = Renderer() + r.mode = OutputMode.JSON + r.machine_stream = _MACHINE + r.pretty_stream = _HUMAN + set_renderer(r) + yield r + reset_renderer_for_testing() + + +def _stdout() -> str: + return _MACHINE.getvalue() + + +def _envelope() -> dict[str, Any]: + """Parse the single envelope the command is allowed to put on stdout.""" + output = _stdout() + objects = [] + for line in output.strip().splitlines(): + if not line.strip(): + continue + try: + objects.append(json.loads(line)) + except json.JSONDecodeError: + pytest.fail(f"non-JSON line on stdout in --json mode: {line!r}\nfull stdout:\n{output}") + assert len(objects) == 1, f"expected exactly one envelope on stdout, got {len(objects)}:\n{output}" + return objects[0] + + +def _assert_error_envelope(result, code: str) -> dict[str, Any]: + assert result.exit_code == 1, f"expected exit 1, got {result.exit_code}\n{result.output}" + env = _envelope() + assert env["schema"] == "envelope/1" + assert env["ok"] is False + assert env["error"]["code"] == code + assert env["error"]["hint"], "every error must carry a navigation hint" + return env + + +def _invoke(args: list[str], **patches): + """Run `model download` with the network/prompt surface stubbed out. + + Defaults: not a CivitAI URL, not a Hugging Face URL, no config values. Each + test overrides only the piece it exercises. + + Note: there is deliberately no ``patch("comfy_cli.tracking.track_command")`` + here. ``download`` is decorated at import time, so patching the factory + afterwards never rebinds the wrapper — the mock would be pure decoration. The + real wrapper runs and no-ops because telemetry consent is off under pytest. + """ + cfg = patches.pop("config_manager", None) + if cfg is None: + cfg = Mock() + cfg.get_or_override.return_value = None + cfg.get.return_value = None + + with ( + patch("comfy_cli.command.models.models.check_civitai_url", return_value=(False, False, None, None)), + patch("comfy_cli.command.models.models.check_huggingface_url", return_value=(False, None, None, None, None)), + patch("comfy_cli.command.models.models.config_manager", cfg), + ): + stack = [] + try: + for target, kwargs in patches.items(): + p = patch(f"comfy_cli.command.models.models.{target}", **kwargs) + stack.append(p) + p.start() + return runner.invoke(app, ["download", *args]) + finally: + for p in reversed(stack): + p.stop() + + +# --------------------------------------------------------------------------- # +# transfer failure — the headline acceptance case +# --------------------------------------------------------------------------- # + + +def test_download_exception_emits_download_failed(tmp_path): + result = _invoke( + ["--url", "https://example.com/missing.safetensors", "--filename", "x.safetensors"], + get_workspace={"return_value": tmp_path}, + download_file={"side_effect": DownloadException("404 Not Found")}, + ) + + env = _assert_error_envelope(result, "download_failed") + assert "404 Not Found" in env["error"]["message"] + assert env["error"]["details"]["url"] == "https://example.com/missing.safetensors" + + +def test_local_write_failure_emits_download_failed(tmp_path): + """`download_file` converts *network* failures to DownloadException, but a local + filesystem failure still surfaces as OSError — which used to end the command + with a traceback and no envelope.""" + result = _invoke( + ["--url", "https://example.com/m.bin", "--filename", "x.bin"], + get_workspace={"return_value": tmp_path}, + download_file={"side_effect": PermissionError(13, "Permission denied")}, + ) + + env = _assert_error_envelope(result, "download_failed") + assert "Permission denied" in env["error"]["message"] + + +def test_download_exception_message_with_markup_survives(tmp_path): + """A server message containing rich-markup metacharacters must reach the + envelope verbatim — the JSON path never runs it through Rich markup.""" + result = _invoke( + ["--url", "https://example.com/m.bin", "--filename", "x.bin"], + get_workspace={"return_value": tmp_path}, + download_file={"side_effect": DownloadException("server said [/] at /path/[id]/resource")}, + ) + + env = _assert_error_envelope(result, "download_failed") + assert env["error"]["message"] == "server said [/] at /path/[id]/resource" + + +# --------------------------------------------------------------------------- # +# file already exists — used to exit 0 with no envelope +# --------------------------------------------------------------------------- # + + +def test_file_exists_emits_error_and_exits_nonzero(tmp_path): + target = tmp_path / "models" / "checkpoints" / "x.safetensors" + target.parent.mkdir(parents=True) + target.write_bytes(b"already here") + + result = _invoke( + [ + "--url", + "https://example.com/x.safetensors", + "--relative-path", + "models/checkpoints", + "--filename", + "x.safetensors", + ], + get_workspace={"return_value": tmp_path}, + download_file={}, + ) + + env = _assert_error_envelope(result, "model_file_exists") + assert env["error"]["details"]["path"] == str(target) + + +def test_file_exists_does_not_download(tmp_path): + target = tmp_path / "x.bin" + target.write_bytes(b"already here") + + with patch("comfy_cli.command.models.models.download_file") as mock_dl: + result = _invoke( + ["--url", "https://example.com/x.bin", "--relative-path", ".", "--filename", "x.bin"], + get_workspace={"return_value": tmp_path}, + ) + + assert result.exit_code == 1 + assert not mock_dl.called + + +# --------------------------------------------------------------------------- # +# Hugging Face unauthorized with no token — used to exit 0 with no envelope +# --------------------------------------------------------------------------- # + + +def test_hf_unauthorized_without_token_emits_error(tmp_path): + cfg = Mock() + cfg.get_or_override.return_value = None # no CivitAI token, no HF token + cfg.get.return_value = None + + with ( + patch("comfy_cli.command.models.models.check_civitai_url", return_value=(False, False, None, None)), + patch( + "comfy_cli.command.models.models.check_huggingface_url", + return_value=(True, "org/repo", "m.safetensors", None, "main"), + ), + patch("comfy_cli.command.models.models.check_unauthorized", return_value=True), + patch("comfy_cli.command.models.models.get_workspace", return_value=tmp_path), + patch("comfy_cli.command.models.models.download_file") as mock_dl, + patch("comfy_cli.command.models.models.config_manager", cfg), + ): + result = runner.invoke( + app, + [ + "download", + "--url", + "https://huggingface.co/org/repo/resolve/main/m.safetensors", + "--relative-path", + "models/checkpoints", + "--filename", + "m.safetensors", + ], + ) + + env = _assert_error_envelope(result, "hf_unauthorized") + assert constants.HF_API_TOKEN_ENV_KEY in env["error"]["hint"] + assert env["error"]["details"]["repo_id"] == "org/repo" + assert not mock_dl.called + + +# --------------------------------------------------------------------------- # +# no resolvable filename — used to be a bare `typer.Exit(1)` / a raw traceback +# --------------------------------------------------------------------------- # + + +def test_unprompted_empty_filename_emits_missing_argument(tmp_path): + """`ui.prompt_input` returns its (empty) default when prompting is skipped — + e.g. for an agentic caller. That used to raise an unhandled + DownloadException("Filename cannot be empty") and print a traceback.""" + result = _invoke( + ["--url", "https://example.com/"], + get_workspace={"return_value": tmp_path}, + download_file={}, + ui={"new": Mock(**{"prompt_input.return_value": ""})}, + ) + + _assert_error_envelope(result, "missing_argument") + + +def test_cancelled_filename_prompt_emits_missing_argument(tmp_path): + """`questionary` returns None when the prompt is cancelled / gets EOF. That + used to be a bare `typer.Exit(1)` with no message and no envelope.""" + result = _invoke( + ["--url", "https://example.com/"], + get_workspace={"return_value": tmp_path}, + download_file={}, + ui={"new": Mock(**{"prompt_input.return_value": None})}, + ) + + _assert_error_envelope(result, "missing_argument") + + +# --------------------------------------------------------------------------- # +# CivitAI metadata resolution — used to escape as an unhandled traceback +# --------------------------------------------------------------------------- # + + +def test_civitai_lookup_failure_emits_download_failed(tmp_path): + with ( + patch("comfy_cli.command.models.models.check_civitai_url", return_value=(True, False, 4242, None)), + patch("comfy_cli.command.models.models.check_huggingface_url", return_value=(False, None, None, None, None)), + patch( + "comfy_cli.command.models.models.request_civitai_model_api", + side_effect=RuntimeError("404 Client Error"), + ), + patch("comfy_cli.command.models.models.get_workspace", return_value=tmp_path), + patch("comfy_cli.command.models.models.config_manager", Mock(**{"get_or_override.return_value": None})), + ): + result = runner.invoke( + app, + ["download", "--url", "https://civitai.com/models/4242", "--filename", "x.safetensors"], + ) + + env = _assert_error_envelope(result, "download_failed") + assert env["error"]["details"]["stage"] == "resolve" + + +def test_civitai_version_without_primary_file_emits_download_failed(tmp_path): + with ( + patch("comfy_cli.command.models.models.check_civitai_url", return_value=(False, True, None, 777)), + patch("comfy_cli.command.models.models.check_huggingface_url", return_value=(False, None, None, None, None)), + patch("comfy_cli.command.models.models.request_civitai_model_version_api", return_value=None), + patch("comfy_cli.command.models.models.get_workspace", return_value=tmp_path), + patch("comfy_cli.command.models.models.config_manager", Mock(**{"get_or_override.return_value": None})), + ): + result = runner.invoke( + app, + ["download", "--url", "https://civitai.com/api/download/models/777", "--filename", "x.safetensors"], + ) + + env = _assert_error_envelope(result, "download_failed") + assert env["error"]["details"]["stage"] == "resolve" + + +# --------------------------------------------------------------------------- # +# the success path is unchanged: no envelope, exit 0 +# --------------------------------------------------------------------------- # + + +def test_direct_url_success_still_emits_no_envelope(tmp_path): + """A plain (non-CivitAI, non-Hugging-Face) file URL is a SUPPORTED source — + it must still download, and must still leave stdout empty on success so the + MCP's exit-0-no-envelope success synthesizer keeps working.""" + with patch("comfy_cli.command.models.models.download_file") as mock_dl: + result = _invoke( + ["--url", "https://example.com/model.bin", "--relative-path", ".", "--filename", "model.bin"], + get_workspace={"return_value": tmp_path}, + ) + + assert result.exit_code == 0, result.output + assert mock_dl.called, "a direct file URL must still be downloaded" + assert _stdout().strip() == "", f"success path must keep stdout clean, got: {_stdout()!r}" + assert "Done in" in _HUMAN.getvalue(), "the human-facing log still goes to stderr" + + +# --------------------------------------------------------------------------- # +# the URL is never echoed back with its credentials attached +# --------------------------------------------------------------------------- # + + +def test_error_details_url_is_scrubbed_of_token(tmp_path): + """CivitAI download links carry the API token as `?token=`. `tracking._scrub_value` + strips it before telemetry; an error envelope is a louder channel than telemetry + (agent/MCP wrappers log stdout verbatim), so it must be stripped there too.""" + result = _invoke( + ["--url", "https://example.com/m.bin?token=SUPERSECRET#frag", "--filename", "x.bin"], + get_workspace={"return_value": tmp_path}, + download_file={"side_effect": DownloadException("boom")}, + ) + + env = _assert_error_envelope(result, "download_failed") + assert env["error"]["details"]["url"] == "https://example.com/m.bin" + assert "SUPERSECRET" not in json.dumps(env) + assert "SUPERSECRET" not in _HUMAN.getvalue(), "the human log leaks the token too" + + +def test_userinfo_is_stripped_from_details_url(tmp_path): + result = _invoke( + ["--url", "https://user:pw@example.com/m.bin", "--filename", "x.bin"], + get_workspace={"return_value": tmp_path}, + download_file={"side_effect": DownloadException("boom")}, + ) + + env = _assert_error_envelope(result, "download_failed") + assert env["error"]["details"]["url"] == "https://example.com/m.bin" + + +# --------------------------------------------------------------------------- # +# a filename must not escape the workspace +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize("bad", ["../../.bashrc", "..", "sub/dir.bin", "C:evil.bin", "a\\b.bin"]) +def test_traversing_filename_is_rejected_before_any_download(tmp_path, bad): + """`--filename` — and, in a non-interactive run, the CivitAI-supplied `file["name"]` + that becomes its default — is joined into the destination with no sanitisation by + `pathlib`, so `..` or an absolute value writes outside the workspace.""" + with patch("comfy_cli.command.models.models.download_file") as mock_dl: + result = _invoke( + ["--url", "https://example.com/m.bin", "--relative-path", ".", "--filename", bad], + get_workspace={"return_value": tmp_path}, + ) + + env = _assert_error_envelope(result, "invalid_argument") + assert "--relative-path" in env["error"]["hint"], "the hint must point at the supported way to pick a directory" + assert not mock_dl.called + + +def test_plain_filename_is_still_accepted(tmp_path): + """The traversal guard must not deny the ordinary case.""" + with patch("comfy_cli.command.models.models.download_file") as mock_dl: + result = _invoke( + ["--url", "https://example.com/m.bin", "--relative-path", "models/checkpoints", "--filename", "v1.5.bin"], + get_workspace={"return_value": tmp_path}, + ) + + assert result.exit_code == 0, result.output + assert mock_dl.called + + +def test_traversing_civitai_basemodel_is_rejected(tmp_path): + """`version["baseModel"]` is remote input joined straight into the destination path.""" + with ( + patch("comfy_cli.command.models.models.check_civitai_url", return_value=(False, True, None, 777)), + patch("comfy_cli.command.models.models.check_huggingface_url", return_value=(False, None, None, None, None)), + patch( + "comfy_cli.command.models.models.request_civitai_model_version_api", + return_value=("m.safetensors", "https://civitai.com/x", "checkpoint", "../../../.."), + ), + patch("comfy_cli.command.models.models.get_workspace", return_value=tmp_path), + patch("comfy_cli.command.models.models.download_file") as mock_dl, + patch("comfy_cli.command.models.models.config_manager", Mock(**{"get_or_override.return_value": None})), + ): + result = runner.invoke(app, ["download", "--url", "https://civitai.com/api/download/models/777"]) + + _assert_error_envelope(result, "invalid_argument") + assert not mock_dl.called + + +# --------------------------------------------------------------------------- # +# rich markup in dynamic values must not become the crash it was meant to catch +# --------------------------------------------------------------------------- # + + +def test_markup_hostile_url_does_not_crash_the_log_line(tmp_path): + """`print` here is Rich's markup-parsing print. A URL holding `[/]` (or an IPv6 + literal) raised MarkupError out of the *logging*, ending the command with a + traceback and no envelope — the exact failure this command promises not to have.""" + with patch("comfy_cli.command.models.models.download_file") as mock_dl: + result = _invoke( + ["--url", "http://[::1]/a/[/]/m.bin", "--relative-path", ".", "--filename", "m.bin"], + get_workspace={"return_value": tmp_path}, + ) + + assert result.exit_code == 0, result.output + assert mock_dl.called + + +# --------------------------------------------------------------------------- # +# nothing escapes without an envelope, including what we didn't name +# --------------------------------------------------------------------------- # + + +def test_unexpected_downloader_exception_still_emits_envelope(tmp_path): + """The downloader can raise something neither `DownloadException` nor `OSError` — + a malformed `Content-Length` used to surface as ValueError out of `int()`. The + `--json` contract is an envelope on *every* failure, so there's a backstop.""" + result = _invoke( + ["--url", "https://example.com/m.bin", "--filename", "x.bin"], + get_workspace={"return_value": tmp_path}, + download_file={"side_effect": ValueError("invalid literal for int()")}, + ) + + env = _assert_error_envelope(result, "download_failed") + assert "ValueError" in env["error"]["message"] + + +# --------------------------------------------------------------------------- # +# --filename is honoured on the gated Hugging Face path too +# --------------------------------------------------------------------------- # + + +def test_hf_download_honours_explicit_filename(tmp_path): + """`hf_hub_download` names the file after the repo path, so `--filename` was + silently ignored on the *authenticated* HF branch only. That made the + `local_filepath.exists()` guard inspect a path this branch never writes, and made + the `model_file_exists` hint ("pass `--filename`") impossible to act on.""" + hf_written = tmp_path / "models" / "checkpoints" / "repo-name.safetensors" + hf_written.parent.mkdir(parents=True) + hf_written.write_bytes(b"weights") + + fake_hub = Mock() + fake_hub.hf_hub_download.return_value = str(hf_written) + + cfg = Mock() + cfg.get_or_override.return_value = "hf_token" + cfg.get.return_value = None + + with ( + patch.dict(sys.modules, {"huggingface_hub": fake_hub}), + patch("comfy_cli.command.models.models.check_civitai_url", return_value=(False, False, None, None)), + patch( + "comfy_cli.command.models.models.check_huggingface_url", + return_value=(True, "org/repo", "repo-name.safetensors", None, "main"), + ), + patch("comfy_cli.command.models.models.check_unauthorized", return_value=True), + patch("comfy_cli.command.models.models.get_workspace", return_value=tmp_path), + patch("comfy_cli.command.models.models.config_manager", cfg), + ): + result = runner.invoke( + app, + [ + "download", + "--url", + "https://huggingface.co/org/repo/resolve/main/repo-name.safetensors", + "--relative-path", + "models/checkpoints", + "--filename", + "mine.safetensors", + ], + ) + + assert result.exit_code == 0, result.output + assert (tmp_path / "models" / "checkpoints" / "mine.safetensors").read_bytes() == b"weights" + assert not hf_written.exists() diff --git a/tests/test_file_utils_network.py b/tests/test_file_utils_network.py index 219e3001..2fe1347c 100644 --- a/tests/test_file_utils_network.py +++ b/tests/test_file_utils_network.py @@ -107,6 +107,25 @@ def test_download_file_success_without_content_length(mock_stream, tmp_path): assert test_file.read_bytes() == b"chunk1chunk2" +@patch("httpx.stream") +def test_download_file_success_with_garbage_content_length(mock_stream, tmp_path): + """A non-numeric Content-Length (broken server/proxy) must degrade to an + indeterminate progress bar, not blow the transfer up with a ValueError out of + ``int()`` — which escaped `model download`'s handlers as a bare traceback.""" + mock_response = Mock() + mock_response.status_code = 200 + mock_response.headers = {"Content-Length": "not-a-number"} + mock_response.iter_bytes.return_value = [b"chunk1", b"chunk2"] + mock_response.__enter__ = Mock(return_value=mock_response) + mock_response.__exit__ = Mock(return_value=None) + mock_stream.return_value = mock_response + + test_file = tmp_path / "test.txt" + download_file("http://example.com", test_file) + + assert test_file.read_bytes() == b"chunk1chunk2" + + @patch("httpx.stream") def test_download_file_failure(mock_stream): mock_response = Mock()