unsloth/studio/backend/hub/tests/test_model_services.py

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

import asyncio
import sys
from pathlib import Path
from types import SimpleNamespace

import pytest
from fastapi import HTTPException

from hub.dependencies import get_hf_token
from hub.storage import scan_folders
from hub.services import download_lifecycle
from hub.services import snapshot_progress
from hub.services.datasets import downloads as dataset_downloads
from hub.services.models import (
    cache_inventory,
    common as model_common,
    deletion,
    downloads,
    folder_browser,
    gguf_variants,
    local_inventory,
    ollama,
)
from hub.utils import (
    download_manifest,
    download_registry,
    gguf,
    hf_cache_state,
    inventory_scan,
    paths,
    state_dir,
)
from hub.workers import hf_download


@pytest.fixture(autouse = True)
def _denylist_inert(monkeypatch):
    # The browse tests here exercise allowlist containment, symlink safety and
    # the sensitive-name filter, not the system-directory denylist (which has
    # its own suite in tests/test_browse_denylist.py). On macOS tmp_path
    # resolves under /private/var, a denied prefix, so _resolve_browse_target
    # would 403 the fixture dirs before that logic runs. Keep the denylist inert
    # so these assertions hold on every platform. folder_browser binds
    # is_denied_system_path at import, so patch it on that module, not on
    # scan_folders. The "rejects" cases still 403 via the allowlist/sensitive
    # checks, and the non-browse tests never call it.
    monkeypatch.setattr(folder_browser, "is_denied_system_path", lambda _p: False)


def _repo(repo_id: str, files: list[SimpleNamespace], repo_path: Path):
    return SimpleNamespace(
        repo_id = repo_id,
        repo_type = "model",
        repo_path = repo_path,
        revisions = [SimpleNamespace(files = files)],
    )


def _file(
    name: str,
    size: int,
    blob_path: str | None = None,
):
    return SimpleNamespace(file_name = name, size_on_disk = size, blob_path = blob_path)


def _sibling(name: str, size: int, sha: str):
    return SimpleNamespace(rfilename = name, size = size, lfs = {"sha256": sha})


class TestExtractQuantToken:
    def test_trailing_precision_is_kept(self):
        assert gguf.extract_quant_token("model-it-F16.gguf") == "F16"
        assert gguf.extract_quant_token("model-BF16.gguf") == "BF16"

    def test_real_quant_wins_over_infix_precision(self):
        assert gguf.extract_quant_token("Foo-BF16-Q4_K_M.gguf") == "Q4_K_M"
        assert gguf.extract_quant_token("Foo-F16-Q8_0.gguf") == "Q8_0"
        assert gguf.extract_quant_token("Foo-F32-IQ4_XS.gguf") == "IQ4_XS"

    def test_ud_prefix_preserved(self):
        assert gguf.extract_quant_token("Foo-BF16-UD-Q4_K_XL.gguf") == "UD-Q4_K_XL"

    def test_precision_infix_variants_do_not_collapse(self):
        labels = {
            gguf.extract_quant_label("Foo-BF16-Q4_K_M.gguf"),
            gguf.extract_quant_label("Foo-BF16-Q8_0.gguf"),
        }
        assert labels == {"Q4_K_M", "Q8_0"}


def test_big_endian_detection_ignores_model_name_be_token():
    assert gguf.is_big_endian_gguf_path("model-Q4_K_M-be.gguf", "Q4_K_M")
    assert gguf.is_big_endian_gguf_path("model-Q4_K_M_be_infill.gguf", "Q4_K_M")
    assert not gguf.is_big_endian_gguf_path("foo-be-Q4_K_M.gguf", "Q4_K_M")
    assert not gguf.is_big_endian_gguf_path("Q4_K_M/foo-be.gguf", "Q4_K_M")
    assert gguf.pick_best_gguf(["model-Q4_K_M-be.gguf", "model-Q4_K_M.gguf"]) == (
        "model-Q4_K_M.gguf"
    )


def test_list_local_gguf_variants_skips_big_endian_sibling(tmp_path):
    (tmp_path / "model-Q4_K_M-be.gguf").write_bytes(b"x" * 100)
    (tmp_path / "model-Q4_K_M.gguf").write_bytes(b"y" * 10)

    variants, has_vision = gguf.list_local_gguf_variants(str(tmp_path))

    assert has_vision is False
    assert [(v.quant, v.filename, v.size_bytes) for v in variants] == [
        ("Q4_K_M", "model-Q4_K_M.gguf", 10)
    ]


@pytest.mark.parametrize("repo_id", ["bert-base-uncased", "owner/repo"])
def test_repo_id_validation_accepts_hf_repo_id_contract(repo_id):
    assert paths.is_valid_repo_id(repo_id)


def test_repo_id_validation_accepts_max_length_namespaced_repo():
    assert paths.is_valid_repo_id(f"{'a' * 96}/{'b' * 96}")


@pytest.mark.parametrize(
    "repo_id",
    [
        "datasets/foo/bar",
        ".repo",
        "repo.git",
        "foo..bar",
        "foo--bar",
        "../repo",
        "owner/../repo",
    ],
)
def test_repo_id_validation_rejects_unsafe_or_invalid_ids(repo_id):
    assert not paths.is_valid_repo_id(repo_id)


def test_download_state_preserves_readable_keys_when_safe(monkeypatch, tmp_path):
    monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path)

    path = state_dir.marker_path("model", "Owner/Repo", "Q4_K_M")

    assert path is not None
    assert path.name == "models--owner--repo--variant--q4_k_m.json"


@pytest.mark.parametrize("variant", ["bad variant with spaces", "q" * 64])
def test_download_state_bounds_long_repo_variant_filenames(monkeypatch, tmp_path, variant):
    monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path)
    repo_id = f"{'a' * 96}/{'b' * 96}"

    assert paths.is_valid_repo_id(repo_id)
    assert download_manifest.write_cancel_marker("model", repo_id, variant, "http")
    assert download_manifest.write_manifest(
        "model",
        repo_id,
        variant,
        [download_manifest.ExpectedFile(path = "model.gguf", size = 1)],
        "http",
    )

    marker_path = state_dir.marker_path("model", repo_id, variant)
    manifest_path = state_dir.manifest_path("model", repo_id, variant)

    assert marker_path is not None
    assert manifest_path is not None
    assert "--sha256-" in marker_path.name
    assert len(marker_path.name.encode("utf-8")) <= 255
    assert len(f".{marker_path.name}.tmp-00000000".encode("utf-8")) <= 255
    assert download_manifest.has_cancel_marker("model", repo_id, variant)
    assert download_manifest.read_manifest("model", repo_id, variant) is not None
    assert list(download_manifest.iter_variant_markers("model", repo_id)) == [
        (variant, marker_path)
    ]
    assert list(download_manifest.iter_variant_manifests("model", repo_id)) == [
        (variant, manifest_path)
    ]


class _RecordingLogger:
    def __init__(self):
        self.warnings = []

    def warning(self, *args, **kwargs):
        self.warnings.append((args, kwargs))


def test_resolve_browse_target_preserves_allowlist_and_symlink_safety(tmp_path):
    home = tmp_path / "home"
    scan = tmp_path / "scan"
    target = scan / "nested"
    home.mkdir()
    target.mkdir(parents = True)
    (home / "scan-link").symlink_to(scan, target_is_directory = True)

    resolved = folder_browser._resolve_browse_target(
        str(home / "scan-link" / "nested"),
        [home, scan],
    )

    assert resolved == target.resolve()


def test_resolve_browse_target_rejects_outside_allowlist(tmp_path):
    allowed = tmp_path / "allowed"
    outside = tmp_path / "outside"
    allowed.mkdir()
    outside.mkdir()

    with pytest.raises(HTTPException) as exc_info:
        folder_browser._resolve_browse_target(str(outside), [allowed])

    assert exc_info.value.status_code == 403


def test_resolve_browse_target_rejects_sensitive_dir(tmp_path):
    home = tmp_path / "home"
    ssh = home / ".ssh"
    ssh.mkdir(parents = True)

    with pytest.raises(HTTPException) as exc_info:
        folder_browser._resolve_browse_target(str(ssh), [home])

    assert exc_info.value.status_code == 403


def test_resolve_browse_target_rejects_sensitive_root(tmp_path):
    ssh = tmp_path / "home" / ".ssh"
    ssh.mkdir(parents = True)

    with pytest.raises(HTTPException) as exc_info:
        folder_browser._resolve_browse_target(str(ssh), [ssh])

    assert exc_info.value.status_code == 403


def test_browse_folders_hides_sensitive_dirs(monkeypatch, tmp_path):
    home = tmp_path / "home"
    (home / ".ssh").mkdir(parents = True)
    (home / "models").mkdir()
    # Accept and ignore the optional (media_roots, drive_roots) args the caller now passes.
    monkeypatch.setattr(folder_browser, "_build_browse_allowlist", lambda *_a, **_k: [home])

    response = folder_browser.browse_folders_response(str(home), show_hidden = True)

    names = {entry.name for entry in response.entries}
    assert "models" in names
    assert ".ssh" not in names


def test_browse_allowlist_includes_linux_run_media_mounts(monkeypatch, tmp_path):
    home = tmp_path / "home"
    media_root = tmp_path / "run" / "media" / "dspofu" / "nvmeB"
    model_dir = media_root / "modelsAI" / "gguf" / "qwen3.6"
    home.mkdir()
    model_dir.mkdir(parents = True)
    monkeypatch.setattr(folder_browser.Path, "home", lambda: home)
    monkeypatch.setattr(folder_browser, "linux_run_media_mount_roots", lambda: [media_root])
    monkeypatch.setattr(folder_browser, "_resolve_hf_cache_dir", lambda: tmp_path / "missing-hf")
    monkeypatch.setattr(scan_folders, "list_scan_folders", lambda: [])
    monkeypatch.setattr(folder_browser, "well_known_model_dirs", lambda: [])

    allowlist = folder_browser._build_browse_allowlist()

    assert media_root.resolve() in allowlist
    assert folder_browser._resolve_browse_target(str(model_dir), allowlist) == model_dir.resolve()


def test_get_models_folder_response_creates_and_returns_dir(monkeypatch, tmp_path):
    # The endpoint creates the cache dir on demand so the desktop "Open folder"
    # action works even before the first download.
    target = tmp_path / "hub"
    monkeypatch.setattr(local_inventory, "_resolve_hf_cache_dir", lambda: target)

    response = local_inventory.get_models_folder_response()

    assert response == {"path": str(target)}
    assert target.is_dir()


def test_get_models_folder_response_reports_create_failure(monkeypatch, tmp_path):
    target = tmp_path / "hub"
    target.write_text("not a directory")
    monkeypatch.setattr(local_inventory, "_resolve_hf_cache_dir", lambda: target)

    with pytest.raises(HTTPException) as exc_info:
        local_inventory.get_models_folder_response()

    assert exc_info.value.status_code == 500
    assert "Failed to create models folder" in exc_info.value.detail


def test_get_models_folder_response_requires_directory(monkeypatch, tmp_path):
    class MissingPath:
        def __init__(self, value: Path):
            self.value = value

        def mkdir(self, *, parents: bool, exist_ok: bool):
            assert parents is True
            assert exist_ok is True

        def is_dir(self):
            return False

        def __str__(self):
            return str(self.value)

    target = MissingPath(tmp_path / "hub")
    monkeypatch.setattr(local_inventory, "_resolve_hf_cache_dir", lambda: target)

    with pytest.raises(HTTPException) as exc_info:
        local_inventory.get_models_folder_response()

    assert exc_info.value.status_code == 500
    assert "not a directory" in exc_info.value.detail


def test_contained_link_path_confines_to_link_dir(tmp_path):
    link_dir = tmp_path / "ollama" / ".studio_links" / "abc123"

    legit = ollama._contained_link_path(link_dir, "llama3-latest-Q4_K_M.gguf")
    assert legit == link_dir / "llama3-latest-Q4_K_M.gguf"

    for unsafe in (
        "",
        ".",
        "..",
        "a/b.gguf",
        "../evil.gguf",
        "/etc/passwd",
        "model-tag-../../../pwned.gguf",
    ):
        assert ollama._contained_link_path(link_dir, unsafe) is None


def test_make_ollama_blob_link_refuses_escaping_name(tmp_path):
    root = tmp_path / "ollama"
    link_dir = root / ".studio_links" / "abc123"
    blob = root / "blobs" / "sha256-deadbeef"
    blob.parent.mkdir(parents = True)
    blob.write_bytes(b"weights")

    escaped = ollama._make_ollama_blob_link(link_dir, "model-tag-../../../pwned.gguf", blob)
    assert escaped is None
    assert not list(tmp_path.rglob("pwned.gguf"))

    safe = ollama._make_ollama_blob_link(link_dir, "model-tag.gguf", blob)
    assert safe == str(link_dir / "model-tag.gguf")
    assert (link_dir / "model-tag.gguf").exists()


def test_cached_gguf_scan_dedupes_and_excludes_mmproj_only(monkeypatch, tmp_path):
    smaller = _repo("Org/Dupe", [_file("Q4_K_M.gguf", 100)], tmp_path / "small")
    larger = _repo(
        "org/dupe",
        [_file("Q4_K_M.gguf", 300), _file("Q8_0.gguf", 200)],
        tmp_path / "large",
    )
    mmproj_only = _repo("Org/VisionAdapter", [_file("mmproj-F16.gguf", 900)], tmp_path / "mmproj")
    monkeypatch.setattr(
        cache_inventory,
        "all_hf_cache_scans",
        lambda: [SimpleNamespace(repos = [smaller, larger, mmproj_only])],
    )
    monkeypatch.setattr(
        cache_inventory.hf_cache_scan,
        "is_gguf_repo_partial",
        lambda _repo_id, _path: False,
    )

    result = {"cached": cache_inventory._scan_cached_gguf()}

    assert [row["repo_id"] for row in result["cached"]] == ["org/dupe"]
    assert result["cached"][0]["size_bytes"] == 500
    assert result["cached"][0]["model_format"] == "gguf"
    assert result["cached"][0]["capabilities"]["requires_variant"] is True


