fix(download): allow downloading additional files of an in-library model version (#1058)

2026-08-20 12:31:27 -03:00 · 2026-08-19 16:29:59 +08:00
parent 0a28500848
commit cef4129fc9
9 changed files with 1092 additions and 83 deletions
@@ -82,14 +82,18 @@ def stub_metadata(monkeypatch):


 class DummyScanner:
-    def __init__(self, exists: bool = False):
+    def __init__(self, exists: bool = False, raw_data=None):
        self.exists = exists
        self.calls = []
+        self._cache = SimpleNamespace(raw_data=list(raw_data or []))

    async def check_model_version_exists(self, version_id):
        self.calls.append(version_id)
        return self.exists

+    async def get_cached_data(self):
+        return self._cache
+

@pytest.fixture
 def scanners(monkeypatch):
@@ -1692,3 +1696,310 @@ async def test_download_proceeds_when_history_skip_disabled(monkeypatch, scanner
    assert result.get("skipped") is not True
    history_service.has_been_downloaded.assert_not_called()
    execute_download.assert_awaited_once()
+
+
+# ---------------------------------------------------------------------------
+# Multi-file downloads within a single model version (#1058)
+# ---------------------------------------------------------------------------
+
+
+def _multi_file_payload(include_hashes: bool = True):
+    """Version payload with two weight files under the same version."""
+    files = [
+        {
+            "id": 1001,
+            "type": "Model",
+            "primary": True,
+            "downloadUrl": "https://example.invalid/file-a.safetensors",
+            "name": "file-a.safetensors",
+        },
+        {
+            "id": 1002,
+            "type": "Model",
+            "primary": False,
+            "downloadUrl": "https://example.invalid/file-b.safetensors",
+            "name": "file-b.safetensors",
+        },
+    ]
+    if include_hashes:
+        files[0]["hashes"] = {"SHA256": "AAA111"}
+        files[1]["hashes"] = {"SHA256": "BBB222"}
+    return {
+        "id": 42,
+        "modelId": 7,
+        "model": {"type": "LoRA", "tags": ["fantasy"]},
+        "baseModel": "BaseModel",
+        "creator": {"username": "Author"},
+        "files": files,
+    }
+
+
+def _local_entry(file_name: str, sha256: str, version_id: int = 42):
+    """A library cache entry for one already-downloaded file of a version."""
+    return {
+        "file_name": file_name,
+        "file_path": f"/tmp/{file_name}.safetensors",
+        "sha256": sha256,
+        "civitai": {"id": version_id, "modelId": 7},
+    }
+
+
+def _stub_history_service(monkeypatch, has_been_downloaded: bool = False):
+    history_service = AsyncMock()
+    history_service.has_been_downloaded = AsyncMock(
+        return_value=has_been_downloaded
+    )
+    history_service.mark_downloaded = AsyncMock()
+    monkeypatch.setattr(
+        ServiceRegistry,
+        "get_downloaded_version_history_service",
+        AsyncMock(return_value=history_service),
+    )
+    return history_service
+
+
+@pytest.mark.asyncio
+async def test_download_allows_other_file_of_in_library_version(
+    monkeypatch, scanners, metadata_provider
+):
+    """A different file of an in-library version must still download (#1058)."""
+    scanners.lora._cache.raw_data.append(_local_entry("file-a", "aaa111"))
+    metadata_provider.get_model_version = AsyncMock(
+        return_value=_multi_file_payload()
+    )
+    _stub_history_service(monkeypatch)
+
+    execute_download = AsyncMock(return_value={"success": True, "download_id": "done"})
+    monkeypatch.setattr(
+        DownloadManager, "_execute_download", execute_download, raising=False
+    )
+
+    manager = DownloadManager()
+    result = await manager.download_from_civitai(
+        model_version_id=42,
+        save_dir="/tmp",
+        file_params={"id": 1002, "type": "Model", "name": "file-b.safetensors"},
+    )
+
+    assert result["success"] is True
+    execute_download.assert_awaited_once()
+    # Version-level gates must not run for an explicit file selection
+    assert scanners.lora.calls == []
+
+
+@pytest.mark.asyncio
+async def test_download_blocks_same_file_of_in_library_version(
+    monkeypatch, scanners, metadata_provider
+):
+    """Re-downloading the SAME file of an in-library version stays blocked."""
+    scanners.lora._cache.raw_data.append(_local_entry("file-a", "aaa111"))
+    metadata_provider.get_model_version = AsyncMock(
+        return_value=_multi_file_payload()
+    )
+    _stub_history_service(monkeypatch)
+
+    execute_download = AsyncMock(return_value={"success": True})
+    monkeypatch.setattr(
+        DownloadManager, "_execute_download", execute_download, raising=False
+    )
+
+    manager = DownloadManager()
+    result = await manager.download_from_civitai(
+        model_version_id=42,
+        save_dir="/tmp",
+        file_params={"id": 1001, "type": "Model", "name": "file-a.safetensors"},
+    )
+
+    assert result["success"] is False
+    assert "file-a.safetensors" in result["error"]
+    assert "already exists in lora library" in result["error"]
+    execute_download.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_download_unresolvable_file_params_falls_back_to_version_gate(
+    monkeypatch, scanners, metadata_provider
+):
+    """file_params that match no file must not bypass version-level gates."""
+    scanners.lora.exists = True
+    metadata_provider.get_model_version = AsyncMock(
+        return_value=_multi_file_payload()
+    )
+    _stub_history_service(monkeypatch)
+
+    execute_download = AsyncMock(return_value={"success": True})
+    monkeypatch.setattr(
+        DownloadManager, "_execute_download", execute_download, raising=False
+    )
+
+    manager = DownloadManager()
+    result = await manager.download_from_civitai(
+        model_version_id=42,
+        save_dir="/tmp",
+        file_params={"id": 9999, "type": "Model"},
+    )
+
+    assert result["success"] is False
+    assert result["error"] == "Model version already exists in lora library"
+    execute_download.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_download_empty_file_params_treated_as_no_selection(
+    monkeypatch, scanners, metadata_provider
+):
+    """An empty file_params dict is normalized to None (version-level gates)."""
+    scanners.lora.exists = True
+
+    execute_download = AsyncMock(return_value={"success": True})
+    monkeypatch.setattr(
+        DownloadManager, "_execute_download", execute_download, raising=False
+    )
+
+    manager = DownloadManager()
+    result = await manager.download_from_civitai(
+        model_version_id=42,
+        save_dir="/tmp",
+        file_params={},
+    )
+
+    assert result["success"] is False
+    assert result["error"] == "Model version already exists in lora library"
+    # Early version-level gate fired (file_params normalized to None)
+    assert scanners.lora.calls == [42]
+    execute_download.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_download_explicit_file_bypasses_history_skip(
+    monkeypatch, scanners, metadata_provider
+):
+    """Explicit file selection bypasses the previously-downloaded skip."""
+    get_settings_manager().settings[
+        "skip_previously_downloaded_model_versions"
+    ] = True
+    scanners.lora._cache.raw_data.append(_local_entry("file-a", "aaa111"))
+    metadata_provider.get_model_version = AsyncMock(
+        return_value=_multi_file_payload()
+    )
+    history_service = _stub_history_service(monkeypatch, has_been_downloaded=True)
+
+    execute_download = AsyncMock(return_value={"success": True, "download_id": "done"})
+    monkeypatch.setattr(
+        DownloadManager, "_execute_download", execute_download, raising=False
+    )
+
+    manager = DownloadManager()
+    result = await manager.download_from_civitai(
+        model_version_id=42,
+        save_dir="/tmp",
+        file_params={"id": 1002, "type": "Model", "name": "file-b.safetensors"},
+    )
+
+    assert result["success"] is True
+    execute_download.assert_awaited_once()
+    # History gate is bypassed before it even queries the service
+    history_service.has_been_downloaded.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_download_file_match_by_name_when_local_hash_missing(
+    monkeypatch, scanners, metadata_provider
+):
+    """Legacy local metadata without sha256 falls back to name matching."""
+    scanners.lora._cache.raw_data.append(_local_entry("file-a", ""))
+    metadata_provider.get_model_version = AsyncMock(
+        return_value=_multi_file_payload()
+    )
+    _stub_history_service(monkeypatch)
+
+    execute_download = AsyncMock(return_value={"success": True})
+    monkeypatch.setattr(
+        DownloadManager, "_execute_download", execute_download, raising=False
+    )
+
+    manager = DownloadManager()
+    result = await manager.download_from_civitai(
+        model_version_id=42,
+        save_dir="/tmp",
+        file_params={"id": 1001, "type": "Model"},
+    )
+
+    assert result["success"] is False
+    assert "already exists in lora library" in result["error"]
+    execute_download.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_download_no_false_positive_when_both_hashes_empty(
+    monkeypatch, scanners, metadata_provider
+):
+    """Two empty hashes must never compare equal; name decides instead."""
+    # Local entry for a DIFFERENT file of the same version, no hash stored
+    scanners.lora._cache.raw_data.append(_local_entry("file-b", ""))
+    metadata_provider.get_model_version = AsyncMock(
+        return_value=_multi_file_payload(include_hashes=False)
+    )
+    _stub_history_service(monkeypatch)
+
+    execute_download = AsyncMock(return_value={"success": True, "download_id": "done"})
+    monkeypatch.setattr(
+        DownloadManager, "_execute_download", execute_download, raising=False
+    )
+
+    manager = DownloadManager()
+    result = await manager.download_from_civitai(
+        model_version_id=42,
+        save_dir="/tmp",
+        file_params={"id": 1001, "type": "Model"},
+    )
+
+    assert result["success"] is True
+    execute_download.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_download_model_id_only_with_file_params_resolves_file(
+    monkeypatch, scanners, metadata_provider
+):
+    """model_id-only requests with file_params resolve the file post-fetch."""
+    scanners.lora.exists = True  # version-level index says "in library"
+    scanners.lora._cache.raw_data.append(_local_entry("file-a", "aaa111"))
+    metadata_provider.get_model_version = AsyncMock(
+        return_value=_multi_file_payload()
+    )
+    _stub_history_service(monkeypatch)
+
+    execute_download = AsyncMock(return_value={"success": True, "download_id": "done"})
+    monkeypatch.setattr(
+        DownloadManager, "_execute_download", execute_download, raising=False
+    )
+
+    manager = DownloadManager()
+    result = await manager.download_from_civitai(
+        model_id=7,
+        save_dir="/tmp",
+        file_params={"id": 1002, "type": "Model"},
+    )
+
+    assert result["success"] is True
+    execute_download.assert_awaited_once()
+
+
+def test_resolve_target_file_by_id():
+    files = _multi_file_payload()["files"]
+    resolved = DownloadManager._resolve_target_file(files, {"id": 1002})
+    assert resolved is files[1]
+
+
+def test_resolve_target_file_by_primary_flag():
+    files = _multi_file_payload()["files"]
+    resolved = DownloadManager._resolve_target_file(files, {"isPrimary": True})
+    assert resolved is files[0]
+
+
+def test_resolve_target_file_returns_none_for_no_match():
+    files = _multi_file_payload()["files"]
+    assert DownloadManager._resolve_target_file(files, {"id": 9999}) is None
+    assert DownloadManager._resolve_target_file(files, None) is None
+    assert DownloadManager._resolve_target_file(files, {}) is None
@@ -5,6 +5,8 @@ Covers the new ``download_id``-based code paths in
 compatibility with ``id``.
 """

+import json
+import sqlite3
 from pathlib import Path

 import pytest
@@ -191,3 +193,131 @@ async def test_retry_unknown_download_id(tmp_path: Path) -> None:

    item = await svc.retry_from_history(download_id="dl-nope")
    assert item is None
+
+
+# ---------------------------------------------------------------------------
+# file_params persistence across queue -> history -> retry (#1058)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_file_params_survive_queue_to_history(tmp_path: Path) -> None:
+    """complete_download copies the queue row's file_params into history."""
+    svc = _make_service(tmp_path)
+    await svc.add_to_queue(
+        download_id="dl-fp",
+        model_id=1,
+        model_version_id=100,
+        file_params={"id": 1002, "type": "Model"},
+    )
+
+    await svc.complete_download("dl-fp", status="failed", error="boom")
+
+    history = await svc.get_history()
+    assert len(history["items"]) == 1
+    assert json.loads(history["items"][0]["file_params"]) == {
+        "id": 1002,
+        "type": "Model",
+    }
+
+
+@pytest.mark.asyncio
+async def test_retry_restores_file_params(tmp_path: Path) -> None:
+    """retry_from_history re-queues with the originally selected file."""
+    svc = _make_service(tmp_path)
+    await svc.add_to_history(
+        download_id="dl-fail-fp",
+        model_id=1,
+        model_version_id=100,
+        status="failed",
+        file_params={"id": 1002, "type": "Model"},
+    )
+
+    item = await svc.retry_from_history(download_id="dl-fail-fp")
+
+    assert item is not None
+    assert item["status"] == "queued"
+    assert json.loads(item["file_params"]) == {"id": 1002, "type": "Model"}
+
+
+@pytest.mark.asyncio
+async def test_retry_all_restores_file_params(tmp_path: Path) -> None:
+    """retry_all_failed preserves file_params for every re-queued item."""
+    svc = _make_service(tmp_path)
+    await svc.add_to_history(
+        download_id="dl-f1", status="failed", file_params={"id": 1001}
+    )
+    await svc.add_to_history(download_id="dl-f2", status="canceled")
+
+    count = await svc.retry_all_failed()
+    assert count == 2
+
+    queue = await svc.get_queue()
+    restored = sorted(
+        (q["file_params"] or "") for q in queue
+    )
+    assert restored[0] == ""  # dl-f2 never had file_params
+    assert json.loads(restored[1]) == {"id": 1001}
+
+
+@pytest.mark.asyncio
+async def test_legacy_history_db_gains_file_params_column(tmp_path: Path) -> None:
+    """Databases created before the file_params column get migrated (#1058)."""
+    db_path = tmp_path / "queue.sqlite"
+    conn = sqlite3.connect(db_path)
+    conn.executescript(
+        """
+        CREATE TABLE download_queue (
+            download_id TEXT PRIMARY KEY,
+            model_id INTEGER,
+            model_version_id INTEGER,
+            model_name TEXT NOT NULL DEFAULT '',
+            version_name TEXT DEFAULT '',
+            thumbnail_url TEXT DEFAULT '',
+            source TEXT,
+            file_params TEXT,
+            status TEXT NOT NULL DEFAULT 'queued',
+            priority INTEGER DEFAULT 0,
+            progress INTEGER DEFAULT 0,
+            bytes_downloaded INTEGER DEFAULT 0,
+            total_bytes INTEGER,
+            bytes_per_second REAL DEFAULT 0.0,
+            error TEXT,
+            file_path TEXT,
+            added_at REAL NOT NULL,
+            started_at REAL,
+            completed_at REAL
+        );
+        CREATE TABLE download_history (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            download_id TEXT,
+            model_id INTEGER,
+            model_version_id INTEGER,
+            model_name TEXT NOT NULL DEFAULT '',
+            version_name TEXT DEFAULT '',
+            thumbnail_url TEXT DEFAULT '',
+            status TEXT NOT NULL,
+            error TEXT,
+            file_path TEXT,
+            bytes_downloaded INTEGER DEFAULT 0,
+            total_bytes INTEGER,
+            completed_at REAL NOT NULL,
+            is_already_exists INTEGER DEFAULT 0
+        );
+        """
+    )
+    conn.close()
+
+    svc = DownloadQueueService(db_path=str(db_path))
+
+    # The migrated table accepts file_params writes
+    await svc.add_to_history(
+        download_id="dl-legacy-fp", status="failed", file_params={"id": 5}
+    )
+    history = await svc.get_history()
+    assert json.loads(history["items"][0]["file_params"]) == {"id": 5}
+
+    # And retry restores them
+    item = await svc.retry_from_history(download_id="dl-legacy-fp")
+    assert item is not None
+    assert json.loads(item["file_params"]) == {"id": 5}