refactor(hash): extract hash route handlers into HashService and add tests
This commit is contained in:
Dariusz L
+15
-623
@@ -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)
|
||||
|
||||
@@ -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...",
|
||||
)
|
||||
@@ -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...",
|
||||
)
|
||||
Reference in New Issue
Block a user