def test_cached_gguf_scan_preserves_partial_flag(monkeypatch, tmp_path):
    partial = _repo("Org/Partial", [_file("Q4_K_M.gguf", 100)], tmp_path / "partial")
    monkeypatch.setattr(
        cache_inventory,
        "all_hf_cache_scans",
        lambda: [SimpleNamespace(repos = [partial])],
    )
    monkeypatch.setattr(
        cache_inventory.hf_cache_scan,
        "is_gguf_repo_partial",
        lambda _repo_id, _path: True,
    )

    result = {"cached": cache_inventory._scan_cached_gguf()}
    row = result["cached"][0]

    assert row["partial"] is True
    assert row["partial_transport"] is None
    assert row["capabilities"]["can_chat"] is False


def test_cached_gguf_scan_includes_variant_state_without_completed_gguf(monkeypatch, tmp_path):
    monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
    repo_path = tmp_path / "hub" / "models--Org--PartialGguf"
    repo_path.mkdir(parents = True)
    partial = _repo(
        "Org/PartialGguf",
        [_file("config.json", 12)],
        repo_path,
    )
    assert download_manifest.write_manifest(
        "model",
        "Org/PartialGguf",
        "Q4_K_M",
        [download_manifest.ExpectedFile(path = "model-Q4_K_M.gguf", size = 4096)],
        "http",
    )
    assert download_manifest.write_cancel_marker("model", "Org/PartialGguf", "Q4_K_M", "http")
    monkeypatch.setattr(
        cache_inventory,
        "all_hf_cache_scans",
        lambda: [SimpleNamespace(repos = [partial])],
    )
    monkeypatch.setattr(
        cache_inventory.hf_cache_scan,
        "is_gguf_repo_partial",
        lambda _repo_id, _path: True,
    )

    result = {"cached": cache_inventory._scan_cached_gguf()}
    row = result["cached"][0]

    assert row["repo_id"] == "Org/PartialGguf"
    assert row["model_format"] == "gguf"
    assert row["size_bytes"] == 4096
    assert row["partial"] is True
    assert row["capabilities"]["requires_variant"] is True


def test_gguf_variant_requirements_include_split_files_and_preferred_mmproj():
    requirements = gguf_variants._build_gguf_variant_requirements(
        [
            _sibling("model-Q4_K_M-00001-of-00002.gguf", 10, "main-a"),
            _sibling("model-Q4_K_M-00002-of-00002.gguf", 20, "main-b"),
            _sibling("mmproj-BF16.gguf", 7, "mm-bf16"),
            _sibling("mmproj-F16.gguf", 5, "mm-f16"),
        ]
    )

    req = requirements["q4_k_m"]

    assert req.main_size_bytes == 30
    assert req.download_size_bytes == 35
    assert req.main_hashes == frozenset({"main-a", "main-b"})
    assert req.required_hashes == frozenset({"main-a", "main-b", "mm-f16"})
    assert req.companion_hashes == frozenset({"mm-f16"})
    assert req.mmproj_hashes == frozenset({"mm-bf16", "mm-f16"})
    assert req.target_filenames == (
        "model-Q4_K_M-00001-of-00002.gguf",
        "model-Q4_K_M-00002-of-00002.gguf",
        "mmproj-F16.gguf",
    )


def test_gguf_variant_requirements_skip_big_endian_sibling():
    requirements = gguf_variants._build_gguf_variant_requirements(
        [
            _sibling("model-Q4_K_M-be.gguf", 100, "main-be"),
            _sibling("model-Q4_K_M.gguf", 10, "main-le"),
        ]
    )

    req = requirements["q4_k_m"]

    assert req.main_size_bytes == 10
    assert req.main_hashes == frozenset({"main-le"})
    assert req.main_filenames == frozenset({"model-Q4_K_M.gguf"})
    assert req.target_filenames == ("model-Q4_K_M.gguf",)


def test_worker_gguf_variant_plan_matches_service_requirement(monkeypatch):
    siblings = [
        _sibling("model-Q4_K_M-00001-of-00002.gguf", 10, "main-a"),
        _sibling("model-Q4_K_M-00002-of-00002.gguf", 20, "main-b"),
        _sibling("mmproj-BF16.gguf", 7, "mm-bf16"),
        _sibling("mmproj-F16.gguf", 5, "mm-f16"),
    ]
    monkeypatch.setattr(
        hf_download,
        "_model_info_with_retry",
        lambda *_args, **_kwargs: SimpleNamespace(siblings = siblings),
    )

    service_req = gguf_variants._build_gguf_variant_requirements(siblings)["q4_k_m"]
    worker_plan = hf_download._gguf_variant_target_plan("Org/Vision", "Q4_K_M", None)

    assert worker_plan == service_req


def test_gguf_variant_blob_hashes_accept_dict_lfs_fallback(monkeypatch):
    with gguf_variants._VARIANT_HASH_LOCK:
        gguf_variants._VARIANT_HASH_CACHE.clear()
        gguf_variants._VARIANT_REQUIREMENT_CACHE.clear()
    monkeypatch.setitem(
        sys.modules,
        "huggingface_hub",
        SimpleNamespace(
            HfApi = lambda *_args, **_kwargs: SimpleNamespace(
                model_info = lambda *_a, **_k: SimpleNamespace(
                    siblings = [
                        _sibling("model-Q4_K_M.gguf", 10, "main-dict"),
                        _sibling("model-Q8_0.gguf", 20, "other"),
                        _sibling("mmproj-F16.gguf", 5, "mmproj"),
                    ]
                )
            )
        ),
    )

    result = gguf_variants.gguf_variant_blob_hashes("Org/DictLfs", "Q4_K_M", None)
    main_only = gguf_variants.gguf_variant_blob_hashes(
        "Org/DictLfs",
        "Q4_K_M",
        None,
        include_companions = False,
    )

    assert result == frozenset({"main-dict", "mmproj"})
    assert main_only == frozenset({"main-dict"})


def test_gguf_variant_blob_hashes_skip_missing_rfilename(monkeypatch):
    with gguf_variants._VARIANT_HASH_LOCK:
        gguf_variants._VARIANT_HASH_CACHE.clear()
        gguf_variants._VARIANT_REQUIREMENT_CACHE.clear()
    siblings = [
        SimpleNamespace(rfilename = None, size = 1, lfs = {"sha256": "bad"}),
        _sibling("model-Q4_K_M.gguf", 10, "main"),
    ]
    monkeypatch.setattr(
        gguf_variants,
        "_fetch_gguf_variant_requirements",
        lambda _repo_id, _hf_token = None: gguf_variants._build_gguf_variant_requirements(siblings),
    )

    result = gguf_variants.gguf_variant_blob_hashes("Org/Malformed", "Q4_K_M", None)

    assert result == frozenset({"main"})


def test_worker_gguf_variant_targets_skip_missing_rfilename(monkeypatch):
    monkeypatch.setattr(
        hf_download,
        "_model_info_with_retry",
        lambda *_args, **_kwargs: SimpleNamespace(
            siblings = [
                SimpleNamespace(rfilename = None, size = 1),
                _sibling("model-Q4_K_M.gguf", 10, "main"),
                _sibling("mmproj-F16.gguf", 5, "mm"),
            ]
        ),
    )

    result = hf_download._gguf_variant_target_plan("Org/Malformed", "Q4_K_M", None)

    assert list(result.target_filenames) == ["model-Q4_K_M.gguf", "mmproj-F16.gguf"]


def test_download_gguf_variant_purges_only_main_quant_hashes(monkeypatch, tmp_path):
    prepare_calls = []
    snapshot_calls = []
    written = []
    verified = []

    monkeypatch.setattr(
        hf_download,
        "_model_info_with_retry",
        lambda *_args, **_kwargs: SimpleNamespace(
            siblings = [
                _sibling("model-Q4_K_M.gguf", 10, "q4-main"),
                _sibling("model-Q8_0.gguf", 20, "q8-main"),
                _sibling("mmproj-F16.gguf", 5, "shared-mmproj"),
            ]
        ),
    )
    monkeypatch.setattr(
        hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args)
    )
    monkeypatch.setattr(
        download_registry,
        "prepare_cache_for_transport",
        lambda *args, **kwargs: prepare_calls.append((args, kwargs)) or 0,
    )
    monkeypatch.setattr(download_manifest, "clear_cancel_marker", lambda *_args: None)
    monkeypatch.setattr(
        download_manifest, "write_manifest", lambda *args: written.append(args) or True
    )
    monkeypatch.setitem(
        sys.modules,
        "huggingface_hub",
        SimpleNamespace(
            snapshot_download = lambda **kwargs: snapshot_calls.append(kwargs) or str(tmp_path)
        ),
    )

    hf_download._download_gguf_variant("Org/Vision", "Q4_K_M", None, "http")

    assert prepare_calls == [
        (
            ("model", "Org/Vision", "http", "Q4_K_M"),
            {
                "only_blob_hashes": frozenset({"q4-main"}),
                "companion_blob_hashes": frozenset({"shared-mmproj"}),
                "protected_blob_hashes": frozenset(),
            },
        )
    ]
    assert [file.path for file in written[0][3]] == ["model-Q4_K_M.gguf", "mmproj-F16.gguf"]
    assert snapshot_calls[0]["allow_patterns"] == ["model-Q4_K_M.gguf", "mmproj-F16.gguf"]
    assert verified == [("model", "Org/Vision", "Q4_K_M", str(tmp_path))]


def test_download_gguf_variant_manifest_resume_purges_only_main_quant_hashes(monkeypatch, tmp_path):
    prepare_calls = []
    snapshot_calls = []

    def _metadata_unavailable(*_args, **_kwargs):
        raise RuntimeError("metadata down")

    manifest = download_manifest.Manifest(
        repo_type = "model",
        repo_id = "Org/Vision",
        variant = "Q4_K_M",
        started_at = "",
        expected_files = (
            download_manifest.ExpectedFile(
                path = "model-Q4_K_M.gguf",
                size = 10,
                sha256 = "q4-main",
            ),
            download_manifest.ExpectedFile(
                path = "mmproj-F16.gguf",
                size = 5,
                sha256 = "shared-mmproj",
            ),
        ),
        transport = "http",
    )
    monkeypatch.setattr(
        hf_download,
        "_gguf_variant_target_plan",
        _metadata_unavailable,
    )
    monkeypatch.setattr(download_manifest, "read_manifest", lambda *_args: manifest)
    monkeypatch.setattr(download_manifest, "clear_cancel_marker", lambda *_args: None)
    monkeypatch.setattr(
        download_registry,
        "prepare_cache_for_transport",
        lambda *args, **kwargs: prepare_calls.append((args, kwargs)) or 0,
    )
    monkeypatch.setattr(hf_download, "_verify_completed_download", lambda *_args, **_kwargs: None)
    monkeypatch.setitem(
        sys.modules,
        "huggingface_hub",
        SimpleNamespace(
            snapshot_download = lambda **kwargs: snapshot_calls.append(kwargs) or str(tmp_path)
        ),
    )

    hf_download._download_gguf_variant("Org/Vision", "Q4_K_M", None, "http")

    assert prepare_calls == [
        (
            ("model", "Org/Vision", "http", "Q4_K_M"),
            {
                "only_blob_hashes": frozenset({"q4-main"}),
                "companion_blob_hashes": frozenset({"shared-mmproj"}),
                "protected_blob_hashes": frozenset(),
            },
        )
    ]
    assert snapshot_calls[0]["allow_patterns"] == ["model-Q4_K_M.gguf", "mmproj-F16.gguf"]


def test_download_snapshot_recovers_manifest_after_metadata_fallback(monkeypatch, tmp_path):
    metadata_calls = []
    written = []
    cleared = []
    verified = []

    def _metadata(*_args, **_kwargs):
        metadata_calls.append(True)
        if len(metadata_calls) == 1:
            raise RuntimeError("metadata down")
        return SimpleNamespace(siblings = [SimpleNamespace(rfilename = "config.json", size = 12)])

    monkeypatch.setattr(hf_download, "_model_info_with_retry", _metadata)
    monkeypatch.setattr(
        hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args)
    )
    monkeypatch.setattr(
        download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0
    )
    monkeypatch.setattr(
        download_manifest, "clear_cancel_marker", lambda *args: cleared.append(args)
    )
    monkeypatch.setattr(
        download_manifest, "write_manifest", lambda *args: written.append(args) or True
    )
    monkeypatch.setitem(
        sys.modules,
        "huggingface_hub",
        SimpleNamespace(snapshot_download = lambda **_kwargs: str(tmp_path)),
    )

    hf_download._download_snapshot("Org/Model", None, "http")

    assert len(metadata_calls) == 2
    assert cleared == [("model", "Org/Model", None)]
    assert written[0][0:3] == ("model", "Org/Model", None)
    assert written[0][3][0].path == "config.json"
    assert verified == [("model", "Org/Model", None, str(tmp_path))]


