diff --git a/docs/plans/issue-1058-multi-file-version-downloads.md b/docs/plans/issue-1058-multi-file-version-downloads.md new file mode 100644 index 00000000..ad8b9788 --- /dev/null +++ b/docs/plans/issue-1058-multi-file-version-downloads.md @@ -0,0 +1,206 @@ +# Plan: Multi-File Downloads Within a Single CivitAI Model Version + +**Issue:** [#1058 — Cannot download multiple file variants from the same model version](https://github.com/willmiao/ComfyUI-Lora-Manager/issues/1058) +**Status:** v2 — revised after adversarial review (backend correctness + frontend/tests) +**Scope:** CivitAI/CivArchive downloads of `lora`, `checkpoint`, `embedding` model types. HuggingFace downloads are out of scope (already per-file). + +> v2 changelog: incorporated 18 review findings. Key changes vs v1: +> shared file resolver + `resolved_version_id` for the gate (R1); `file_params` normalization at API boundary (R2); D2 hash-matching rule fixed for empty-hash cases (R6/R7); D3 extended to re-point `version_index` on removal (R4); D4 replaced with a child table (R3); `delete_model_version` interaction documented (R5); `ModelVersionsTab` surface added to phase 2 (F6); phase-2 multi-file loop requires a reload-deferred download variant (F7); queue-retry `file_params=NULL` known issue recorded (R9); test-fixture gaps and revised estimates (F10). + +--- + +## 1. Problem Statement + +A CivitAI model version can contain multiple downloadable weight files (e.g. fp16/fp32, safetensors/ckpt, different sizes). LoRA Manager already has a working file-selection pipeline (frontend file dialog → `fileParams` → backend file matching), but downloaded state is tracked at the **model-version** level. After any single file of a version is downloaded: + +1. The version is marked **In Library** and the file-selection entry point disappears. +2. The backend rejects further download attempts for that version. + +There is no way to download the remaining files of the same version through LoRA Manager. + +## 2. Current State (verified against code; all references confirmed by review) + +### 2.1 Download gating — backend (`py/services/download_manager.py`) + +`_execute_original_download` enforces two version-level gates: + +- **Library gate, early** (lines 1157–1184, before metadata fetch, fires when `model_version_id` given) and **late** (lines 1350–1376, fires only when `model_version_id is None`): `scanner.check_model_version_exists(version_id)` across lora/checkpoint/embedding scanners → hard error `"Model version already exists in ... library"`. +- **History gate** (lines 1238–1279): when `skip_previously_downloaded_model_versions` setting is on, `_has_been_downloaded(model_type, version_id)` → silent skip. History DB primary key is `(model_type, version_id)` (`py/services/downloaded_version_history_service.py:61`). + +File selection works: `file_params {id, type, format, size, fp}` is matched against `version_info.files` (lines 1498–1569), **but only under `if file_params and model_version_id:` (line 1499)** — with `model_id`-only requests the selection silently falls back to the primary file (1571–1619). `file_params` currently carries no file `name` or hash. + +### 2.2 Downloaded-state surfacing — backend (`py/routes/handlers/model_handlers.py`) + +`get_civitai_versions` (lines 2148–2188) sets per-version `existsLocally` via `cache.version_index.get(version_id)` (plus a single `localPath` from that entry) and `hasBeenDownloaded` via the history service. No per-file granularity. + +### 2.3 Frontend blockers (`static/js/managers/DownloadManager.js`) + +Three independent gates prevent re-entering the file dialog: + +1. **Line 598:** file-select badge rendered only when `modelFiles.length > 1 && !existsLocally`. +2. **Lines 666–681 (`updateNextButtonState`):** Next button disabled with "Already in Library" when `currentVersion.existsLocally`. +3. **Lines 784–787 (`proceedToLocation`):** toast + abort when `currentVersion.existsLocally`. + +The badge path (`confirmFileSelection` lines 737–759 → `proceedToLocationContent` → `startDownload` single mode → `executeDownloadWithProgress` → POST `file_params`, `static/js/api/baseModelApi.js:1236–1250`) has **zero** `existsLocally` guards (all 12 occurrences enumerated; none on this path; `import/DownloadManager.js` has none either). The `.exists-locally` CSS class is purely visual (`download-modal.css:496–499`). **Making the badge visible again is sufficient to unlock the flow** for phase 1. + +Post-download refresh is clean: the modal closes and `resetAndReload(true)` performs a full library refetch (`DownloadManager.js:1063`); dialog reopen resets state and refetches versions with no client-side cache. No same-session staleness. + +### 2.4 Local identity of the downloaded file + +`LoraMetadata/CheckpointMetadata/EmbeddingMetadata.from_civitai_info(version_info, file_info, ...)` (`py/utils/models.py:245–369`) persists: + +- `sha256` = `file_info.hashes.SHA256` (lowercased, defaults to `""`) — a stable per-file identity; +- `civitai` = the full `version_info` payload (including the `files` list). + +Metadata refresh (`metadata_sync_service.py:104–105`) replaces the `civitai` blob wholesale but never overwrites top-level `sha256`; `verify_duplicate_hashes` (481–526) corrects it to the on-disk hash. Top-level-sha256 matching is refresh-robust. + +**Caveats (review R6/R7):** +- SHA256 is not guaranteed: CivArchive's transform only sets `hashes` when source data carries it (`civarchive_client.py:185–189`); `from_civitai_info` defaults to `""`. +- Name fallback is unreliable exactly when it matters: local `file_name` is extension-less (`models.py:264`) and `generate_unique_filename` rewrites it with a hash suffix on conflict (`download_manager.py:1125–1136`); checkpoints with `hash_status='pending'` keep empty sha256 until on-demand hashing (`model_scanner.py:1232–1240`). + +### 2.5 Version index collision (pre-existing hazard) + +`ModelCache.version_index` is single-valued (`model_cache.py:133`: `version_index[version_id] = item`). Two files of the same version in the library → second entry overwrites the first; `remove_from_version_index` (lines 151–181) drops the whole version key when the indexed entry is removed, even if a sibling file remains. ~10 read sites depend on this index (48 grep touch points total; readers include `recipe_scanner.py:2682–2726`, `recipe_format.py:37–40`, `misc_handlers.py:2440–2444`, `model_handlers.py`, `model_scanner.check_model_version_exists:2444`). + +Review correction (F3): bulk paths `remove_models` (`model_scanner.py:2376`) and `update_single_model_cache` (`:1689`) call `rebuild_version_index()` right after, so a sibling re-enters the index in those flows — the hazard is narrower than v1 stated, but direct `remove_from_version_index` callers (e.g. `model_scanner.py:1018`) still drop the key, and the user-visible artifact in phase 1 is real: `localPath` in the dialog flips to whichever file was indexed last. + +### 2.6 Entry points that send / don't send `file_params` (fully enumerated by review) + +**Send `file_params` (user-initiated dialog flows only):** `DownloadManager.js:1611–1639` (single mode). API surface accepting arbitrary JSON `file_params`: GET `/api/lm/download-model-get` (`model_handlers.py:1634–1686`), POST `/api/lm/downloads/queue/add` (`model_handlers.py:1799–1832`). + +**Never send `file_params` (keep version-level semantics):** batch download (`DownloadManager.js:1756–1766`; batch also filters out in-library versions at `:1648`), `downloadVersionWithDefaults` (`:1810–1830`), recipe import (`import/DownloadManager.js:269–276`), bulk missing-LoRA (`BulkMissingLoraDownloadManager.js:292–299`), `RecipeModal.js:1728–1736`, `ModelVersionsTab.js:1427`. `web/comfyui/` and `vue-widgets/src` contain **no** download triggers at all (grep-verified). `py/services/use_cases/` has only `download_model_use_case.py` (pass-through). + +### 2.7 Paths that do NOT need changes (verified) + +- **aria2 pause/resume** (`_resume_restored_aria2_download`, line 754+): resumes from persisted `resume_context`; never re-runs existence gates. +- **`download_coordinator.py:90`**: pure pass-through of `file_params`. +- **Update checker / plugin self-update** (`update_routes.py:496–501`): only closes the history DB handle. +- **History delete semantics**: `mark_as_deleted` sets `is_deleted_override=1` and `has_been_downloaded` then returns False (`downloaded_version_history_service.py:276`) — LM-initiated deletes already reset the history skip. + +### 2.8 Related pre-existing issues (record, not necessarily fix) + +- **Queue retry drops file selection** (R9): `download_queue_service.retry_from_history` / `retry_all_failed` re-queue with `file_params=NULL` (`download_queue_service.py:705, 758`) although the queue table has a `file_params` column (`:43`) — a retried non-primary download silently reverts to the primary file. Fix alongside phase 1 (small: persist and reuse the column). +- **`delete_model_version`** (`misc_handlers.py:2410–2487`): resolves the file via the single-valued `version_index` (2440–2444), deletes only that one file, and `mark_as_deleted` flags the **entire version** as deleted in history (2479) even when a sibling file remains in the library. See phase 2 item 6.1.5. + +## 3. Goals / Non-Goals + +**Goals** + +- G1: A user can download any not-yet-downloaded file of a version already partially in the library (issue repro steps 6–8). +- G2: True duplicates stay blocked: downloading the *same* file of the same version twice is rejected. +- G3: Per-file downloaded state visible in the file dialog; multiple files selectable and downloadable in one pass. +- G4: No regression for version-level semantics relied on by batch download, recipe missing-LoRA detection, and `skip_previously_downloaded_model_versions`. + +**Non-Goals** + +- No change to recipe `inLibrary` semantics ("any file of the version present" remains sufficient). +- No change to the update-checker (version-level comparison). +- No primary-key rebuild of the history database. +- HuggingFace download flow untouched. + +## 4. Design Decisions + +- **D1 — Explicit file selection bypasses the history gate, version-level gates stay for everyone else.** The history skip exists to dedupe automated flows. A user explicitly picking a file is unambiguous intent; the file-level library gate (G2) still prevents real duplicates. **Guard conditions use normalized truthiness** (see D1a). All confirmed `file_params` senders are user-initiated dialog flows (2.6), and LM-initiated deletes already reset history (2.7), so the bypass only affects "downloaded but not LM-deleted" versions with the setting on — intended. +- **D1a — `file_params` normalization at the boundary (R2).** `download-model-get` and `downloads/queue/add` accept arbitrary JSON; `{}` is `not None` but falsy and would bypass gates while downloading the primary file. Normalize `file_params = file_params or None` in the coordinator/handlers, and treat the bypass as active only when a target file id is resolvable. +- **D2 — File identity matching rule (R6/R7):** hash-compare **only when both sides are non-empty** (lowercase SHA256 equality); name-compare when either side is empty. Never let `"" == ""` match. Name fallback caveats from 2.4 apply (renamed files, pending checkpoint hashes) — acceptable residual risk, worst case is a blocked re-download the user can retry after hashing completes. +- **D3 — Cache indexes: additive multi-index + removal re-pointing (R4).** Add `version_files_index: Dict[int, List[dict]]` maintained alongside `version_index` by the same add/remove/rebuild methods; existing readers of `version_index` untouched. Additionally fix `remove_from_version_index`: when the popped entry has a surviving sibling (per the multi-index), re-point `version_index[version_id]` to the sibling instead of dropping the key; same for the `model_id_index` descriptor. This closes the 2.5 hazard for existing readers (`check_model_version_exists`, `existsLocally`, recipe matching) without restructuring anything. +- **D4 — Per-file history via a child table (R3).** v1's additive-column approach is structurally impossible on a `(model_type, version_id)` PK (`ON CONFLICT DO UPDATE` would keep only the last file). Instead add `downloaded_version_files(model_type, version_id, file_id, file_name, downloaded_at, PRIMARY KEY(model_type, version_id, file_id))` — additive, no PK rebuild, honors the Non-Goal. Existing version-level table and queries unchanged. New per-file queries are opt-in. `_initialize_schema` uses `CREATE TABLE IF NOT EXISTS`, so the new table is created for existing DBs without any ALTER. +- **D5 — UI flow reuse, with an extracted inner download function for multi-file (F7).** Phase 1 unlocks the existing badge → file dialog → location → download pipeline. Phase 2 upgrades the dialog to multi-select; iterating `executeDownloadWithProgress` as-is would produce N full library reloads, N toasts, and competing failure-summary modals — so phase 2 extracts a reload-deferred, failure-aggregating inner variant and runs one reload + one summary at the end. + +## 5. Implementation — Phase 1 (fix the issue; independently shippable) + +### 5.1 Backend — `py/services/download_manager.py` + +1. **Normalize `file_params`** at the boundary (D1a): `download_coordinator.schedule_download` and the two API handlers (`model_handlers.py:1649–1666`, `1810–1832`) apply `file_params = file_params or None`. +2. **Extract a shared file resolver** (R1): pull the matching logic at 1498–1569 into `_resolve_target_file(version_info, file_params) -> Optional[dict]`, used by **both** the new gate and the download-selection path. The selection path's condition (line 1499) switches from `model_version_id` to `resolved_version_id` (already computed at 1230–1236 from `version_info.id`), so gate and download always agree on the target file — including the `model_id`-only case. +3. **New helper** `_find_local_file_entry(version_id, target_file) -> Optional[dict]`: iterate the three scanners' cached `raw_data` (NOT `version_index` — single-valued); candidates = entries whose `civitai.id` normalizes to `version_id`; match per D2. +4. **Gate restructure in `_execute_original_download`**: + - Early scanner gate (1157–1184): add `file_params is None` guard; with normalized `file_params`, defer (file identity not resolvable before metadata fetch). + - After `version_info` fetch + `resolved_version_id` (~1229): when `file_params` present, resolve target file via the shared resolver; unresolvable → hard error "No matching file" (fail closed, prevents empty-dict bypass). Resolvable → `_find_local_file_entry`; hit → same hard error shape as today with the file name in the message. + - History gate (1238–1279): add `file_params is None` (D1). Base-model skip (1281–1324) unchanged — still applies. + - Late gate (1350–1376): add `file_params is None` guard (F2) — the post-fetch file-level check above already covers this case. + - Nothing between the early gate and the post-fetch point assumes the version is absent (review task 6: only provider selection + metadata fetch; no DB writes; `_persist_aria2_state` runs only when actually downloading at 1659). +5. **Queue retry fix** (2.8, small): persist `file_params` into the queue table on enqueue and reuse it in `retry_from_history` / `retry_all_failed`. +6. Logging: `[download]` lines for file-level allow/block, consistent with existing style. + +**Estimated:** ~150–220 LOC + resolver extraction. + +### 5.2 Frontend — `static/js/managers/DownloadManager.js` + +1. Line 598: drop `&& !existsLocally` from the badge condition (badge shows whenever `modelFiles.length > 1`). +2. `fileParams` construction (1611–1616): add `name: this.selectedFile.name`. +3. Surface the backend "file already in library" hard error as a toast instead of only the batch-summary modal (R10/F12 nit; reuse existing error message field). +4. No changes to `updateNextButtonState` / `proceedToLocation` in phase 1; no template or CSS changes. + +**Known phase-1 UX limitations (acknowledged, fixed in phase 2):** with all files downloaded the badge still renders and re-picking a downloaded file fails late (backend error after the location step); `localPath` may point at a sibling file; batch-preview "In Library" badge stays version-level and gives no hint of remaining files. + +**Estimated:** ~10–30 LOC (confirmed realistic by review). + +### 5.3 Phase 1 tests + +Backend — extend `tests/services/test_download_manager_basic.py` (1694 lines; all fixture patterns exist): + +- **Fixture gaps to add (F10):** `DummyScanner.get_cached_data()`/`raw_data` stub (~10 lines); `hashes.SHA256` in the metadata-provider payload's `files`. +- Cases: same version + different SHA256 in library + `file_params` → proceeds; same SHA256 → hard error; `file_params=None` + version in library → hard error (unchanged); history-skip on + `file_params` → not skipped; without → skipped (unchanged); empty-dict `file_params` normalized → version-level behavior; `model_id`-only + `file_params` → gate and selection resolve the same file; legacy metadata (empty local sha256) matched by name; target file with empty SHA256 → name fallback, no `""==""` false positive. +- Queue retry: `file_params` survives retry. +- Assert proceed/abort via the existing `_execute_download` mock pattern. + +Frontend (`tests/frontend/`): badge renders for multi-file version with `existsLocally=true` (pattern from `downloadManager.history.test.js`). + +**Estimated:** ~150–250 LOC (confirmed realistic). + +## 6. Implementation — Phase 2 (per-file status + multi-select + index hardening) + +### 6.1 Backend + +1. **`py/services/model_cache.py`** (D3): add `version_files_index`; maintain in `add_to_version_index` / `remove_from_version_index` / `rebuild_version_index`; removal re-points `version_index[version_id]` (and the `model_id_index` descriptor) to a surviving sibling instead of dropping the key. +2. **`py/services/model_scanner.py`**: expose `get_files_for_version(version_id) -> List[dict]`. +3. **`py/routes/handlers/model_handlers.py` `get_civitai_versions`**: annotate each version with `downloadedFiles: [{fileId, fileName, filePath}]` via `version_files_index` + D2 matching against `version.files`. +4. **`py/services/downloaded_version_history_service.py`** (D4): new child table `downloaded_version_files`; `mark_downloaded` also upserts the child row when `file_id` known; `mark_as_deleted` clears the version's child rows only when no sibling remains in the library; new `get_downloaded_file_ids(model_type, version_id) -> set[int]`. `_record_downloaded_version_history` passes `file_info` through. +5. **`delete_model_version`** (`misc_handlers.py:2410–2487`, R5): resolve **all** local files of the version via `version_files_index`; delete all (current endpoint semantics are version-level) or — if kept per-file — only `mark_as_deleted` when no sibling remains. Decide at implementation time; minimum is documenting current behavior. +6. **`ModelVersionsTab` backend support**: none needed beyond item 3 (`downloadedFiles`); the tab consumes the same versions payload. + +### 6.2 Frontend + +1. **File dialog multi-select** — change surface (F8): option markup (`DownloadManager.js:712–724`), the single-select click handler (`727–734`), the `input[type="radio"]:checked` selector in `confirmFileSelection` (`738`); template `templates/components/modals/download_modal.html:48–60` (confirm-button label only); CSS `download-modal.css` — checkbox variant of `.file-option-radio input` (595–604) and a **new** `.file-option.disabled` style (does not exist). Files whose id ∈ `downloadedFiles` render disabled with an "In Library" tag. +2. **Mixed-type guard (F8):** multi-select is restricted to files sharing the same routing target (`_isDiffusionModel` is computed once from a single `selectedFile` at 798–803; e.g. "Model" + "UNet" files route to different roots). Disallow mixed-type multi-select (simplest, predictable); single-file selection unchanged. +3. **Multi-file download loop (D5/F7):** extract from `executeDownloadWithProgress` a reload-deferred, no-toast inner function; iterate per selected file with per-file progress; one `resetAndReload(true)` + one aggregated success/failure summary at the end (reuse `showDownloadBatchSummary`). +4. **`updateNextButtonState` / `proceedToLocation`:** for multi-file versions, Next routes into the file dialog; hard block only when *every* weight file is downloaded. +5. **`ModelVersionsTab.js` (F6):** the Download action (`:576` hidden when `isInLibrary`) — for multi-file versions with remaining files, show it and route into the download modal's file dialog; keep hidden when all files present. +6. **Batch preview (F5):** `batch-preview-local-badge` (`:1320`) gains a "partially downloaded" hint for multi-file versions with remaining files. +7. New i18n keys (`modals.download.fileSelection.inLibrary`, `downloadSelected`, partial-download tooltip, etc.) → run `python scripts/sync_translation_keys.py`. + +### 6.3 Phase 2 tests + +- `model_cache` (`tests/services/test_model_cache.py` already covers add/remove at 44–55): multi-valued index; sibling re-point on removal; rebuild. +- `get_civitai_versions`: `downloadedFiles` correctness (hash match, name fallback, no match, CivArchive no-hash payload). +- History service (`tests/services/test_downloaded_version_history_service.py` uses real SQLite on tmp_path): child-table creation on a legacy DB; per-file record/query; `mark_as_deleted` sibling semantics. +- Frontend: dialog checkbox rendering/disabled state and multi-file confirm — **greenfield behavior coverage** (F10: no existing test exercises `showFileSelectionStep`/`confirmFileSelection`; infra exists, patterns must be built). + +## 7. Risks and Mitigations + +| Risk | Impact | Mitigation | +|---|---|---| +| History-gate bypass (D1) causes unwanted re-downloads in automated flows | Large checkpoint files re-downloaded | Bypass only with normalized, resolvable `file_params` (D1a); all such senders are user-initiated dialog flows (2.6, verified); tests pin batch/recipe/bulk behavior. | +| Empty-hash matching edge cases (R6) | Duplicate download of the same file, or false block | D2 rule: hash only when both non-empty; name otherwise; never `""==""`. Residual risk documented (2.4). | +| Phase-1 late-failure UX (F12) | User picks a downloaded file, fails only after location step | Toast surfacing (5.2.3); phase 2 disables downloaded files up front. | +| Phase-2 index change corrupts existing behavior | Recipe matching, delete flows | Additive index + re-point only; `version_index` read semantics unchanged; `remove_models`/`update_single_model_cache` already rebuild (F3); tests. | +| `delete_model_version` marks whole version deleted while sibling remains (R5) | History wrongly suppresses re-download of the surviving sibling's version | Phase 2 item 6.1.5; documented until then. | +| History child-table migration failure on user installs | Service init crash | `CREATE TABLE IF NOT EXISTS` in `_initialize_schema`; failure degrades to version-level behavior (per-file queries return empty). | +| Batch-preview badge misleading for partial versions (F5) | Minor UX confusion | Acknowledged in phase 1; fixed in phase 2 item 6.2.6. | +| UI confusion: version shows "In Library" while files remain downloadable | Support burden | Phase 2: per-file disabled state + partial-download tooltip. | +| Hash-identical sibling files (repacked content) | Second file blocked | Acceptable: scanner hash dedup already collapses them. | + +## 8. Rollout + +1. **Commit 1** — `fix(download): allow downloading additional files of an in-library model version (#1058)` → Phase 1 (5.1–5.3). +2. **Commit 2** — `feat(download): per-file download status and multi-file selection (#1058)` → Phase 2 (6.1–6.3). + +Phase 1 alone resolves the issue as reported; phase 2 can ship in a later release if review prefers smaller increments. + +## 9. Effort Estimate (revised after review) + +| Phase | Backend | Frontend | Tests | Risk | +|---|---|---|---|---| +| 1 | ~150–220 LOC (+ queue-retry fix ~30) | ~10–30 LOC | ~150–250 LOC | Low | +| 2 | ~250–350 LOC | ~250–350 LOC (multi-file loop refactor + ModelVersionsTab + batch badge) | ~250–350 LOC (dialog tests greenfield) | Medium | diff --git a/py/routes/handlers/model_handlers.py b/py/routes/handlers/model_handlers.py index 0ef47c21..a22501f9 100644 --- a/py/routes/handlers/model_handlers.py +++ b/py/routes/handlers/model_handlers.py @@ -1659,7 +1659,8 @@ class ModelDownloadHandler: import json try: - data["file_params"] = json.loads(file_params_json) + # Normalize falsy payloads (e.g. {}) to None (#1058) + data["file_params"] = json.loads(file_params_json) or None except json.JSONDecodeError: self._logger.warning( "Invalid file_params JSON: %s", file_params_json @@ -1811,7 +1812,8 @@ class ModelDownloadHandler: model_id = int(model_id_str) if model_id_str else None model_version_id = int(model_version_id_str) if model_version_id_str else None - file_params = json.loads(file_params_json) if file_params_json else None + # Normalize falsy payloads (e.g. {}) to None (#1058) + file_params = (json.loads(file_params_json) if file_params_json else None) or None service = await DownloadQueueService.get_instance() item = await service.add_to_queue( diff --git a/py/services/download_coordinator.py b/py/services/download_coordinator.py index 2b0f6466..4577ba57 100644 --- a/py/services/download_coordinator.py +++ b/py/services/download_coordinator.py @@ -87,7 +87,9 @@ class DownloadCoordinator: progress_callback=progress_callback, download_id=download_id, source=payload.get("source"), - file_params=payload.get("file_params"), + # Normalize falsy file_params (e.g. {}) to None so download gates + # treat it as "no explicit file selection" (#1058). + file_params=payload.get("file_params") or None, ) result["download_id"] = download_id diff --git a/py/services/download_manager.py b/py/services/download_manager.py index ff4a2ca6..73e43f99 100644 --- a/py/services/download_manager.py +++ b/py/services/download_manager.py @@ -213,6 +213,162 @@ class DownloadManager: ) return False + async def _get_scanner_for_model_type(self, model_type: str): + """Return the scanner responsible for the given model type.""" + if model_type == "checkpoint": + return await self._get_checkpoint_scanner() + if model_type == "embedding": + return await ServiceRegistry.get_embedding_scanner() + return await self._get_lora_scanner() + + @staticmethod + def _resolve_target_file( + files: Any, file_params: Dict[str, Any] | None + ) -> Optional[Dict[str, Any]]: + """Resolve the target file within a version's file list from file_params. + + Shared by the existence gate and the actual file selection so both + always agree on which file a download refers to (#1058). Returns None + when file_params is None or no file matches. + """ + if not file_params or not isinstance(files, list): + return None + + target_file_id = file_params.get("id") + target_type = file_params.get("type", "Model") + target_format = file_params.get("format") + target_size = file_params.get("size") + target_fp = file_params.get("fp") + is_primary = file_params.get("isPrimary", False) + + logger.debug( + "[download] file_params received: id=%s, type=%s, format=%s, size=%s, fp=%s, " + "isPrimary=%s, total_files=%d", + target_file_id, target_type, target_format, target_size, target_fp, + is_primary, len(files), + ) + + file_info: Optional[Dict[str, Any]] = None + + if target_file_id: + target_id_str = str(target_file_id) + for f in files: + if not isinstance(f, dict): + continue + f_id = f.get("id") + if str(f_id) == target_id_str: + file_info = f + logger.debug( + "[download] MATCH by ID: id=%s name='%s'", + f_id, f.get("name"), + ) + break + if not file_info: + logger.debug("[download] No file found with id=%s", target_file_id) + + elif is_primary: + file_info = next( + ( + f + for f in files + if isinstance(f, dict) + and f.get("primary") + and f.get("type") in MODEL_WEIGHT_FILE_TYPES + ), + None, + ) + else: + # Lenient metadata match: only compare fields present on both sides + for f in files: + if not isinstance(f, dict): + continue + f_type = f.get("type", "") + if f_type != target_type: + continue + + f_meta = f.get("metadata", {}) + f_format = f_meta.get("format") or f.get("format") + f_size = f_meta.get("size") or f.get("size") + f_fp = f_meta.get("fp") or f.get("fp") + + if target_format and f_format != target_format: + continue + if target_size and f_size and f_size != target_size: + continue + if target_fp and f_fp and f_fp != target_fp: + continue + + file_info = f + break + + return file_info + + async def _find_local_file_entry( + self, + model_type: str, + model_version_id: int, + target_file: Dict[str, Any], + ) -> Optional[Dict[str, Any]]: + """Find a local library entry for a specific file of a model version. + + Matches per design rule D2 (#1058): SHA256 is only compared when both + sides carry a non-empty hash; otherwise fall back to (extension-less) + file name equality. Never let two empty hashes compare equal. + """ + try: + normalized_version_id = int(model_version_id) + except (TypeError, ValueError): + return None + + try: + scanner = await self._get_scanner_for_model_type(model_type) + cache = await scanner.get_cached_data() + except Exception as exc: + logger.debug( + "Failed to scan local entries for version %s file check: %s", + model_version_id, + exc, + ) + return None + + raw_data = getattr(cache, "raw_data", None) if cache else None + if not raw_data: + return None + + target_hash = str( + (target_file.get("hashes") or {}).get("SHA256") or "" + ).strip().lower() + target_name = str(target_file.get("name") or "").strip() + target_base = os.path.splitext(target_name)[0] if target_name else "" + + for item in raw_data: + if not isinstance(item, dict): + continue + civitai_data = item.get("civitai") + if not isinstance(civitai_data, dict): + continue + try: + item_version_id = int(civitai_data.get("id")) + except (TypeError, ValueError): + continue + if item_version_id != normalized_version_id: + continue + + local_hash = str(item.get("sha256") or "").strip().lower() + if target_hash and local_hash: + if local_hash == target_hash: + return item + # Both sides carry hashes that differ: this is a different + # file of the same version — do not fall back to name match. + continue + + if target_base: + local_name = str(item.get("file_name") or "").strip() + if local_name == target_base: + return item + + return None + async def download_from_civitai( self, model_id: int | None = None, @@ -242,6 +398,10 @@ class DownloadManager: Returns: Dict with download result """ + # Normalize falsy file_params (e.g. an empty dict from API JSON + # parsing) to None so gate conditions behave consistently (#1058). + file_params = file_params or None + logger.debug( "[download] download_from_civitai called: model_id=%s, model_version_id=%s, " "source=%s, file_params=%s", @@ -1152,9 +1312,13 @@ class DownloadManager: use_save_dir_as_root: bool = False, ) -> Dict[str, Any]: """Wrapper for original download_from_civitai implementation""" + file_params = file_params or None try: - # Check if model version already exists in library - if model_version_id is not None: + # Check if model version already exists in library. + # With an explicit file selection (file_params) the version-level + # check is deferred until after the metadata fetch, when the target + # file can be resolved and checked individually (#1058). + if model_version_id is not None and file_params is None: # Check both scanners lora_scanner = await self._get_lora_scanner() checkpoint_scanner = await self._get_checkpoint_scanner() @@ -1235,8 +1399,26 @@ class DownloadManager: except (TypeError, ValueError): resolved_version_id = None + # Resolve the explicitly selected file (if any) up front so the + # existence gates and the actual file selection below always agree + # on the target file (#1058). + target_file: Optional[Dict[str, Any]] = None + if file_params is not None: + target_file = self._resolve_target_file( + version_info.get("files") or [], file_params + ) + if target_file is None: + logger.warning( + "[download] file_params provided but no file matched; " + "falling back to version-level checks and primary file " + "selection (model_version_id=%s)", + resolved_version_id, + ) + explicit_file = target_file is not None + if ( - get_settings_manager().get_skip_previously_downloaded_model_versions() + not explicit_file + and get_settings_manager().get_skip_previously_downloaded_model_versions() and resolved_version_id is not None and await self._has_been_downloaded(model_type, resolved_version_id) ): @@ -1346,9 +1528,38 @@ class DownloadManager: f"baseModel '{base_model_value}' is a known diffusion model, routing to unet folder" ) - # Case 2: model_version_id was None, check after getting version_info - if model_version_id is None: - version_id = version_info.get("id") + # Existence check after the metadata fetch (#1058): + # - An explicit file selection only blocks when THIS file is + # already in the library; other files of the same version + # remain downloadable. + # - Without file_params (or when file_params failed to resolve), + # keep version-level protection. The case "model_version_id + # given + no file_params" was already covered by the early + # gate above. + if explicit_file and resolved_version_id is not None: + existing_entry = await self._find_local_file_entry( + model_type, resolved_version_id, target_file + ) + if existing_entry is not None: + error_message = ( + f"File '{target_file.get('name')}' from model version " + f"{resolved_version_id} already exists in {model_type} library" + ) + logger.info("[download] %s", error_message) + return {"success": False, "error": error_message} + logger.info( + "[download] File '%s' of model version %s not in %s library — " + "download allowed (other files of this version may exist locally)", + target_file.get("name"), resolved_version_id, model_type, + ) + elif file_params is not None or model_version_id is None: + # Case 2: model_version_id was None, or file_params did not + # resolve to a concrete file — check at version level. + version_id = ( + resolved_version_id + if resolved_version_id is not None + else version_info.get("id") + ) if model_type == "lora": # Check lora scanner @@ -1495,73 +1706,16 @@ class DownloadManager: files = version_info.get("files", []) file_info = None - # If file_params is provided, try to find matching file - if file_params and model_version_id: - target_file_id = file_params.get("id") - target_type = file_params.get("type", "Model") - target_format = file_params.get("format") - target_size = file_params.get("size") - target_fp = file_params.get("fp") - is_primary = file_params.get("isPrimary", False) - - logger.debug( - "[download] file_params received: id=%s, type=%s, format=%s, size=%s, fp=%s, isPrimary=%s, " - "model_version_id=%s, total_files=%d", - target_file_id, target_type, target_format, target_size, target_fp, is_primary, - model_version_id, len(files), - ) - - if target_file_id: - target_id_str = str(target_file_id) - for f in files: - f_id = f.get("id") - if str(f_id) == target_id_str: - file_info = f - logger.debug( - "[download] MATCH by ID: id=%s name='%s'", - f_id, f.get("name"), - ) - break - if not file_info: - logger.debug("[download] No file found with id=%s", target_file_id) - - elif is_primary: - file_info = next( - ( - f - for f in files - if f.get("primary") - and f.get("type") in MODEL_WEIGHT_FILE_TYPES - ), - None, - ) - else: - # Lenient metadata match: only compare fields present on both sides - for f in files: - f_type = f.get("type", "") - if f_type != target_type: - continue - - f_meta = f.get("metadata", {}) - f_format = f_meta.get("format") or f.get("format") - f_size = f_meta.get("size") or f.get("size") - f_fp = f_meta.get("fp") or f.get("fp") - - if target_format and f_format != target_format: - continue - if target_size and f_size and f_size != target_size: - continue - if target_fp and f_fp and f_fp != target_fp: - continue - - file_info = f - break - + # If file_params is provided, reuse the file resolved right after + # the metadata fetch so the existence gate and this selection + # always agree on the target file (#1058). + if file_params is not None: + file_info = target_file if not file_info: logger.debug( "[download] No match found via file_params — falling back to primary file lookup", ) - elif not file_params: + else: logger.debug( "[download] No file_params provided (null/None) — will use primary file lookup. " "model_version_id=%s, total_files=%d", diff --git a/py/services/download_queue_service.py b/py/services/download_queue_service.py index 09ae068f..4c87699d 100644 --- a/py/services/download_queue_service.py +++ b/py/services/download_queue_service.py @@ -64,6 +64,7 @@ class DownloadQueueService: model_name TEXT NOT NULL DEFAULT '', version_name TEXT DEFAULT '', thumbnail_url TEXT DEFAULT '', + file_params TEXT, status TEXT NOT NULL, error TEXT, file_path TEXT, @@ -120,6 +121,18 @@ class DownloadQueueService: with self._connect() as conn: conn.executescript(self._SCHEMA_TABLES) + # Databases created by older versions lack + # download_history.file_params; add it so retry-from-history can + # restore the originally selected file (#1058). + history_columns = { + row["name"] + for row in conn.execute("PRAGMA table_info(download_history)") + } + if "file_params" not in history_columns: + conn.execute( + "ALTER TABLE download_history ADD COLUMN file_params TEXT" + ) + # Creating the unique index on download_history.download_id can # fail if pre-existing rows have duplicate values (e.g. from a # previous version that lacked the index). Deduplicate first so @@ -418,6 +431,12 @@ class DownloadQueueService: return None now = completed_at if completed_at is not None else time.time() + # Guard against legacy databases whose download_queue table + # predates the file_params column. + queue_columns = set(row.keys()) + file_params_json = ( + row["file_params"] if "file_params" in queue_columns else None + ) conn.execute( "DELETE FROM download_queue WHERE download_id = ?", (download_id,), @@ -426,9 +445,9 @@ class DownloadQueueService: """ INSERT OR IGNORE INTO download_history ( download_id, model_id, model_version_id, model_name, - version_name, thumbnail_url, status, error, file_path, - bytes_downloaded, total_bytes, completed_at - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + version_name, thumbnail_url, file_params, status, error, + file_path, bytes_downloaded, total_bytes, completed_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( row["download_id"], @@ -437,6 +456,7 @@ class DownloadQueueService: row["model_name"], row["version_name"], row["thumbnail_url"], + file_params_json, status, error, file_path, @@ -503,6 +523,7 @@ class DownloadQueueService: bytes_downloaded: int = 0, total_bytes: Optional[int] = None, is_already_exists: int = 0, + file_params: Optional[dict[str, Any]] = None, ) -> int: """Insert a record into the download history. @@ -510,6 +531,7 @@ class DownloadQueueService: inserted row. """ now = time.time() + file_params_json = json.dumps(file_params) if file_params is not None else None async with self._lock: conn = self._get_conn() @@ -517,9 +539,10 @@ class DownloadQueueService: """ INSERT INTO download_history ( download_id, model_id, model_version_id, model_name, - version_name, thumbnail_url, status, error, file_path, - bytes_downloaded, total_bytes, completed_at, is_already_exists - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + version_name, thumbnail_url, file_params, status, error, + file_path, bytes_downloaded, total_bytes, completed_at, + is_already_exists + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( download_id, @@ -528,6 +551,7 @@ class DownloadQueueService: model_name, version_name, thumbnail_url, + file_params_json, status, error, file_path, @@ -702,7 +726,7 @@ class DownloadQueueService: download_id, model_id, model_version_id, model_name, version_name, thumbnail_url, source, file_params, status, priority, added_at - ) VALUES (?, ?, ?, ?, ?, ?, ?, NULL, 'queued', 0, ?) + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'queued', 0, ?) """, ( new_id, @@ -712,6 +736,7 @@ class DownloadQueueService: row["version_name"], row["thumbnail_url"], "retry", + row["file_params"], now, ), ) @@ -755,7 +780,7 @@ class DownloadQueueService: download_id, model_id, model_version_id, model_name, version_name, thumbnail_url, source, file_params, status, priority, added_at - ) VALUES (?, ?, ?, ?, ?, ?, ?, NULL, 'queued', 0, ?) + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'queued', 0, ?) """, ( new_id, @@ -765,6 +790,7 @@ class DownloadQueueService: row["version_name"], row["thumbnail_url"], "retry", + row["file_params"], now, ), ) diff --git a/static/js/managers/DownloadManager.js b/static/js/managers/DownloadManager.js index 2cb5b50b..c8d5c2a0 100644 --- a/static/js/managers/DownloadManager.js +++ b/static/js/managers/DownloadManager.js @@ -595,7 +595,10 @@ export class DownloadManager { `; } - const fileBadge = modelFiles.length > 1 && !existsLocally + // Always offer the file-selection entry for multi-file versions, + // even when the version is already (partially) in the library, so + // remaining files can still be downloaded (#1058). + const fileBadge = modelFiles.length > 1 ? ` ${modelFiles.length} ${translate('modals.download.fileSelection.files')} ` @@ -1016,6 +1019,16 @@ export class DownloadManager { if (!response?.success) { this.loadingManager.setStatus(translate('modals.download.status.finalizing')); + const errorMessage = response?.error || 'Unknown error'; + // A file-level "already in library" rejection is an expected + // outcome when browsing files of a partially downloaded + // version — surface it as a lightweight toast instead of the + // failure summary modal so the user can simply go back and + // pick another file (#1058). + if (typeof errorMessage === 'string' && errorMessage.includes('already exists in')) { + showToast(errorMessage, {}, 'info'); + return false; + } showDownloadBatchSummary({ total: 1, completed: 0, @@ -1026,7 +1039,7 @@ export class DownloadManager { source, url: this._buildSingleItemUrl({ modelId, versionId, source }), }, - error: response?.error || 'Unknown error', + error: errorMessage, name: displayName, }], onRetry: () => this.executeDownloadWithProgress(retryParams), @@ -1610,6 +1623,7 @@ export class DownloadManager { const fileParams = this.selectedFile ? { id: this.selectedFile.id, + name: this.selectedFile.name || null, type: this.selectedFile.type || 'Model', format: this.selectedFile.metadata?.format || null, size: this.selectedFile.metadata?.size || null, diff --git a/tests/frontend/managers/downloadManager.multiFile.test.js b/tests/frontend/managers/downloadManager.multiFile.test.js new file mode 100644 index 00000000..41a9c5f3 --- /dev/null +++ b/tests/frontend/managers/downloadManager.multiFile.test.js @@ -0,0 +1,164 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; + +const { + DOWNLOAD_MANAGER_MODULE, + MODAL_MANAGER_MODULE, + UI_HELPERS_MODULE, + STATE_MODULE, + LOADING_MANAGER_MODULE, + API_FACTORY_MODULE, + STORAGE_HELPERS_MODULE, + FOLDER_TREE_MANAGER_MODULE, + I18N_HELPERS_MODULE, +} = vi.hoisted(() => ({ + DOWNLOAD_MANAGER_MODULE: new URL('../../../static/js/managers/DownloadManager.js', import.meta.url).pathname, + MODAL_MANAGER_MODULE: new URL('../../../static/js/managers/ModalManager.js', import.meta.url).pathname, + UI_HELPERS_MODULE: new URL('../../../static/js/utils/uiHelpers.js', import.meta.url).pathname, + STATE_MODULE: new URL('../../../static/js/state/index.js', import.meta.url).pathname, + LOADING_MANAGER_MODULE: new URL('../../../static/js/managers/LoadingManager.js', import.meta.url).pathname, + API_FACTORY_MODULE: new URL('../../../static/js/api/modelApiFactory.js', import.meta.url).pathname, + STORAGE_HELPERS_MODULE: new URL('../../../static/js/utils/storageHelpers.js', import.meta.url).pathname, + FOLDER_TREE_MANAGER_MODULE: new URL('../../../static/js/components/FolderTreeManager.js', import.meta.url).pathname, + I18N_HELPERS_MODULE: new URL('../../../static/js/utils/i18nHelpers.js', import.meta.url).pathname, +})); + +vi.mock(MODAL_MANAGER_MODULE, () => ({ + modalManager: { + showModal: vi.fn(), + closeModal: vi.fn(), + }, +})); + +vi.mock(UI_HELPERS_MODULE, () => ({ + showToast: vi.fn(), +})); + +vi.mock(STATE_MODULE, () => ({ + state: { + global: { + settings: {}, + }, + }, +})); + +vi.mock(LOADING_MANAGER_MODULE, () => ({ + LoadingManager: vi.fn(() => ({ + showSimpleLoading: vi.fn(), + hide: vi.fn(), + restoreProgressBar: vi.fn(), + showDownloadProgress: vi.fn(() => vi.fn()), + setStatus: vi.fn(), + })), +})); + +vi.mock(API_FACTORY_MODULE, () => ({ + getModelApiClient: vi.fn(() => ({ + apiConfig: { + config: { + displayName: 'LoRA', + singularName: 'lora', + }, + }, + })), + resetAndReload: vi.fn(), +})); + +vi.mock(STORAGE_HELPERS_MODULE, () => ({ + getStorageItem: vi.fn((_key, defaultValue) => defaultValue), + setStorageItem: vi.fn(), +})); + +vi.mock(FOLDER_TREE_MANAGER_MODULE, () => ({ + FolderTreeManager: vi.fn(() => ({ + clearSelection: vi.fn(), + init: vi.fn(), + })), +})); + +vi.mock(I18N_HELPERS_MODULE, () => ({ + translate: vi.fn((_, __, fallback) => fallback ?? ''), +})); + +const MULTI_FILE_VERSION = { + id: 201, + name: 'Multi-file version', + images: [], + files: [ + { id: 1001, type: 'Model', sizeKB: 2048, name: 'file-a.safetensors' }, + { id: 1002, type: 'Model', sizeKB: 2048, name: 'file-b.safetensors' }, + ], + createdAt: '2026-01-01T00:00:00Z', +}; + +const SINGLE_FILE_VERSION = { + id: 202, + name: 'Single-file version', + images: [], + files: [{ id: 1003, type: 'Model', sizeKB: 2048, name: 'file-c.safetensors' }], + createdAt: '2026-01-01T00:00:00Z', +}; + +describe('DownloadManager multi-file version badge (#1058)', () => { + let DownloadManager; + + beforeEach(async () => { + vi.resetModules(); + document.body.innerHTML = ` +
+
+
+ + `; + ({ DownloadManager } = await import(DOWNLOAD_MANAGER_MODULE)); + }); + + afterEach(() => { + document.body.innerHTML = ''; + }); + + it('shows the file-select badge for a multi-file version already in library', () => { + const manager = new DownloadManager(); + manager.versions = [ + { + ...MULTI_FILE_VERSION, + existsLocally: true, + localPath: '/models/loras/file-a.safetensors', + }, + ]; + + manager.showVersionStep(); + + const badge = document.querySelector('.file-select-badge'); + expect(badge).not.toBeNull(); + expect(badge.dataset.versionId).toBe('201'); + expect(badge.textContent).toContain('2'); + // The in-library badge is still shown alongside + expect(document.querySelector('.local-badge')).not.toBeNull(); + }); + + it('still shows the file-select badge for a multi-file version not in library', () => { + const manager = new DownloadManager(); + manager.versions = [{ ...MULTI_FILE_VERSION, existsLocally: false }]; + + manager.showVersionStep(); + + const badge = document.querySelector('.file-select-badge'); + expect(badge).not.toBeNull(); + expect(badge.dataset.versionId).toBe('201'); + }); + + it('does not show the file-select badge for a single-file version in library', () => { + const manager = new DownloadManager(); + manager.versions = [ + { + ...SINGLE_FILE_VERSION, + existsLocally: true, + localPath: '/models/loras/file-c.safetensors', + }, + ]; + + manager.showVersionStep(); + + expect(document.querySelector('.file-select-badge')).toBeNull(); + }); +}); diff --git a/tests/services/test_download_manager_basic.py b/tests/services/test_download_manager_basic.py index cc98794a..ae807463 100644 --- a/tests/services/test_download_manager_basic.py +++ b/tests/services/test_download_manager_basic.py @@ -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 diff --git a/tests/services/test_download_queue_service.py b/tests/services/test_download_queue_service.py index c46c7794..7b8c50b6 100644 --- a/tests/services/test_download_queue_service.py +++ b/tests/services/test_download_queue_service.py @@ -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}