From bf898cc552ffb415d58c57d0326604703dd122a6 Mon Sep 17 00:00:00 2001 From: Simon Pinfold Date: Sat, 20 Jun 2026 12:41:10 +1200 Subject: [PATCH] fix(assets): move model asset on model_type: edit to stay loader-coherent Under supports_model_type_tags, a model_type: edit on a filesystem-backed model asset now relocates and re-registers the file to the folder matching the new model_type:, so the file location stays coherent with the label instead of a label-only relabel that left the file loader-orphaned. Triggered off the POST /tags add (apply_tags): on a registered folder change the file moves (preserving subfolder + filename), path-derived system tags are re-derived, and filename metadata refreshed, all in one transaction with file rollback on failure. Collisions are rejected (never clobber); an unregistered model_type: stays a permissive label; non-model and hash-only references are untouched. Shared on-disk folders keep their plural twins. --- app/assets/api/routes.py | 6 + app/assets/services/__init__.py | 4 + app/assets/services/ingest.py | 185 ++++++++++ app/assets/services/path_utils.py | 47 ++- app/assets/services/tagging.py | 7 + .../services/test_model_type_move.py | 338 ++++++++++++++++++ tests-unit/assets_test/test_tags_api.py | 89 +++++ 7 files changed, 668 insertions(+), 8 deletions(-) create mode 100644 tests-unit/assets_test/services/test_model_type_move.py diff --git a/app/assets/api/routes.py b/app/assets/api/routes.py index 2a3f4a858..6e79558a0 100644 --- a/app/assets/api/routes.py +++ b/app/assets/api/routes.py @@ -26,6 +26,7 @@ from app.assets.seeder import ScanInProgressError, asset_seeder from app.assets.services import ( DependencyMissingError, HashMismatchError, + ModelMoveError, apply_tags, asset_exists, create_from_hash, @@ -623,6 +624,11 @@ async def add_asset_tags(request: web.Request) -> web.Response: already_present=result.already_present, total_tags=result.total_tags, ) + except ModelMoveError as me: + # A model_type: edit that couldn't be applied coherently (unknown + # folder, or destination collision). The FE re-adds the prior + # model_type: tag on a non-2xx, so the asset stays coherent. + return _build_error_response(me.status, me.code, me.message, {"id": reference_id}) except PermissionError as pe: return _build_error_response(403, "FORBIDDEN", str(pe), {"id": reference_id}) except ValueError as ve: diff --git a/app/assets/services/__init__.py b/app/assets/services/__init__.py index 03990966b..5f5464266 100644 --- a/app/assets/services/__init__.py +++ b/app/assets/services/__init__.py @@ -22,9 +22,11 @@ from app.assets.services.file_utils import ( from app.assets.services.ingest import ( DependencyMissingError, HashMismatchError, + ModelMoveError, create_from_hash, ingest_existing_file, register_output_files, + relocate_model_asset_for_model_type_tags, upload_from_temp_path, ) from app.assets.database.queries import ( @@ -62,8 +64,10 @@ __all__ = [ "HashMismatchError", "IngestResult", "ListAssetsResult", + "ModelMoveError", "RegisterAssetResult", "RemoveTagsResult", + "relocate_model_asset_for_model_type_tags", "TagUsage", "UploadResult", "UserMetadata", diff --git a/app/assets/services/ingest.py b/app/assets/services/ingest.py index 28aac33d5..b0e285741 100644 --- a/app/assets/services/ingest.py +++ b/app/assets/services/ingest.py @@ -2,11 +2,13 @@ import contextlib import logging import mimetypes import os +import shutil from typing import Any, Sequence from sqlalchemy.orm import Session import app.assets.services.hashing as hashing +from app.assets.database.models import AssetReference from app.assets.database.queries import ( add_tags_to_reference, count_active_siblings, @@ -20,6 +22,7 @@ from app.assets.database.queries import ( list_references_by_asset_id, reference_exists, remove_missing_tag_for_asset_id, + remove_tags_from_reference, set_reference_metadata, set_reference_system_metadata, set_reference_tags, @@ -35,7 +38,9 @@ from app.assets.services.image_dimensions import extract_image_dimensions from app.assets.services.path_utils import ( compute_relative_filename, get_backend_system_tags_from_path, + get_model_base_for_folder, get_name_and_tags_from_asset_path, + model_folders_for_path, resolve_destination_from_tags, validate_path_within_base, ) @@ -453,6 +458,186 @@ class DependencyMissingError(Exception): super().__init__(message) +class ModelMoveError(Exception): + """A model_type: edit could not be applied coherently (BE-1641). + + Carries an HTTP-ish ``status``/``code`` so the route can surface a precise + 4xx (rather than the generic 404 the bare ValueError path produces). The FE + edit-type flow compensates on any non-2xx by re-adding the prior + ``model_type:`` tag, so a reject here leaves the asset coherent. + """ + + def __init__(self, code: str, message: str, status: int = 409): + self.code = code + self.message = message + self.status = status + super().__init__(message) + + +def _move_file(src: str, dst: str) -> None: + """Relocate a file, falling back to a cross-device copy+unlink. + + ``os.replace`` is atomic but fails with ``EXDEV`` when src and dst live on + different filesystems (``extra_model_paths`` may point at another mount), so + fall back to ``shutil.move`` there. + """ + try: + os.replace(src, dst) + except OSError: + shutil.move(src, dst) + + +def relocate_model_asset_for_model_type_tags( + session: Session, + ref: AssetReference, + requested_tags: Sequence[str], + origin: str = "automatic", +) -> bool: + """Move a filesystem-backed model asset to match an added ``model_type:`` tag. + + BE-1641 / spec-drift §2: under the ``supports_model_type_tags`` contract a + ``model_type:`` edit must stay coherent on *edit*, not just upload. When a + ``model_type:`` tag is applied to a filesystem-backed model asset + whose file is not already under that folder, move the file to the folder's + base and re-derive the path-based system tags so location and label agree. + + Mutates ``ref`` and its tags in-place within ``session`` (the caller owns + the commit). Returns True if a physical move happened, False otherwise + (non-filesystem / hash-only / non-model asset, no ``model_type:`` added, or + the target folder already covers the current path — the shared-dir case in + spec-drift §1). + + Raises: + ModelMoveError: the target folder is unknown, or the destination is + already occupied (collision) — never clobbers (Q2). + """ + if not ref.file_path: + # API-created / hash-only reference: nothing on disk to move. Labels + # stay labels (matches AC scope: "filesystem-backed model asset"). + return False + + requested_folders = [ + t.split(":", 1)[1] + for t in normalize_tags(list(requested_tags)) + if t.startswith("model_type:") and t.split(":", 1)[1] + ] + if not requested_folders: + return False + + old_path = os.path.abspath(ref.file_path) + current_folders = set(model_folders_for_path(old_path)) + if not current_folders: + # Not under any model base (e.g. an input/output asset). A model_type: + # label here is meaningless for placement; leave it as a plain label. + return False + + # The FE emits exactly one model_type: per edit; if several are requested, + # the last one wins deterministically. + target_folder = requested_folders[-1] + + # Shared on-disk dir (spec-drift §1): the path already covers the target + # folder, so re-deriving would keep both twins — no physical move needed. + if target_folder in current_folders: + return False + + # An unregistered folder_name cannot correspond to any real on-disk + # location, so there is nothing to relocate. Keep Core's established + # permissive model_type: labeling (spec-drift §3: Core is local/trusted and + # does not reject model_type: labels) — store it as a plain label, don't + # reject. A genuine edit-type action always targets a registered folder_name + # from the discovery vocabulary, so this only affects manual label adds. + try: + new_base = get_model_base_for_folder(target_folder) + except ValueError: + return False + + rel = compute_relative_filename(old_path) + if not rel: + raise ModelMoveError( + "INVALID_MODEL_PATH", + f"cannot determine relative path for model asset: {old_path}", + status=400, + ) + new_path = os.path.abspath(os.path.join(new_base, rel)) + try: + validate_path_within_base(new_path, new_base) + except ValueError as e: + raise ModelMoveError("INVALID_MODEL_PATH", str(e), status=400) + + if new_path == old_path: + return False + + # Q2: collision -> reject, never clobber. Cover both an on-disk file and a + # reference that already owns the destination path. + if os.path.exists(new_path): + raise ModelMoveError( + "DESTINATION_EXISTS", f"destination already exists: {new_path}" + ) + if get_reference_by_file_path(session, new_path) is not None: + raise ModelMoveError( + "DESTINATION_EXISTS", f"destination already registered: {new_path}" + ) + + os.makedirs(os.path.dirname(new_path), exist_ok=True) + _move_file(old_path, new_path) + try: + _reregister_moved_reference(session, ref, new_path, origin=origin) + except Exception: + # Never half-move: roll the file back before surfacing the failure. + with contextlib.suppress(Exception): + _move_file(new_path, old_path) + raise + return True + + +def _reregister_moved_reference( + session: Session, + ref: AssetReference, + new_path: str, + origin: str = "automatic", +) -> None: + """Point ``ref`` at ``new_path`` and reconcile path-derived system tags. + + Re-derives ``models`` + ``model_type:*`` from the new location, drops any + stale ``model_type:`` no longer justified by the path, and refreshes the + relative ``filename`` metadata. User labels are left untouched. + """ + # Bytes are unchanged by a move; only the location and mtime can shift + # (shutil.move's cross-device fallback re-stats). + _size_bytes, mtime_ns = get_size_and_mtime_ns(new_path) + ref.file_path = new_path + ref.mtime_ns = mtime_ns + ref.updated_at = get_utc_now() + session.flush() + + derived = get_backend_system_tags_from_path(new_path) + derived_model_types = {t for t in derived if t.startswith("model_type:")} + + current = set(get_reference_tags(session, reference_id=ref.id)) + stale = { + t for t in current if t.startswith("model_type:") + } - derived_model_types + if stale: + remove_tags_from_reference( + session, reference_id=ref.id, tags=sorted(stale) + ) + add_tags_to_reference( + session, + reference_id=ref.id, + tags=derived, + origin=origin, + create_if_missing=True, + ) + + _update_metadata_with_filename( + session, + reference_id=ref.id, + file_path=new_path, + current_metadata=ref.user_metadata, + user_metadata={}, + ) + + def upload_from_temp_path( temp_path: str, name: str | None = None, diff --git a/app/assets/services/path_utils.py b/app/assets/services/path_utils.py index 8b2eba37f..17f9d7356 100644 --- a/app/assets/services/path_utils.py +++ b/app/assets/services/path_utils.py @@ -25,6 +25,44 @@ def get_comfy_models_folders() -> list[tuple[str, list[str]]]: return targets +def get_model_base_for_folder(folder_name: str) -> str: + """Resolve a registered model ``folder_name`` to its canonical base path. + + This is the single-valued reverse lookup (model_type -> on-disk base) the + upload destination and the edit-type move (BE-1641) share: resolve the one + named folder to its first configured base. Never fans out. + + Raises: + ValueError: the folder_name is not registered, or has no base path. + """ + model_folder_paths = dict(get_comfy_models_folders()) + try: + bases = model_folder_paths[folder_name] + except KeyError: + raise ValueError(f"unknown model category '{folder_name}'") + if not bases: + raise ValueError(f"no base path configured for category '{folder_name}'") + return os.path.abspath(bases[0]) + + +def model_folders_for_path(path: str) -> list[str]: + """Return every model folder_name whose registered base covers ``path``. + + Set-valued (spec-drift §1): a path shared by >=2 registered folders (e.g. + ``diffusion_models`` and ``unet_gguf`` both registering ``models/unet``) + belongs to all of them. Purely lexical — does not require the file to exist. + Empty when ``path`` is not under any model base (e.g. an input/output file). + """ + fp_path = Path(os.path.abspath(path)) + out: list[str] = [] + for folder_name, bases in get_comfy_models_folders(): + for base in bases: + if fp_path.is_relative_to(os.path.abspath(base)): + out.append(folder_name) + break + return out + + def _validate_subfolder(subfolder: str | None) -> list[str]: if not subfolder: return [] @@ -65,14 +103,7 @@ def resolve_destination_from_tags( folder_name = model_type_tags[0].split(":", 1)[1] if not folder_name: raise ValueError("models uploads require exactly one model_type: tag") - model_folder_paths = dict(get_comfy_models_folders()) - try: - bases = model_folder_paths[folder_name] - except KeyError: - raise ValueError(f"unknown model category '{folder_name}'") - if not bases: - raise ValueError(f"no base path configured for category '{folder_name}'") - base_dir = os.path.abspath(bases[0]) + base_dir = get_model_base_for_folder(folder_name) elif root == "input": base_dir = os.path.abspath(folder_paths.get_input_directory()) else: diff --git a/app/assets/services/tagging.py b/app/assets/services/tagging.py index 5fa39d26a..8e45f7188 100644 --- a/app/assets/services/tagging.py +++ b/app/assets/services/tagging.py @@ -9,6 +9,7 @@ from app.assets.database.queries import ( remove_tags_from_reference, ) from app.assets.database.queries.tags import list_tag_counts_for_filtered_assets +from app.assets.services.ingest import relocate_model_asset_for_model_type_tags from app.assets.services.schemas import TagUsage from app.database.db import create_session @@ -22,6 +23,12 @@ def apply_tags( with create_session() as session: ref_row = get_reference_with_owner_check(session, reference_id, owner_id) + # BE-1641: a flag-on model_type: edit on a filesystem-backed model asset + # must MOVE/re-register the file so its location stays coherent with the + # label (not a label-only relabel). Runs before the label add so the + # path-derived system tags are reconciled first; may raise ModelMoveError. + relocate_model_asset_for_model_type_tags(session, ref_row, tags) + result = add_tags_to_reference( session, reference_id=reference_id, diff --git a/tests-unit/assets_test/services/test_model_type_move.py b/tests-unit/assets_test/services/test_model_type_move.py new file mode 100644 index 000000000..4ff3155d5 --- /dev/null +++ b/tests-unit/assets_test/services/test_model_type_move.py @@ -0,0 +1,338 @@ +"""Tests for the BE-1641 edit-type MOVE/re-register behavior. + +A flag-on ``model_type:`` edit on a filesystem-backed model asset must move the +file to the folder that matches the new ``model_type:`` so the file +location stays coherent with the label (not a label-only relabel). The move runs +off the ``POST /tags`` add path (``apply_tags``). +""" +import os +import tempfile +from pathlib import Path +from unittest.mock import patch + +import pytest +from sqlalchemy.orm import Session + +from app.assets.database.models import Asset, AssetReference +from app.assets.database.queries import ( + add_tags_to_reference, + ensure_tags_exist, + get_reference_tags, +) +from app.assets.helpers import get_utc_now +from app.assets.services import ModelMoveError, apply_tags +from app.assets.services.ingest import relocate_model_asset_for_model_type_tags + + +@pytest.fixture +def model_dirs(): + """Temp model/input/output/temp dirs with a shared on-disk model folder. + + ``diffusion_models`` and ``unet_gguf`` both register the same ``models/unet`` + dir (spec-drift §1 plural membership), so a file there belongs to both. + """ + with tempfile.TemporaryDirectory() as root: + root_path = Path(root) + checkpoints = root_path / "models" / "checkpoints" + loras = root_path / "models" / "loras" + unet = root_path / "models" / "unet" + input_dir = root_path / "input" + output_dir = root_path / "output" + temp_dir = root_path / "temp" + for d in (checkpoints, loras, unet, input_dir, output_dir, temp_dir): + d.mkdir(parents=True) + + folders = [ + ("checkpoints", [str(checkpoints)]), + ("loras", [str(loras)]), + ("diffusion_models", [str(unet)]), + ("unet_gguf", [str(unet)]), + ] + + with patch("app.assets.services.path_utils.folder_paths") as mock_fp, patch( + "app.assets.services.path_utils.get_comfy_models_folders", + return_value=folders, + ): + mock_fp.get_input_directory.return_value = str(input_dir) + mock_fp.get_output_directory.return_value = str(output_dir) + mock_fp.get_temp_directory.return_value = str(temp_dir) + yield { + "checkpoints": checkpoints, + "loras": loras, + "unet": unet, + "input": input_dir, + "output": output_dir, + "temp": temp_dir, + } + + +def _make_fs_ref( + session: Session, + file_path: Path, + tags: list[str], + *, + contents: bytes = b"weights", + name: str | None = None, + owner_id: str = "", +) -> AssetReference: + """Create an on-disk file plus a filesystem-backed reference carrying ``tags``.""" + file_path.parent.mkdir(parents=True, exist_ok=True) + file_path.write_bytes(contents) + + asset = Asset(hash=f"blake3:{file_path.name}", size_bytes=len(contents)) + session.add(asset) + session.flush() + + now = get_utc_now() + ref = AssetReference( + owner_id=owner_id, + name=name or file_path.name, + asset_id=asset.id, + file_path=str(file_path), + mtime_ns=os.stat(file_path).st_mtime_ns, + created_at=now, + updated_at=now, + last_access_time=now, + ) + session.add(ref) + session.flush() + + if tags: + ensure_tags_exist(session, tags) + add_tags_to_reference(session, reference_id=ref.id, tags=tags) + session.commit() + return ref + + +def _tags_after(session: Session, reference_id: str) -> set[str]: + session.expire_all() + return set(get_reference_tags(session, reference_id)) + + +class TestMoveHappyPath: + def test_checkpoint_to_lora_moves_file_and_reconciles_tags( + self, mock_create_session, session: Session, model_dirs + ): + src = model_dirs["checkpoints"] / "model.safetensors" + ref = _make_fs_ref(session, src, ["models", "model_type:checkpoints"]) + + apply_tags(reference_id=ref.id, tags=["model_type:loras"]) + + dst = model_dirs["loras"] / "model.safetensors" + assert not src.exists() + assert dst.exists() + assert dst.read_bytes() == b"weights" + + session.expire_all() + moved = session.get(AssetReference, ref.id) + assert moved.file_path == str(dst) + + tags = _tags_after(session, ref.id) + assert "model_type:loras" in tags + assert "model_type:checkpoints" not in tags + assert "models" in tags + + def test_preserves_subfolder_structure( + self, mock_create_session, session: Session, model_dirs + ): + src = model_dirs["checkpoints"] / "sub" / "nested" / "m.safetensors" + ref = _make_fs_ref(session, src, ["models", "model_type:checkpoints"]) + + apply_tags(reference_id=ref.id, tags=["model_type:loras"]) + + dst = model_dirs["loras"] / "sub" / "nested" / "m.safetensors" + assert not src.exists() + assert dst.exists() + session.expire_all() + assert session.get(AssetReference, ref.id).file_path == str(dst) + + def test_preserves_user_labels( + self, mock_create_session, session: Session, model_dirs + ): + src = model_dirs["checkpoints"] / "labelled.safetensors" + ref = _make_fs_ref( + session, src, ["models", "model_type:checkpoints", "favorite", "sdxl"] + ) + + apply_tags(reference_id=ref.id, tags=["model_type:loras"]) + + tags = _tags_after(session, ref.id) + assert {"favorite", "sdxl"} <= tags + + def test_refreshes_filename_metadata( + self, mock_create_session, session: Session, model_dirs + ): + src = model_dirs["checkpoints"] / "deep" / "m.safetensors" + ref = _make_fs_ref(session, src, ["models", "model_type:checkpoints"]) + + apply_tags(reference_id=ref.id, tags=["model_type:loras"]) + + session.expire_all() + moved = session.get(AssetReference, ref.id) + # filename is relative to the category root, so the subfolder survives. + assert moved.user_metadata["filename"] == "deep/m.safetensors" + + +class TestNoMoveCases: + def test_same_type_is_noop( + self, mock_create_session, session: Session, model_dirs + ): + src = model_dirs["checkpoints"] / "m.safetensors" + ref = _make_fs_ref(session, src, ["models", "model_type:checkpoints"]) + + moved = relocate_model_asset_for_model_type_tags( + session, ref, ["model_type:checkpoints"] + ) + assert moved is False + assert src.exists() + + def test_shared_dir_does_not_move( + self, mock_create_session, session: Session, model_dirs + ): + # File in the shared unet dir belongs to BOTH diffusion_models and + # unet_gguf; editing to the sibling that shares the dir must not move. + src = model_dirs["unet"] / "g.gguf" + ref = _make_fs_ref( + session, + src, + ["models", "model_type:diffusion_models", "model_type:unet_gguf"], + ) + + moved = relocate_model_asset_for_model_type_tags( + session, ref, ["model_type:unet_gguf"] + ) + assert moved is False + assert src.exists() + + def test_hash_only_reference_is_label_only( + self, mock_create_session, session: Session, model_dirs + ): + asset = Asset(hash="blake3:hashonly", size_bytes=10) + session.add(asset) + session.flush() + now = get_utc_now() + ref = AssetReference( + name="hashonly", + asset_id=asset.id, + file_path=None, + created_at=now, + updated_at=now, + last_access_time=now, + ) + session.add(ref) + session.commit() + + moved = relocate_model_asset_for_model_type_tags( + session, ref, ["model_type:loras"] + ) + assert moved is False + + def test_non_model_filesystem_asset_is_label_only( + self, mock_create_session, session: Session, model_dirs + ): + src = model_dirs["input"] / "photo.png" + ref = _make_fs_ref(session, src, ["input"]) + + moved = relocate_model_asset_for_model_type_tags( + session, ref, ["model_type:loras"] + ) + assert moved is False + assert src.exists() + + def test_no_model_type_tag_is_noop( + self, mock_create_session, session: Session, model_dirs + ): + src = model_dirs["checkpoints"] / "m.safetensors" + ref = _make_fs_ref(session, src, ["models", "model_type:checkpoints"]) + + moved = relocate_model_asset_for_model_type_tags( + session, ref, ["favorite"] + ) + assert moved is False + assert src.exists() + + +class TestUnknownFolderIsLabelOnly: + def test_unknown_folder_is_stored_as_label_not_rejected( + self, mock_create_session, session: Session, model_dirs + ): + # Core stays permissive about model_type: LABELS (spec-drift §3): an + # unregistered folder_name can't map to a real location, so it's a plain + # label, not a move and not a reject. + src = model_dirs["checkpoints"] / "m.safetensors" + ref = _make_fs_ref(session, src, ["models", "model_type:checkpoints"]) + + result = apply_tags(reference_id=ref.id, tags=["model_type:bogus"]) + + assert src.exists() # not moved + assert "model_type:bogus" in result.total_tags + # The real path-derived type is untouched. + assert "model_type:checkpoints" in result.total_tags + + +class TestRejects: + def test_collision_rejected_409_and_not_clobbered( + self, mock_create_session, session: Session, model_dirs + ): + src = model_dirs["checkpoints"] / "m.safetensors" + ref = _make_fs_ref(session, src, ["models", "model_type:checkpoints"]) + + # Pre-existing file at the destination must not be overwritten. + dst = model_dirs["loras"] / "m.safetensors" + dst.write_bytes(b"existing-lora") + + with pytest.raises(ModelMoveError) as ei: + apply_tags(reference_id=ref.id, tags=["model_type:loras"]) + assert ei.value.status == 409 + assert ei.value.code == "DESTINATION_EXISTS" + assert src.exists() + assert dst.read_bytes() == b"existing-lora" + + def test_collision_with_registered_reference_rejected( + self, mock_create_session, session: Session, model_dirs + ): + src = model_dirs["checkpoints"] / "m.safetensors" + ref = _make_fs_ref(session, src, ["models", "model_type:checkpoints"]) + + # Another reference already owns the destination path (no on-disk file). + dst = model_dirs["loras"] / "m.safetensors" + other = Asset(hash="blake3:other", size_bytes=5) + session.add(other) + session.flush() + now = get_utc_now() + session.add( + AssetReference( + name="m.safetensors", + asset_id=other.id, + file_path=str(dst), + created_at=now, + updated_at=now, + last_access_time=now, + ) + ) + session.commit() + + with pytest.raises(ModelMoveError) as ei: + apply_tags(reference_id=ref.id, tags=["model_type:loras"]) + assert ei.value.code == "DESTINATION_EXISTS" + assert src.exists() + + +class TestRollback: + def test_file_rolled_back_when_db_update_fails( + self, mock_create_session, session: Session, model_dirs + ): + src = model_dirs["checkpoints"] / "m.safetensors" + ref = _make_fs_ref(session, src, ["models", "model_type:checkpoints"]) + dst = model_dirs["loras"] / "m.safetensors" + + with patch( + "app.assets.services.ingest.get_size_and_mtime_ns", + side_effect=OSError("boom"), + ): + with pytest.raises(OSError, match="boom"): + apply_tags(reference_id=ref.id, tags=["model_type:loras"]) + + # The half-move must be undone: source restored, destination clean. + assert src.exists() + assert not dst.exists() diff --git a/tests-unit/assets_test/test_tags_api.py b/tests-unit/assets_test/test_tags_api.py index 93786696f..2df318188 100644 --- a/tests-unit/assets_test/test_tags_api.py +++ b/tests-unit/assets_test/test_tags_api.py @@ -259,3 +259,92 @@ def test_tags_prefix_treats_underscore_literal( assert tag_ok in names, f"Expected {tag_ok} to be returned for prefix '{base}_'" assert tag_bad not in names, f"'{tag_bad}' must not match — '_' is not a wildcard" assert body["total"] == 1 + + +def _upload_model(http, api_base, name, tags): + """Upload a fresh filesystem-backed model asset, returning the response dict.""" + content = uuid.uuid4().bytes + b"W" * (4096 - 16) + files = {"file": (name, content, "application/octet-stream")} + form = {"tags": json.dumps(tags), "name": name, "user_metadata": json.dumps({})} + r = http.post(api_base + "/api/assets", files=files, data=form, timeout=120) + body = r.json() + assert r.status_code == 201, body + return body + + +def test_edit_type_moves_file_and_reregisters( + http: requests.Session, + api_base: str, + comfy_tmp_base_dir, +): + """BE-1641: a flag-on model_type: edit MOVEs the file to the new folder. + + Mirrors the FE edit-type flow (remove old model_type:, add new) and asserts + both the tag set and the on-disk location end up coherent. + """ + + name = f"edit_type_{uuid.uuid4().hex[:8]}.safetensors" + asset = _upload_model(http, api_base, name, ["models", "model_type:checkpoints", "unit-tests"]) + aid = asset["id"] + + models_dir = comfy_tmp_base_dir / "models" + rel = asset["user_metadata"]["filename"] + src_path = models_dir / "checkpoints" / rel + dst_path = models_dir / "loras" / rel + assert src_path.is_file(), f"uploaded model should be under checkpoints: {src_path}" + + # FE edit-type: DELETE old model_type:, then POST the new one. + rd = http.delete( + f"{api_base}/api/assets/{aid}/tags", + json={"tags": ["model_type:checkpoints"]}, + timeout=120, + ) + assert rd.status_code == 200, rd.json() + + ra = http.post( + f"{api_base}/api/assets/{aid}/tags", + json={"tags": ["model_type:loras"]}, + timeout=120, + ) + assert ra.status_code == 200, ra.json() + + # File relocated on disk, no stale copy left behind. + assert dst_path.is_file(), f"file should have moved to loras: {dst_path}" + assert not src_path.exists(), "stale checkpoints copy must be gone" + + # Tag set is coherent with the new location. + rg = http.get(f"{api_base}/api/assets/{aid}", timeout=120) + detail = rg.json() + assert rg.status_code == 200, detail + tags = set(detail["tags"]) + assert "model_type:loras" in tags + assert "model_type:checkpoints" not in tags + assert "models" in tags + + http.delete(f"{api_base}/api/assets/{aid}", timeout=30) + + +def test_edit_type_unknown_folder_is_label_only( + http: requests.Session, + api_base: str, +): + """An unregistered model_type: target stays a plain label (no move/reject). + + Core stays permissive about model_type: labels; only a registered folder + triggers the edit-type move. The file must not move. + """ + name = f"edit_bad_{uuid.uuid4().hex[:8]}.safetensors" + asset = _upload_model(http, api_base, name, ["models", "model_type:checkpoints", "unit-tests"]) + aid = asset["id"] + + r = http.post( + f"{api_base}/api/assets/{aid}/tags", + json={"tags": ["model_type:does_not_exist"]}, + timeout=120, + ) + body = r.json() + assert r.status_code == 200, body + assert "model_type:does_not_exist" in body["total_tags"] + assert "model_type:checkpoints" in body["total_tags"] + + http.delete(f"{api_base}/api/assets/{aid}", timeout=30)