def test_download_dataset_continues_without_metadata_manifest(monkeypatch, tmp_path):
    metadata_calls = []
    snapshot_calls = []
    written = []
    cleared = []
    verified = []

    def _metadata(*_args, **_kwargs):
        metadata_calls.append(True)
        raise RuntimeError("metadata down")

    monkeypatch.setattr(hf_download, "_dataset_info_with_retry", _metadata)
    monkeypatch.setattr(
        hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args)
    )
    monkeypatch.setattr(
        download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0
    )
    monkeypatch.setattr(
        download_manifest, "clear_cancel_marker", lambda *args: cleared.append(args)
    )
    monkeypatch.setattr(
        download_manifest, "write_manifest", lambda *args: written.append(args) or True
    )
    monkeypatch.setattr(
        hf_cache_state, "has_active_incomplete_blobs", lambda *_args, **_kwargs: False
    )
    monkeypatch.setitem(
        sys.modules,
        "huggingface_hub",
        SimpleNamespace(
            snapshot_download = lambda **kwargs: snapshot_calls.append(kwargs) or str(tmp_path)
        ),
    )

    hf_download._download_dataset("Org/Data", None, "http")

    assert len(metadata_calls) == 2
    assert cleared == [("dataset", "Org/Data", None)]
    assert written == []
    assert snapshot_calls == [
        {
            "repo_id": "Org/Data",
            "token": False,
            "repo_type": "dataset",
            "max_workers": 1,
        }
    ]
    assert verified == [("dataset", "Org/Data", None, str(tmp_path))]


def test_download_snapshot_fails_when_metadata_unavailable_and_partial_remains(
    monkeypatch, tmp_path
):
    """No prior manifest + metadata unavailable + leftover .incomplete blobs means
    a cached partial was returned without downloading: the worker must exit 1, not
    derive a self-certifying manifest from the finalized subset."""
    written = []
    verified = []

    def _metadata(*_args, **_kwargs):
        raise RuntimeError("metadata down")

    monkeypatch.setattr(hf_download, "_model_info_with_retry", _metadata)
    monkeypatch.setattr(
        hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args)
    )
    monkeypatch.setattr(
        download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0
    )
    monkeypatch.setattr(download_manifest, "clear_cancel_marker", lambda *_args: None)
    monkeypatch.setattr(download_manifest, "read_manifest", lambda *_args, **_kwargs: None)
    monkeypatch.setattr(
        download_manifest, "write_manifest", lambda *args: written.append(args) or True
    )
    monkeypatch.setattr(
        hf_cache_state, "has_active_incomplete_blobs", lambda *_args, **_kwargs: True
    )
    monkeypatch.setitem(
        sys.modules,
        "huggingface_hub",
        SimpleNamespace(snapshot_download = lambda **_kwargs: str(tmp_path)),
    )

    with pytest.raises(SystemExit) as excinfo:
        hf_download._download_snapshot("Org/Model", None, "http")

    assert excinfo.value.code == 1
    assert written == []
    assert verified == []


def test_purge_repo_cache_dirs_skips_top_level_symlink(monkeypatch, tmp_path):
    root = tmp_path / "hub"
    target = tmp_path / "target"
    root.mkdir()
    target.mkdir()
    link = root / "models--Org--Repo"
    link.symlink_to(target, target_is_directory = True)
    monkeypatch.setattr(hf_cache_state, "hf_cache_roots", lambda: [root])

    removed = hf_cache_state.purge_repo_cache_dirs("model", "Org/Repo")

    assert removed is False
    assert link.is_symlink()
    assert target.is_dir()


def test_gguf_download_progress_fallback_logs_warning(monkeypatch):
    token = "hf_12345678901234567890"
    logger = _RecordingLogger()

    async def _run_inline(fn, *args, **kwargs):
        return fn(*args, **kwargs)

    def _raise_permission_error(*_args, **_kwargs):
        raise PermissionError(f"denied {token}")

    monkeypatch.setattr(snapshot_progress, "logger", logger)
    monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline)
    monkeypatch.setattr(
        downloads.gguf_variants,
        "gguf_variant_requirements",
        lambda *_args, **_kwargs: None,
    )
    monkeypatch.setattr(
        downloads.gguf_variants,
        "gguf_variant_blob_hashes",
        lambda *_args, **_kwargs: frozenset(),
    )
    monkeypatch.setattr(
        snapshot_progress,
        "preferred_repo_cache_dirs",
        _raise_permission_error,
    )
    monkeypatch.setattr(
        downloads,
        "_registry",
        SimpleNamespace(get_job = lambda _key: SimpleNamespace(state = "running")),
    )

    result = asyncio.run(
        downloads.get_gguf_download_progress_response(
            "Org/Model",
            variant = "Q4_K_M",
            expected_bytes = -1,
            hf_token = token,
        )
    )

    assert result == {
        "downloaded_bytes": 0,
        "completed_bytes": 0,
        "complete_on_disk": False,
        "expected_bytes": 0,
        "progress": 0,
        "cache_path": None,
    }
    assert logger.warnings
    args, kwargs = logger.warnings[0]
    assert args[:4] == (
        "Error checking %s download progress for %s: %s: %s",
        "model",
        "Org/Model",
        "PermissionError",
    )
    assert token not in args[4]
    assert "***" in args[4]
    assert kwargs == {}


def test_gguf_progress_counts_completed_mmproj_with_expected_bytes(monkeypatch, tmp_path):
    """A finished mmproj companion keeps counting toward progress once the caller
    supplies expected bytes; resolving the variant requirement credits it."""
    entry = tmp_path / "models--Org--Model-GGUF"
    snap = entry / "snapshots" / "rev0"
    blobs = entry / "blobs"
    snap.mkdir(parents = True)
    blobs.mkdir(parents = True)
    (snap / "model-Q4_K_M.gguf").write_bytes(b"x" * 100)
    (snap / "mmproj-F16.gguf").write_bytes(b"y" * 30)
    (blobs / "mainhash").write_bytes(b"x" * 100)
    (blobs / "mmprojhash").write_bytes(b"y" * 30)
    monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
    assert download_manifest.write_manifest(
        "model",
        "Org/Model-GGUF",
        "Q4_K_M",
        [
            download_manifest.ExpectedFile(
                path = "model-Q4_K_M.gguf",
                size = 100,
                sha256 = "mainhash",
            ),
            download_manifest.ExpectedFile(
                path = "mmproj-F16.gguf",
                size = 30,
                sha256 = "mmprojhash",
            ),
        ],
        "http",
    )

    requirement = gguf_variants._GgufVariantRequirement(
        main_filenames = frozenset({"model-Q4_K_M.gguf"}),
        target_filenames = ("model-Q4_K_M.gguf", "mmproj-F16.gguf"),
        main_hashes = frozenset({"mainhash"}),
        required_hashes = frozenset({"mainhash", "mmprojhash"}),
        companion_hashes = frozenset({"mmprojhash"}),
        mmproj_filenames = frozenset({"mmproj-F16.gguf"}),
        mmproj_hashes = frozenset({"mmprojhash"}),
        expected_files = (
            download_manifest.ExpectedFile(
                path = "model-Q4_K_M.gguf",
                size = 100,
                sha256 = "mainhash",
            ),
            download_manifest.ExpectedFile(
                path = "mmproj-F16.gguf",
                size = 30,
                sha256 = "mmprojhash",
            ),
        ),
        main_size_bytes = 100,
        download_size_bytes = 130,
    )

    async def _run_inline(fn, *args, **kwargs):
        return fn(*args, **kwargs)

    monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline)
    monkeypatch.setattr(
        downloads.gguf_variants,
        "gguf_variant_requirements",
        lambda *_args, **_kwargs: requirement,
    )
    monkeypatch.setattr(
        snapshot_progress,
        "preferred_repo_cache_dirs",
        lambda *_args, **_kwargs: [entry],
    )
    monkeypatch.setattr(
        downloads,
        "_registry",
        SimpleNamespace(get_job = lambda _key: SimpleNamespace(state = "idle")),
    )

    result = asyncio.run(
        downloads.get_gguf_download_progress_response(
            "Org/Model-GGUF",
            variant = "Q4_K_M",
            expected_bytes = 130,
        )
    )

    assert result["completed_bytes"] == 130
    assert result["downloaded_bytes"] == 130
    assert result["complete_on_disk"] is True
    assert result["progress"] == 1.0


def test_gguf_progress_subtracts_new_job_completed_baseline(monkeypatch, tmp_path):
    entry = tmp_path / "models--Org--Model-GGUF"
    snap = entry / "snapshots" / "rev0"
    blobs = entry / "blobs"
    snap.mkdir(parents = True)
    blobs.mkdir(parents = True)
    (snap / "mmproj-F16.gguf").write_bytes(b"y" * 30)
    (blobs / "mmprojhash").write_bytes(b"y" * 30)
    monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
    assert download_manifest.write_manifest(
        "model",
        "Org/Model-GGUF",
        "Q4_K_M",
        [
            download_manifest.ExpectedFile(
                path = "model-Q4_K_M.gguf",
                size = 100,
                sha256 = "mainhash",
            ),
            download_manifest.ExpectedFile(
                path = "mmproj-F16.gguf",
                size = 30,
                sha256 = "mmprojhash",
            ),
        ],
        "http",
    )

    requirement = gguf_variants._GgufVariantRequirement(
        main_filenames = frozenset({"model-Q4_K_M.gguf"}),
        target_filenames = ("model-Q4_K_M.gguf", "mmproj-F16.gguf"),
        main_hashes = frozenset({"mainhash"}),
        required_hashes = frozenset({"mainhash", "mmprojhash"}),
        companion_hashes = frozenset({"mmprojhash"}),
        mmproj_filenames = frozenset({"mmproj-F16.gguf"}),
        mmproj_hashes = frozenset({"mmprojhash"}),
        expected_files = (
            download_manifest.ExpectedFile(
                path = "model-Q4_K_M.gguf",
                size = 100,
                sha256 = "mainhash",
            ),
            download_manifest.ExpectedFile(
                path = "mmproj-F16.gguf",
                size = 30,
                sha256 = "mmprojhash",
            ),
        ),
        main_size_bytes = 100,
        download_size_bytes = 130,
    )

    async def _run_inline(fn, *args, **kwargs):
        return fn(*args, **kwargs)

    monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline)
    monkeypatch.setattr(
        downloads.gguf_variants,
        "gguf_variant_requirements",
        lambda *_args, **_kwargs: requirement,
    )
    monkeypatch.setattr(
        snapshot_progress,
        "preferred_repo_cache_dirs",
        lambda *_args, **_kwargs: [entry],
    )
    monkeypatch.setattr(
        downloads,
        "_registry",
        SimpleNamespace(
            get_job = lambda _key: SimpleNamespace(state = "running"),
            get_job_metadata = lambda _key: SimpleNamespace(
                completed_baseline_bytes = 30,
            ),
        ),
    )

    result = asyncio.run(
        downloads.get_gguf_download_progress_response(
            "Org/Model-GGUF",
            variant = "Q4_K_M",
            expected_bytes = 130,
        )
    )

    assert result["completed_bytes"] == 0
    assert result["downloaded_bytes"] == 0
    assert result["expected_bytes"] == 100
    assert result["complete_on_disk"] is False
    assert result["progress"] == 0


def test_gguf_progress_shows_main_when_companion_left_the_count(monkeypatch, tmp_path):
    # The mmproj companion that seeded the baseline is gone, so completed_bytes
    # is main-only and below the baseline; it must not be subtracted to 0.
    entry = tmp_path / "models--Org--Model-GGUF"
    blobs = entry / "blobs"
    blobs.mkdir(parents = True)
    (blobs / "mainhash").write_bytes(b"x" * 20)
    monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")

    requirement = gguf_variants._GgufVariantRequirement(
        main_filenames = frozenset({"model-Q4_K_M.gguf"}),
        target_filenames = ("model-Q4_K_M.gguf", "mmproj-F16.gguf"),
        main_hashes = frozenset({"mainhash"}),
        required_hashes = frozenset({"mainhash", "mmprojhash"}),
        companion_hashes = frozenset({"mmprojhash"}),
        mmproj_filenames = frozenset({"mmproj-F16.gguf"}),
        mmproj_hashes = frozenset({"mmprojhash"}),
        expected_files = (
            download_manifest.ExpectedFile(
                path = "model-Q4_K_M.gguf",
                size = 100,
                sha256 = "mainhash",
            ),
            download_manifest.ExpectedFile(
                path = "mmproj-F16.gguf",
                size = 30,
                sha256 = "mmprojhash",
            ),
        ),
        main_size_bytes = 100,
        download_size_bytes = 130,
    )

    async def _run_inline(fn, *args, **kwargs):
        return fn(*args, **kwargs)

    monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline)
    monkeypatch.setattr(
        downloads.gguf_variants,
        "gguf_variant_requirements",
        lambda *_args, **_kwargs: requirement,
    )
    monkeypatch.setattr(
        snapshot_progress,
        "preferred_repo_cache_dirs",
        lambda *_args, **_kwargs: [entry],
    )
    monkeypatch.setattr(
        downloads,
        "_registry",
        SimpleNamespace(
            get_job = lambda _key: SimpleNamespace(state = "running"),
            get_job_metadata = lambda _key: SimpleNamespace(
                completed_baseline_bytes = 30,
            ),
        ),
    )

    result = asyncio.run(
        downloads.get_gguf_download_progress_response(
            "Org/Model-GGUF",
            variant = "Q4_K_M",
            expected_bytes = 130,
        )
    )

    assert result["completed_bytes"] == 20
    assert result["downloaded_bytes"] == 20
    assert result["expected_bytes"] == 130
    assert result["complete_on_disk"] is False


