From 259e08e47ca2a168893295979b1ae94261e8be63 Mon Sep 17 00:00:00 2001 From: Will Miao Date: Thu, 20 Aug 2026 20:59:19 +0800 Subject: [PATCH] feat(download): expose per-file downloadedFiles in check-model-exists (#1058) The version branch of check-model-exists now returns downloadedFiles: [{fileId, fileName, filePath}] so clients (e.g. the browser extension) can tell a partially downloaded version apart from a fully downloaded one. Reuses ModelCivitaiHandler._match_downloaded_files (D2 rule) against the local cache; unmatchable local files are reported with fileId: None. No CivitAI API call added. --- py/routes/handlers/misc_handlers.py | 66 +++++++++++++ tests/routes/test_misc_routes.py | 143 ++++++++++++++++++++++++++++ 2 files changed, 209 insertions(+) diff --git a/py/routes/handlers/misc_handlers.py b/py/routes/handlers/misc_handlers.py index ec912a5d..5680d68e 100644 --- a/py/routes/handlers/misc_handlers.py +++ b/py/routes/handlers/misc_handlers.py @@ -56,6 +56,7 @@ from ...utils.constants import ( ) from .hf_handlers import HfHandler from .agent_handlers import AgentHandler +from .model_handlers import ModelCivitaiHandler from ...utils.civitai_utils import rewrite_preview_url from ...utils.example_images_paths import ( find_non_compliant_items_in_example_images_root, @@ -2061,6 +2062,63 @@ class ModelLibraryHandler: enriched.append(entry) return enriched + @staticmethod + async def _get_downloaded_files( + scanner: Any, model_version_id: int + ) -> list[dict[str, Any]]: + """Return per-file downloaded state for a version in the library. + + This handler has no CivitAI version payload, so the remote file list + is taken from the local entries' cached ``civitai`` metadata (the + full version payload persisted at download time, see + ``BaseModelMetadata.from_civitai_info``) and matched with the same + D2 rule used by ``get_civitai_versions`` (#1058). Local entries that + cannot be matched to a known remote file (e.g. missing metadata or + renamed files) are still reported with ``fileId`` set to None. + Returns ``[{fileId, fileName, filePath}]``. + """ + try: + cache = await scanner.get_cached_data() + except Exception: # pragma: no cover - defensive fallback + logger.debug( + "Failed to read cache for downloaded files of version %s", + model_version_id, + exc_info=True, + ) + return [] + + files_getter = getattr(cache, "get_files_by_version_id", None) + local_entries = files_getter(model_version_id) if files_getter else [] + if not local_entries: + return [] + + version_payload: Mapping[str, Any] = {} + for entry in local_entries: + civitai = entry.get("civitai") if isinstance(entry, Mapping) else None + if isinstance(civitai, Mapping) and isinstance(civitai.get("files"), list): + version_payload = civitai + break + + downloaded = ModelCivitaiHandler._match_downloaded_files( + version_payload, local_entries + ) + + # Surface local files that D2 could not map to a known remote file + matched_paths = {item.get("filePath") for item in downloaded} + for entry in local_entries: + if not isinstance(entry, Mapping): + continue + if entry.get("file_path") in matched_paths: + continue + downloaded.append( + { + "fileId": None, + "fileName": entry.get("file_name"), + "filePath": entry.get("file_path"), + } + ) + return downloaded + async def check_model_exists(self, request: web.Request) -> web.Response: try: model_id_str = request.query.get("modelId") @@ -2096,9 +2154,11 @@ class ModelLibraryHandler: exists = False model_type = None + matched_scanner = None if await lora_scanner.check_model_version_exists(model_version_id): exists = True model_type = "lora" + matched_scanner = lora_scanner elif ( checkpoint_scanner and await checkpoint_scanner.check_model_version_exists( @@ -2107,6 +2167,7 @@ class ModelLibraryHandler: ): exists = True model_type = "checkpoint" + matched_scanner = checkpoint_scanner elif ( embedding_scanner and await embedding_scanner.check_model_version_exists( @@ -2115,6 +2176,7 @@ class ModelLibraryHandler: ): exists = True model_type = "embedding" + matched_scanner = embedding_scanner if exists: return web.json_response( @@ -2123,6 +2185,9 @@ class ModelLibraryHandler: "exists": True, "modelType": model_type, "hasBeenDownloaded": False, + "downloadedFiles": await self._get_downloaded_files( + matched_scanner, model_version_id + ), } ) @@ -2144,6 +2209,7 @@ class ModelLibraryHandler: "exists": False, "modelType": history_type, "hasBeenDownloaded": has_been_downloaded, + "downloadedFiles": [], } ) diff --git a/tests/routes/test_misc_routes.py b/tests/routes/test_misc_routes.py index eed69553..132d1e6f 100644 --- a/tests/routes/test_misc_routes.py +++ b/tests/routes/test_misc_routes.py @@ -908,6 +908,33 @@ class FakeExistenceScanner: return [] +class FakeVersionFilesCache: + """Cache stub exposing the multi-valued version->files index (#1058).""" + + def __init__(self, entries): + self._entries = entries + + def get_files_by_version_id(self, _version_id): + return list(self._entries) + + +class FakeDownloadedFilesScanner: + """Scanner stub with one known version and per-file cache entries.""" + + def __init__(self, version_id, entries): + self._version_id = version_id + self._cache = FakeVersionFilesCache(entries) + + async def check_model_version_exists(self, version_id): + return version_id == self._version_id + + async def get_model_versions_by_id(self, _model_id): + return [] + + async def get_cached_data(self): + return self._cache + + class FakeMetadataProvider: async def get_model_versions(self, _model_id): return {"modelVersions": [], "name": "", "type": "lora"} @@ -1524,9 +1551,125 @@ async def test_check_model_exists_returns_download_history_when_file_missing(): "exists": False, "modelType": "checkpoint", "hasBeenDownloaded": True, + "downloadedFiles": [], } +@pytest.mark.asyncio +async def test_check_model_exists_returns_downloaded_files(): + """The version branch exposes per-file downloaded state (#1058).""" + civitai_payload = { + "id": 100, + "files": [ + { + "id": 1001, + "name": "model-fp16.safetensors", + "hashes": {"SHA256": "A" * 64}, + }, + { + "id": 1002, + "name": "model-fp32.safetensors", + "hashes": {"SHA256": "B" * 64}, + }, + ], + } + entries = [ + { + "file_name": "model-fp16", + "file_path": "/models/loras/model-fp16.safetensors", + "sha256": "a" * 64, + "civitai": civitai_payload, + }, + { + # Legacy entry without a hash: matched by extension-less name (D2) + "file_name": "model-fp32", + "file_path": "/models/loras/model-fp32.safetensors", + "sha256": "", + "civitai": civitai_payload, + }, + ] + lora_scanner = FakeDownloadedFilesScanner(100, entries) + + async def lora_factory(): + return lora_scanner + + handler = ModelLibraryHandler( + ServiceRegistryAdapter( + get_lora_scanner=lora_factory, + get_checkpoint_scanner=fake_scanner_factory, + get_embedding_scanner=fake_scanner_factory, + get_downloaded_version_history_service=fake_download_history_service_factory, + ), + metadata_provider_factory=fake_metadata_provider_factory, + ) + + response = await handler.check_model_exists( + FakeRequest(query={"modelId": "5", "modelVersionId": "100"}) # pyright: ignore[reportArgumentType] + ) + payload = _json_payload(response) + + assert payload == { + "success": True, + "exists": True, + "modelType": "lora", + "hasBeenDownloaded": False, + "downloadedFiles": [ + { + "fileId": 1001, + "fileName": "model-fp16.safetensors", + "filePath": "/models/loras/model-fp16.safetensors", + }, + { + "fileId": 1002, + "fileName": "model-fp32.safetensors", + "filePath": "/models/loras/model-fp32.safetensors", + }, + ], + } + + +@pytest.mark.asyncio +async def test_check_model_exists_downloaded_files_without_remote_metadata(): + """Unmatchable local files are still reported with fileId None (#1058).""" + entries = [ + { + "file_name": "renamed-model", + "file_path": "/models/loras/renamed-model.safetensors", + "sha256": "c" * 64, + "civitai": {"id": 100}, # no remote files list cached + }, + ] + lora_scanner = FakeDownloadedFilesScanner(100, entries) + + async def lora_factory(): + return lora_scanner + + handler = ModelLibraryHandler( + ServiceRegistryAdapter( + get_lora_scanner=lora_factory, + get_checkpoint_scanner=fake_scanner_factory, + get_embedding_scanner=fake_scanner_factory, + get_downloaded_version_history_service=fake_download_history_service_factory, + ), + metadata_provider_factory=fake_metadata_provider_factory, + ) + + response = await handler.check_model_exists( + FakeRequest(query={"modelId": "5", "modelVersionId": "100"}) # pyright: ignore[reportArgumentType] + ) + payload = _json_payload(response) + + assert payload["success"] is True + assert payload["exists"] is True + assert payload["downloadedFiles"] == [ + { + "fileId": None, + "fileName": "renamed-model", + "filePath": "/models/loras/renamed-model.safetensors", + } + ] + + @pytest.mark.asyncio async def test_model_version_download_status_endpoints(): history_service = FakeDownloadHistoryService({"lora": {123}})