diff --git a/core/routes/hashes.py b/core/routes/hashes.py index 53e344f..4e8892d 100644 --- a/core/routes/hashes.py +++ b/core/routes/hashes.py @@ -1,663 +1,55 @@ -"""Local model and hash route registration.""" +"""HTTP adapters for local model and hash operations.""" +from ..services.hash_service import HashService from .context import RouteContext def register_hash_routes(context: RouteContext): - FileManagerError = context.get('FileManagerError') - FileManagerUnavailableError = context.get('FileManagerUnavailableError') - MODEL_RESOLVER_METADATA_SCHEMA = context.get('MODEL_RESOLVER_METADATA_SCHEMA') - MODEL_RESOLVER_METADATA_SCHEMA_VERSION = context.get('MODEL_RESOLVER_METADATA_SCHEMA_VERSION') - UnsupportedFileManagerPlatformError = context.get('UnsupportedFileManagerPlatformError') - asyncio = context.get('asyncio') - cancel_progress_response = context.get('cancel_progress_response') - get_existing_model_preview_path = context.get('get_existing_model_preview_path') - get_filename_from_path = context.get('get_filename_from_path') - get_local_model_hash_metadata = context.get('get_local_model_hash_metadata') - get_progress_response = context.get('get_progress_response') - get_safe_model_resolver_sidecar_path = context.get('get_safe_model_resolver_sidecar_path') - get_workflow_model_inventory = context.get('get_workflow_model_inventory') - is_path_in_configured_model_roots = context.get('is_path_in_configured_model_roots') - json_api_endpoint = context.get('json_api_endpoint') - load_resolver_settings = context.get('load_resolver_settings') - normalize_file_manager_path = context.get('normalize_file_manager_path') - normalize_sha256 = context.get('normalize_sha256') - open_in_file_manager = context.get('open_in_file_manager') - os = context.get('os') - read_json_safe = context.get('read_json_safe') - resolver_bool_setting = context.get('resolver_bool_setting') - routes = context.get('routes') - run_in_background_thread = context.get('run_in_background_thread') - search_local_matches_by_hash = context.get('search_local_matches_by_hash') - self = context.get('self') - time = context.get('time') - to_bool = context.get('to_bool') - to_int = context.get('to_int') - web = context.get('web') - write_json_atomic = context.get('write_json_atomic') + """Register hash routes and delegate behavior to :class:`HashService`.""" + service = HashService(context) + json_api_endpoint = context.get("json_api_endpoint") + routes = context.get("routes") @routes.post("/model_resolver/local-model-hashes") @json_api_endpoint("local-model-hashes") async def local_model_hashes(request): - """Return SHA256 hashes already stored in local sidecar metadata.""" - data = await request.json() - model = data.get("model") if isinstance(data.get("model"), dict) else {} - path = ( - data.get("path") - or data.get("file_path") - or data.get("resolved_path") - or model.get("path") - or model.get("resolved_path") - or "" - ) - - if not path: - return web.json_response( - {"error": "path is required"}, status=400 - ) - - normalized_path = os.path.realpath( - os.path.abspath(os.path.normpath(str(path))) - ) - if not is_path_in_configured_model_roots(normalized_path): - return web.json_response( - {"error": "path is outside configured model directories"}, - status=403, - ) - - return web.json_response( - get_local_model_hash_metadata(normalized_path, model=model) - ) + return await service.local_model_hashes(request) @routes.get("/model_resolver/model-preview") async def get_model_preview(request): - """Serve an adjacent model preview from a configured model directory.""" - model_path = str(request.query.get("path") or "").strip() - if not model_path: - return web.Response(text="path is required", status=400) - - try: - normalized_path = os.path.realpath( - os.path.abspath(os.path.normpath(model_path)) - ) - except (OSError, TypeError, ValueError): - return web.Response(text="invalid model path", status=400) - if not os.path.isfile(normalized_path): - return web.Response(text="model file does not exist", status=404) - if not is_path_in_configured_model_roots(normalized_path): - return web.Response( - text="path is outside configured model directories", - status=403, - ) - - preview_path = get_existing_model_preview_path(normalized_path) - if not preview_path: - return web.Response(text="preview not found", status=404) - if not is_path_in_configured_model_roots(preview_path): - return web.Response( - text="preview is outside configured model directories", - status=403, - ) - - return web.FileResponse( - preview_path, - headers={"Cache-Control": "private, no-cache"}, - ) + return await service.get_model_preview(request) @routes.post("/model_resolver/workflow-model-hashes") @json_api_endpoint("workflow-model-hashes") async def workflow_model_hashes(request): - """Return hash metadata for existing local models used by a workflow.""" - data = await request.json() - workflow_json = data.get("workflow") - if not isinstance(workflow_json, dict): - return web.json_response( - {"error": "Workflow JSON must be an object"}, status=400 - ) - - settings = await asyncio.to_thread(load_resolver_settings) - if not resolver_bool_setting( - settings.get("workflow_hash_metadata_enabled"), True - ): - return web.json_response( - { - "success": True, - "enabled": False, - "models": [], - "by_node": {}, - "by_path": {}, - "count": 0, - } - ) - - inventory = await asyncio.to_thread( - get_workflow_model_inventory, - workflow_json, - ) - refs = inventory["model_refs"] - - by_node = {} - by_path = {} - models = [] - seen = set() - - for ref in refs: - if not isinstance(ref, dict) or not ref.get("exists"): - continue - full_path = str(ref.get("full_path") or "").strip() - if not full_path: - continue - model_info = { - "path": full_path, - "filename": get_filename_from_path(full_path), - "relative_path": ref.get("original_path") or "", - "category": ref.get("category") or "", - } - metadata = get_local_model_hash_metadata( - full_path, model=model_info - ) - sha256 = normalize_sha256(metadata.get("sha256")) - if not sha256: - continue - - entry = { - "node_id": ref.get("node_id"), - "node_type": ref.get("node_type") or "", - "widget_index": ref.get("widget_index"), - "widget_name": ref.get("widget_name") or "", - "path": ref.get("original_path") or "", - "filename": get_filename_from_path( - ref.get("original_path") or full_path - ), - "category": ref.get("category") or "", - "sha256": sha256, - "size": metadata.get("size") or 0, - } - entry_key = ( - str(entry.get("node_id")), - str(entry.get("widget_index")), - entry.get("path"), - sha256, - ) - if entry_key in seen: - continue - seen.add(entry_key) - models.append(entry) - - path_key = str(entry.get("path") or entry.get("filename") or "") - if path_key: - by_path[path_key] = entry - by_path[get_filename_from_path(path_key)] = entry - node_key = f"{entry.get('node_id')}:{entry.get('widget_index')}" - by_node[node_key] = entry - - return web.json_response( - { - "success": True, - "enabled": True, - "models": models, - "by_node": by_node, - "by_path": by_path, - "count": len(models), - } - ) + return await service.workflow_model_hashes(request) @routes.post("/model_resolver/local-matches-by-hash") @json_api_endpoint("local-matches-by-hash") async def local_matches_by_hash(request): - """Search local model metadata sidecars for a remote SHA256.""" - data = await request.json() - sha256 = normalize_sha256( - data.get("sha256") - or data.get("hash") - or data.get("SHA256") - or "" - ) - if not sha256: - return web.json_response( - {"error": "sha256 is required"}, status=400 - ) - - category = data.get("category", "") - source = str( - data.get("source") - or data.get("hash_lookup_source") - or "download_source" - ).strip().lower().replace("-", "_") - filename = ( - data.get("filename") - or data.get("path") - or data.get("model_name") - or "" - ) - max_matches = to_int(data.get("max_matches"), 20) - force_rescan = to_bool(data.get("force_rescan"), False) - - matches = search_local_matches_by_hash( - sha256, - category=category or None, - max_matches=max_matches, - force_rescan=force_rescan, - ) - enriched_matches = [ - { - **match, - "hash_lookup_source": source or "download_source", - "hash_lookup_filename": filename, - "hash_lookup_sha256": sha256, - } - for match in matches - ] - return web.json_response( - { - "sha256": sha256, - "local_hash_matches": enriched_matches, - "matches": enriched_matches, - } - ) + return await service.local_matches_by_hash(request) @routes.post("/model_resolver/open-containing-folder") @json_api_endpoint("open-containing-folder") async def open_containing_folder(request): - """Reveal a file or open a directory in the host file manager.""" - try: - data = await request.json() - except Exception as exc: - return web.json_response( - {"success": False, "error": f"invalid JSON body: {exc}"}, - status=400, - ) - if not isinstance(data, dict): - return web.json_response( - {"success": False, "error": "JSON body must be an object"}, - status=400, - ) - target_path = data.get("path", "") - - try: - normalized_path = normalize_file_manager_path(target_path) - except ValueError as exc: - return web.json_response( - {"success": False, "error": str(exc)}, status=400 - ) - - if not os.path.exists(normalized_path): - return web.json_response( - {"success": False, "error": "path does not exist"}, status=404 - ) - if not is_path_in_configured_model_roots(normalized_path): - return web.json_response( - { - "success": False, - "error": "path is outside configured model directories", - }, - status=403, - ) - - try: - result = await asyncio.to_thread( - open_in_file_manager, normalized_path - ) - except FileNotFoundError: - return web.json_response( - {"success": False, "error": "path no longer exists"}, - status=404, - ) - except UnsupportedFileManagerPlatformError as exc: - return web.json_response( - {"success": False, "error": str(exc)}, status=501 - ) - except FileManagerUnavailableError as exc: - return web.json_response( - { - "success": False, - "error": ( - f"{exc} The ComfyUI host may be running without " - "a graphical desktop session." - ), - }, - status=503, - ) - except FileManagerError as exc: - return web.json_response( - {"success": False, "error": str(exc)}, status=500 - ) - - return web.json_response({"success": True, **result}) - - # cleanup_hash_progress, update_hash_progress, is_hash_progress_cancelled, and mark_hash_progress_cancelled removed - - from ..path_utils import HashCalculationCancelled - - def resolve_hash_file_request(data): - import os as _os - - file_path = ( - data.get("file_path") - or data.get("resolved_path") - or data.get("path") - or "" - ) - if not file_path: - return "", "file_path is required" - - normalized_path = _os.path.realpath( - _os.path.abspath(_os.path.normpath(file_path)) - ) - if not _os.path.exists(normalized_path) or not _os.path.isfile(normalized_path): - return "", "file does not exist" - if not is_path_in_configured_model_roots(normalized_path): - return "", "file is outside configured model directories" - return normalized_path, "" - - def write_calculated_hash_metadata( - normalized_path, - sha256, - sha256_source="file", - ): - import os as _os - - resolved_metadata_path = get_safe_model_resolver_sidecar_path( - normalized_path - ) - - metadata_updated = False - try: - metadata = read_json_safe(resolved_metadata_path, {}) - if not isinstance(metadata, dict): - metadata = {} - - filename = get_filename_from_path(normalized_path) - stem, _ext = _os.path.splitext(filename) - hashes = metadata.get("hashes") - if not isinstance(hashes, dict): - hashes = {} - hashes["SHA256"] = sha256 - - metadata["schema"] = MODEL_RESOLVER_METADATA_SCHEMA - metadata["schema_version"] = ( - MODEL_RESOLVER_METADATA_SCHEMA_VERSION - ) - metadata["managed_by"] = MODEL_RESOLVER_METADATA_SCHEMA - metadata["sha256"] = sha256 - metadata["hashes"] = hashes - metadata["hash_status"] = "completed" - metadata["sha256_source"] = sha256_source or "file" - metadata["last_checked_at"] = time.time() - metadata.setdefault("file_name", stem) - metadata.setdefault("model_name", stem) - metadata.setdefault("file_path", normalized_path.replace("\\", "/")) - try: - metadata.setdefault("size", _os.path.getsize(normalized_path)) - except Exception: - pass - - write_json_atomic(resolved_metadata_path, metadata, indent=2) - metadata_updated = True - self.logger.info( - f"Stored SHA256 and updated metadata: {resolved_metadata_path}" - ) - except Exception as metadata_error: - self.logger.warning( - f"Could not update metadata with calculated SHA256 for {normalized_path}: {metadata_error}" - ) - - return resolved_metadata_path, metadata_updated - - def calculate_sha256_with_progress(normalized_path, progress_id=""): - import os as _os - - from ..path_utils import calculate_file_sha256 as _calculate_file_sha256_core - - total_bytes = max(0, _os.path.getsize(normalized_path)) - - self.hash_tracker.update( - progress_id, - status="running", - stage="header", - message="Checking safetensors header...", - percent=0, - bytes_read=0, - total_bytes=total_bytes, - ) - - if self.hash_tracker.is_cancelled(progress_id): - raise HashCalculationCancelled() - - hash_source = ["file"] - stage_transitioned = [False] - _bytes_read = [0] - _last_update = [0.0] - - def on_progress(bytes_read, total_bytes): - _bytes_read[0] = bytes_read - if not stage_transitioned[0]: - stage_transitioned[0] = True - self.logger.info( - "No SHA256 found in safetensors header; calculating full " - f"file SHA256: {normalized_path}" - ) - self.hash_tracker.update( - progress_id, - status="running", - stage="hashing", - message="Calculating SHA256...", - percent=0, - bytes_read=0, - total_bytes=total_bytes, - ) - - now = time.time() - if now - _last_update[0] >= 0.15 or bytes_read >= total_bytes: - percent = 98 if total_bytes <= 0 else min( - 98, - (bytes_read / total_bytes) * 98, - ) - self.hash_tracker.update( - progress_id, - status="running", - stage="hashing", - message="Calculating SHA256...", - percent=percent, - bytes_read=bytes_read, - total_bytes=total_bytes, - ) - _last_update[0] = now - - def is_cancelled(): - if self.hash_tracker.is_cancelled(progress_id): - percent = 0 if total_bytes <= 0 else min( - 98, - (_bytes_read[0] / total_bytes) * 98, - ) - self.hash_tracker.update( - progress_id, - status="cancelled", - stage="cancelled", - message="Hash calculation cancelled", - percent=percent, - bytes_read=_bytes_read[0], - total_bytes=total_bytes, - ) - return True - return False - - def on_hash_source(source_name): - hash_source[0] = source_name - if source_name == "safetensors_header": - self.logger.info( - f"Found SHA256 in safetensors header: {normalized_path}" - ) - self.hash_tracker.update( - progress_id, - status="running", - stage="header", - message="Found SHA256 in safetensors header", - percent=98, - bytes_read=0, - total_bytes=total_bytes, - sha256_source="safetensors_header", - ) - - sha256 = _calculate_file_sha256_core( - normalized_path, - chunk_size=1024 * 1024 * 4, - on_progress=on_progress, - is_cancelled=is_cancelled, - use_safetensors_header=True, - on_hash_source=on_hash_source, - ) - - if is_cancelled(): - raise HashCalculationCancelled() - - return sha256, hash_source[0] + return await service.open_containing_folder(request) @routes.post("/model_resolver/calculate-file-hash") @json_api_endpoint("calculate-file-hash") async def calculate_file_hash_route(request): - """Calculate SHA256 for a local model and persist it to sidecar metadata.""" - data = await request.json() - - normalized_path, error = resolve_hash_file_request(data) - if error == "file_path is required": - return web.json_response( - {"error": "file_path is required"}, status=400 - ) - if error: - status = 403 if error == "file is outside configured model directories" else 404 - return web.json_response( - {"error": error}, status=status - ) - - sha256, sha256_source = calculate_sha256_with_progress(normalized_path) - if not sha256: - return web.json_response( - {"error": "could not calculate hash"}, status=500 - ) - resolved_metadata_path, metadata_updated = write_calculated_hash_metadata( - normalized_path, - sha256, - sha256_source, - ) - - return web.json_response( - { - "success": True, - "sha256": sha256, - "hash": sha256, - "sha256_source": sha256_source, - "file_path": normalized_path, - "metadata_path": resolved_metadata_path, - "metadata_updated": metadata_updated, - } - ) + return await service.calculate_file_hash(request) @routes.post("/model_resolver/calculate-file-hash/start") @json_api_endpoint("calculate-file-hash-start") async def calculate_file_hash_start_route(request): - """Start SHA256 calculation in a background thread.""" - import uuid - - data = await request.json() - normalized_path, error = resolve_hash_file_request(data) - if error == "file_path is required": - return web.json_response( - {"error": "file_path is required"}, status=400 - ) - if error: - status = 403 if error == "file is outside configured model directories" else 404 - return web.json_response( - {"error": error}, status=status - ) - - self.hash_tracker.cleanup() - progress_id = f"hash_{uuid.uuid4().hex}" - self.hash_tracker.update( - progress_id, - status="queued", - stage="queued", - message="Preparing hash calculation...", - percent=0, - file_path=normalized_path, - ) - - def run_hash_task(): - sha256, sha256_source = calculate_sha256_with_progress( - normalized_path, - progress_id=progress_id, - ) - if self.hash_tracker.is_cancelled(progress_id): - raise HashCalculationCancelled() - self.hash_tracker.update( - progress_id, - status="running", - stage="metadata", - message="Saving metadata...", - percent=99, - ) - resolved_metadata_path, metadata_updated = write_calculated_hash_metadata( - normalized_path, - sha256, - sha256_source, - ) - return sha256, sha256_source, resolved_metadata_path, metadata_updated - - def on_success(res): - sha256, sha256_source, resolved_metadata_path, metadata_updated = res - self.hash_tracker.update( - progress_id, - status="done", - stage="done", - message="Hash calculated", - percent=100, - sha256=sha256, - hash=sha256, - sha256_source=sha256_source, - file_path=normalized_path, - metadata_path=resolved_metadata_path, - metadata_updated=metadata_updated, - ) - - def on_cancel(res=None): - self.hash_tracker.update( - progress_id, - status="cancelled", - stage="cancelled", - message="Hash calculation cancelled", - ) - - run_in_background_thread( - self.hash_tracker, - progress_id, - run_hash_task, - on_success, - on_cancel, - error_log_msg=f"Hash calculation failed for {normalized_path}", - ) - return web.json_response( - { - "success": True, - "progress_id": progress_id, - } - ) + return await service.calculate_file_hash_start(request) @routes.get("/model_resolver/calculate-file-hash/progress/{progress_id}") @json_api_endpoint("calculate-file-hash-progress") async def calculate_file_hash_progress_route(request): - """Return progress for a background SHA256 calculation.""" - return get_progress_response( - self.hash_tracker, - request, - not_found_status=404 - ) + return await service.calculate_file_hash_progress(request) @routes.post("/model_resolver/calculate-file-hash/cancel/{progress_id}") @json_api_endpoint("calculate-file-hash-cancel") async def calculate_file_hash_cancel_route(request): - """Cancel a background SHA256 calculation.""" - return cancel_progress_response( - self.hash_tracker, - request, - cancel_message="Stopping hash calculation..." - ) + return await service.calculate_file_hash_cancel(request) diff --git a/core/services/hash_service.py b/core/services/hash_service.py new file mode 100644 index 0000000..2a48f88 --- /dev/null +++ b/core/services/hash_service.py @@ -0,0 +1,696 @@ +"""Local model metadata and hash calculation service.""" + +import os +import time +import uuid + +from .. import path_utils +from ..routes.context import RouteContext + + +class HashService: + """Implement local model hash operations used by HTTP routes.""" + + def __init__(self, context: RouteContext): + self.file_manager_error = context.get("FileManagerError") + self.file_manager_unavailable_error = context.get( + "FileManagerUnavailableError" + ) + self.metadata_schema = context.get("MODEL_RESOLVER_METADATA_SCHEMA") + self.metadata_schema_version = context.get( + "MODEL_RESOLVER_METADATA_SCHEMA_VERSION" + ) + self.unsupported_file_manager_error = context.get( + "UnsupportedFileManagerPlatformError" + ) + self.asyncio = context.get("asyncio") + self.cancel_progress_response = context.get("cancel_progress_response") + self.get_existing_model_preview_path = context.get( + "get_existing_model_preview_path" + ) + self.get_filename_from_path = context.get("get_filename_from_path") + self.get_local_model_hash_metadata = context.get( + "get_local_model_hash_metadata" + ) + self.get_progress_response = context.get("get_progress_response") + self.get_safe_model_resolver_sidecar_path = context.get( + "get_safe_model_resolver_sidecar_path" + ) + self.get_workflow_model_inventory = context.get( + "get_workflow_model_inventory" + ) + self.is_path_in_configured_model_roots = context.get( + "is_path_in_configured_model_roots" + ) + self.load_resolver_settings = context.get("load_resolver_settings") + self.normalize_file_manager_path = context.get( + "normalize_file_manager_path" + ) + self.normalize_sha256 = context.get("normalize_sha256") + self.open_in_file_manager = context.get("open_in_file_manager") + self.os = context.get("os") or os + self.read_json_safe = context.get("read_json_safe") + self.resolver_bool_setting = context.get("resolver_bool_setting") + self.run_in_background_thread = context.get("run_in_background_thread") + self.search_local_matches_by_hash = context.get( + "search_local_matches_by_hash" + ) + self.extension = context.get("self") + self.time = context.get("time") or time + self.to_bool = context.get("to_bool") + self.to_int = context.get("to_int") + self.web = context.get("web") + self.write_json_atomic = context.get("write_json_atomic") + + @property + def hash_tracker(self): + return self.extension.hash_tracker + + @property + def logger(self): + return self.extension.logger + + def _normalize_path(self, path): + return self.os.path.realpath( + self.os.path.abspath(self.os.path.normpath(str(path))) + ) + + async def local_model_hashes(self, request): + """Return SHA256 hashes already stored in local sidecar metadata.""" + data = await request.json() + model = data.get("model") if isinstance(data.get("model"), dict) else {} + path = ( + data.get("path") + or data.get("file_path") + or data.get("resolved_path") + or model.get("path") + or model.get("resolved_path") + or "" + ) + + if not path: + return self.web.json_response( + {"error": "path is required"}, status=400 + ) + + normalized_path = self._normalize_path(path) + if not self.is_path_in_configured_model_roots(normalized_path): + return self.web.json_response( + {"error": "path is outside configured model directories"}, + status=403, + ) + + return self.web.json_response( + self.get_local_model_hash_metadata(normalized_path, model=model) + ) + + async def get_model_preview(self, request): + """Serve an adjacent model preview from a configured model directory.""" + model_path = str(request.query.get("path") or "").strip() + if not model_path: + return self.web.Response(text="path is required", status=400) + + try: + normalized_path = self._normalize_path(model_path) + except (OSError, TypeError, ValueError): + return self.web.Response(text="invalid model path", status=400) + if not self.os.path.isfile(normalized_path): + return self.web.Response(text="model file does not exist", status=404) + if not self.is_path_in_configured_model_roots(normalized_path): + return self.web.Response( + text="path is outside configured model directories", + status=403, + ) + + preview_path = self.get_existing_model_preview_path(normalized_path) + if not preview_path: + return self.web.Response(text="preview not found", status=404) + if not self.is_path_in_configured_model_roots(preview_path): + return self.web.Response( + text="preview is outside configured model directories", + status=403, + ) + + return self.web.FileResponse( + preview_path, + headers={"Cache-Control": "private, no-cache"}, + ) + + async def workflow_model_hashes(self, request): + """Return hash metadata for existing local models used by a workflow.""" + data = await request.json() + workflow_json = data.get("workflow") + if not isinstance(workflow_json, dict): + return self.web.json_response( + {"error": "Workflow JSON must be an object"}, status=400 + ) + + settings = await self.asyncio.to_thread(self.load_resolver_settings) + if not self.resolver_bool_setting( + settings.get("workflow_hash_metadata_enabled"), True + ): + return self.web.json_response( + { + "success": True, + "enabled": False, + "models": [], + "by_node": {}, + "by_path": {}, + "count": 0, + } + ) + + inventory = await self.asyncio.to_thread( + self.get_workflow_model_inventory, + workflow_json, + ) + refs = inventory["model_refs"] + + by_node = {} + by_path = {} + models = [] + seen = set() + + for ref in refs: + if not isinstance(ref, dict) or not ref.get("exists"): + continue + full_path = str(ref.get("full_path") or "").strip() + if not full_path: + continue + model_info = { + "path": full_path, + "filename": self.get_filename_from_path(full_path), + "relative_path": ref.get("original_path") or "", + "category": ref.get("category") or "", + } + metadata = self.get_local_model_hash_metadata( + full_path, + model=model_info, + ) + sha256 = self.normalize_sha256(metadata.get("sha256")) + if not sha256: + continue + + entry = { + "node_id": ref.get("node_id"), + "node_type": ref.get("node_type") or "", + "widget_index": ref.get("widget_index"), + "widget_name": ref.get("widget_name") or "", + "path": ref.get("original_path") or "", + "filename": self.get_filename_from_path( + ref.get("original_path") or full_path + ), + "category": ref.get("category") or "", + "sha256": sha256, + "size": metadata.get("size") or 0, + } + entry_key = ( + str(entry.get("node_id")), + str(entry.get("widget_index")), + entry.get("path"), + sha256, + ) + if entry_key in seen: + continue + seen.add(entry_key) + models.append(entry) + + path_key = str(entry.get("path") or entry.get("filename") or "") + if path_key: + by_path[path_key] = entry + by_path[self.get_filename_from_path(path_key)] = entry + node_key = f"{entry.get('node_id')}:{entry.get('widget_index')}" + by_node[node_key] = entry + + return self.web.json_response( + { + "success": True, + "enabled": True, + "models": models, + "by_node": by_node, + "by_path": by_path, + "count": len(models), + } + ) + + async def local_matches_by_hash(self, request): + """Search local model metadata sidecars for a remote SHA256.""" + data = await request.json() + sha256 = self.normalize_sha256( + data.get("sha256") + or data.get("hash") + or data.get("SHA256") + or "" + ) + if not sha256: + return self.web.json_response( + {"error": "sha256 is required"}, status=400 + ) + + category = data.get("category", "") + source = str( + data.get("source") + or data.get("hash_lookup_source") + or "download_source" + ).strip().lower().replace("-", "_") + filename = ( + data.get("filename") + or data.get("path") + or data.get("model_name") + or "" + ) + max_matches = self.to_int(data.get("max_matches"), 20) + force_rescan = self.to_bool(data.get("force_rescan"), False) + + matches = self.search_local_matches_by_hash( + sha256, + category=category or None, + max_matches=max_matches, + force_rescan=force_rescan, + ) + enriched_matches = [ + { + **match, + "hash_lookup_source": source or "download_source", + "hash_lookup_filename": filename, + "hash_lookup_sha256": sha256, + } + for match in matches + ] + return self.web.json_response( + { + "sha256": sha256, + "local_hash_matches": enriched_matches, + "matches": enriched_matches, + } + ) + + async def open_containing_folder(self, request): + """Reveal a file or open a directory in the host file manager.""" + try: + data = await request.json() + except Exception as exc: + return self.web.json_response( + {"success": False, "error": f"invalid JSON body: {exc}"}, + status=400, + ) + if not isinstance(data, dict): + return self.web.json_response( + {"success": False, "error": "JSON body must be an object"}, + status=400, + ) + target_path = data.get("path", "") + + try: + normalized_path = self.normalize_file_manager_path(target_path) + except ValueError as exc: + return self.web.json_response( + {"success": False, "error": str(exc)}, status=400 + ) + + if not self.os.path.exists(normalized_path): + return self.web.json_response( + {"success": False, "error": "path does not exist"}, + status=404, + ) + if not self.is_path_in_configured_model_roots(normalized_path): + return self.web.json_response( + { + "success": False, + "error": "path is outside configured model directories", + }, + status=403, + ) + + try: + result = await self.asyncio.to_thread( + self.open_in_file_manager, + normalized_path, + ) + except FileNotFoundError: + return self.web.json_response( + {"success": False, "error": "path no longer exists"}, + status=404, + ) + except self.unsupported_file_manager_error as exc: + return self.web.json_response( + {"success": False, "error": str(exc)}, status=501 + ) + except self.file_manager_unavailable_error as exc: + return self.web.json_response( + { + "success": False, + "error": ( + f"{exc} The ComfyUI host may be running without " + "a graphical desktop session." + ), + }, + status=503, + ) + except self.file_manager_error as exc: + return self.web.json_response( + {"success": False, "error": str(exc)}, status=500 + ) + + return self.web.json_response({"success": True, **result}) + + def resolve_hash_file_request(self, data): + """Validate and normalize a file path accepted by hash endpoints.""" + file_path = ( + data.get("file_path") + or data.get("resolved_path") + or data.get("path") + or "" + ) + if not file_path: + return "", "file_path is required" + + normalized_path = self._normalize_path(file_path) + if ( + not self.os.path.exists(normalized_path) + or not self.os.path.isfile(normalized_path) + ): + return "", "file does not exist" + if not self.is_path_in_configured_model_roots(normalized_path): + return "", "file is outside configured model directories" + return normalized_path, "" + + def write_calculated_hash_metadata( + self, + normalized_path, + sha256, + sha256_source="file", + ): + """Persist calculated SHA256 metadata next to a local model.""" + resolved_metadata_path = self.get_safe_model_resolver_sidecar_path( + normalized_path + ) + + metadata_updated = False + try: + metadata = self.read_json_safe(resolved_metadata_path, {}) + if not isinstance(metadata, dict): + metadata = {} + + filename = self.get_filename_from_path(normalized_path) + stem, _ext = self.os.path.splitext(filename) + hashes = metadata.get("hashes") + if not isinstance(hashes, dict): + hashes = {} + hashes["SHA256"] = sha256 + + metadata["schema"] = self.metadata_schema + metadata["schema_version"] = self.metadata_schema_version + metadata["managed_by"] = self.metadata_schema + metadata["sha256"] = sha256 + metadata["hashes"] = hashes + metadata["hash_status"] = "completed" + metadata["sha256_source"] = sha256_source or "file" + metadata["last_checked_at"] = self.time.time() + metadata.setdefault("file_name", stem) + metadata.setdefault("model_name", stem) + metadata.setdefault( + "file_path", + normalized_path.replace("\\", "/"), + ) + try: + metadata.setdefault("size", self.os.path.getsize(normalized_path)) + except Exception: + pass + + self.write_json_atomic(resolved_metadata_path, metadata, indent=2) + metadata_updated = True + self.logger.info( + f"Stored SHA256 and updated metadata: {resolved_metadata_path}" + ) + except Exception as metadata_error: + self.logger.warning( + "Could not update metadata with calculated SHA256 for " + f"{normalized_path}: {metadata_error}" + ) + + return resolved_metadata_path, metadata_updated + + def calculate_sha256_with_progress(self, normalized_path, progress_id=""): + """Calculate a SHA256 while updating and honoring hash progress.""" + total_bytes = max(0, self.os.path.getsize(normalized_path)) + + self.hash_tracker.update( + progress_id, + status="running", + stage="header", + message="Checking safetensors header...", + percent=0, + bytes_read=0, + total_bytes=total_bytes, + ) + + if self.hash_tracker.is_cancelled(progress_id): + raise path_utils.HashCalculationCancelled() + + hash_source = ["file"] + stage_transitioned = [False] + bytes_read_state = [0] + last_update = [0.0] + + def on_progress(bytes_read, callback_total_bytes): + bytes_read_state[0] = bytes_read + if not stage_transitioned[0]: + stage_transitioned[0] = True + self.logger.info( + "No SHA256 found in safetensors header; calculating full " + f"file SHA256: {normalized_path}" + ) + self.hash_tracker.update( + progress_id, + status="running", + stage="hashing", + message="Calculating SHA256...", + percent=0, + bytes_read=0, + total_bytes=callback_total_bytes, + ) + + now = self.time.time() + if ( + now - last_update[0] >= 0.15 + or bytes_read >= callback_total_bytes + ): + percent = ( + 98 + if callback_total_bytes <= 0 + else min(98, (bytes_read / callback_total_bytes) * 98) + ) + self.hash_tracker.update( + progress_id, + status="running", + stage="hashing", + message="Calculating SHA256...", + percent=percent, + bytes_read=bytes_read, + total_bytes=callback_total_bytes, + ) + last_update[0] = now + + def is_cancelled(): + if self.hash_tracker.is_cancelled(progress_id): + percent = ( + 0 + if total_bytes <= 0 + else min(98, (bytes_read_state[0] / total_bytes) * 98) + ) + self.hash_tracker.update( + progress_id, + status="cancelled", + stage="cancelled", + message="Hash calculation cancelled", + percent=percent, + bytes_read=bytes_read_state[0], + total_bytes=total_bytes, + ) + return True + return False + + def on_hash_source(source_name): + hash_source[0] = source_name + if source_name == "safetensors_header": + self.logger.info( + f"Found SHA256 in safetensors header: {normalized_path}" + ) + self.hash_tracker.update( + progress_id, + status="running", + stage="header", + message="Found SHA256 in safetensors header", + percent=98, + bytes_read=0, + total_bytes=total_bytes, + sha256_source="safetensors_header", + ) + + sha256 = path_utils.calculate_file_sha256( + normalized_path, + chunk_size=1024 * 1024 * 4, + on_progress=on_progress, + is_cancelled=is_cancelled, + use_safetensors_header=True, + on_hash_source=on_hash_source, + ) + + if is_cancelled(): + raise path_utils.HashCalculationCancelled() + + return sha256, hash_source[0] + + async def calculate_file_hash(self, request): + """Calculate SHA256 for a local model and persist sidecar metadata.""" + data = await request.json() + + normalized_path, error = self.resolve_hash_file_request(data) + if error == "file_path is required": + return self.web.json_response( + {"error": "file_path is required"}, status=400 + ) + if error: + status = ( + 403 + if error == "file is outside configured model directories" + else 404 + ) + return self.web.json_response({"error": error}, status=status) + + sha256, sha256_source = self.calculate_sha256_with_progress(normalized_path) + if not sha256: + return self.web.json_response( + {"error": "could not calculate hash"}, status=500 + ) + resolved_metadata_path, metadata_updated = ( + self.write_calculated_hash_metadata( + normalized_path, + sha256, + sha256_source, + ) + ) + + return self.web.json_response( + { + "success": True, + "sha256": sha256, + "hash": sha256, + "sha256_source": sha256_source, + "file_path": normalized_path, + "metadata_path": resolved_metadata_path, + "metadata_updated": metadata_updated, + } + ) + + async def calculate_file_hash_start(self, request): + """Start SHA256 calculation in a background thread.""" + data = await request.json() + normalized_path, error = self.resolve_hash_file_request(data) + if error == "file_path is required": + return self.web.json_response( + {"error": "file_path is required"}, status=400 + ) + if error: + status = ( + 403 + if error == "file is outside configured model directories" + else 404 + ) + return self.web.json_response({"error": error}, status=status) + + self.hash_tracker.cleanup() + progress_id = f"hash_{uuid.uuid4().hex}" + self.hash_tracker.update( + progress_id, + status="queued", + stage="queued", + message="Preparing hash calculation...", + percent=0, + file_path=normalized_path, + ) + + def run_hash_task(): + sha256, sha256_source = self.calculate_sha256_with_progress( + normalized_path, + progress_id=progress_id, + ) + if self.hash_tracker.is_cancelled(progress_id): + raise path_utils.HashCalculationCancelled() + self.hash_tracker.update( + progress_id, + status="running", + stage="metadata", + message="Saving metadata...", + percent=99, + ) + resolved_metadata_path, metadata_updated = ( + self.write_calculated_hash_metadata( + normalized_path, + sha256, + sha256_source, + ) + ) + return sha256, sha256_source, resolved_metadata_path, metadata_updated + + def on_success(result): + ( + sha256, + sha256_source, + resolved_metadata_path, + metadata_updated, + ) = result + self.hash_tracker.update( + progress_id, + status="done", + stage="done", + message="Hash calculated", + percent=100, + sha256=sha256, + hash=sha256, + sha256_source=sha256_source, + file_path=normalized_path, + metadata_path=resolved_metadata_path, + metadata_updated=metadata_updated, + ) + + def on_cancel(result=None): + del result + self.hash_tracker.update( + progress_id, + status="cancelled", + stage="cancelled", + message="Hash calculation cancelled", + ) + + self.run_in_background_thread( + self.hash_tracker, + progress_id, + run_hash_task, + on_success, + on_cancel, + error_log_msg=f"Hash calculation failed for {normalized_path}", + ) + return self.web.json_response( + { + "success": True, + "progress_id": progress_id, + } + ) + + async def calculate_file_hash_progress(self, request): + """Return progress for a background SHA256 calculation.""" + return self.get_progress_response( + self.hash_tracker, + request, + not_found_status=404, + ) + + async def calculate_file_hash_cancel(self, request): + """Cancel a background SHA256 calculation.""" + return self.cancel_progress_response( + self.hash_tracker, + request, + cancel_message="Stopping hash calculation...", + ) diff --git a/tests/test_hash_routes.py b/tests/test_hash_routes.py new file mode 100644 index 0000000..ae1719f --- /dev/null +++ b/tests/test_hash_routes.py @@ -0,0 +1,522 @@ +import asyncio +import json +import os +import tempfile +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from aiohttp import web + +from core.path_utils import HashCalculationCancelled +from core.routes.context import RouteContext +from core.routes.hashes import register_hash_routes +from core.services.hash_service import HashService +from core.type_utils import normalize_sha256, to_bool, to_int + + +class _Routes: + def __init__(self): + self.handlers = {} + + def post(self, path): + def decorator(handler): + self.handlers[("POST", path)] = handler + return handler + + return decorator + + def get(self, path): + def decorator(handler): + self.handlers[("GET", path)] = handler + return handler + + return decorator + + +class _UnsupportedFileManagerPlatformError(Exception): + pass + + +class _FileManagerUnavailableError(Exception): + pass + + +class _FileManagerError(Exception): + pass + + +def _build_hash_routes(overrides=None): + routes = _Routes() + hash_tracker = MagicMock() + hash_tracker.is_cancelled.return_value = False + extension = SimpleNamespace( + hash_tracker=hash_tracker, + logger=MagicMock(), + ) + + def run_in_background_thread(tracker, progress_id, task, on_success, on_cancel, **kwargs): + try: + on_success(task()) + except HashCalculationCancelled: + on_cancel() + + values = { + "FileManagerError": _FileManagerError, + "FileManagerUnavailableError": _FileManagerUnavailableError, + "MODEL_RESOLVER_METADATA_SCHEMA": "comfyui-model-resolver", + "MODEL_RESOLVER_METADATA_SCHEMA_VERSION": 2, + "UnsupportedFileManagerPlatformError": _UnsupportedFileManagerPlatformError, + "asyncio": asyncio, + "cancel_progress_response": MagicMock( + return_value=web.json_response({"cancelled": True}) + ), + "get_existing_model_preview_path": MagicMock(return_value=None), + "get_filename_from_path": os.path.basename, + "get_local_model_hash_metadata": MagicMock( + return_value={"sha256": "a" * 64, "size": 123} + ), + "get_progress_response": MagicMock( + return_value=web.json_response({"status": "running"}) + ), + "get_safe_model_resolver_sidecar_path": lambda path: ( + f"{path}.modelresolver.json" + ), + "get_workflow_model_inventory": MagicMock( + return_value={"model_refs": []} + ), + "is_path_in_configured_model_roots": MagicMock(return_value=True), + "json_api_endpoint": lambda _name: lambda handler: handler, + "load_resolver_settings": MagicMock(return_value={}), + "normalize_file_manager_path": lambda path: str(path), + "normalize_sha256": normalize_sha256, + "open_in_file_manager": MagicMock(return_value={"opened": True}), + "os": os, + "read_json_safe": MagicMock(return_value={}), + "resolver_bool_setting": lambda value, default=True: ( + default if value is None else bool(value) + ), + "routes": routes, + "run_in_background_thread": run_in_background_thread, + "search_local_matches_by_hash": MagicMock(return_value=[]), + "self": extension, + "time": __import__("time"), + "to_bool": to_bool, + "to_int": to_int, + "web": web, + "write_json_atomic": MagicMock(), + } + if overrides: + values.update(overrides) + + register_hash_routes(RouteContext(values)) + return routes.handlers, values + + +def _request(payload): + request = SimpleNamespace(json=AsyncMock(return_value=payload)) + return request + + +@pytest.mark.asyncio +async def test_local_model_hashes_uses_model_path_and_returns_metadata(): + handlers, values = _build_hash_routes() + handler = handlers[("POST", "/model_resolver/local-model-hashes")] + with tempfile.TemporaryDirectory() as temp_dir: + model_path = os.path.join(temp_dir, "model.safetensors") + response = await handler(_request({"model": {"path": model_path}})) + + assert response.status == 200 + assert json.loads(response.text) == {"sha256": "a" * 64, "size": 123} + values["get_local_model_hash_metadata"].assert_called_once_with( + os.path.realpath(os.path.abspath(os.path.normpath(model_path))), + model={"path": model_path}, + ) + + +@pytest.mark.asyncio +async def test_hash_routes_validate_local_metadata_and_preview_paths(): + handlers, values = _build_hash_routes() + local_handler = handlers[("POST", "/model_resolver/local-model-hashes")] + preview_handler = handlers[("GET", "/model_resolver/model-preview")] + + response = await local_handler(_request({})) + assert response.status == 400 + + values["is_path_in_configured_model_roots"].return_value = False + response = await local_handler(_request({"path": "model.safetensors"})) + assert response.status == 403 + + response = await preview_handler(SimpleNamespace(query={})) + assert response.status == 400 + + with tempfile.TemporaryDirectory() as temp_dir: + model_path = os.path.join(temp_dir, "model.safetensors") + preview_path = os.path.join(temp_dir, "model.png") + with open(model_path, "wb") as model_file: + model_file.write(b"model") + with open(preview_path, "wb") as preview_file: + preview_file.write(b"preview") + + values["is_path_in_configured_model_roots"].return_value = True + response = await preview_handler( + SimpleNamespace(query={"path": os.path.join(temp_dir, "missing.safetensors")}) + ) + assert response.status == 404 + + values["get_existing_model_preview_path"].return_value = None + response = await preview_handler(SimpleNamespace(query={"path": model_path})) + assert response.status == 404 + + values["get_existing_model_preview_path"].return_value = preview_path + values["is_path_in_configured_model_roots"].side_effect = [True, False] + response = await preview_handler(SimpleNamespace(query={"path": model_path})) + assert response.status == 403 + + +@pytest.mark.asyncio +async def test_workflow_hashes_can_be_disabled_and_local_matches_require_hash(): + handlers, values = _build_hash_routes() + workflow_handler = handlers[("POST", "/model_resolver/workflow-model-hashes")] + matches_handler = handlers[("POST", "/model_resolver/local-matches-by-hash")] + + response = await workflow_handler(_request({"workflow": []})) + assert response.status == 400 + + values["load_resolver_settings"].return_value = { + "workflow_hash_metadata_enabled": False + } + response = await workflow_handler(_request({"workflow": {}})) + assert json.loads(response.text) == { + "success": True, + "enabled": False, + "models": [], + "by_node": {}, + "by_path": {}, + "count": 0, + } + + response = await matches_handler(_request({"filename": "model.safetensors"})) + assert response.status == 400 + + +@pytest.mark.asyncio +async def test_workflow_model_hashes_skips_invalid_refs_and_deduplicates_models(): + refs = [ + {"exists": False, "full_path": "ignored.safetensors"}, + { + "exists": True, + "full_path": r"C:\models\one.safetensors", + "original_path": "one.safetensors", + "category": "checkpoints", + "node_id": 1, + "widget_index": 0, + "node_type": "CheckpointLoaderSimple", + }, + { + "exists": True, + "full_path": r"C:\models\one.safetensors", + "original_path": "one.safetensors", + "category": "checkpoints", + "node_id": 1, + "widget_index": 0, + "node_type": "CheckpointLoaderSimple", + }, + {"exists": True, "full_path": ""}, + { + "exists": True, + "full_path": r"C:\models\without-hash.safetensors", + "original_path": "without-hash.safetensors", + }, + ] + get_metadata = MagicMock( + side_effect=[ + {"sha256": "a" * 64, "size": 456}, + {"sha256": "a" * 64, "size": 456}, + {}, + ] + ) + handlers, values = _build_hash_routes( + { + "get_workflow_model_inventory": MagicMock( + return_value={"model_refs": refs} + ), + "get_local_model_hash_metadata": get_metadata, + } + ) + handler = handlers[("POST", "/model_resolver/workflow-model-hashes")] + + response = await handler(_request({"workflow": {"nodes": []}})) + + body = json.loads(response.text) + assert response.status == 200 + assert body["count"] == 1 + assert body["models"][0]["sha256"] == "a" * 64 + assert body["by_node"]["1:0"]["filename"] == "one.safetensors" + assert body["by_path"]["one.safetensors"]["size"] == 456 + assert values["get_workflow_model_inventory"].called + + +@pytest.mark.asyncio +async def test_local_matches_by_hash_normalizes_and_enriches_results(): + matches = [{"path": r"C:\models\local.safetensors"}] + search_matches = MagicMock(return_value=matches) + handlers, values = _build_hash_routes( + {"search_local_matches_by_hash": search_matches} + ) + handler = handlers[("POST", "/model_resolver/local-matches-by-hash")] + sha256 = "b" * 64 + + response = await handler( + _request( + { + "hash": "sha256:" + sha256.upper(), + "category": "checkpoints", + "source": "Civit-AI", + "filename": "remote.safetensors", + "max_matches": "3", + "force_rescan": "true", + } + ) + ) + + body = json.loads(response.text) + assert body["sha256"] == sha256 + assert body["local_hash_matches"] == [ + { + "path": r"C:\models\local.safetensors", + "hash_lookup_source": "civit_ai", + "hash_lookup_filename": "remote.safetensors", + "hash_lookup_sha256": sha256, + } + ] + values["search_local_matches_by_hash"].assert_called_once_with( + sha256, + category="checkpoints", + max_matches=3, + force_rescan=True, + ) + + +@pytest.mark.asyncio +async def test_open_containing_folder_maps_file_manager_errors(): + handlers, values = _build_hash_routes() + handler = handlers[("POST", "/model_resolver/open-containing-folder")] + with tempfile.TemporaryDirectory() as temp_dir: + target = os.path.join(temp_dir, "model.safetensors") + with open(target, "wb") as model_file: + model_file.write(b"model") + + error_cases = [ + (FileNotFoundError("gone"), 404), + (_UnsupportedFileManagerPlatformError("unsupported"), 501), + (_FileManagerUnavailableError("unavailable"), 503), + (_FileManagerError("failed"), 500), + ] + for error, expected_status in error_cases: + values["open_in_file_manager"].side_effect = error + response = await handler(_request({"path": target})) + assert response.status == expected_status + + values["open_in_file_manager"].side_effect = None + values["open_in_file_manager"].return_value = {"opened": True} + response = await handler(_request({"path": target})) + + assert response.status == 200 + assert json.loads(response.text) == {"success": True, "opened": True} + + +@pytest.mark.asyncio +async def test_open_containing_folder_validates_request_and_path(): + normalizer = MagicMock(side_effect=ValueError("invalid path")) + handlers, values = _build_hash_routes( + {"normalize_file_manager_path": normalizer} + ) + handler = handlers[("POST", "/model_resolver/open-containing-folder")] + + invalid_json_request = SimpleNamespace( + json=AsyncMock(side_effect=ValueError("malformed")) + ) + response = await handler(invalid_json_request) + assert response.status == 400 + + response = await handler(_request([])) + assert response.status == 400 + + response = await handler(_request({"path": "bad"})) + assert response.status == 400 + + normalizer.side_effect = lambda path: path + values["is_path_in_configured_model_roots"].return_value = False + response = await handler(_request({"path": "missing"})) + assert response.status == 404 + + with tempfile.TemporaryDirectory() as temp_dir: + target = os.path.join(temp_dir, "model.safetensors") + with open(target, "wb") as model_file: + model_file.write(b"model") + response = await handler(_request({"path": target})) + assert response.status == 403 + + +@pytest.mark.asyncio +async def test_calculate_file_hash_writes_metadata_and_reports_source(): + sha256 = "c" * 64 + written = {} + + def write_metadata(path, metadata, indent=2): + written["path"] = path + written["metadata"] = metadata + written["indent"] = indent + + handlers, values = _build_hash_routes( + { + "read_json_safe": MagicMock(return_value={"file_name": "old"}), + "write_json_atomic": write_metadata, + } + ) + handler = handlers[("POST", "/model_resolver/calculate-file-hash")] + with tempfile.TemporaryDirectory() as temp_dir: + model_path = os.path.join(temp_dir, "model.safetensors") + with open(model_path, "wb") as model_file: + model_file.write(b"model") + with patch( + "core.path_utils.calculate_file_sha256", + return_value=sha256, + ): + response = await handler(_request({"file_path": model_path})) + + body = json.loads(response.text) + assert body["success"] is True + assert body["sha256"] == sha256 + assert body["sha256_source"] == "file" + assert body["metadata_updated"] is True + assert written["metadata"]["sha256"] == sha256 + assert written["metadata"]["hashes"]["SHA256"] == sha256 + assert written["metadata"]["file_name"] == "old" + + +def test_hash_service_reports_header_source_and_cancellation(): + sha256 = "e" * 64 + _, values = _build_hash_routes() + service = HashService(RouteContext(values)) + + def calculate_with_header(path, **callbacks): + callbacks["on_hash_source"]("safetensors_header") + callbacks["on_progress"](1, 1) + return sha256 + + with tempfile.TemporaryDirectory() as temp_dir: + model_path = os.path.join(temp_dir, "model.safetensors") + with open(model_path, "wb") as model_file: + model_file.write(b"model") + with patch( + "core.path_utils.calculate_file_sha256", + side_effect=calculate_with_header, + ): + calculated_hash, source = service.calculate_sha256_with_progress( + model_path, + progress_id="header-test", + ) + + assert calculated_hash == sha256 + assert source == "safetensors_header" + assert any( + call.kwargs.get("sha256_source") == "safetensors_header" + for call in values["self"].hash_tracker.update.call_args_list + ) + + values["self"].hash_tracker.reset_mock() + values["self"].hash_tracker.is_cancelled.side_effect = [False, True] + with tempfile.TemporaryDirectory() as temp_dir: + model_path = os.path.join(temp_dir, "model.safetensors") + with open(model_path, "wb") as model_file: + model_file.write(b"model") + with ( + patch( + "core.path_utils.calculate_file_sha256", + return_value=sha256, + ), + pytest.raises(HashCalculationCancelled), + ): + service.calculate_sha256_with_progress(model_path, "cancel-test") + + assert any( + call.kwargs.get("status") == "cancelled" + for call in values["self"].hash_tracker.update.call_args_list + ) + + +def test_hash_service_returns_false_when_metadata_write_fails(): + _, values = _build_hash_routes( + {"read_json_safe": MagicMock(side_effect=OSError("read failed"))} + ) + service = HashService(RouteContext(values)) + + with tempfile.TemporaryDirectory() as temp_dir: + model_path = os.path.join(temp_dir, "model.safetensors") + with open(model_path, "wb") as model_file: + model_file.write(b"model") + metadata_path, updated = service.write_calculated_hash_metadata( + model_path, + "f" * 64, + ) + + assert metadata_path.endswith(".modelresolver.json") + assert updated is False + values["self"].logger.warning.assert_called_once() + + +@pytest.mark.asyncio +async def test_calculate_file_hash_start_runs_background_success_callback(): + sha256 = "d" * 64 + handlers, values = _build_hash_routes() + handler = handlers[("POST", "/model_resolver/calculate-file-hash/start")] + with tempfile.TemporaryDirectory() as temp_dir: + model_path = os.path.join(temp_dir, "model.safetensors") + with open(model_path, "wb") as model_file: + model_file.write(b"model") + with patch( + "core.path_utils.calculate_file_sha256", + return_value=sha256, + ): + response = await handler(_request({"file_path": model_path})) + + body = json.loads(response.text) + assert response.status == 200 + assert body["success"] is True + assert body["progress_id"].startswith("hash_") + progress_calls = values["self"].hash_tracker.update.call_args_list + assert any( + call.kwargs.get("status") == "done" + and call.kwargs.get("sha256") == sha256 + for call in progress_calls + ) + + +@pytest.mark.asyncio +async def test_hash_progress_and_cancel_routes_delegate_to_helpers(): + handlers, values = _build_hash_routes() + progress_handler = handlers[ + ("GET", "/model_resolver/calculate-file-hash/progress/{progress_id}") + ] + cancel_handler = handlers[ + ("POST", "/model_resolver/calculate-file-hash/cancel/{progress_id}") + ] + request = SimpleNamespace() + + progress_response = await progress_handler(request) + cancel_response = await cancel_handler(request) + + assert progress_response.status == 200 + assert cancel_response.status == 200 + values["get_progress_response"].assert_called_once_with( + values["self"].hash_tracker, + request, + not_found_status=404, + ) + values["cancel_progress_response"].assert_called_once_with( + values["self"].hash_tracker, + request, + cancel_message="Stopping hash calculation...", + )