def test_gguf_progress_complete_on_disk_ignores_full_baseline(monkeypatch, tmp_path):
    # A variant already complete on disk carries a baseline equal to its full
    # size; subtracting it would report 0/0 for a finished variant (frontend
    # evicts it as gone). Once complete_on_disk is verified, the full figures
    # must survive.
    entry = tmp_path / "models--Org--Model-GGUF"
    snap = entry / "snapshots" / "rev0"
    blobs = entry / "blobs"
    snap.mkdir(parents = True)
    blobs.mkdir(parents = True)
    (snap / "model-Q4_K_M.gguf").write_bytes(b"x" * 100)
    (snap / "mmproj-F16.gguf").write_bytes(b"y" * 30)
    (blobs / "mainhash").write_bytes(b"x" * 100)
    (blobs / "mmprojhash").write_bytes(b"y" * 30)
    monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
    assert download_manifest.write_manifest(
        "model",
        "Org/Model-GGUF",
        "Q4_K_M",
        [
            download_manifest.ExpectedFile(
                path = "model-Q4_K_M.gguf",
                size = 100,
                sha256 = "mainhash",
            ),
            download_manifest.ExpectedFile(
                path = "mmproj-F16.gguf",
                size = 30,
                sha256 = "mmprojhash",
            ),
        ],
        "http",
    )

    requirement = gguf_variants._GgufVariantRequirement(
        main_filenames = frozenset({"model-Q4_K_M.gguf"}),
        target_filenames = ("model-Q4_K_M.gguf", "mmproj-F16.gguf"),
        main_hashes = frozenset({"mainhash"}),
        required_hashes = frozenset({"mainhash", "mmprojhash"}),
        companion_hashes = frozenset({"mmprojhash"}),
        mmproj_filenames = frozenset({"mmproj-F16.gguf"}),
        mmproj_hashes = frozenset({"mmprojhash"}),
        expected_files = (
            download_manifest.ExpectedFile(
                path = "model-Q4_K_M.gguf",
                size = 100,
                sha256 = "mainhash",
            ),
            download_manifest.ExpectedFile(
                path = "mmproj-F16.gguf",
                size = 30,
                sha256 = "mmprojhash",
            ),
        ),
        main_size_bytes = 100,
        download_size_bytes = 130,
    )

    async def _run_inline(fn, *args, **kwargs):
        return fn(*args, **kwargs)

    monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline)
    monkeypatch.setattr(
        downloads.gguf_variants,
        "gguf_variant_requirements",
        lambda *_args, **_kwargs: requirement,
    )
    monkeypatch.setattr(
        snapshot_progress,
        "preferred_repo_cache_dirs",
        lambda *_args, **_kwargs: [entry],
    )
    monkeypatch.setattr(
        downloads,
        "_registry",
        SimpleNamespace(
            get_job = lambda _key: SimpleNamespace(state = "running"),
            get_job_metadata = lambda _key: SimpleNamespace(
                completed_baseline_bytes = 130,
            ),
        ),
    )

    result = asyncio.run(
        downloads.get_gguf_download_progress_response(
            "Org/Model-GGUF",
            variant = "Q4_K_M",
            expected_bytes = 130,
        )
    )

    assert result["complete_on_disk"] is True
    assert result["completed_bytes"] == 130
    assert result["downloaded_bytes"] == 130
    assert result["expected_bytes"] == 130
    assert result["progress"] == 1.0


def test_gguf_progress_scoped_hashes_exclude_sibling_quant(monkeypatch, tmp_path):
    # The "instant ~900 MB" bug: a sibling quant is fully cached when a different
    # variant starts. With this variant's hashes resolved, progress counts ONLY
    # its in-progress blob, never the sibling's finalized bytes in the shared
    # blobs/ dir.
    entry = tmp_path / "models--Org--Model-GGUF"
    blobs = entry / "blobs"
    blobs.mkdir(parents = True)
    (blobs / "siblinghash").write_bytes(b"z" * 900)  # other quant, complete
    (blobs / "mainhash.incomplete").write_bytes(b"x" * 5)  # this variant, started

    requirement = gguf_variants._GgufVariantRequirement(
        main_filenames = frozenset({"model-Q4_K_M.gguf"}),
        target_filenames = ("model-Q4_K_M.gguf",),
        main_hashes = frozenset({"mainhash"}),
        required_hashes = frozenset({"mainhash"}),
        companion_hashes = frozenset(),
        mmproj_filenames = frozenset(),
        mmproj_hashes = frozenset(),
        expected_files = (
            download_manifest.ExpectedFile(
                path = "model-Q4_K_M.gguf",
                size = 100,
                sha256 = "mainhash",
            ),
        ),
        main_size_bytes = 100,
        download_size_bytes = 100,
    )

    async def _run_inline(fn, *args, **kwargs):
        return fn(*args, **kwargs)

    monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline)
    monkeypatch.setattr(
        downloads.gguf_variants,
        "gguf_variant_requirements",
        lambda *_args, **_kwargs: requirement,
    )
    monkeypatch.setattr(
        snapshot_progress,
        "preferred_repo_cache_dirs",
        lambda *_args, **_kwargs: [entry],
    )
    monkeypatch.setattr(
        downloads,
        "_registry",
        SimpleNamespace(
            get_job = lambda _key: SimpleNamespace(state = "running"),
            get_job_metadata = lambda _key: SimpleNamespace(
                completed_baseline_bytes = 0,
            ),
        ),
    )

    result = asyncio.run(
        downloads.get_gguf_download_progress_response(
            "Org/Model-GGUF",
            variant = "Q4_K_M",
            expected_bytes = 100,
        )
    )

    assert result["completed_bytes"] == 0
    assert result["downloaded_bytes"] == 5


def test_gguf_progress_unknown_hashes_does_not_count_foreign_blobs(monkeypatch, tmp_path):
    # With a variant's hashes unresolved (metadata flaked, no manifest), the
    # shared blobs/ dir's FINALIZED blobs must NOT be counted wholesale: a cached
    # sibling quant (``siblinghash``) alongside is the "instant ~900 MB" bug.
    # With no .incomplete present, downloaded must be 0.
    entry = tmp_path / "models--Org--Model-GGUF"
    snap = entry / "snapshots" / "rev0"
    blobs = entry / "blobs"
    snap.mkdir(parents = True)
    blobs.mkdir(parents = True)
    (blobs / "mainhash").write_bytes(b"x" * 100)
    (blobs / "mmprojhash").write_bytes(b"y" * 30)
    (blobs / "siblinghash").write_bytes(b"z" * 900)

    async def _run_inline(fn, *args, **kwargs):
        return fn(*args, **kwargs)

    monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline)
    monkeypatch.setattr(
        downloads.gguf_variants,
        "gguf_variant_requirements",
        lambda *_args, **_kwargs: None,
    )
    monkeypatch.setattr(
        downloads.gguf_variants,
        "gguf_variant_blob_hashes",
        lambda *_args, **_kwargs: frozenset(),
    )
    monkeypatch.setattr(
        snapshot_progress,
        "preferred_repo_cache_dirs",
        lambda *_args, **_kwargs: [entry],
    )
    monkeypatch.setattr(
        downloads,
        "_registry",
        SimpleNamespace(get_job = lambda _key: SimpleNamespace(state = "running")),
    )

    result = asyncio.run(
        downloads.get_gguf_download_progress_response(
            "Org/Model-GGUF",
            variant = "Q4_K_M",
            expected_bytes = 130,
        )
    )

    assert result["completed_bytes"] == 0
    assert result["downloaded_bytes"] == 0
    assert result["complete_on_disk"] is False


def test_gguf_progress_unknown_hashes_drops_unscoped_incomplete_blob(monkeypatch, tmp_path):
    # With hashes unresolved, an .incomplete in the shared blobs/ dir can't be
    # attributed to this variant (it may be a concurrent sibling's active write),
    # so it is dropped, mirroring the finalized-blob guard. In production the
    # worker writes the manifest before any .incomplete exists, so hashes resolve
    # via the manifest backstop and this window never suppresses real progress.
    entry = tmp_path / "models--Org--Model-GGUF"
    blobs = entry / "blobs"
    blobs.mkdir(parents = True)
    (blobs / "activehash.incomplete").write_bytes(b"x" * 50)  # unattributable
    (blobs / "siblinghash").write_bytes(b"z" * 900)  # finalized sibling

    async def _run_inline(fn, *args, **kwargs):
        return fn(*args, **kwargs)

    monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline)
    monkeypatch.setattr(
        downloads.gguf_variants,
        "gguf_variant_requirements",
        lambda *_args, **_kwargs: None,
    )
    monkeypatch.setattr(
        downloads.gguf_variants,
        "gguf_variant_blob_hashes",
        lambda *_args, **_kwargs: frozenset(),
    )
    monkeypatch.setattr(
        snapshot_progress,
        "preferred_repo_cache_dirs",
        lambda *_args, **_kwargs: [entry],
    )
    monkeypatch.setattr(
        downloads,
        "_registry",
        SimpleNamespace(get_job = lambda _key: SimpleNamespace(state = "running")),
    )

    result = asyncio.run(
        downloads.get_gguf_download_progress_response(
            "Org/Model-GGUF",
            variant = "Q4_K_M",
            expected_bytes = 1000,
        )
    )

    assert result["downloaded_bytes"] == 0  # unscoped .incomplete not leaked
    assert result["completed_bytes"] == 0  # finalized sibling still ignored


def test_gguf_progress_unknown_hashes_no_backward_dip_when_variant_finalizes(monkeypatch, tmp_path):
    # Regression for the two-variant dip: with hashes unresolved, the first quant
    # finalizes while the sibling still writes its .incomplete. The sibling's
    # bytes used to leak into this numerator, dipping the bar ~99% -> ~78% for
    # one poll. The unscoped .incomplete must be dropped so the reading stays 0.
    entry = tmp_path / "models--unsloth--SmolLM2-360M-Instruct-GGUF"
    blobs = entry / "blobs"
    snap = entry / "snapshots" / "rev0"
    blobs.mkdir(parents = True)
    snap.mkdir(parents = True)
    own_total = 218_673_760  # Q2_K finished blob size (denominator)
    sibling_total = 234_686_560  # Q3_K_M total

    def _sparse_file(path: Path, size: int) -> None:
        with path.open("wb") as handle:
            handle.truncate(size)

    own_finalized = blobs / "q2hash"
    _sparse_file(own_finalized, own_total)
    # ~72.7% of the sibling => sibling_partial / own_total == 0.78 pre-fix.
    sibling_incomplete = blobs / "q3hash.incomplete"
    _sparse_file(sibling_incomplete, int(sibling_total * 0.727))

    async def _run_inline(fn, *args, **kwargs):
        return fn(*args, **kwargs)

    monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline)
    monkeypatch.setattr(
        downloads.gguf_variants,
        "gguf_variant_requirements",
        lambda *_args, **_kwargs: None,
    )
    monkeypatch.setattr(
        downloads.gguf_variants,
        "gguf_variant_blob_hashes",
        lambda *_args, **_kwargs: frozenset(),
    )
    monkeypatch.setattr(
        snapshot_progress,
        "preferred_repo_cache_dirs",
        lambda *_args, **_kwargs: [entry],
    )
    monkeypatch.setattr(
        downloads,
        "_registry",
        SimpleNamespace(get_job = lambda _key: SimpleNamespace(state = "running")),
    )

    result = asyncio.run(
        downloads.get_gguf_download_progress_response(
            "unsloth/SmolLM2-360M-Instruct-GGUF",
            variant = "Q2_K",
            expected_bytes = own_total,
        )
    )

    assert result["downloaded_bytes"] == 0  # sibling .incomplete did not leak
    assert result["progress"] == 0  # no ~0.78 backward dip


def test_hf_cache_model_file_probe_is_bounded(monkeypatch, tmp_path):
    snapshot = tmp_path / "snapshot"
    snapshot.mkdir()
    first = tmp_path / "README.md"
    second = tmp_path / "notes.txt"
    model = tmp_path / "model.safetensors"
    first.write_text("readme", encoding = "utf-8")
    second.write_text("notes", encoding = "utf-8")
    model.write_bytes(b"weights")
    entries = [first, second, model]

    monkeypatch.setattr(model_common.Path, "rglob", lambda _self, _pattern: iter(entries))
    monkeypatch.setattr(model_common, "_HF_CACHE_MODEL_FILE_PROBE_LIMIT", 2)

    bounded = model_common._iter_hf_cache_model_files(snapshot)

    assert bounded == [first, second]

    monkeypatch.setattr(model_common, "_HF_CACHE_MODEL_FILE_PROBE_LIMIT", 3)

    unbounded = model_common._iter_hf_cache_model_files(snapshot)

    assert unbounded == [first, second, model]


def test_download_state_lookup_is_repo_case_insensitive(monkeypatch, tmp_path):
    monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path)

    assert download_manifest.write_manifest(
        "model",
        "Owner/Repo",
        None,
        [download_manifest.ExpectedFile(path = "config.json", size = 12)],
    )
    assert download_manifest.write_cancel_marker("model", "Owner/Repo", "Q4_K_M", "http")

    manifest = download_manifest.read_manifest("model", "owner/repo", None)

    assert manifest is not None
    assert manifest.repo_id == "Owner/Repo"
    assert manifest.expected_files[0].path == "config.json"
    assert download_manifest.has_cancel_marker("model", "owner/repo", "Q4_K_M")
    assert (
        download_manifest.read_cancel_marker_transport(
            "model",
            "owner/repo",
            "Q4_K_M",
        )
        == "http"
    )
    assert [
        variant
        for variant, _path in download_manifest.iter_variant_markers(
            "model",
            "owner/repo",
        )
    ] == ["Q4_K_M"]
    assert download_manifest.purge_all_state_for_repo("model", "owner/repo") == 2
    assert download_manifest.read_manifest("model", "owner/repo", None) is None


def test_hf_cache_scan_fallback_row_uses_local_model_info_alias(monkeypatch, tmp_path):
    cache_dir = tmp_path / "hub"
    repo_dir = cache_dir / "models--Org--Broken"
    blobs_dir = repo_dir / "blobs"
    blobs_dir.mkdir(parents = True)
    (blobs_dir / "blob").write_bytes(b"content")
    monkeypatch.setattr(local_inventory, "_classify_local_path", lambda *_args, **_kwargs: [])
    monkeypatch.setattr(
        local_inventory.hf_cache_scan,
        "is_snapshot_partial",
        lambda *_args, **_kwargs: False,
    )
    monkeypatch.setattr(
        local_inventory.hf_cache_scan,
        "is_gguf_repo_partial",
        lambda *_args, **_kwargs: False,
    )
    monkeypatch.setattr(
        local_inventory.hf_cache_scan,
        "resolve_hf_cache_realpath",
        lambda *_args, **_kwargs: None,
    )

    rows = local_inventory._scan_hf_cache(cache_dir)

    assert len(rows) == 1
    assert rows[0].model_id == "Org/Broken"
    assert rows[0].source == "hf_cache"
    assert rows[0].model_format == "unknown"


def test_hf_cache_scan_uses_gguf_partial_row_for_variant_state(monkeypatch, tmp_path):
    monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
    cache_dir = tmp_path / "hub"
    repo_dir = cache_dir / "models--Org--PartialGguf"
    blobs_dir = repo_dir / "blobs"
    blobs_dir.mkdir(parents = True)
    (blobs_dir / "partial").write_bytes(b"content")
    assert download_manifest.write_manifest(
        "model",
        "Org/PartialGguf",
        "Q4_K_M",
        [download_manifest.ExpectedFile(path = "model-Q4_K_M.gguf", size = 8192)],
        "http",
    )
    assert download_manifest.write_cancel_marker("model", "Org/PartialGguf", "Q4_K_M", "http")
    monkeypatch.setattr(local_inventory, "_classify_local_path", lambda *_args, **_kwargs: [])
    monkeypatch.setattr(
        local_inventory.hf_cache_scan,
        "is_snapshot_partial",
        lambda *_args, **_kwargs: False,
    )
    monkeypatch.setattr(
        local_inventory.hf_cache_scan,
        "is_gguf_repo_partial",
        lambda *_args, **_kwargs: True,
    )
    monkeypatch.setattr(
        local_inventory.hf_cache_scan,
        "resolve_hf_cache_realpath",
        lambda *_args, **_kwargs: None,
    )

    rows = local_inventory._scan_hf_cache(cache_dir)

    assert len(rows) == 1
    assert rows[0].model_id == "Org/PartialGguf"
    assert rows[0].source == "hf_cache"
    assert rows[0].model_format == "gguf"
    assert rows[0].partial is True
    assert rows[0].size_bytes == 8192
    assert rows[0].capabilities.requires_variant is True


def test_model_download_job_helpers_preserve_idle_shape():
    key = downloads._download_job_key("Org/Model", None)
    status = downloads._job_status(key)

    assert key == "org/model::"
    assert status.state == "idle"
    assert status.error is None


def test_gguf_repo_partial_treats_completed_disk_variant_as_clean(monkeypatch, tmp_path):
    monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
    snapshot = tmp_path / "cache" / "models--Org--Repo" / "snapshots" / "abc"
    snapshot.mkdir(parents = True)
    (snapshot / "model-Q8_0.gguf").write_bytes(b"complete")
    assert download_manifest.write_cancel_marker("model", "Org/Repo", "Q4_K_M", "xet")
    monkeypatch.setattr(
        inventory_scan,
        "resolve_snapshot_dir_for_scan",
        lambda *_args: snapshot,
    )

    assert inventory_scan.is_gguf_repo_partial("Org/Repo", snapshot.parents[1]) is False


def test_gguf_repo_partial_flags_vision_variant_missing_mmproj(monkeypatch, tmp_path):
    monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
    snapshot = tmp_path / "cache" / "models--Org--Vision" / "snapshots" / "abc"
    snapshot.mkdir(parents = True)
    (snapshot / "model-Q4_K_M.gguf").write_bytes(b"complete-weight")
    assert download_manifest.write_manifest(
        "model",
        "Org/Vision",
        "Q4_K_M",
        [
            download_manifest.ExpectedFile(path = "model-Q4_K_M.gguf", size = 15),
            download_manifest.ExpectedFile(path = "mmproj-F16.gguf", size = 8),
        ],
        "http",
    )
    monkeypatch.setattr(
        inventory_scan,
        "resolve_snapshot_dir_for_scan",
        lambda *_args: snapshot,
    )

    assert inventory_scan.is_gguf_repo_partial("Org/Vision") is True


def test_cancel_worker_leaves_exited_process_to_watcher():
    calls: list = []

    class _Registry:
        def get_process(self, _key):
            return SimpleNamespace(poll = lambda: 1)

        def get_job(self, _key):
            return SimpleNamespace(state = "running")

        def mark_pending_cancel(self, key, generation):
            calls.append(("pending", key, generation))
            return True

        def request_cancel(self, key, proc, generation):
            calls.append(("request", key, generation))
            return True

        def cancel_requested(self, _key):
            return False

    state = download_lifecycle.cancel_worker(
        _Registry(),
        "org/model::",
        generation = 3,
        label = "Org/Model",
        logger = SimpleNamespace(warning = lambda *_a, **_k: None),
    )

    assert state == "running"
    assert calls == []


def test_completed_gguf_split_variant_requires_all_shards(tmp_path):
    snapshot = tmp_path / "snapshot"
    snapshot.mkdir()
    first = snapshot / "model-Q8_0-00001-of-00002.gguf"
    second = snapshot / "model-Q8_0-00002-of-00002.gguf"
    first.write_bytes(b"first")

    assert "Q8_0" not in inventory_scan._completed_gguf_variants(snapshot)

    second.write_bytes(b"second")
    assert "Q8_0" in inventory_scan._completed_gguf_variants(snapshot)


def test_variant_partial_accepts_variant_filtered_legacy_hashes(monkeypatch, tmp_path):
    monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path)

    assert inventory_scan.is_variant_partial(
        "Org/Repo",
        "Q4_K_M",
        incomplete_blob_hashes = {"main-q4", "main-q8"},
        variant_blob_hashes = frozenset({"main-q4"}),
    )
    assert not inventory_scan.is_variant_partial(
        "Org/Repo",
        "Q5_K_M",
        incomplete_blob_hashes = {"main-q4"},
        variant_blob_hashes = frozenset({"main-q5"}),
    )


def test_variant_partial_accepts_completed_variant_in_non_latest_snapshot(monkeypatch, tmp_path):
    """A verified GGUF update can prune an older snapshot and make that old
    directory the newest by mtime. The variant is still complete when another
    snapshot satisfies its manifest."""
    monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
    repo_dir = tmp_path / "cache" / "models--Org--Repo"
    old_snapshot = repo_dir / "snapshots" / "old"
    new_snapshot = repo_dir / "snapshots" / "new"
    old_snapshot.mkdir(parents = True)
    new_snapshot.mkdir(parents = True)
    (old_snapshot / "model-Q8_0.gguf").write_bytes(b"sibling")
    (new_snapshot / "model-Q4_K_M.gguf").write_bytes(b"new")
    assert download_manifest.write_manifest(
        "model",
        "Org/Repo",
        "Q4_K_M",
        [download_manifest.ExpectedFile(path = "model-Q4_K_M.gguf", size = 3)],
        "http",
    )

    assert not inventory_scan.is_variant_partial(
        "Org/Repo",
        "Q4_K_M",
        snapshot_dir = old_snapshot,
        repo_cache_dir = repo_dir,
    )


def test_gguf_variants_partial_marker_overrides_size_only_downloaded(monkeypatch, tmp_path):
    async def _run_inline(fn, *args, **kwargs):
        return fn(*args, **kwargs)

    monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
    monkeypatch.setattr(gguf_variants.asyncio, "to_thread", _run_inline)
    assert download_manifest.write_cancel_marker("model", "Org/PartialRepo", "Q4_K_M", "http")
    snapshot = tmp_path / "cache" / "models--Org--PartialRepo" / "snapshots" / "rev0"
    snapshot.mkdir(parents = True)
    (snapshot / "model-Q4_K_M.gguf").write_bytes(b"x" * 100)

    monkeypatch.setattr(
        gguf_variants,
        "list_gguf_variants",
        lambda *_args, **_kwargs: (
            [
                SimpleNamespace(
                    filename = "model-Q4_K_M.gguf",
                    quant = "Q4_K_M",
                    display_label = None,
                    size_bytes = 100,
                )
            ],
            False,
            None,
        ),
    )
    monkeypatch.setattr(
        gguf_variants,
        "iter_hf_cache_snapshots",
        lambda _repo_id: [snapshot],
    )
    monkeypatch.setattr(
        gguf_variants,
        "_gguf_all_variant_requirements",
        lambda *_args, **_kwargs: {},
    )
    monkeypatch.setattr(
        gguf_variants.download_registry,
        "incomplete_blob_hashes",
        lambda *_args, **_kwargs: set(),
    )

    result = asyncio.run(gguf_variants.get_gguf_variants_response("Org/PartialRepo"))

    assert result.variants[0].downloaded is False
    assert result.variants[0].partial is True


def test_download_registry_repo_keys_are_case_insensitive():
    registry = download_registry.DownloadRegistry()

    claimed, state = registry.claim(
        "Org/Repo::Q8_0",
        download_registry.TRANSPORT_HTTP,
        repo_type = "model",
        repo_id = "Org/Repo",
        variant = "Q8_0",
    )
    # The same variant under a different-cased repo id resolves to the same
    # job, so the second claim attaches to the running one instead of starting
    # a duplicate.
    duplicate_claimed, duplicate_state = registry.claim(
        "org/repo::Q8_0",
        download_registry.TRANSPORT_HTTP,
        repo_type = "model",
        repo_id = "org/repo",
        variant = "Q8_0",
    )

    assert claimed is True
    assert state == "running"
    assert duplicate_claimed is False
    assert duplicate_state == "running"
    assert registry.active_jobs("ORG/REPO") == {"org/repo::Q8_0": "running"}


def test_download_registry_allows_disjoint_gguf_variant_downloads():
    registry = download_registry.DownloadRegistry()

    claimed, state = registry.claim(
        "Org/Repo::Q8_0",
        download_registry.TRANSPORT_HTTP,
        repo_type = "model",
        repo_id = "Org/Repo",
        variant = "Q8_0",
        blob_hashes = frozenset({"q8-main"}),
        progress_blob_hashes = frozenset({"q8-main", "shared-mmproj"}),
    )
    second_claimed, second_state = registry.claim(
        "Org/Repo::Q4_K_M",
        download_registry.TRANSPORT_HTTP,
        repo_type = "model",
        repo_id = "Org/Repo",
        variant = "Q4_K_M",
        blob_hashes = frozenset({"q4-main"}),
        progress_blob_hashes = frozenset({"q4-main", "shared-mmproj"}),
    )

    assert claimed is True
    assert state == "running"
    assert second_claimed is True
    assert second_state == "running"
    assert registry.active_jobs("org/repo") == {
        "org/repo::Q8_0": "running",
        "org/repo::Q4_K_M": "running",
    }


def test_download_registry_allows_overlapping_same_transport_variant_downloads():
    # Two variants sharing one mmproj blob still download together on one
    # transport: huggingface_hub's per-blob lock serializes the shared write and
    # prepare_cache_for_transport never purges a blob a peer is writing.
    registry = download_registry.DownloadRegistry()

    claimed, state = registry.claim(
        "Org/Repo::Q8_0",
        download_registry.TRANSPORT_HTTP,
        repo_type = "model",
        repo_id = "Org/Repo",
        variant = "Q8_0",
        blob_hashes = frozenset({"q8-main"}),
        progress_blob_hashes = frozenset({"q8-main", "shared-mmproj"}),
    )
    second_claimed, second_state = registry.claim(
        "Org/Repo::Q4_K_M",
        download_registry.TRANSPORT_HTTP,
        repo_type = "model",
        repo_id = "Org/Repo",
        variant = "Q4_K_M",
        blob_hashes = frozenset({"q4-main"}),
        progress_blob_hashes = frozenset({"q4-main", "shared-mmproj"}),
    )

    assert claimed is True
    assert state == "running"
    assert second_claimed is True
    assert second_state == "running"


def test_download_registry_variant_delete_does_not_block_sibling_download():
    # Deleting one quant's partial must be allowed while a different quant of the
    # same repo is downloading, and must protect every blob the live sibling is
    # writing (including a shared mmproj companion).
    registry = download_registry.DownloadRegistry()
    registry.claim(
        "Org/Repo::Q8_0",
        download_registry.TRANSPORT_HTTP,
        repo_type = "model",
        repo_id = "Org/Repo",
        variant = "Q8_0",
        blob_hashes = frozenset({"q8-main"}),
        progress_blob_hashes = frozenset({"q8-main", "shared-mmproj"}),
    )

    # A sibling variant delete is allowed; deleting the in-flight variant is not.
    assert registry.begin_delete("Org/Repo", "Q4_K_M") is True
    assert registry.begin_delete("Org/Repo", "Q8_0") is False
    # A whole-repo delete still waits for every active download.
    assert registry.begin_delete("Org/Repo") is False

    # The live sibling is detected so the delete keeps the shared companion.
    assert registry.has_active_peer_variant("Org/Repo", "Q4_K_M") is True
    assert registry.has_active_peer_variant("Org/Repo", "Q8_0") is False

    # While Q4_K_M is being deleted, re-downloading it is blocked but an
    # untouched third variant may still start.
    blocked, blocked_state = registry.claim(
        "Org/Repo::Q4_K_M",
        download_registry.TRANSPORT_HTTP,
        repo_type = "model",
        repo_id = "Org/Repo",
        variant = "Q4_K_M",
    )
    assert blocked is False
    assert blocked_state == "deleting"
    started, started_state = registry.claim(
        "Org/Repo::Q5_K_M",
        download_registry.TRANSPORT_HTTP,
        repo_type = "model",
        repo_id = "Org/Repo",
        variant = "Q5_K_M",
    )
    assert started is True
    assert started_state == "running"

    registry.end_delete("Org/Repo", "Q4_K_M")
    assert registry.begin_delete("Org/Repo", "Q4_K_M") is True


def test_partial_gguf_reconstruction_dedupes_variant_casing(monkeypatch):
    # The manifest keeps original casing while the marker is lowercased; offline
    # reconstruction must collapse them to ONE entry (manifest's casing), not two.
    monkeypatch.setattr(
        download_manifest,
        "iter_variant_manifests",
        lambda _repo_type, _repo_id: iter([("Q4_K_M", Path("manifest.json"))]),
    )
    monkeypatch.setattr(
        download_manifest,
        "iter_variant_markers",
        lambda _repo_type, _repo_id: iter([("q4_k_m", Path("marker.json"))]),
    )
    monkeypatch.setattr(download_manifest, "read_manifest", lambda *_a, **_k: None)

    result = gguf.list_partial_gguf_variants_from_state("Org/Repo")

    assert result is not None
    variants, _has_vision = result
    assert [variant.quant for variant in variants] == ["Q4_K_M"]


def test_download_registry_serializes_cross_transport_variant_downloads():
    # An HTTP append-resume and an XET rewrite of the same shared blob would
    # corrupt each other, so different-transport variants are serialized.
    registry = download_registry.DownloadRegistry()

    claimed, state = registry.claim(
        "Org/Repo::Q8_0",
        download_registry.TRANSPORT_HTTP,
        repo_type = "model",
        repo_id = "Org/Repo",
        variant = "Q8_0",
        blob_hashes = frozenset({"q8-main"}),
        progress_blob_hashes = frozenset({"q8-main", "shared-mmproj"}),
    )
    second_claimed, second_state = registry.claim(
        "Org/Repo::Q4_K_M",
        download_registry.TRANSPORT_XET,
        repo_type = "model",
        repo_id = "Org/Repo",
        variant = "Q4_K_M",
        blob_hashes = frozenset({"q4-main"}),
        progress_blob_hashes = frozenset({"q4-main", "shared-mmproj"}),
    )

    assert claimed is True
    assert state == "running"
    assert second_claimed is False
    assert second_state == "running"


def test_download_registry_allows_unknown_hash_gguf_variant_downloads():
    # Resolved blob hashes are NOT required to run two same-transport variants
    # concurrently: on-disk safety comes from each worker purging only its own
    # main-quant blobs plus huggingface_hub's per-etag lock. Requiring them here
    # used to reject the second variant whenever a metadata fetch flaked.
    registry = download_registry.DownloadRegistry()

    claimed, state = registry.claim(
        "Org/Repo::Q8_0",
        download_registry.TRANSPORT_HTTP,
        repo_type = "model",
        repo_id = "Org/Repo",
        variant = "Q8_0",
    )
    second_claimed, second_state = registry.claim(
        "Org/Repo::Q4_K_M",
        download_registry.TRANSPORT_HTTP,
        repo_type = "model",
        repo_id = "Org/Repo",
        variant = "Q4_K_M",
        blob_hashes = frozenset({"q4-main"}),
        progress_blob_hashes = frozenset({"q4-main", "shared-mmproj"}),
    )

    assert claimed is True
    assert state == "running"
    assert second_claimed is True
    assert second_state == "running"
    assert registry.active_jobs("org/repo") == {
        "org/repo::Q8_0": "running",
        "org/repo::Q4_K_M": "running",
    }


def test_finalize_worker_exit_never_kills_a_healthy_worker(monkeypatch, tmp_path):
    # finalize_worker_exit relies solely on the worker's exit code and never kills
    # a live process: huggingface_hub already bounds reads with timeouts, so a
    # liveness kill could only false-cancel a healthy download.
    import inspect
    import io
    import logging

    monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")

    class _Proc:
        pid = 4242

        def __init__(self):
            self.killed = False
            self.stderr = io.BytesIO(b"")

        def poll(self):
            return 0

        def wait(self, timeout = None):
            return 0

        def kill(self):
            self.killed = True

    registry = download_registry.DownloadRegistry()
    proc = _Proc()
    key = "Org/Repo::Q4_K_M"
    registry.claim(
        key,
        download_registry.TRANSPORT_HTTP,
        repo_type = "model",
        repo_id = "Org/Repo",
        variant = "Q4_K_M",
    )
    registry.register_process(key, proc)

    download_lifecycle.finalize_worker_exit(
        registry,
        key,
        proc,
        hf_token = None,
        label = "Org/Repo [Q4_K_M]",
        log_prefix = "Download",
        logger = logging.getLogger("test"),
        repo_type = "model",
        repo_id = "Org/Repo",
        transport = "http",
    )

    assert proc.killed is False
    assert registry.get_job(key).state == "complete"
    # The stall-watchdog knob is gone entirely; no caller may re-enable it.
    assert (
        "enable_stall_watchdog"
        not in inspect.signature(download_lifecycle.finalize_worker_exit).parameters
    )


def test_prepare_cache_for_transport_purges_only_requested_hashes(monkeypatch, tmp_path):
    root = tmp_path / "hub"
    blobs = root / "models--Org--Repo" / "blobs"
    blobs.mkdir(parents = True)
    (blobs / "variant-main.incomplete").write_bytes(b"x")
    (blobs / "shared-mmproj.incomplete").write_bytes(b"y")
    monkeypatch.setattr(download_registry, "hf_cache_root", lambda create = False: root)

    purged = download_registry.prepare_cache_for_transport(
        "model",
        "Org/Repo",
        download_registry.TRANSPORT_XET,
        "Q4_K_M",
        frozenset({"variant-main"}),
    )

    assert purged == 1
    assert not (blobs / "variant-main.incomplete").exists()
    assert (blobs / "shared-mmproj.incomplete").exists()


def _vision_cache_root(monkeypatch, tmp_path):
    root = tmp_path / "hub"
    blobs = root / "models--Org--Vision" / "blobs"
    blobs.mkdir(parents = True)
    monkeypatch.setattr(download_registry, "hf_cache_root", lambda create = False: root)
    return blobs


def test_prepare_cache_for_transport_purges_cross_transport_companion(monkeypatch, tmp_path):
    blobs = _vision_cache_root(monkeypatch, tmp_path)
    companion = frozenset({"shared-mmproj"})

    # An interrupted XET download stamps the companion marker "xet" and leaves a
    # sparse partial. A later HTTP download of a different variant must purge it,
    # else the HTTP resumer appends to the sparse bytes and corrupts the blob.
    download_registry.prepare_cache_for_transport(
        "model",
        "Org/Vision",
        download_registry.TRANSPORT_XET,
        "Q4_K_M",
        only_blob_hashes = frozenset({"q4-main"}),
        companion_blob_hashes = companion,
    )
    (blobs / "shared-mmproj.incomplete").write_bytes(b"sparse")

    purged = download_registry.prepare_cache_for_transport(
        "model",
        "Org/Vision",
        download_registry.TRANSPORT_HTTP,
        "Q8_0",
        only_blob_hashes = frozenset({"q8-main"}),
        companion_blob_hashes = companion,
    )

    assert purged == 1
    assert not (blobs / "shared-mmproj.incomplete").exists()


def test_prepare_cache_for_transport_preserves_same_transport_companion(monkeypatch, tmp_path):
    blobs = _vision_cache_root(monkeypatch, tmp_path)
    companion = frozenset({"shared-mmproj"})

    download_registry.prepare_cache_for_transport(
        "model",
        "Org/Vision",
        download_registry.TRANSPORT_HTTP,
        "Q4_K_M",
        only_blob_hashes = frozenset({"q4-main"}),
        companion_blob_hashes = companion,
    )
    (blobs / "shared-mmproj.incomplete").write_bytes(b"resumable")

    purged = download_registry.prepare_cache_for_transport(
        "model",
        "Org/Vision",
        download_registry.TRANSPORT_HTTP,
        "Q4_K_M",
        only_blob_hashes = frozenset({"q4-main"}),
        companion_blob_hashes = companion,
    )

    assert purged == 0
    assert (blobs / "shared-mmproj.incomplete").exists()


def test_prepare_cache_for_transport_protects_peer_companion(monkeypatch, tmp_path):
    blobs = _vision_cache_root(monkeypatch, tmp_path)
    companion = frozenset({"shared-mmproj"})

    download_registry.prepare_cache_for_transport(
        "model",
        "Org/Vision",
        download_registry.TRANSPORT_XET,
        "Q4_K_M",
        only_blob_hashes = frozenset({"q4-main"}),
        companion_blob_hashes = companion,
    )
    (blobs / "shared-mmproj.incomplete").write_bytes(b"sparse")

    purged = download_registry.prepare_cache_for_transport(
        "model",
        "Org/Vision",
        download_registry.TRANSPORT_HTTP,
        "Q8_0",
        only_blob_hashes = frozenset({"q8-main"}),
        companion_blob_hashes = companion,
        protected_blob_hashes = companion,
    )

    assert purged == 0
    assert (blobs / "shared-mmproj.incomplete").exists()


def test_model_download_records_completed_baseline_for_new_gguf_variant(monkeypatch, tmp_path):
    async def _run_inline(fn, *args, **kwargs):
        return fn(*args, **kwargs)

    monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline)
    monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
    monkeypatch.setattr(
        downloads,
        "resolve_cached_repo_id_case",
        lambda repo_id, repo_type = "model": repo_id,
    )
    monkeypatch.setattr(
        downloads.gguf_variants,
        "gguf_variant_blob_hashes",
        lambda _repo, _variant, _token = None, include_companions = True, **_kwargs: (
            frozenset({"mainhash", "mmprojhash"}) if include_companions else frozenset({"mainhash"})
        ),
    )
    monkeypatch.setattr(
        downloads.download_registry,
        "completed_blob_bytes",
        lambda *_args, **_kwargs: 30,
    )

    class _Registry:
        claim_kwargs = None

        def claim(self, _key, _transport, **kwargs):
            self.claim_kwargs = kwargs
            return True, "running"

        def current_generation(self, _key):
            return 1

        def get_job(self, _key):
            return SimpleNamespace(state = "running")

        def register_process(self, _key, _proc):
            return False

        def peer_blob_hashes(self, _key):
            return frozenset()

    class _Proc:
        pid = 123
        stderr = None

        def poll(self):
            return None

        def kill(self):
            return None

        def wait(self, timeout = None):
            return 0

    registry = _Registry()
    monkeypatch.setattr(downloads, "_registry", registry)
    monkeypatch.setattr(downloads, "_spawn_download_worker", lambda *_args, **_kwargs: _Proc())

    asyncio.run(
        downloads.download_model_response(
            SimpleNamespace(repo_id = "Org/Model", gguf_variant = "Q4_K_M", use_xet = False)
        )
    )

    assert registry.claim_kwargs["blob_hashes"] == frozenset({"mainhash"})
    assert registry.claim_kwargs["progress_blob_hashes"] == frozenset({"mainhash", "mmprojhash"})
    assert registry.claim_kwargs["completed_baseline_bytes"] == 30


def test_gguf_model_download_skips_completed_baseline_for_variant_resume_state(
    monkeypatch, tmp_path
):
    async def _run_inline(fn, *args, **kwargs):
        return fn(*args, **kwargs)

    monkeypatch.setattr(downloads.asyncio, "to_thread", _run_inline)
    monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
    assert download_manifest.write_manifest(
        "model",
        "Org/Model",
        "Q4_K_M",
        [
            download_manifest.ExpectedFile(
                path = "model-Q4_K_M.gguf",
                size = 100,
                sha256 = "mainhash",
            )
        ],
        "http",
    )
    monkeypatch.setattr(
        downloads,
        "resolve_cached_repo_id_case",
        lambda repo_id, repo_type = "model": repo_id,
    )
    monkeypatch.setattr(
        downloads.gguf_variants,
        "gguf_variant_blob_hashes",
        lambda _repo, _variant, _token = None, include_companions = True, **_kwargs: (
            frozenset({"mainhash", "mmprojhash"}) if include_companions else frozenset({"mainhash"})
        ),
    )
    monkeypatch.setattr(
        downloads.download_registry,
        "completed_blob_bytes",
        lambda *_args, **_kwargs: 30,
    )

    class _Registry:
        claim_kwargs = None

        def claim(self, _key, _transport, **kwargs):
            self.claim_kwargs = kwargs
            return True, "running"

        def current_generation(self, _key):
            return 1

        def get_job(self, _key):
            return SimpleNamespace(state = "running")

        def register_process(self, _key, _proc):
            return False

        def peer_blob_hashes(self, _key):
            return frozenset()

    class _Proc:
        pid = 123
        stderr = None

        def poll(self):
            return None

        def kill(self):
            return None

        def wait(self, timeout = None):
            return 0

    registry = _Registry()
    monkeypatch.setattr(downloads, "_registry", registry)
    monkeypatch.setattr(downloads, "_spawn_download_worker", lambda *_args, **_kwargs: _Proc())

    asyncio.run(
        downloads.download_model_response(
            SimpleNamespace(repo_id = "Org/Model", gguf_variant = "Q4_K_M", use_xet = False)
        )
    )

    assert registry.claim_kwargs["completed_baseline_bytes"] == 0


def test_model_idle_status_uses_cancel_marker_after_restart(monkeypatch, tmp_path):
    monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path)
    monkeypatch.setattr(downloads, "_registry", download_registry.DownloadRegistry())
    assert download_manifest.write_cancel_marker("model", "Owner/Repo", "Q4_K_M", "http")

    status = asyncio.run(downloads.get_download_status_response("owner/repo", "Q4_K_M"))

    assert status.state == "cancelled"
    assert status.error is None


def test_shutdown_kills_all_workers_before_shared_deadline_reap(monkeypatch):
    events = []
    now = [100.0]

    class _Proc:
        def __init__(self, name):
            self.name = name

        def poll(self):
            return None

        def kill(self):
            events.append(("kill", self.name))

        def wait(self, timeout):
            events.append(("wait", self.name, timeout))
            now[0] += 7.0

    registry = download_registry.DownloadRegistry()
    proc_a = _Proc("a")
    proc_b = _Proc("b")
    registry.claim(
        "Org/A",
        download_registry.TRANSPORT_HTTP,
        repo_type = "dataset",
        repo_id = "Org/A",
    )
    registry.claim(
        "Org/B",
        download_registry.TRANSPORT_HTTP,
        repo_type = "dataset",
        repo_id = "Org/B",
    )
    assert registry.register_process("org/a", proc_a)
    assert registry.register_process("org/b", proc_b)
    monkeypatch.setattr(
        download_registry,
        "persist_cancel_marker",
        lambda *args, **kwargs: events.append(("marker", args[1])),
    )
    monkeypatch.setattr(download_registry.time, "monotonic", lambda: now[0])

    registry.terminate_all("dataset download")

    assert events == [
        ("kill", "a"),
        ("kill", "b"),
        ("wait", "a", 10.0),
        ("marker", "Org/A"),
        ("wait", "b", 3.0),
        ("marker", "Org/B"),
    ]


def test_shutdown_skips_marker_for_worker_that_exits_cleanly(monkeypatch):
    markers = []

    class _Proc:
        def __init__(self, final_rc):
            self._final_rc = final_rc
            self._exited = False

        def poll(self):
            return self._final_rc if self._exited else None

        def kill(self):
            pass

        def wait(self, timeout):
            self._exited = True

    registry = download_registry.DownloadRegistry()
    clean = _Proc(0)
    interrupted = _Proc(-9)
    registry.claim(
        "Org/Clean",
        download_registry.TRANSPORT_HTTP,
        repo_type = "dataset",
        repo_id = "Org/Clean",
    )
    registry.claim(
        "Org/Cut",
        download_registry.TRANSPORT_HTTP,
        repo_type = "dataset",
        repo_id = "Org/Cut",
    )
    assert registry.register_process("org/clean", clean)
    assert registry.register_process("org/cut", interrupted)
    monkeypatch.setattr(
        download_registry,
        "persist_cancel_marker",
        lambda *args, **kwargs: markers.append(args[1]),
    )

    registry.terminate_all("dataset download")

    assert markers == ["Org/Cut"]


def test_model_claim_register_cancel_uses_registry_marker_owner(monkeypatch):
    killed = []

    class _Registry:
        def claim(self, *_args, **_kwargs):
            return True, "running"

        def current_generation(self, _key):
            return 1

        def register_process(self, _key, _proc):
            return False

        def persist_cancel_for_key(self, *_args, **_kwargs):
            raise AssertionError("register_process owns pending-cancel markers")

        def get_job(self, _key):
            return SimpleNamespace(state = "cancelled", error = None)

    monkeypatch.setattr(downloads, "_registry", _Registry())
    monkeypatch.setattr(
        downloads,
        "resolve_cached_repo_id_case",
        lambda repo_id, **_kwargs: repo_id,
    )
    monkeypatch.setattr(
        downloads.download_registry,
        "download_transport_unavailable_reason",
        lambda _transport: None,
    )
    monkeypatch.setattr(
        downloads,
        "_spawn_download_worker",
        lambda *_args, **_kwargs: object(),
    )
    monkeypatch.setattr(
        downloads.download_lifecycle,
        "kill_and_reap_process",
        lambda proc, **_kwargs: killed.append(proc),
    )

    result = asyncio.run(
        downloads.download_model_response(
            SimpleNamespace(repo_id = "Org/Model", gguf_variant = None, use_xet = False)
        )
    )

    assert result["state"] == "cancelled"
    assert killed


def test_model_cancel_registered_worker_requests_and_kills(monkeypatch):
    events = []

    class _Proc:
        def poll(self):
            return None

        def kill(self):
            events.append(("kill",))

    class _Registry:
        def get_process(self, _key):
            return _Proc()

        def request_cancel(self, key, _proc, generation):
            events.append(("request", key, generation))
            return True

        def persist_cancel_for_key(self, *_args, **_kwargs):
            raise AssertionError(
                "cancel_worker must leave marker persistence to the exit watcher; "
                "an eager persist races a clean completion and strands a stale marker"
            )

        def get_job(self, _key):
            return SimpleNamespace(state = "running")

    monkeypatch.setattr(downloads, "_registry", _Registry())
    monkeypatch.setattr(
        downloads,
        "resolve_cached_repo_id_case",
        lambda repo_id, **_kwargs: repo_id,
    )

    result = asyncio.run(
        downloads.cancel_download_model_response(
            SimpleNamespace(repo_id = "Org/Model", gguf_variant = "Q4_K_M", generation = 7)
        )
    )

    assert result == {
        "job_key": downloads._download_job_key("Org/Model", "Q4_K_M"),
        "state": "cancelling",
    }
    assert events == [
        ("request", downloads._download_job_key("Org/Model", "Q4_K_M"), 7),
        ("kill",),
    ]


def test_model_download_watcher_invalidates_hf_cache_scan(monkeypatch):
    invalidated = []

    class _Registry:
        def claim(self, *_args, **_kwargs):
            return True, "running"

        def current_generation(self, _key):
            return 1

        def register_process(self, _key, _proc):
            return True

        def get_job(self, _key):
            return SimpleNamespace(state = "complete", error = None)

    class _ImmediateThread:
        def __init__(self, *, target, **_kwargs):
            self._target = target

        def start(self):
            self._target()

    monkeypatch.setattr(downloads, "_registry", _Registry())
    monkeypatch.setattr(
        downloads,
        "resolve_cached_repo_id_case",
        lambda repo_id, **_kwargs: repo_id,
    )
    monkeypatch.setattr(
        downloads.download_registry,
        "download_transport_unavailable_reason",
        lambda _transport: None,
    )
    monkeypatch.setattr(
        downloads.download_lifecycle,
        "finalize_worker_exit",
        lambda *_args, **_kwargs: None,
    )
    monkeypatch.setattr(
        downloads,
        "_spawn_download_worker",
        lambda *_args, **_kwargs: object(),
    )
    monkeypatch.setattr(downloads.download_lifecycle.threading, "Thread", _ImmediateThread)
    monkeypatch.setattr(
        downloads.hf_cache_scan,
        "invalidate_hf_cache_scans",
        lambda: invalidated.append(True),
    )

    async def _inline_to_thread(func, *args, **kwargs):
        return func(*args, **kwargs)

    monkeypatch.setattr(downloads.asyncio, "to_thread", _inline_to_thread)

    result = asyncio.run(
        downloads.download_model_response(
            SimpleNamespace(repo_id = "Org/Model", gguf_variant = None, use_xet = False)
        )
    )

    assert result["accepted"] is True
    assert invalidated == [True]


def test_two_concurrent_same_repo_variants_both_complete(monkeypatch, tmp_path):
    # End-to-end proof that two GGUF variants of ONE repo download concurrently
    # without cancelling each other, with real registry/finalize/subprocess/watch
    # threads exercising the claim gate, register, finalize funnel, and
    # classify_exit under true concurrency.
    import subprocess
    import time

    monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
    monkeypatch.setattr(
        downloads,
        "_registry",
        download_registry.DownloadRegistry(),
    )
    monkeypatch.setattr(
        downloads,
        "resolve_cached_repo_id_case",
        lambda repo_id, **_k: repo_id,
    )
    monkeypatch.setattr(
        downloads.download_registry,
        "download_transport_unavailable_reason",
        lambda _transport: None,
    )
    # Per-variant blob hashes (distinct main shard, shared mmproj companion).
    monkeypatch.setattr(
        downloads.gguf_variants,
        "gguf_variant_blob_hashes",
        lambda _repo, variant, _token = None, include_companions = True, **_k: (
            frozenset({f"{variant.lower()}-main", "shared-mmproj"})
            if include_companions
            else frozenset({f"{variant.lower()}-main"})
        ),
    )
    monkeypatch.setattr(
        downloads.download_registry,
        "completed_blob_bytes",
        lambda *_a, **_k: 0,
    )
    monkeypatch.setattr(
        downloads.hf_cache_scan,
        "invalidate_hf_cache_scans",
        lambda: None,
    )
    # Real subprocess that exits 0 immediately, with a stderr pipe to drain.
    spawned: list[subprocess.Popen] = []

    def _fake_spawn(*_args, **_kwargs):
        proc = subprocess.Popen(
            [sys.executable, "-c", "import sys; sys.exit(0)"],
            stderr = subprocess.PIPE,
        )
        spawned.append(proc)
        return proc

    monkeypatch.setattr(downloads, "_spawn_download_worker", _fake_spawn)

    async def _run_both():
        return await asyncio.gather(
            downloads.download_model_response(
                SimpleNamespace(
                    repo_id = "Org/Model",
                    gguf_variant = "Q4_K_M",
                    use_xet = False,
                )
            ),
            downloads.download_model_response(
                SimpleNamespace(
                    repo_id = "Org/Model",
                    gguf_variant = "Q8_0",
                    use_xet = False,
                )
            ),
        )

    results = asyncio.run(_run_both())
    assert all(r["accepted"] is True for r in results), results

    registry = downloads._registry
    key_q4 = downloads._download_job_key("Org/Model", "Q4_K_M")
    key_q8 = downloads._download_job_key("Org/Model", "Q8_0")
    deadline = time.monotonic() + 10.0
    while time.monotonic() < deadline:
        s4 = registry.get_job(key_q4).state
        s8 = registry.get_job(key_q8).state
        if s4 in download_registry.TERMINAL_STATES and s8 in download_registry.TERMINAL_STATES:
            break
        time.sleep(0.02)

    for p in spawned:
        try:
            p.wait(timeout = 5)
        except Exception:
            pass

    assert registry.get_job(key_q4).state == "complete"
    assert registry.get_job(key_q8).state == "complete"


def test_download_registry_factories_reuse_service_singletons():
    registry_module = downloads.download_registry
    before_count = len(registry_module._REGISTRIES)

    assert registry_module.get_models_registry() is downloads.registry
    assert registry_module.get_models_registry() is downloads.registry
    assert registry_module.get_datasets_registry() is dataset_downloads.registry
    assert registry_module.get_datasets_registry() is dataset_downloads.registry
    assert len(registry_module._REGISTRIES) == before_count


def test_hub_hf_token_header_uses_namespaced_header_only():
    assert get_hf_token("new-token") == "new-token"
    assert get_hf_token(None) is None


def test_scan_folder_rejects_credential_directories(tmp_path):
    sensitive_dir = tmp_path / ".ssh" / "models"
    sensitive_dir.mkdir(parents = True)

    with pytest.raises(ValueError, match = "Credential or configuration"):
        scan_folders.add_scan_folder(str(sensitive_dir))


def _build_variant_cache_repo(repo_dir, blob_specs, snapshot_links):
    """Build a HF cache repo dir with blobs + snapshot symlinks for the
    per-variant deletion path. blob_specs: {blob_name: bytes_payload};
    snapshot_links: list of (revision, filename, blob_name)."""
    blobs_dir = repo_dir / "blobs"
    blobs_dir.mkdir(parents = True)
    for blob_name, payload in blob_specs.items():
        (blobs_dir / blob_name).write_bytes(payload)

    files = []
    for revision, filename, blob_name in snapshot_links:
        snap_dir = repo_dir / "snapshots" / revision
        snap_dir.mkdir(parents = True, exist_ok = True)
        blob = blobs_dir / blob_name
        link = snap_dir / filename
        link.symlink_to(blob)
        files.append(
            SimpleNamespace(
                file_name = filename,
                file_path = str(link),
                blob_path = str(blob),
                size_on_disk = blob.stat().st_size,
            )
        )
    repo = SimpleNamespace(
        repo_id = "Org/Repo-GGUF",
        repo_type = "model",
        repo_path = repo_dir,
        revisions = [SimpleNamespace(commit_hash = "rev1", files = files)],
    )
    return repo


def _patch_variant_delete_side_effects(monkeypatch):
    monkeypatch.setattr(
        deletion.download_manifest,
        "purge_state",
        lambda *_args, **_kwargs: False,
    )


def test_snapshot_progress_filters_stale_blobs(monkeypatch, tmp_path):
    """Exclude superseded-revision blobs; count an in-progress blob only when its
    hash belongs to the target."""
    entry = tmp_path / "datasets--Org--Data"
    blobs = entry / "blobs"
    blobs.mkdir(parents = True)
    (blobs / "keep1").write_bytes(b"a" * 100)
    (blobs / "stale").write_bytes(b"b" * 500)
    (blobs / "keep2.incomplete").write_bytes(b"c" * 40)

    monkeypatch.setattr(
        snapshot_progress,
        "preferred_repo_cache_dirs",
        lambda _repo_type, _repo_id, force_active = False: [entry],
    )

    result = snapshot_progress.compute_snapshot_progress(
        repo_type = "dataset",
        repo_id = "Org/Data",
        job_key = "org/data",
        expected_bytes = 0,
        hf_token = None,
        registry = SimpleNamespace(
            get_job = lambda _key: SimpleNamespace(state = "running"),
        ),
        metadata_resolver = lambda _repo_id, _hf_token: (
            140,
            frozenset({"keep1", "keep2"}),
        ),
    )

    assert result["completed_bytes"] == 100
    assert result["downloaded_bytes"] == 140
    assert result["complete_on_disk"] is False
    assert result["expected_bytes"] == 140


def test_snapshot_progress_confirms_complete_only_with_verified_snapshot(monkeypatch, tmp_path):
    entry = tmp_path / "models--Org--Model"
    blobs = entry / "blobs"
    snap = entry / "snapshots" / "rev0"
    blobs.mkdir(parents = True)
    snap.mkdir(parents = True)
    (blobs / "keep1").write_bytes(b"a" * 100)
    (snap / "model.safetensors").write_bytes(b"a" * 100)

    monkeypatch.setattr(
        snapshot_progress,
        "preferred_repo_cache_dirs",
        lambda _repo_type, _repo_id, force_active = False: [entry],
    )
    monkeypatch.setattr(
        snapshot_progress.download_manifest,
        "has_cancel_marker",
        lambda *_args, **_kwargs: False,
    )
    monkeypatch.setattr(
        snapshot_progress.download_manifest,
        "read_manifest",
        lambda *_args, **_kwargs: SimpleNamespace(),
    )
    monkeypatch.setattr(
        snapshot_progress.download_manifest,
        "verify_against_disk",
        lambda *_args, **_kwargs: SimpleNamespace(ok = True),
    )

    result = snapshot_progress.compute_snapshot_progress(
        repo_type = "model",
        repo_id = "Org/Model",
        job_key = "org/model::",
        expected_bytes = 100,
        hf_token = None,
        registry = SimpleNamespace(
            get_job = lambda _key: SimpleNamespace(state = "idle"),
        ),
        metadata_resolver = lambda _repo_id, _hf_token: (
            100,
            frozenset({"keep1"}),
        ),
    )

    assert result["completed_bytes"] == 100
    assert result["complete_on_disk"] is True


def test_expected_files_from_snapshot_dir_records_relative_paths_and_sizes(tmp_path):
    snap = tmp_path / "snapshots" / "rev0"
    (snap / "nested").mkdir(parents = True)
    (snap / "model.safetensors").write_bytes(b"a" * 12)
    (snap / "nested" / "config.json").write_bytes(b"b" * 3)

    files = download_manifest.expected_files_from_snapshot_dir(snap)

    by_path = {f.path: f for f in files}
    assert by_path["model.safetensors"].size == 12
    assert by_path["nested/config.json"].size == 3
    assert all(f.sha256 is None for f in files)


def test_snapshot_progress_complete_with_manifest_synthesized_from_disk(monkeypatch, tmp_path):
    """A finished snapshot whose only manifest was synthesized from on-disk files
    still verifies as complete, so a refresh finalizes it instead of capping at
    99% and evicting it as gone."""
    entry = tmp_path / "models--Org--Model"
    blobs = entry / "blobs"
    snap = entry / "snapshots" / "rev0"
    blobs.mkdir(parents = True)
    snap.mkdir(parents = True)
    (blobs / "keep1").write_bytes(b"a" * 100)
    (snap / "model.safetensors").write_bytes(b"a" * 100)

    synthesized = download_manifest.expected_files_from_snapshot_dir(snap)
    manifest = download_manifest.Manifest(
        repo_type = "model",
        repo_id = "Org/Model",
        variant = None,
        started_at = "",
        expected_files = tuple(synthesized),
    )

    monkeypatch.setattr(
        snapshot_progress,
        "preferred_repo_cache_dirs",
        lambda _repo_type, _repo_id, force_active = False: [entry],
    )
    monkeypatch.setattr(
        snapshot_progress.download_manifest,
        "has_cancel_marker",
        lambda *_args, **_kwargs: False,
    )
    monkeypatch.setattr(
        snapshot_progress.download_manifest,
        "read_manifest",
        lambda *_args, **_kwargs: manifest,
    )

    result = snapshot_progress.compute_snapshot_progress(
        repo_type = "model",
        repo_id = "Org/Model",
        job_key = "org/model::",
        expected_bytes = 100,
        hf_token = None,
        registry = SimpleNamespace(
            get_job = lambda _key: SimpleNamespace(state = "idle"),
        ),
        metadata_resolver = lambda _repo_id, _hf_token: (
            100,
            frozenset({"keep1"}),
        ),
    )

    assert result["complete_on_disk"] is True
    assert result["progress"] == 1.0


def test_delete_variant_keeps_blob_shared_with_other_snapshot(monkeypatch, tmp_path):
    """A blob still referenced by a non-target snapshot symlink survives so that
    symlink doesn't dangle (which the scanner reports as partial)."""
    repo_dir = tmp_path / "models--Org--Repo-GGUF"
    repo = _build_variant_cache_repo(
        repo_dir,
        blob_specs = {"sharedblob": b"x" * 200, "q8blob": b"y" * 300},
        snapshot_links = [
            ("rev1", "model-Q4_K_M.gguf", "sharedblob"),
            ("rev1", "model-Q8_0.gguf", "q8blob"),
            # An unrelated file that happens to share Q4's blob content.
            ("rev1", "extra-copy.gguf", "sharedblob"),
        ],
    )
    monkeypatch.setattr(
        deletion.cache_inventory,
        "all_hf_cache_scans",
        lambda: [SimpleNamespace(repos = [repo])],
    )
    _patch_variant_delete_side_effects(monkeypatch)

    result = deletion._delete_cached_model_blocking("Org/Repo-GGUF", "Q4_K_M", None)

    assert result["status"] == "deleted"
    # Q4 snapshot link gone, but its blob survives (extra-copy still links it).
    assert not (repo_dir / "snapshots" / "rev1" / "model-Q4_K_M.gguf").exists()
    assert (repo_dir / "blobs" / "sharedblob").exists()
    extra = repo_dir / "snapshots" / "rev1" / "extra-copy.gguf"
    assert extra.is_symlink() and extra.exists()  # not dangling


def test_delete_variant_unlinks_unshared_blob(monkeypatch, tmp_path):
    repo_dir = tmp_path / "models--Org--Repo-GGUF"
    repo = _build_variant_cache_repo(
        repo_dir,
        blob_specs = {"q4blob": b"x" * 200, "q8blob": b"y" * 300},
        snapshot_links = [
            ("rev1", "model-Q4_K_M.gguf", "q4blob"),
            ("rev1", "model-Q8_0.gguf", "q8blob"),
        ],
    )
    monkeypatch.setattr(
        deletion.cache_inventory,
        "all_hf_cache_scans",
        lambda: [SimpleNamespace(repos = [repo])],
    )
    _patch_variant_delete_side_effects(monkeypatch)

    result = deletion._delete_cached_model_blocking("Org/Repo-GGUF", "Q4_K_M", None)

    assert result["status"] == "deleted"
    assert not (repo_dir / "blobs" / "q4blob").exists()
    # Untouched sibling variant remains fully intact.
    assert (repo_dir / "blobs" / "q8blob").exists()
    q8 = repo_dir / "snapshots" / "rev1" / "model-Q8_0.gguf"
    assert q8.is_symlink() and q8.exists()


def test_delete_variant_surfaces_locked_file_as_conflict(monkeypatch, tmp_path):
    """A blob unlink that fails (e.g. a Windows file lock on a loaded model)
    must raise a clear 409, not report a misleading success."""
    repo_dir = tmp_path / "models--Org--Repo-GGUF"
    repo = _build_variant_cache_repo(
        repo_dir,
        blob_specs = {"lockedblob": b"x" * 200},
        snapshot_links = [("rev1", "model-Q4_K_M.gguf", "lockedblob")],
    )
    monkeypatch.setattr(
        deletion.cache_inventory,
        "all_hf_cache_scans",
        lambda: [SimpleNamespace(repos = [repo])],
    )
    _patch_variant_delete_side_effects(monkeypatch)

    real_unlink = Path.unlink

    def fake_unlink(self, *args, **kwargs):
        if self.name == "lockedblob":
            raise PermissionError("file in use")
        return real_unlink(self, *args, **kwargs)

    monkeypatch.setattr(Path, "unlink", fake_unlink)

    with pytest.raises(HTTPException) as exc_info:
        deletion._delete_cached_model_blocking("Org/Repo-GGUF", "Q4_K_M", None)

    assert exc_info.value.status_code == 409


def test_download_snapshot_writes_manifest_for_xet(monkeypatch, tmp_path):
    written = []
    verified = []

    monkeypatch.setattr(
        hf_download,
        "_model_info_with_retry",
        lambda *_args, **_kwargs: SimpleNamespace(
            siblings = [SimpleNamespace(rfilename = "config.json", size = 12)]
        ),
    )
    monkeypatch.setattr(
        hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args)
    )
    monkeypatch.setattr(
        download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0
    )
    monkeypatch.setattr(download_manifest, "clear_cancel_marker", lambda *_args: None)
    monkeypatch.setattr(
        download_manifest, "write_manifest", lambda *args: written.append(args) or True
    )
    monkeypatch.setitem(
        sys.modules,
        "huggingface_hub",
        SimpleNamespace(snapshot_download = lambda **_kwargs: str(tmp_path)),
    )

    hf_download._download_snapshot("Org/Model", None, "xet")

    assert written, "XET snapshot download must still record a manifest"
    assert written[0][0:3] == ("model", "Org/Model", None)
    assert written[0][3][0].path == "config.json"
    assert verified == [("model", "Org/Model", None, str(tmp_path))]


def test_download_gguf_variant_writes_manifest_for_xet(monkeypatch, tmp_path):
    written = []
    verified = []

    monkeypatch.setattr(
        hf_download,
        "_model_info_with_retry",
        lambda *_args, **_kwargs: SimpleNamespace(
            siblings = [_sibling("model-Q4_K_M.gguf", 10, "main")]
        ),
    )
    monkeypatch.setattr(
        hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args)
    )
    monkeypatch.setattr(
        download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0
    )
    monkeypatch.setattr(download_manifest, "clear_cancel_marker", lambda *_args: None)
    monkeypatch.setattr(
        download_manifest, "write_manifest", lambda *args: written.append(args) or True
    )
    monkeypatch.setitem(
        sys.modules,
        "huggingface_hub",
        SimpleNamespace(snapshot_download = lambda **_kwargs: str(tmp_path)),
    )

    hf_download._download_gguf_variant("Org/Model", "Q4_K_M", None, "xet")

    assert written, "XET GGUF variant download must still record a manifest"
    assert written[0][0:3] == ("model", "Org/Model", "Q4_K_M")
    assert written[0][3][0].path == "model-Q4_K_M.gguf"
    assert verified == [("model", "Org/Model", "Q4_K_M", str(tmp_path))]


def test_download_dataset_writes_manifest_for_xet(monkeypatch, tmp_path):
    written = []
    verified = []

    monkeypatch.setattr(
        hf_download,
        "_dataset_info_with_retry",
        lambda *_args, **_kwargs: SimpleNamespace(
            siblings = [SimpleNamespace(rfilename = "data.parquet", size = 30)]
        ),
    )
    monkeypatch.setattr(
        hf_download, "_verify_completed_download", lambda *args, **kwargs: verified.append(args)
    )
    monkeypatch.setattr(
        download_registry, "prepare_cache_for_transport", lambda *_args, **_kwargs: 0
    )
    monkeypatch.setattr(download_manifest, "clear_cancel_marker", lambda *_args: None)
    monkeypatch.setattr(
        download_manifest, "write_manifest", lambda *args: written.append(args) or True
    )
    monkeypatch.setitem(
        sys.modules,
        "huggingface_hub",
        SimpleNamespace(snapshot_download = lambda **_kwargs: str(tmp_path)),
    )

    hf_download._download_dataset("Org/Data", None, "xet")

    assert written, "XET dataset download must still record a manifest"
    assert written[0][0:3] == ("dataset", "Org/Data", None)
    assert written[0][3][0].path == "data.parquet"
    assert verified == [("dataset", "Org/Data", None, str(tmp_path))]


def test_dataset_status_includes_generation(monkeypatch):
    class _Registry:
        def get_job(self, _key):
            return SimpleNamespace(state = "running", error = None)

        def current_generation(self, _key):
            return 4

    monkeypatch.setattr(dataset_downloads, "_registry", _Registry())
    monkeypatch.setattr(
        dataset_downloads,
        "resolve_cached_repo_id_case",
        lambda repo_id, **_kwargs: repo_id,
    )

    result = asyncio.run(dataset_downloads.get_dataset_download_status_response("Org/Data"))

    assert result.state == "running"
    assert result.generation == 4