From 2d6cf545b935efa0176d7bd55a5d841278ab3813 Mon Sep 17 00:00:00 2001 From: Will Miao Date: Tue, 11 Aug 2026 14:08:15 +0800 Subject: [PATCH] feat(delete): stage model and recipe deletes for 30s undo --- py/services/checkpoint_scanner.py | 5 +- py/services/model_lifecycle_service.py | 26 +- py/services/model_scanner.py | 78 +- py/services/pending_delete_service.py | 974 +++++++++++ py/services/recipes/persistence_service.py | 66 +- py/services/settings_manager.py | 1 + py/utils/usage_stats.py | 4 +- .../services/test_model_lifecycle_service.py | 169 ++ tests/services/test_model_scanner.py | 264 +++ tests/services/test_pending_delete_service.py | 1535 +++++++++++++++++ tests/services/test_recipe_persistence.py | 338 ++++ tests/services/test_settings_manager.py | 5 + 12 files changed, 3444 insertions(+), 21 deletions(-) create mode 100644 py/services/pending_delete_service.py create mode 100644 tests/services/test_pending_delete_service.py create mode 100644 tests/services/test_recipe_persistence.py diff --git a/py/services/checkpoint_scanner.py b/py/services/checkpoint_scanner.py index 707a0d0d..af4318e5 100644 --- a/py/services/checkpoint_scanner.py +++ b/py/services/checkpoint_scanner.py @@ -13,7 +13,7 @@ from ..utils.models import CheckpointMetadata from ..utils.file_utils import find_preview_file, normalize_path, calculate_autov3 from ..utils.metadata_manager import MetadataManager from ..config import config -from .model_scanner import ModelScanner +from .model_scanner import ModelScanner, _is_excluded_dir from .model_hash_index import ModelHashIndex logger = logging.getLogger(__name__) @@ -328,7 +328,8 @@ class CheckpointScanner(ModelScanner): if not os.path.exists(root_path): continue - for dirpath, _dirnames, filenames in os.walk(root_path): + for dirpath, dirnames, filenames in os.walk(root_path): + dirnames[:] = [d for d in dirnames if not _is_excluded_dir(d)] for filename in filenames: if not filename.endswith(".metadata.json"): continue diff --git a/py/services/model_lifecycle_service.py b/py/services/model_lifecycle_service.py index 3f7e6318..be24cc6f 100644 --- a/py/services/model_lifecycle_service.py +++ b/py/services/model_lifecycle_service.py @@ -7,6 +7,7 @@ import os from typing import Any, Awaitable, Callable, Dict, Iterable, List, Mapping, Optional, TYPE_CHECKING, cast from ..services.service_registry import ServiceRegistry +from ..services.pending_delete_service import get_pending_delete_service from ..utils.constants import PREVIEW_EXTENSIONS from ..utils.metadata_manager import MetadataManager @@ -129,9 +130,24 @@ class ModelLifecycleService: target_dir = os.path.dirname(file_path) base_name = os.path.basename(file_path) file_name, main_extension = os.path.splitext(base_name) - deleted_files = await delete_model_artifacts( - target_dir, file_name, main_extension=main_extension + + # Stage the delete into the pending-delete service when undo is + # enabled; a successful stage renames the artifacts away, otherwise + # fall back to the direct hard delete. + pending_delete_service = await get_pending_delete_service() + batch_id = await pending_delete_service.stage_model_delete( + scanner=self._scanner, + target_dir=target_dir, + file_name=file_name, + main_extension=main_extension, + original_file_path=file_path, + cached_entry=cached_entry, ) + deleted_files: List[str] = [] + if batch_id is None: + deleted_files = await delete_model_artifacts( + target_dir, file_name, main_extension=main_extension + ) if cache: cache.raw_data = [ @@ -151,7 +167,11 @@ class ModelLifecycleService: if callable(persist_current_cache): await cast(Awaitable[Any], persist_current_cache()) - return {"success": True, "deleted_files": deleted_files} + return { + "success": True, + "deleted_files": deleted_files, + "batch_id": batch_id, + } @staticmethod def _extract_model_id_from_payload(payload: Any) -> Optional[int]: diff --git a/py/services/model_scanner.py b/py/services/model_scanner.py index 8409abcc..e3f1226f 100644 --- a/py/services/model_scanner.py +++ b/py/services/model_scanner.py @@ -19,12 +19,28 @@ from .service_registry import ServiceRegistry from .websocket_manager import ws_manager from .persistent_model_cache import get_persistent_cache from .settings_manager import get_settings_manager +from .pending_delete_service import PENDING_DELETE_DIR_NAME, get_pending_delete_service from .cache_entry_validator import CacheEntryValidator from .cache_health_monitor import CacheHealthMonitor, CacheHealthStatus logger = logging.getLogger(__name__) +def _is_excluded_dir(name: str) -> bool: + """Return True when a directory entry must be skipped during model walks. + + The pending-delete staging directory is excluded so staged files never + appear in the library as ghost model entries. + """ + return name == PENDING_DELETE_DIR_NAME + + +def _is_pending_delete_path(path: str) -> bool: + """Return True when any path component is the pending-delete staging dir.""" + normalized = str(path).replace(os.sep, "/") + return any(part == PENDING_DELETE_DIR_NAME for part in normalized.split("/")) + + @dataclass class CacheBuildResult: """Represents the outcome of scanning model files for cache building.""" @@ -711,6 +727,8 @@ class ModelScanner: if ext in self.file_extensions: total_files += 1 elif entry.is_dir(follow_symlinks=True): + if _is_excluded_dir(entry.name): + continue count_recursive(entry.path) except Exception as e: logger.error(f"Error counting files in entry {entry.path}: {e}") @@ -864,7 +882,8 @@ class ModelScanner: continue # Recursively scan directory - for root, _, files in os.walk(root_path, followlinks=True): + for root, dirnames, files in os.walk(root_path, followlinks=True): + dirnames[:] = [d for d in dirnames if not _is_excluded_dir(d)] real_root = os.path.realpath(root) if real_root in visited_real_paths: continue @@ -1137,6 +1156,11 @@ class ModelScanner: hash_index = hash_index or self._hash_index excluded_models = excluded_models if excluded_models is not None else self._excluded_models + # Belt-and-braces: staged files must never become library entries even + # if a caller invokes this method directly with a staging path. + if _is_pending_delete_path(file_path): + return None + metadata, should_skip = await MetadataManager.load_metadata(file_path, self.model_class) if should_skip: @@ -1456,6 +1480,8 @@ class ModelScanner: if self.is_cancelled(): return elif entry.is_dir(follow_symlinks=True): + if _is_excluded_dir(entry.name): + continue await scan_recursive(entry.path, root_path, visited_paths) except Exception as entry_error: logger.error(f"Error processing entry {entry.path}: {entry_error}") @@ -2206,6 +2232,11 @@ class ModelScanner: # Track deleted models to update cache once deleted_models = [] + # Stage each file into the pending-delete staging area and merge + # all per-file batches into ONE batch for the whole bulk action. + pending_delete_service = await get_pending_delete_service() + batch_ids: List[str] = [] + for file_path in file_paths: if self.is_cancelled(): logger.info(f"{self.model_type.capitalize()} Scanner: Bulk delete cancelled by user") @@ -2218,11 +2249,35 @@ class ModelScanner: base_name = os.path.basename(file_path) file_name, main_extension = os.path.splitext(base_name) - deleted_files = await delete_model_artifacts( - target_dir, - file_name, + # Snapshot the cache entry BEFORE the cache mutation that + # runs after the loop - the manifest needs it for undo. + cached_entry = None + if cache is not None: + cached_entry = next( + (item for item in cache.raw_data if item.get('file_path') == file_path), + None, + ) + + batch_id = await pending_delete_service.stage_model_delete( + scanner=self, + target_dir=target_dir, + file_name=file_name, main_extension=main_extension, + original_file_path=file_path, + cached_entry=cached_entry, ) + + if batch_id is not None: + # Artifacts were renamed into staging: the main file is + # gone from its original location. + batch_ids.append(batch_id) + deleted_files = [file_path] + else: + deleted_files = await delete_model_artifacts( + target_dir, + file_name, + main_extension=main_extension, + ) if deleted_files: deleted_models.append(file_path) @@ -2246,6 +2301,18 @@ class ModelScanner: 'error': str(e) }) + # Merge every staged per-file batch into ONE undoable batch. On a + # merge failure (cross-volume EXDEV etc.) the response falls back + # to the constituent batch_ids array so the frontend can undo them + # sequentially. + batch_field: Dict[str, Any] = {} + if batch_ids: + merged_id = await pending_delete_service.merge_batches(batch_ids) + if merged_id is not None: + batch_field['batch_id'] = merged_id + else: + batch_field['batch_ids'] = list(batch_ids) + # Batch update cache if any models were deleted if deleted_models: # Update the cache in a batch operation @@ -2257,7 +2324,8 @@ class ModelScanner: 'total_deleted': total_deleted, 'total_attempted': len(file_paths), 'cache_updated': cache_updated, - 'results': results + 'results': results, + **batch_field } except Exception as e: diff --git a/py/services/pending_delete_service.py b/py/services/pending_delete_service.py new file mode 100644 index 00000000..91fc620f --- /dev/null +++ b/py/services/pending_delete_service.py @@ -0,0 +1,974 @@ +"""Pending-delete staging service. + +Stages model/recipe deletes into hidden per-root staging directories so a +30-second undo window can restore them before the physical purge runs. The +service is the foundation for the delete-undo feature: every staged batch is +described by a ``manifest.json`` which is the ONLY source of truth. + +LOCK HIERARCHY (critical - asyncio.Lock is NOT re-entrant): +``_ops_lock`` is acquired ONLY by stage_model_delete, stage_recipe_delete, +merge_batches, undo and purge_batch. ``purge_expired()`` NEVER acquires it - +it enumerates staging dirs and delegates each batch to ``purge_batch`` (which +locks). The opportunistic ``await self.purge_expired()`` at the start of +stage_*/undo MUST therefore run BEFORE those methods acquire the lock. +""" + +from __future__ import annotations + +import asyncio +import errno +import json +import logging +import os +import shutil +import tempfile +import time +import uuid +from typing import ( + Any, + Awaitable, + Callable, + Dict, + List, + Optional, + Sequence, + Set, + Tuple, + cast, +) + +from ..utils.constants import PREVIEW_EXTENSIONS +from ..utils import settings_paths +from .settings_manager import get_settings_manager + +logger = logging.getLogger(__name__) + +# Undo window in seconds before a staged batch becomes purge-eligible. +PENDING_DELETE_TTL_SECONDS = 30 +# Hidden staging directory name placed under each model root (and the settings +# dir for recipes). +PENDING_DELETE_DIR_NAME = ".lm-pending-delete" +# Manifest file name inside every batch directory. +MANIFEST_FILE_NAME = "manifest.json" +# Suffix appended when quarantining malformed/manifest-less batch dirs. The +# quarantine is terminal: never re-renamed, never re-quarantined, never +# deleted by the sweep. +ORPHANED_SUFFIX = ".orphaned" + +# Map scanner.model_type (singular) to the manifest page type values. +_MODEL_TYPE_PAGE_MAP = { + "lora": "loras", + "checkpoint": "checkpoints", + "embedding": "embeddings", +} + +# Module-level alias so tests can spy on timer task creation without patching +# the global asyncio module. +_create_task = asyncio.create_task + + +class PendingDeleteService: + """Stage, undo and purge pending model/recipe deletions. + + Singleton + asyncio.Lock pattern (mirrors py/services/model_scanner.py). + """ + + _instance: Optional["PendingDeleteService"] = None + _lock: asyncio.Lock = asyncio.Lock() + + @classmethod + async def get_instance(cls) -> "PendingDeleteService": + """Return the lazily initialised singleton instance.""" + async with cls._lock: + if cls._instance is None: + cls._instance = cls() + return cls._instance + + def __init__(self) -> None: + if hasattr(self, "_initialized"): + return + self._initialized = True + # Serialises stage/merge/undo/purge_batch. purge_expired never locks. + self._ops_lock = asyncio.Lock() + # Track fire-and-forget purge timer tasks to keep them alive and to + # cancel them on shutdown / singleton reset. + self._purge_tasks: Set[Any] = set() + # Roots the service has staged into (in-process). Combined with the + # ServiceRegistry roots during sweeps so undo/purge work even before + # every scanner is registered. + self._known_roots: List[str] = [] + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + async def stage_model_delete( + self, + *, + scanner: Any, + target_dir: str, + file_name: str, + main_extension: Optional[str], + original_file_path: str, + cached_entry: Optional[Dict[str, Any]], + ) -> Optional[str]: + """Rename a model's artifacts into a per-root staging batch. + + Returns the batch id, or ``None`` when undo is disabled, the staging + root cannot be resolved, or staging failed (caller falls back to a + hard delete). + """ + # LOCK-FREE section: opportunistic purge must never run while holding + # the ops lock (the lock is not re-entrant). + await self._opportunistic_purge() + + if not self._undo_enabled(): + return None + + async with self._ops_lock: + batch_dir: Optional[str] = None + staged_pairs: List[Dict[str, Any]] = [] + try: + root = self._find_model_root(scanner, original_file_path) + if not root: + logger.warning( + "No model root contains %s; skipping staging", + original_file_path, + ) + return None + + artifacts = self._enumerate_model_artifacts( + target_dir, file_name, main_extension + ) + if not artifacts: + logger.warning( + "No existing artifacts for %s; skipping staging", + original_file_path, + ) + return None + + batch_id = self._new_batch_id() + batch_dir = os.path.join( + os.path.join(root, PENDING_DELETE_DIR_NAME), batch_id + ) + os.makedirs(batch_dir, exist_ok=True) + + staged_pairs = self._rename_artifacts_into_batch( + batch_dir, artifacts, staged_pairs + ) + # Attach the model snapshot to the MAIN-file entry (the one + # whose original path is the model file itself, not the + # metadata/preview sidecars). Merged bulk manifests therefore + # carry EVERY deleted model's snapshot on its entry; the + # top-level model_snapshot is kept for backward compatibility + # and the single-delete path. + main_abs = os.path.abspath(original_file_path) + for entry in staged_pairs: + if entry.get("original") == main_abs: + entry["snapshot"] = cached_entry + break + manifest = self._build_manifest( + batch_id=batch_id, + kind="model", + model_type=self._resolve_model_type(scanner), + expires_at=int(time.time()) + PENDING_DELETE_TTL_SECONDS, + entries=staged_pairs, + model_snapshot=cached_entry, + ) + self._write_manifest_atomic(batch_dir, manifest) + self._remember_root(root) + # Arm the per-batch purge timer. Safe inside the lock: task + # creation does not await, and purge_batch re-reads the + # manifest's expires_at at fire time, so stale timers no-op. + self._arm_purge_timer(batch_id) + logger.info( + "Staged model delete batch %s with %d file(s)", + batch_id, + len(staged_pairs), + ) + return batch_id + except OSError as exc: + logger.warning( + "Staging model %s failed: %s; rolling back", original_file_path, exc + ) + if batch_dir: + self._rollback_model_staging(batch_dir, staged_pairs) + self._remove_empty_dir(batch_dir) + return None + except Exception as exc: # defensive - never block the delete flow + logger.warning( + "Unexpected error staging model %s: %s", original_file_path, exc + ) + return None + + async def stage_recipe_delete( + self, + *, + recipe_json_path: str, + image_path: Optional[str], + recipe_data: Optional[Dict[str, Any]], + ) -> Optional[str]: + """Copy a recipe JSON (and, when it exists, its image) into staging. + + Returns the batch id, or ``None`` when undo is disabled / staging + failed. Missing or shared preview images are skipped. + """ + await self._opportunistic_purge() + + if not self._undo_enabled(): + return None + + async with self._ops_lock: + batch_dir: Optional[str] = None + staged_pairs: List[Dict[str, Any]] = [] + try: + json_path = os.path.abspath(os.path.normpath(recipe_json_path)) + if not os.path.exists(json_path): + logger.warning( + "Recipe JSON %s does not exist; skipping staging", json_path + ) + return None + + batch_id = self._new_batch_id() + batch_dir = os.path.join(self._recipe_staging_parent(), batch_id) + os.makedirs(batch_dir, exist_ok=True) + + staged_pairs = self._copy_recipe_artifacts( + batch_dir, json_path, image_path, staged_pairs + ) + manifest = self._build_manifest( + batch_id=batch_id, + kind="recipe", + model_type=None, + expires_at=int(time.time()) + PENDING_DELETE_TTL_SECONDS, + entries=staged_pairs, + recipe_snapshot=recipe_data, + ) + self._write_manifest_atomic(batch_dir, manifest) + self._arm_purge_timer(batch_id) + logger.info( + "Staged recipe delete batch %s with %d file(s)", + batch_id, + len(staged_pairs), + ) + return batch_id + except OSError as exc: + logger.warning( + "Staging recipe %s failed: %s; rolling back", + recipe_json_path, + exc, + ) + if batch_dir: + self._rollback_recipe_staging(batch_dir, staged_pairs) + self._remove_empty_dir(batch_dir) + return None + except Exception as exc: # defensive - never block the delete flow + logger.warning( + "Unexpected error staging recipe %s: %s", recipe_json_path, exc + ) + return None + + async def merge_batches(self, batch_ids: Sequence[str]) -> Optional[str]: + """Merge several batches into the first batch's manifest. + + Winner is ``batch_ids[0]``. The staged files of losing batches are + MOVED (os.rename) into the winner's batch dir and their ``staged`` + paths rewritten in the merged manifest BEFORE any loser dir is + removed. ``expires_at`` is re-anchored to ``now + TTL`` at merge time + and a FRESH purge timer is armed for the winner. + + On any move failure every already-moved file is moved BACK and the + original batch dirs/manifests are left intact; ``None`` is returned so + callers fall back to the ``batch_ids`` array contract. Cross-volume + merges hit EXDEV here - expected and fine (the fallback is the normal + path for those bulks). + """ + if not batch_ids: + return None + + async with self._ops_lock: + winner_id = batch_ids[0] + winner_dir = await self._find_batch_dir(winner_id) + if not winner_dir: + return None + winner_manifest = self._read_manifest(winner_dir) + if winner_manifest is None: + return None + + # Track (entry, original_staged_path, loser_dir) for rollback. + moved: List[Tuple[Dict[str, Any], str, str]] = [] + processed_losers: List[str] = [] + + try: + for loser_id in batch_ids[1:]: + loser_dir = await self._find_batch_dir(loser_id) + if not loser_dir or os.path.normpath(loser_dir) == os.path.normpath( + winner_dir + ): + continue + loser_manifest = self._read_manifest(loser_dir) + if loser_manifest is None: + # Corrupted loser: leave it for the sweep to quarantine. + continue + for entry in loser_manifest.get("entries") or []: + if entry.get("restored"): + continue + staged_path = entry.get("staged") + if not staged_path or not os.path.exists(staged_path): + continue + new_staged = os.path.join( + winner_dir, os.path.basename(staged_path) + ) + if os.path.exists(new_staged): + # os.rename would silently overwrite the existing + # staged file on POSIX - never drop a staged file. + # Abort the merge so callers fall back to the + # batch_ids array contract. + raise OSError( + f"Merge collision: {os.path.basename(staged_path)} " + "already staged in winner batch" + ) + os.rename(staged_path, new_staged) + original_staged = entry["staged"] + entry["staged"] = os.path.abspath(new_staged) + winner_manifest["entries"].append(entry) + moved.append((entry, original_staged, loser_dir)) + processed_losers.append(loser_dir) + except OSError as exc: + logger.warning( + "Merge of %s failed after moving files: %s; rolling back", + list(batch_ids), + exc, + ) + self._rollback_merge_moves(moved) + return None + + # Re-anchor expiry and persist the merged manifest atomically. + winner_manifest["expires_at"] = ( + int(time.time()) + PENDING_DELETE_TTL_SECONDS + ) + try: + self._write_manifest_atomic(winner_dir, winner_manifest) + except OSError as exc: + logger.warning( + "Failed to write merged manifest for %s: %s; rolling back", + winner_id, + exc, + ) + self._rollback_merge_moves(moved) + return None + + # All moves committed: remove loser dirs (must be empty by now). + for loser_dir in processed_losers: + self._remove_manifest(loser_dir) + self._remove_empty_dir(loser_dir) + + # Arm a fresh purge timer for the winner with the re-anchored + # expiry (the winner's original timer fires at the OLD expiry and + # no-ops after re-reading the manifest - without this fresh timer + # an idle server would never purge the merged batch). + self._arm_purge_timer(winner_id) + logger.info("Merged batches %s into %s", list(batch_ids), winner_id) + return winner_id + + async def undo(self, batch_id: str) -> Dict[str, Any]: + """Restore every staged file of a batch to its original path. + + Raises ``ValueError`` for unknown batches, expired batches ("Undo + window expired") and occupied target paths ("Target path occupied"). + Restores entries one at a time, persisting the manifest after each, so + a mid-undo failure leaves a retry-able state. + """ + await self._opportunistic_purge() + + async with self._ops_lock: + batch_dir = await self._find_batch_dir(batch_id) + if not batch_dir: + raise ValueError(f"Unknown batch id: {batch_id}") + manifest = self._read_manifest(batch_dir) + if manifest is None: + raise ValueError(f"Manifest missing for batch {batch_id}") + + if manifest.get("state") == "restored": + return self._undo_result(manifest) + + now = time.time() + expires_at = manifest.get("expires_at") + if isinstance(expires_at, (int, float)) and expires_at < now: + raise ValueError("Undo window expired") + + entries = manifest.get("entries") or [] + + # Pre-check ALL target paths (except already-restored entries) so + # an occupied original path protects the new file and leaves the + # whole batch intact. + for entry in entries: + if entry.get("restored"): + continue + original_path = entry.get("original") + if original_path and os.path.exists(original_path): + raise ValueError("Target path occupied") + + for entry in entries: + if entry.get("restored"): + continue + staged_path = entry.get("staged") + original_path = entry.get("original") + if not staged_path or not original_path: + entry["restored"] = True + continue + if not os.path.exists(staged_path): + # Staged file already gone (purged or manually removed): + # treat as restored and finish the rest of the batch. + entry["restored"] = True + self._write_manifest_atomic(batch_dir, manifest) + continue + self._restore_file(staged_path, original_path) + entry["restored"] = True + # Persist after each entry so a mid-undo failure is retry-able. + self._write_manifest_atomic(batch_dir, manifest) + + manifest["state"] = "restored" + try: + self._write_manifest_atomic(batch_dir, manifest) + except OSError as exc: + logger.warning( + "Failed to mark manifest restored for %s: %s", batch_id, exc + ) + + # Remove the manifest + batch dir only after all entries restored. + self._remove_manifest(batch_dir) + self._remove_empty_dir(batch_dir) + + logger.info("Restored pending-delete batch %s", batch_id) + return self._undo_result(manifest) + + async def purge_expired(self) -> int: + """Purge every expired batch across ALL model roots and the recipe dir. + + Lock-free by design: enumerates staging parents (all scanner types via + the ServiceRegistry plus the global recipe staging dir) and delegates + each batch to :meth:`purge_batch`, which acquires the ops lock. Never + call this while holding the ops lock. + """ + purged = 0 + for parent in await self._get_all_staging_parents(): + if not os.path.isdir(parent): + continue + for name in self._list_dir_names(parent): + if name.endswith(ORPHANED_SUFFIX): + # Quarantine is terminal - never re-rename or delete. + continue + try: + await self.purge_batch(name) + purged += 1 + except Exception as exc: # defensive - sweep must not crash + logger.warning("Failed to purge batch %s: %s", name, exc) + return purged + + async def purge_batch(self, batch_id: str) -> None: + """Purge one batch. Silent no-op for missing/undone/not-yet-expired. + + Missing staged files (already-restored / partially-restored batches) + are treated as already-purged. A per-file purge failure (locked file) + skips only that file and keeps the batch dir for the next round. + """ + async with self._ops_lock: + batch_dir = await self._find_batch_dir(batch_id) + if not batch_dir: + return + self._purge_batch_dir(batch_dir) + + # ------------------------------------------------------------------ + # Internals + # ------------------------------------------------------------------ + async def _opportunistic_purge(self) -> None: + """Fire the opportunistic sweep. Cheap when empty; never locked.""" + try: + await self.purge_expired() + except Exception as exc: # defensive - staging/undo must still proceed + logger.warning("Opportunistic pending-delete purge failed: %s", exc) + + def _undo_enabled(self) -> bool: + try: + return bool(get_settings_manager().get("delete_undo_enabled", True)) + except Exception as exc: # defensive - default to enabled + logger.warning("Failed to read delete_undo_enabled setting: %s", exc) + return True + + def _remember_root(self, root: str) -> None: + """Record a root the service has staged into (in-process registry).""" + if root and root not in self._known_roots: + self._known_roots.append(root) + + def _find_model_root(self, scanner: Any, original_file_path: Optional[str]) -> Optional[str]: + """Return the configured root containing ``original_file_path``.""" + finder = getattr(scanner, "_find_root_for_file", None) + if callable(finder): + try: + root = cast(Optional[str], finder(original_file_path)) + if root: + return os.path.abspath(root) + except Exception as exc: # defensive - fall back to roots scan + logger.debug("_find_root_for_file failed: %s", exc) + + if not original_file_path: + return None + roots_getter = getattr(scanner, "get_model_roots", None) + if not callable(roots_getter): + return None + try: + normalized = os.path.abspath(os.path.normpath(original_file_path)) + for root in cast(Sequence[str], roots_getter()) or []: + root_abs = os.path.abspath(os.path.normpath(root)) + if normalized == root_abs or normalized.startswith(root_abs + os.sep): + return root_abs + except Exception as exc: # defensive - never block the delete flow + logger.debug("get_model_roots fallback failed: %s", exc) + return None + + def _resolve_model_type(self, scanner: Any) -> Optional[str]: + raw = getattr(scanner, "model_type", None) + if not raw: + return None + return _MODEL_TYPE_PAGE_MAP.get(raw, raw) + + def _enumerate_model_artifacts( + self, target_dir: str, file_name: str, main_extension: Optional[str] + ) -> List[str]: + """Enumerate existing artifacts exactly like delete_model_artifacts.""" + main_extension = ".safetensors" if main_extension is None else main_extension + main_file = f"{file_name}{main_extension}" if main_extension else file_name + patterns = [main_file, f"{file_name}.metadata.json"] + for ext in PREVIEW_EXTENSIONS: + patterns.append(f"{file_name}{ext}") + + artifacts: List[str] = [] + for pattern in patterns: + path = os.path.abspath(os.path.join(target_dir, pattern)) + if os.path.exists(path): + artifacts.append(path) + return artifacts + + def _rename_artifacts_into_batch( + self, + batch_dir: str, + artifacts: Sequence[str], + staged_pairs: List[Dict[str, Any]], + ) -> List[Dict[str, Any]]: + """Rename artifacts into the batch dir, recording progress per file. + + Progress is appended to ``staged_pairs`` before the next move so a + mid-way OSError leaves the caller with the already-moved files for + rollback. + """ + for original_path in artifacts: + staged_path = os.path.join(batch_dir, os.path.basename(original_path)) + os.rename(original_path, staged_path) + staged_pairs.append( + { + "staged": os.path.abspath(staged_path), + "original": os.path.abspath(original_path), + "restored": False, + } + ) + return staged_pairs + + def _copy_recipe_artifacts( + self, + batch_dir: str, + json_path: str, + image_path: Optional[str], + staged_pairs: List[Dict[str, Any]], + ) -> List[Dict[str, Any]]: + """Copy the recipe JSON and, when it exists, the image into the batch.""" + staged_json = os.path.join(batch_dir, os.path.basename(json_path)) + shutil.copy2(json_path, staged_json) + staged_pairs.append( + { + "staged": os.path.abspath(staged_json), + "original": json_path, + "restored": False, + } + ) + if image_path: + image_abs = os.path.abspath(os.path.normpath(image_path)) + if os.path.exists(image_abs): + staged_image = os.path.join(batch_dir, os.path.basename(image_abs)) + shutil.copy2(image_abs, staged_image) + staged_pairs.append( + { + "staged": os.path.abspath(staged_image), + "original": image_abs, + "restored": False, + } + ) + return staged_pairs + + def _restore_file(self, staged_path: str, original_path: str) -> None: + """Restore a staged file to its original path, tolerating EXDEV. + + ``os.rename`` is atomic and preferred (model staging and most recipe + restores are same-volume). Recipe staging copies into the settings-dir + staging parent, which may live on a DIFFERENT filesystem than the + recipes dir; rename then raises EXDEV. Fall back to ``shutil.copy2`` + + ``os.remove`` so the bytes are restored and the staged copy removed. + """ + try: + os.rename(staged_path, original_path) + except OSError as exc: + if exc.errno != errno.EXDEV: + raise + shutil.copy2(staged_path, original_path) + os.remove(staged_path) + + def _rollback_model_staging( + self, batch_dir: str, staged_pairs: Sequence[Dict[str, Any]] + ) -> None: + """Rename already-staged files back to their originals.""" + for pair in reversed(list(staged_pairs)): + staged_path = pair.get("staged") + original_path = pair.get("original") + if not staged_path or not original_path: + continue + if not os.path.exists(staged_path): + continue + try: + os.rename(staged_path, original_path) + except OSError as exc: # pragma: no cover - best-effort rollback + logger.warning( + "Failed to roll back staged file %s -> %s: %s", + staged_path, + original_path, + exc, + ) + + def _rollback_recipe_staging( + self, batch_dir: str, staged_pairs: Sequence[Dict[str, Any]] + ) -> None: + """Remove staged copies (recipe originals were never moved).""" + for pair in staged_pairs: + staged_path = pair.get("staged") + if not staged_path: + continue + try: + if os.path.exists(staged_path): + os.remove(staged_path) + except OSError as exc: # pragma: no cover - best-effort rollback + logger.warning( + "Failed to remove staged copy %s: %s", staged_path, exc + ) + + def _rollback_merge_moves( + self, moved: Sequence[Tuple[Dict[str, Any], str, str]] + ) -> None: + """Move already-merged files back to their original loser batch dirs.""" + for _entry, original_staged, _loser_dir in reversed(list(moved)): + current = _entry.get("staged") + if not current or not original_staged: + continue + if not os.path.exists(current): + continue + try: + os.rename(current, original_staged) + except OSError as exc: # pragma: no cover - best-effort rollback + logger.warning( + "Failed to roll back merge move %s -> %s: %s", + current, + original_staged, + exc, + ) + + def _purge_batch_dir(self, batch_dir: str) -> bool: + """Purge one batch dir. Returns True when the batch was purged/removed.""" + if not os.path.isdir(batch_dir): + return False + + manifest = self._read_manifest(batch_dir) + if manifest is None: + # Corrupted or manifest-less batch: quarantine, NEVER delete the + # staged files (they may be the only copy of the user's data). + self._quarantine_batch_dir(batch_dir) + return True + + if manifest.get("state") == "restored": + return False + + expires_at = manifest.get("expires_at") + if not isinstance(expires_at, (int, float)) or expires_at >= time.time(): + # Not yet expired - stale timers from merged-away/undone batches + # are harmless. + return False + + entries = manifest.get("entries") or [] + remaining: List[Dict[str, Any]] = [] + for entry in entries: + staged_path = entry.get("staged") + if not staged_path or not os.path.exists(staged_path): + # Missing staged file = already restored / already purged. + continue + try: + os.remove(staged_path) + except OSError as exc: + logger.warning( + "Skipping locked staged file %s: %s", staged_path, exc + ) + remaining.append(entry) + + if remaining: + # Never remove the batch dir past per-file errors; the batch is + # retried by the next opportunistic purge. + return False + + self._remove_manifest(batch_dir) + self._remove_empty_dir(batch_dir) + return True + + def _quarantine_batch_dir(self, batch_dir: str) -> str: + """Rename a malformed batch dir to ``.orphaned`` (terminal).""" + orphaned_dir = f"{batch_dir}{ORPHANED_SUFFIX}" + if os.path.exists(orphaned_dir): + orphaned_dir = f"{batch_dir}-{int(time.time())}{ORPHANED_SUFFIX}" + try: + os.rename(batch_dir, orphaned_dir) + except OSError as exc: # pragma: no cover - defensive + logger.warning("Failed to quarantine %s: %s", batch_dir, exc) + return batch_dir + logger.warning( + "Quarantined malformed/manifest-less pending-delete batch %s", + os.path.basename(batch_dir), + ) + return orphaned_dir + + def _build_manifest( + self, + *, + batch_id: str, + kind: str, + model_type: Optional[str], + expires_at: int, + entries: Sequence[Dict[str, Any]], + model_snapshot: Any = None, + recipe_snapshot: Any = None, + ) -> Dict[str, Any]: + is_model = kind == "model" + return { + "batch_id": batch_id, + "kind": kind, + "model_type": model_type if is_model else None, + "state": "staged", + "expires_at": int(expires_at), + "entries": list(entries), + "model_snapshot": model_snapshot if is_model else None, + "recipe_snapshot": recipe_snapshot if not is_model else None, + } + + def _undo_result(self, manifest: Dict[str, Any]) -> Dict[str, Any]: + restored_paths = [ + entry["original"] + for entry in manifest.get("entries") or [] + if entry.get("restored") and entry.get("original") + ] + return { + "batch_id": manifest.get("batch_id"), + "kind": manifest.get("kind"), + "model_type": manifest.get("model_type"), + "restored": restored_paths, + } + + def _write_manifest_atomic( + self, batch_dir: str, manifest: Dict[str, Any] + ) -> None: + """Write manifest.json atomically (temp file + os.replace).""" + manifest_path = os.path.join(batch_dir, MANIFEST_FILE_NAME) + fd, temp_path = tempfile.mkstemp( + dir=batch_dir, prefix=".manifest-", suffix=".tmp" + ) + try: + with os.fdopen(fd, "w", encoding="utf-8") as handle: + json.dump(manifest, handle, indent=2, ensure_ascii=False) + os.replace(temp_path, manifest_path) + except BaseException: + try: + os.remove(temp_path) + except OSError: + pass + raise + + def _read_manifest(self, batch_dir: str) -> Optional[Dict[str, Any]]: + manifest_path = os.path.join(batch_dir, MANIFEST_FILE_NAME) + try: + with open(manifest_path, "r", encoding="utf-8") as handle: + payload = json.load(handle) + except FileNotFoundError: + return None + except (json.JSONDecodeError, OSError) as exc: + logger.warning( + "Corrupted pending-delete manifest at %s: %s", manifest_path, exc + ) + return None + if not isinstance(payload, dict): + logger.warning("Invalid pending-delete manifest at %s", manifest_path) + return None + return payload + + def _recipe_staging_parent(self) -> str: + # Resolve through the module namespace so the conftest settings-dir + # isolation patch takes effect at call time. + return os.path.join( + settings_paths.get_settings_dir(create=True), PENDING_DELETE_DIR_NAME + ) + + async def _get_all_staging_parents(self) -> List[str]: + """Model staging parents for every scanner type + the recipe parent.""" + parents: List[str] = [] + for root in await self._get_all_model_roots(): + parent = os.path.join(root, PENDING_DELETE_DIR_NAME) + if parent not in parents: + parents.append(parent) + recipe_parent = self._recipe_staging_parent() + if recipe_parent not in parents: + parents.append(recipe_parent) + return parents + + async def _get_all_model_roots(self) -> List[str]: + """Collect every configured model root across all scanner types. + + Combines the in-process staging roots with the ServiceRegistry's + per-type scanners so sweeps cover every scanner type while undo/purge + still resolve batches staged before the registry was populated. + """ + from .service_registry import ServiceRegistry + + roots: List[str] = [] + for root in self._known_roots: + if root and root not in roots: + roots.append(root) + for getter_name in ( + "get_lora_scanner", + "get_checkpoint_scanner", + "get_embedding_scanner", + ): + getter = getattr(ServiceRegistry, getter_name, None) + if not callable(getter): + continue + try: + scanner = await cast(Callable[[], Awaitable[Any]], getter)() + except Exception as exc: # defensive - keep sweeping other types + logger.debug( + "Failed to resolve %s for purge enumeration: %s", + getter_name, + exc, + ) + continue + if scanner is None: + continue + get_roots = getattr(scanner, "get_model_roots", None) + if not callable(get_roots): + continue + try: + scanner_roots = cast(Sequence[Any], get_roots()) + except Exception as exc: # defensive + logger.debug( + "get_model_roots failed for %s: %s", getter_name, exc + ) + continue + for root in scanner_roots or []: + if root and root not in roots: + roots.append(root) + return roots + + async def _find_batch_dir(self, batch_id: str) -> Optional[str]: + """Locate a batch directory across every staging parent.""" + if not batch_id: + return None + for parent in await self._get_all_staging_parents(): + candidate = os.path.join(parent, batch_id) + if os.path.isdir(candidate): + return candidate + return None + + def _list_dir_names(self, parent: str) -> List[str]: + try: + return [ + name + for name in os.listdir(parent) + if os.path.isdir(os.path.join(parent, name)) + ] + except OSError as exc: # pragma: no cover - defensive + logger.debug("Failed to list staging parent %s: %s", parent, exc) + return [] + + def _remove_manifest(self, batch_dir: str) -> None: + try: + os.remove(os.path.join(batch_dir, MANIFEST_FILE_NAME)) + except OSError as exc: # pragma: no cover - best-effort + logger.debug("Failed to remove manifest in %s: %s", batch_dir, exc) + + def _remove_empty_dir(self, directory: str) -> None: + try: + os.rmdir(directory) + except OSError as exc: + logger.debug("Directory %s not empty or missing: %s", directory, exc) + + def _new_batch_id(self) -> str: + return uuid.uuid4().hex + + def _arm_purge_timer(self, batch_id: str) -> None: + """Spawn a fire-and-forget purge timer for a batch. + + The timer sleeps until the batch's current expiry and then calls + purge_batch, which re-reads the manifest's ``expires_at`` at fire time + so merged-away/undone/not-yet-expired batches are silent no-ops. + """ + try: + asyncio.get_running_loop() + except RuntimeError: + return + task = _create_task( + self._purge_batch_after_ttl(batch_id), + name=f"pending_delete_purge_{batch_id}", + ) + self._purge_tasks.add(task) + task.add_done_callback(self._purge_tasks.discard) + + async def _purge_batch_after_ttl(self, batch_id: str) -> None: + try: + delay = await self._seconds_until_expiry(batch_id) + if delay is None: + return + await asyncio.sleep(max(0.0, delay)) + await self.purge_batch(batch_id) + except asyncio.CancelledError: + raise + except Exception as exc: # defensive - a timer must never crash the loop + logger.warning("Pending-delete purge timer for %s failed: %s", batch_id, exc) + + async def _seconds_until_expiry(self, batch_id: str) -> Optional[float]: + batch_dir = await self._find_batch_dir(batch_id) + if not batch_dir: + return None + manifest = self._read_manifest(batch_dir) + if manifest is None: + return None + expires_at = manifest.get("expires_at") + if not isinstance(expires_at, (int, float)): + return None + return float(expires_at) - time.time() + + def _cancel_purge_tasks(self) -> None: + for task in list(self._purge_tasks): + task.cancel() + self._purge_tasks.clear() + + +def _reset_pending_delete_service() -> None: + """Reset the singleton and cancel in-flight purge timers (tests/shutdown).""" + instance = PendingDeleteService._instance + if instance is not None: + instance._cancel_purge_tasks() + PendingDeleteService._instance = None + + +async def get_pending_delete_service() -> PendingDeleteService: + """Return the lazily initialised global :class:`PendingDeleteService`.""" + return await PendingDeleteService.get_instance() diff --git a/py/services/recipes/persistence_service.py b/py/services/recipes/persistence_service.py index 10381422..75d81e49 100644 --- a/py/services/recipes/persistence_service.py +++ b/py/services/recipes/persistence_service.py @@ -14,6 +14,7 @@ from typing import Any, Awaitable, Dict, Iterable, Optional, cast from ...config import config from ...recipes.constants import GEN_PARAM_KEYS from ...utils.utils import calculate_recipe_fingerprint +from ..pending_delete_service import get_pending_delete_service from .errors import RecipeNotFoundError, RecipeValidationError @@ -201,12 +202,31 @@ class RecipePersistenceService: recipe_data = json.load(file_obj) image_path = recipe_data.get("file_path") + + # Stage the delete so the recipe can be undone within the undo window. + # The staging service COPIES the JSON (and existing image) into the + # global staging dir and stores recipe_data as the manifest snapshot; + # the originals are removed below as before. When staging is skipped + # (undo disabled / staging failure) the existing hard delete runs. + pending_delete_service = await get_pending_delete_service() + batch_id = await pending_delete_service.stage_recipe_delete( + recipe_json_path=recipe_json_path, + image_path=image_path, + recipe_data=recipe_data, + ) + os.remove(recipe_json_path) if image_path and os.path.exists(image_path): os.remove(image_path) await recipe_scanner.remove_recipe(recipe_id) - return PersistenceResult({"success": True, "message": "Recipe deleted successfully"}) + return PersistenceResult( + { + "success": True, + "message": "Recipe deleted successfully", + "batch_id": batch_id, + } + ) async def update_recipe(self, *, recipe_scanner, recipe_id: str, updates: dict[str, Any]) -> PersistenceResult: """Update persisted metadata for a recipe.""" @@ -450,6 +470,9 @@ class RecipePersistenceService: deleted_recipes: list[str] = [] failed_recipes: list[dict[str, Any]] = [] + batch_ids: list[str] = [] + + pending_delete_service = await get_pending_delete_service() for recipe_id in recipe_ids: recipe_json_path = await recipe_scanner.get_recipe_json_path(recipe_id) @@ -461,6 +484,17 @@ class RecipePersistenceService: with open(recipe_json_path, "r", encoding="utf-8") as file_obj: recipe_data = json.load(file_obj) image_path = recipe_data.get("file_path") + + # Stage each recipe into its own batch; collect the ids so the + # whole bulk action can be merged into ONE undoable batch. + batch_id = await pending_delete_service.stage_recipe_delete( + recipe_json_path=recipe_json_path, + image_path=image_path, + recipe_data=recipe_data, + ) + if batch_id: + batch_ids.append(batch_id) + os.remove(recipe_json_path) if image_path and os.path.exists(image_path): os.remove(image_path) @@ -471,15 +505,27 @@ class RecipePersistenceService: if deleted_recipes: await recipe_scanner.bulk_remove(deleted_recipes) - return PersistenceResult( - { - "success": True, - "deleted": deleted_recipes, - "failed": failed_recipes, - "total_deleted": len(deleted_recipes), - "total_failed": len(failed_recipes), - } - ) + payload: dict[str, Any] = { + "success": True, + "deleted": deleted_recipes, + "failed": failed_recipes, + "total_deleted": len(deleted_recipes), + "total_failed": len(failed_recipes), + } + + if batch_ids: + merged_batch_id = await pending_delete_service.merge_batches(batch_ids) + if merged_batch_id: + # Merge succeeded: one undo action covers the whole bulk. + payload["batch_id"] = merged_batch_id + else: + # Merge failure (e.g. cross-volume move): expose the constituent + # batches so the caller can undo them one at a time. + payload["batch_ids"] = batch_ids + else: + payload["batch_id"] = None + + return PersistenceResult(payload) async def save_recipe_from_widget( self, diff --git a/py/services/settings_manager.py b/py/services/settings_manager.py index 86278fe8..e266739f 100644 --- a/py/services/settings_manager.py +++ b/py/services/settings_manager.py @@ -111,6 +111,7 @@ DEFAULT_SETTINGS: Dict[str, Any] = { "backup_retention_count": 5, "use_new_license_icons": True, "group_by_model": False, + "delete_undo_enabled": True, # AI / LLM provider configuration (BYOK) "llm_provider": "openai", # "openai" | "ollama" | "custom" "llm_api_key": "", diff --git a/py/utils/usage_stats.py b/py/utils/usage_stats.py index 577849db..2a25cf41 100644 --- a/py/utils/usage_stats.py +++ b/py/utils/usage_stats.py @@ -10,6 +10,7 @@ from typing import Any, Awaitable, Dict, Set, cast from ..config import config from ..services.service_registry import ServiceRegistry +from ..services.model_scanner import _is_excluded_dir from ..utils.settings_paths import get_settings_dir # Check if running in standalone mode @@ -421,7 +422,8 @@ class UsageStats: if not os.path.exists(root_path): continue - for dirpath, _dirnames, filenames in os.walk(root_path): + for dirpath, dirnames, filenames in os.walk(root_path): + dirnames[:] = [d for d in dirnames if not _is_excluded_dir(d)] for filename in filenames: extension = os.path.splitext(filename)[1].lower() if extension not in supported_extensions: diff --git a/tests/services/test_model_lifecycle_service.py b/tests/services/test_model_lifecycle_service.py index df76ae47..2c37b3b5 100644 --- a/tests/services/test_model_lifecycle_service.py +++ b/tests/services/test_model_lifecycle_service.py @@ -1,11 +1,16 @@ +from __future__ import annotations + import json import os +from collections.abc import Iterator from pathlib import Path from typing import Any, Dict, cast import pytest from py.services.model_lifecycle_service import ModelLifecycleService, _require_path_in_library_roots +from py.services.pending_delete_service import PENDING_DELETE_DIR_NAME, _reset_pending_delete_service +from py.services.settings_manager import get_settings_manager from py.utils.metadata_manager import MetadataManager from py.utils.models import LoraMetadata @@ -901,3 +906,167 @@ async def test_extract_model_id_handles_string_values(): payload = {"civitai": {"modelId": "54321"}} assert service._extract_model_id_from_payload(payload) == 54321 + + +# ============================================================================= +# Tests for delete_model undo staging +# ============================================================================= + + +@pytest.fixture(autouse=True) +def _reset_pending_delete_singleton() -> Iterator[None]: + """Reset the pending-delete singleton around every test in this module. + + The singleton keeps an in-process list of staging roots across tests; + resetting avoids cross-test pollution (a stale root from one tmp_path + leaking into the next test's opportunistic purge enumeration). + """ + _reset_pending_delete_service() + yield + _reset_pending_delete_service() + + +@pytest.fixture(autouse=True) +def _stub_scanner_registry_getters(monkeypatch: pytest.MonkeyPatch) -> None: + """Prevent purge enumeration from instantiating real scanner singletons. + + ``stage_model_delete`` triggers an opportunistic purge whose root + enumeration queries the ServiceRegistry scanner getters; stubbing them to + ``None`` keeps tests fast and isolated (mirrors test_pending_delete_service). + """ + from py.services.service_registry import ServiceRegistry + + async def _none(*_args: Any, **_kwargs: Any) -> None: + return None + + monkeypatch.setattr(ServiceRegistry, "get_lora_scanner", _none) + monkeypatch.setattr(ServiceRegistry, "get_checkpoint_scanner", _none) + monkeypatch.setattr(ServiceRegistry, "get_embedding_scanner", _none) + + +def _make_delete_service(scanner: Any) -> ModelLifecycleService: + return ModelLifecycleService( + scanner=scanner, + metadata_manager=DummyMetadataManager({"civitai": {"modelId": 1}}), + metadata_loader=_empty_metadata_loader, + ) + + +@pytest.mark.asyncio +async def test_delete_model_stages_file_when_undo_enabled(tmp_path: Path): + """Undo enabled (the default): artifacts are renamed into a + ``.lm-pending-delete//`` staging dir under the model root, the + response carries the batch_id, the cache entry is removed and the cache + is persisted (``_persist_calls`` tracked by ``ScannerForDelete``).""" + root = tmp_path / "loras" + root.mkdir() + model_path = root / "model.safetensors" + model_path.write_bytes(b"content") + metadata_path = root / "model.metadata.json" + metadata_path.write_text(json.dumps({})) + preview_path = root / "model.preview.png" + preview_path.write_bytes(b"preview") + + scanner = ScannerForDelete( + raw_data=[ + { + "file_path": str(model_path), + "civitai": {"modelId": 1, "id": 10}, + "sha256": "abc123", + } + ], + roots=[str(root)], + ) + service = _make_delete_service(scanner) + + result = await service.delete_model(str(model_path)) + + assert result["success"] is True + batch_id = result["batch_id"] + assert isinstance(batch_id, str) + + assert not model_path.exists() + assert not metadata_path.exists() + assert not preview_path.exists() + batch_dir = root / PENDING_DELETE_DIR_NAME / batch_id + assert batch_dir.is_dir() + assert (batch_dir / "model.safetensors").read_bytes() == b"content" + assert (batch_dir / "model.metadata.json").exists() + assert (batch_dir / "model.preview.png").exists() + + assert scanner.cache.raw_data == [] + assert scanner._hash_index.removed == [str(model_path)] + assert scanner._persist_calls == [True] + + +@pytest.mark.asyncio +async def test_delete_model_hard_deletes_when_undo_disabled( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + """delete_undo_enabled=false: old os.remove behavior, batch_id is None and + no staging directory is ever created.""" + root = tmp_path / "loras" + root.mkdir() + model_path = root / "model.safetensors" + model_path.write_bytes(b"content") + + settings_manager = get_settings_manager() + monkeypatch.setattr( + settings_manager, + "get", + lambda key, default=None: False + if key == "delete_undo_enabled" + else default, + ) + + scanner = ScannerForDelete( + raw_data=[{"file_path": str(model_path)}], + roots=[str(root)], + ) + service = _make_delete_service(scanner) + + result = await service.delete_model(str(model_path)) + + assert result["success"] is True + assert result["batch_id"] is None + assert result["deleted_files"] + assert not model_path.exists() + assert not (root / PENDING_DELETE_DIR_NAME).exists() + + +@pytest.mark.asyncio +async def test_delete_model_falls_back_when_staging_fails( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + """Staging rename failure: delete_model_artifacts fallback removes the + files, batch_id is None and no staged data is left behind.""" + root = tmp_path / "loras" + root.mkdir() + model_path = root / "model.safetensors" + model_path.write_bytes(b"content") + + real_rename = os.rename + + def _failing_rename(src: str, dst: str) -> None: + if PENDING_DELETE_DIR_NAME in dst: + raise OSError("simulated staging failure") + real_rename(src, dst) + + monkeypatch.setattr(os, "rename", _failing_rename) + + scanner = ScannerForDelete( + raw_data=[{"file_path": str(model_path)}], + roots=[str(root)], + ) + service = _make_delete_service(scanner) + + result = await service.delete_model(str(model_path)) + + assert result["success"] is True + assert result["batch_id"] is None + assert result["deleted_files"] + assert not model_path.exists() + # No staged batch files remain (the batch dir is rolled back; an empty + # staging parent, if left behind by the rollback, holds no data). + staging_parent = root / PENDING_DELETE_DIR_NAME + assert not staging_parent.exists() or not any(staging_parent.iterdir()) diff --git a/tests/services/test_model_scanner.py b/tests/services/test_model_scanner.py index ee27c004..4fdd71ec 100644 --- a/tests/services/test_model_scanner.py +++ b/tests/services/test_model_scanner.py @@ -1,6 +1,11 @@ +from __future__ import annotations + import asyncio +import json import os import sqlite3 +import time +from collections.abc import Iterator from pathlib import Path from typing import Any, Dict, List, Optional from types import MethodType @@ -11,7 +16,13 @@ from py.services import model_scanner from py.services.model_cache import ModelCache from py.services.model_hash_index import ModelHashIndex from py.services.model_scanner import CacheBuildResult, ModelScanner +from py.services.pending_delete_service import ( + PENDING_DELETE_DIR_NAME, + PENDING_DELETE_TTL_SECONDS, + _reset_pending_delete_service, +) from py.services.persistent_model_cache import PersistentModelCache, DEFAULT_LICENSE_FLAGS +from py.services.settings_manager import get_settings_manager from py.utils.civitai_utils import build_license_flags from py.utils.models import BaseModelMetadata @@ -104,6 +115,27 @@ def stub_register_service(monkeypatch): monkeypatch.setattr(model_scanner.ServiceRegistry, "register_service", noop) +@pytest.fixture(autouse=True) +def _reset_pending_delete_singleton() -> Iterator[None]: + """Reset the pending-delete singleton before and after each test.""" + _reset_pending_delete_service() + yield + _reset_pending_delete_service() + + +@pytest.fixture(autouse=True) +def _stub_service_registry_getters(monkeypatch) -> None: + """Prevent pending-delete purge enumeration from building real scanners.""" + from py.services.service_registry import ServiceRegistry + + async def _none(*_args, **_kwargs) -> None: + return None + + monkeypatch.setattr(ServiceRegistry, "get_lora_scanner", _none) + monkeypatch.setattr(ServiceRegistry, "get_checkpoint_scanner", _none) + monkeypatch.setattr(ServiceRegistry, "get_embedding_scanner", _none) + + def _create_files(root: Path) -> tuple[Path, Path, Path]: first = root / "one.txt" first.write_text("one", encoding="utf-8") @@ -974,3 +1006,235 @@ async def test_sync_cache_conditional_resort_triggered(tmp_path: Path, monkeypat ) assert changed is True assert resort_calls == 1 + + +# ── bulk_delete_models staging (undo-delete feature, todo 3) ─────────────── + + +def _make_bulk_scanner(root: Path, file_paths: List[Path]) -> DummyScanner: + """Build a DummyScanner whose cache mirrors the given files on disk.""" + scanner = DummyScanner(root) + raw_data = [] + for path in file_paths: + name = os.path.splitext(os.path.basename(path))[0] + raw_data.append( + { + "file_path": str(path), + "folder": "", + "sha256": f"hash-{name}", + "tags": ["alpha"] if "one" in name else ["beta"], + "model_name": name, + "file_name": name, + "size": 1, + "modified": 1.0, + } + ) + scanner._cache = ModelCache( + raw_data=raw_data, folders=[], name_display_mode="model_name" + ) + scanner._tags_count = {"alpha": 1, "beta": 1} + for entry in raw_data: + scanner._hash_index.add_entry(entry["sha256"], entry["file_path"]) + return scanner + + +@pytest.mark.asyncio +async def test_bulk_delete_stages_two_files_into_single_batch(tmp_path: Path): + """Two-file bulk delete -> one merged batch id with both files staged.""" + root = tmp_path / "loras" + root.mkdir() + first = root / "one.txt" + first.write_text("one", encoding="utf-8") + second = root / "two.txt" + second.write_text("two", encoding="utf-8") + scanner = _make_bulk_scanner(root, [first, second]) + + result = await scanner.bulk_delete_models([str(first), str(second)]) + + assert result["success"] is True + assert result["status"] == "success" + assert result["total_deleted"] == 2 + assert result["cache_updated"] is True + + # ONE batch id, no batch_ids array, and both files staged in its dir. + assert "batch_id" in result + assert "batch_ids" not in result + batch_id = result["batch_id"] + assert batch_id is not None + staging = root / PENDING_DELETE_DIR_NAME + batch_dir = staging / batch_id + assert batch_dir.is_dir() + assert (batch_dir / "one.txt").read_bytes() == b"one" + assert (batch_dir / "two.txt").read_bytes() == b"two" + + # Loser batch dirs are removed by the merge - exactly one batch remains. + batch_dirs = [d.name for d in staging.iterdir() if d.is_dir()] + assert batch_dirs == [batch_id] + + # The manifest carries the winner's cache snapshot for later undo. + manifest = json.loads((batch_dir / "manifest.json").read_text(encoding="utf-8")) + assert manifest["model_snapshot"]["file_path"] == str(first) + + # Originals gone; cache entries removed. + assert not first.exists() + assert not second.exists() + cached_paths = {item["file_path"] for item in scanner._cache.raw_data} + assert str(first) not in cached_paths + assert str(second) not in cached_paths + + +@pytest.mark.asyncio +async def test_bulk_delete_merged_manifest_reanchors_expiry(tmp_path: Path): + """Merged manifest expires_at is re-anchored to now+TTL at merge time.""" + root = tmp_path / "loras" + root.mkdir() + first = root / "one.txt" + first.write_text("one", encoding="utf-8") + second = root / "two.txt" + second.write_text("two", encoding="utf-8") + scanner = _make_bulk_scanner(root, [first, second]) + + before = int(time.time()) + result = await scanner.bulk_delete_models([str(first), str(second)]) + after = int(time.time()) + + batch_dir = root / PENDING_DELETE_DIR_NAME / result["batch_id"] + manifest = json.loads((batch_dir / "manifest.json").read_text(encoding="utf-8")) + # expires_at >= staging completion time + TTL (re-anchor assertion). + assert manifest["expires_at"] >= after + PENDING_DELETE_TTL_SECONDS - 2 + assert manifest["expires_at"] >= before + PENDING_DELETE_TTL_SECONDS + # Both files are entries of the merged manifest. + assert len(manifest["entries"]) == 2 + assert (batch_dir / "one.txt").exists() + assert (batch_dir / "two.txt").exists() + + +@pytest.mark.asyncio +async def test_bulk_delete_merge_failure_falls_back_to_batch_ids( + tmp_path: Path, monkeypatch +): + """Merge move failure -> batch_ids array of the intact constituent batches.""" + root = tmp_path / "loras" + root.mkdir() + first = root / "one.txt" + first.write_text("one", encoding="utf-8") + second = root / "two.txt" + second.write_text("two", encoding="utf-8") + scanner = _make_bulk_scanner(root, [first, second]) + + real_rename = os.rename + fail_next = {"enabled": True} + + def flaky_merge_rename(src: str, dst: str) -> None: + # Fail only when moving between batch dirs (merge), never during + # staging (src is then the original path, outside .lm-pending-delete). + if ( + fail_next["enabled"] + and PENDING_DELETE_DIR_NAME in src + and PENDING_DELETE_DIR_NAME in dst + ): + fail_next["enabled"] = False + raise OSError("simulated merge failure") + return real_rename(src, dst) + + monkeypatch.setattr( + "py.services.pending_delete_service.os.rename", flaky_merge_rename + ) + + result = await scanner.bulk_delete_models([str(first), str(second)]) + + assert result["success"] is True + assert result["total_deleted"] == 2 + # No single batch id - the constituent ids are returned instead. + assert "batch_id" not in result + assert "batch_ids" in result + assert len(result["batch_ids"]) == 2 + + # Both constituent batches are intact: dirs + manifests + staged files. + staging = root / PENDING_DELETE_DIR_NAME + batch_dirs = sorted(d.name for d in staging.iterdir() if d.is_dir()) + assert sorted(result["batch_ids"]) == batch_dirs + for bid in result["batch_ids"]: + batch_dir = staging / bid + assert (batch_dir / "manifest.json").exists() + staged_files = [ + f.name + for bid in result["batch_ids"] + for f in (staging / bid).iterdir() + if f.is_file() and f.name != "manifest.json" + ] + assert sorted(staged_files) == ["one.txt", "two.txt"] + + +@pytest.mark.asyncio +async def test_bulk_delete_undo_disabled_hard_deletes(tmp_path: Path): + """delete_undo_enabled=false -> old hard delete, no batch, no staging dirs.""" + root = tmp_path / "loras" + root.mkdir() + first = root / "one.txt" + first.write_text("one", encoding="utf-8") + second = root / "two.txt" + second.write_text("two", encoding="utf-8") + scanner = _make_bulk_scanner(root, [first, second]) + + get_settings_manager().settings["delete_undo_enabled"] = False + + result = await scanner.bulk_delete_models([str(first), str(second)]) + + assert result["success"] is True + assert result["status"] == "success" + assert result["total_deleted"] == 2 + assert result.get("batch_id") is None + assert "batch_ids" not in result + + # Old hard-delete behavior: files removed, zero staging dirs created. + assert not first.exists() + assert not second.exists() + assert not (root / PENDING_DELETE_DIR_NAME).exists() + + +@pytest.mark.asyncio +async def test_bulk_delete_cancelled_after_one_staged_batch_present( + tmp_path: Path, monkeypatch +): + """Cancelled mid-way -> status='cancelled' AND the staged subset undoable.""" + root = tmp_path / "loras" + root.mkdir() + first = root / "one.txt" + first.write_text("one", encoding="utf-8") + second = root / "two.txt" + second.write_text("two", encoding="utf-8") + scanner = _make_bulk_scanner(root, [first, second]) + + real_rename = os.rename + rename_count = {"n": 0} + + def cancelling_rename(src: str, dst: str) -> None: + rename_count["n"] += 1 + result = real_rename(src, dst) + # After the first file is staged, request cancellation so the loop + # stops before the second file is processed. + if rename_count["n"] == 1: + scanner.cancel_task() + return result + + monkeypatch.setattr( + "py.services.pending_delete_service.os.rename", cancelling_rename + ) + + result = await scanner.bulk_delete_models([str(first), str(second)]) + + assert result["success"] is True + assert result["status"] == "cancelled" + assert result["total_deleted"] == 1 + assert "batch_id" in result + assert result["batch_id"] is not None + assert "batch_ids" not in result + + # The staged subset is merged into one undoable batch. + batch_dir = root / PENDING_DELETE_DIR_NAME / result["batch_id"] + assert batch_dir.is_dir() + assert (batch_dir / "one.txt").read_bytes() == b"one" + assert not first.exists() + # The second file was never touched. + assert second.exists() diff --git a/tests/services/test_pending_delete_service.py b/tests/services/test_pending_delete_service.py new file mode 100644 index 00000000..8cacca3d --- /dev/null +++ b/tests/services/test_pending_delete_service.py @@ -0,0 +1,1535 @@ +"""Tests for :mod:`py.services.pending_delete_service`. + +Covers the staging service contract: stage (model + recipe), undo (with +partial-undo retry and occupied-path protection), merge (with rollback and a +fresh purge timer), purge (expired-only, quarantine of malformed batches, +per-file lock tolerance) and the scanner exclusion of the staging directory. + +Deterministic time control: no real sleeps - tests rewrite ``expires_at`` in +the manifest or monkeypatch time functions instead. +""" + +from __future__ import annotations + +import asyncio +import errno +import json +import os +import time +from collections.abc import Iterator +from pathlib import Path +from typing import Any, Dict, List, Optional, Sequence + +import pytest + +from py.services.pending_delete_service import ( + PENDING_DELETE_DIR_NAME, + PENDING_DELETE_TTL_SECONDS, + PendingDeleteService, + _reset_pending_delete_service, +) +from py.services.model_hash_index import ModelHashIndex +from py.services.model_scanner import ModelScanner +from py.services.settings_manager import DEFAULT_SETTINGS, get_settings_manager +from py.utils import settings_paths +from py.utils.models import LoraMetadata + + +class ScannerForStage: + """Scanner double exposing the attributes the staging service uses.""" + + def __init__(self, roots: Sequence[Path], model_type: str = "lora") -> None: + self._roots: List[str] = [os.path.abspath(str(r)) for r in roots] + self.model_type = model_type + + def get_model_roots(self) -> List[str]: + return list(self._roots) + + def _find_root_for_file(self, file_path: Optional[str]) -> Optional[str]: + if not file_path: + return None + normalized = os.path.abspath(os.path.normpath(file_path)) + for root in self._roots: + if normalized == root or normalized.startswith(root + os.sep): + return root + return None + + +class CheckpointScannerStub: + """Minimal double for usage-tracking lookups.""" + + def __init__(self, root: Path) -> None: + self._root = str(root) + self.file_extensions = {".safetensors", ".ckpt", ".pt", ".gguf"} + + def get_model_roots(self) -> List[str]: + return [self._root] + + +@pytest.fixture(autouse=True) +def _reset_service_singleton() -> Iterator[None]: + """Reset the pending-delete singleton before and after each test.""" + _reset_pending_delete_service() + yield + _reset_pending_delete_service() + + +@pytest.fixture(autouse=True) +def _stub_scanner_registry(monkeypatch) -> None: + """Prevent purge enumeration from instantiating real scanner singletons.""" + from py.services.service_registry import ServiceRegistry + + async def _none(*_args: Any, **_kwargs: Any) -> None: + return None + + monkeypatch.setattr(ServiceRegistry, "get_lora_scanner", _none) + monkeypatch.setattr(ServiceRegistry, "get_checkpoint_scanner", _none) + monkeypatch.setattr(ServiceRegistry, "get_embedding_scanner", _none) + + +async def _register_model_root( + monkeypatch: pytest.MonkeyPatch, + *, + lora_roots: Sequence[Path] = (), + checkpoint_roots: Sequence[Path] = (), + embedding_roots: Sequence[Path] = (), +) -> None: + """Point the ServiceRegistry scanner getters at tmp-root fakes.""" + from py.services.service_registry import ServiceRegistry + + def _make(roots: Sequence[Path]): + async def _getter(*_args: Any, **_kwargs: Any) -> Optional[ScannerForStage]: + return ScannerForStage(roots) if roots else None + + return _getter + + monkeypatch.setattr(ServiceRegistry, "get_lora_scanner", _make(lora_roots)) + monkeypatch.setattr(ServiceRegistry, "get_checkpoint_scanner", _make(checkpoint_roots)) + monkeypatch.setattr(ServiceRegistry, "get_embedding_scanner", _make(embedding_roots)) + + +def _write_batch_manifest( + batch_dir: Path, + *, + batch_id: str, + kind: str, + expires_at: int, + entries: Sequence[Dict[str, Any]], + model_type: Optional[str] = None, + state: str = "staged", + model_snapshot: Any = None, + recipe_snapshot: Any = None, +) -> None: + """Write a manifest.json with the shape the service reads.""" + manifest: Dict[str, Any] = { + "batch_id": batch_id, + "kind": kind, + "model_type": model_type if kind == "model" else None, + "state": state, + "expires_at": int(expires_at), + "entries": list(entries), + "model_snapshot": model_snapshot if kind == "model" else None, + "recipe_snapshot": recipe_snapshot if kind == "recipe" else None, + } + (batch_dir / "manifest.json").write_text(json.dumps(manifest)) + + +def _spy_purge_timers(monkeypatch: pytest.MonkeyPatch) -> List[Optional[str]]: + """Replace the timer task factory with a recorder (no real tasks).""" + created: List[Optional[str]] = [] + + class DummyTask: + def add_done_callback(self, _cb: Any) -> None: # pragma: no cover - stub + pass + + def cancel(self) -> None: # pragma: no cover - stub + pass + + def done(self) -> bool: # pragma: no cover - stub + return False + + def fake_create_task(coro: Any, *args: Any, **kwargs: Any) -> DummyTask: + created.append(kwargs.get("name")) + coro.close() # never awaited - avoid a "coroutine was never awaited" warning + return DummyTask() + + monkeypatch.setattr("py.services.pending_delete_service._create_task", fake_create_task) + return created + + +async def _stage_simple( + service: PendingDeleteService, + root: Path, + file_name: str, + *, + model_type: str = "lora", + cached_entry: Any = None, +) -> str: + """Stage a single-artifact model delete and return its batch id.""" + model = root / f"{file_name}.safetensors" + model.write_bytes(f"{file_name}-data".encode()) + batch_id = await service.stage_model_delete( + scanner=ScannerForStage([root], model_type=model_type), + target_dir=str(root), + file_name=file_name, + main_extension=".safetensors", + original_file_path=str(model), + cached_entry=cached_entry, + ) + assert batch_id is not None + return batch_id + + +# --------------------------------------------------------------------------- +# (a) stage_model_delete renames all existing artifact patterns + manifest +# --------------------------------------------------------------------------- +async def test_a_stage_model_renames_artifacts_and_writes_manifest(tmp_path: Path) -> None: + root = tmp_path / "loras" + root.mkdir() + model = root / "model.safetensors" + model.write_bytes(b"model-bytes") + metadata = root / "model.metadata.json" + metadata.write_bytes(b'{"key": "value"}') + preview = root / "model.preview.webp" + preview.write_bytes(b"preview-bytes") + + service = await PendingDeleteService.get_instance() + cached_entry = {"file_path": str(model), "sha256": "abc", "tags": ["a"]} + before = int(time.time()) + + batch_id = await service.stage_model_delete( + scanner=ScannerForStage([root]), + target_dir=str(root), + file_name="model", + main_extension=".safetensors", + original_file_path=str(model), + cached_entry=cached_entry, + ) + + assert batch_id is not None + batch_dir = root / PENDING_DELETE_DIR_NAME / batch_id + assert batch_dir.is_dir() + + # Originals are renamed away. + assert not model.exists() + assert not metadata.exists() + assert not preview.exists() + + # Every existing artifact is staged with identical bytes. + assert (batch_dir / "model.safetensors").read_bytes() == b"model-bytes" + assert (batch_dir / "model.metadata.json").read_bytes() == b'{"key": "value"}' + assert (batch_dir / "model.preview.webp").read_bytes() == b"preview-bytes" + + manifest = json.loads((batch_dir / "manifest.json").read_text(encoding="utf-8")) + assert manifest["batch_id"] == batch_id + assert manifest["kind"] == "model" + assert manifest["model_type"] == "loras" + assert manifest["state"] == "staged" + assert manifest["model_snapshot"] == cached_entry + assert manifest["recipe_snapshot"] is None + assert before + PENDING_DELETE_TTL_SECONDS - 2 <= manifest["expires_at"] <= before + PENDING_DELETE_TTL_SECONDS + 2 + + assert len(manifest["entries"]) == 3 + originals = {entry["original"] for entry in manifest["entries"]} + assert originals == {str(model), str(metadata), str(preview)} + for entry in manifest["entries"]: + assert os.path.isabs(entry["staged"]) + assert os.path.isabs(entry["original"]) + assert entry["restored"] is False + + +# --------------------------------------------------------------------------- +# (b) undo() restores all files to original paths and removes batch dir +# --------------------------------------------------------------------------- +async def test_b_undo_restores_files_and_removes_batch_dir(tmp_path: Path) -> None: + root = tmp_path / "loras" + root.mkdir() + model = root / "model.safetensors" + model.write_bytes(b"model-bytes") + metadata = root / "model.metadata.json" + metadata.write_bytes(b"meta-bytes") + + service = await PendingDeleteService.get_instance() + batch_id = await service.stage_model_delete( + scanner=ScannerForStage([root]), + target_dir=str(root), + file_name="model", + main_extension=".safetensors", + original_file_path=str(model), + cached_entry={"file_path": str(model)}, + ) + assert batch_id is not None + batch_dir = root / PENDING_DELETE_DIR_NAME / batch_id + assert batch_dir.is_dir() + + result = await service.undo(batch_id) + + assert result["batch_id"] == batch_id + assert model.read_bytes() == b"model-bytes" + assert metadata.read_bytes() == b"meta-bytes" + assert not batch_dir.exists() + + +# --------------------------------------------------------------------------- +# (c) undo() on expired batch raises ValueError +# --------------------------------------------------------------------------- +async def test_c_undo_expired_raises_value_error(tmp_path: Path, monkeypatch) -> None: + root = tmp_path / "loras" + root.mkdir() + service = await PendingDeleteService.get_instance() + + # Isolate the expiry check: keep the opportunistic purge from consuming it. + async def _no_purge() -> int: + return 0 + + monkeypatch.setattr(service, "purge_expired", _no_purge) + + batch_id = await _stage_simple(service, root, "model") + batch_dir = root / PENDING_DELETE_DIR_NAME / batch_id + manifest_path = batch_dir / "manifest.json" + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + manifest["expires_at"] = int(time.time()) - 10 + manifest_path.write_text(json.dumps(manifest)) + + with pytest.raises(ValueError, match="expired"): + await service.undo(batch_id) + + +# --------------------------------------------------------------------------- +# (d) undo() when original path occupied raises ValueError, batch stays intact +# --------------------------------------------------------------------------- +async def test_d_undo_occupied_path_raises_and_leaves_batch_intact(tmp_path: Path) -> None: + root = tmp_path / "loras" + root.mkdir() + model = root / "model.safetensors" + model.write_bytes(b"original") + metadata = root / "model.metadata.json" + metadata.write_bytes(b"meta") + + service = await PendingDeleteService.get_instance() + batch_id = await service.stage_model_delete( + scanner=ScannerForStage([root]), + target_dir=str(root), + file_name="model", + main_extension=".safetensors", + original_file_path=str(model), + cached_entry=None, + ) + assert batch_id is not None + batch_dir = root / PENDING_DELETE_DIR_NAME / batch_id + manifest_before = (batch_dir / "manifest.json").read_bytes() + + # Simulate a re-download occupying the original path. + model.write_bytes(b"new-file") + + with pytest.raises(ValueError, match="occupied"): + await service.undo(batch_id) + + # Batch dir + manifest untouched, staged file still present, new file safe. + assert (batch_dir / "manifest.json").read_bytes() == manifest_before + assert (batch_dir / "model.safetensors").exists() + assert model.read_bytes() == b"new-file" + + +# --------------------------------------------------------------------------- +# (e) PARTIAL-UNDO RETRY +# --------------------------------------------------------------------------- +async def test_e_partial_undo_retry_completes_on_second_attempt(tmp_path: Path, monkeypatch) -> None: + root = tmp_path / "loras" + root.mkdir() + for name in ("model.safetensors", "model.metadata.json", "model.preview.webp"): + (root / name).write_bytes(name.encode()) + + service = await PendingDeleteService.get_instance() + batch_id = await service.stage_model_delete( + scanner=ScannerForStage([root]), + target_dir=str(root), + file_name="model", + main_extension=".safetensors", + original_file_path=str(root / "model.safetensors"), + cached_entry=None, + ) + assert batch_id is not None + batch_dir = root / PENDING_DELETE_DIR_NAME / batch_id + + real_rename = os.rename + calls = {"n": 0} + fail_next = {"enabled": True} + + def flaky_rename(src: str, dst: str) -> None: + calls["n"] += 1 + if fail_next["enabled"] and calls["n"] == 2: + raise OSError("simulated locked file") + return real_rename(src, dst) + + monkeypatch.setattr("py.services.pending_delete_service.os.rename", flaky_rename) + + with pytest.raises(OSError): + await service.undo(batch_id) + + manifest = json.loads((batch_dir / "manifest.json").read_text(encoding="utf-8")) + assert [entry["restored"] for entry in manifest["entries"]] == [True, False, False] + assert (root / "model.safetensors").exists() + + # Second undo (rename no longer failing) completes the remainder. + fail_next["enabled"] = False + await service.undo(batch_id) + + assert (root / "model.metadata.json").read_bytes() == b"model.metadata.json" + assert (root / "model.preview.webp").read_bytes() == b"model.preview.webp" + assert not batch_dir.exists() + + +# --------------------------------------------------------------------------- +# (e2) UNDO SKIPS A STAGED FILE THAT IS ALREADY GONE and finishes the rest +# --------------------------------------------------------------------------- +async def test_e2_undo_skips_missing_staged_file(tmp_path: Path) -> None: + root = tmp_path / "loras" + root.mkdir() + for name in ("model.safetensors", "model.metadata.json", "model.preview.webp"): + (root / name).write_bytes(name.encode()) + + service = await PendingDeleteService.get_instance() + batch_id = await service.stage_model_delete( + scanner=ScannerForStage([root]), + target_dir=str(root), + file_name="model", + main_extension=".safetensors", + original_file_path=str(root / "model.safetensors"), + cached_entry=None, + ) + assert batch_id is not None + batch_dir = root / PENDING_DELETE_DIR_NAME / batch_id + + # Simulate one staged artifact being removed out-of-band (e.g. an earlier + # purge/manual cleanup) before undo runs. + (batch_dir / "model.metadata.json").unlink() + + await service.undo(batch_id) + + assert (root / "model.safetensors").read_bytes() == b"model.safetensors" + assert (root / "model.preview.webp").read_bytes() == b"model.preview.webp" + assert not (root / "model.metadata.json").exists() + assert not batch_dir.exists() + + +# --------------------------------------------------------------------------- +# (f) purge_expired() removes only expired batches +# --------------------------------------------------------------------------- +async def test_f_purge_expired_removes_only_expired(tmp_path: Path, monkeypatch) -> None: + root = tmp_path / "loras" + root.mkdir() + await _register_model_root(monkeypatch, lora_roots=[root]) + + service = await PendingDeleteService.get_instance() + expired_id = await _stage_simple(service, root, "expired") + fresh_id = await _stage_simple(service, root, "fresh") + + expired_manifest_path = root / PENDING_DELETE_DIR_NAME / expired_id / "manifest.json" + expired_manifest = json.loads(expired_manifest_path.read_text(encoding="utf-8")) + expired_manifest["expires_at"] = int(time.time()) - 10 + expired_manifest_path.write_text(json.dumps(expired_manifest)) + + await service.purge_expired() + + assert not (root / PENDING_DELETE_DIR_NAME / expired_id).exists() + assert (root / PENDING_DELETE_DIR_NAME / fresh_id).exists() + assert not (root / "expired.safetensors").exists() + assert not (root / "fresh.safetensors").exists() + + +# --------------------------------------------------------------------------- +# (g) MANIFEST-LESS dir -> quarantined, files kept +# --------------------------------------------------------------------------- +async def test_g_manifestless_dir_quarantined(tmp_path: Path, monkeypatch) -> None: + root = tmp_path / "loras" + root.mkdir() + staging = root / PENDING_DELETE_DIR_NAME + batch_dir = staging / "batch1" + batch_dir.mkdir(parents=True) + (batch_dir / "model.safetensors").write_bytes(b"user-data") + await _register_model_root(monkeypatch, lora_roots=[root]) + + service = await PendingDeleteService.get_instance() + await service.purge_expired() + + orphaned = staging / "batch1.orphaned" + assert orphaned.is_dir() + assert not batch_dir.exists() + assert (orphaned / "model.safetensors").read_bytes() == b"user-data" + + +# --------------------------------------------------------------------------- +# (h) CORRUPTED manifest -> quarantined, sweep completes +# --------------------------------------------------------------------------- +async def test_h_corrupted_manifest_quarantined(tmp_path: Path, monkeypatch) -> None: + root = tmp_path / "loras" + root.mkdir() + staging = root / PENDING_DELETE_DIR_NAME + batch_dir = staging / "batch2" + batch_dir.mkdir(parents=True) + (batch_dir / "model.safetensors").write_bytes(b"data") + (batch_dir / "manifest.json").write_text("{ not valid json !!!") + + await _register_model_root(monkeypatch, lora_roots=[root]) + service = await PendingDeleteService.get_instance() + + await service.purge_expired() # must not crash + + orphaned = staging / "batch2.orphaned" + assert orphaned.is_dir() + assert (orphaned / "model.safetensors").read_bytes() == b"data" + + +# --------------------------------------------------------------------------- +# (i) PURGE LOCKED FILE -> skip file, keep batch dir, no exception +# --------------------------------------------------------------------------- +async def test_i_purge_locked_file_skips_and_keeps_batch_dir(tmp_path: Path, monkeypatch) -> None: + root = tmp_path / "loras" + root.mkdir() + model = root / "model.safetensors" + model.write_bytes(b"data") + + service = await PendingDeleteService.get_instance() + batch_id = await _stage_simple(service, root, "model") + batch_dir = root / PENDING_DELETE_DIR_NAME / batch_id + manifest_path = batch_dir / "manifest.json" + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + manifest["expires_at"] = int(time.time()) - 10 + manifest_path.write_text(json.dumps(manifest)) + + await _register_model_root(monkeypatch, lora_roots=[root]) + + real_remove = os.remove + + def flaky_remove(path: str, *args: Any, **kwargs: Any) -> None: + if str(path).endswith("model.safetensors"): + raise OSError("simulated locked file") + return real_remove(path, *args, **kwargs) + + monkeypatch.setattr("py.services.pending_delete_service.os.remove", flaky_remove) + + await service.purge_expired() # no exception + + assert batch_dir.is_dir() + assert (batch_dir / "model.safetensors").exists() + assert (batch_dir / "manifest.json").exists() + + +# --------------------------------------------------------------------------- +# (j) STALE TIMER -> purge_batch on missing/undone ids is a silent no-op +# --------------------------------------------------------------------------- +async def test_j_stale_timer_purge_batch_noop(tmp_path: Path, monkeypatch) -> None: + await _register_model_root(monkeypatch, lora_roots=[tmp_path / "nonexistent"]) + + service = await PendingDeleteService.get_instance() + await service.purge_batch("does-not-exist") # silent no-op + + root = tmp_path / "loras" + root.mkdir() + batch_id = await _stage_simple(service, root, "model") + await service.undo(batch_id) + + await service.purge_batch(batch_id) # undone -> silent no-op + assert (root / "model.safetensors").read_bytes() == b"model-data" + + +# --------------------------------------------------------------------------- +# (k) MERGE -> single manifest, re-anchored expiry, all files under winner +# --------------------------------------------------------------------------- +async def test_k_merge_produces_single_manifest_and_moves_all_files(tmp_path: Path, monkeypatch) -> None: + root = tmp_path / "loras" + root.mkdir() + _spy_purge_timers(monkeypatch) + + service = await PendingDeleteService.get_instance() + a1 = root / "alpha.safetensors" + a1.write_bytes(b"alpha-data") + a2 = root / "alpha.metadata.json" + a2.write_bytes(b"alpha-meta") + bid_a = await service.stage_model_delete( + scanner=ScannerForStage([root]), + target_dir=str(root), + file_name="alpha", + main_extension=".safetensors", + original_file_path=str(a1), + cached_entry={"file_path": str(a1)}, + ) + assert bid_a is not None + b1 = root / "beta.safetensors" + b1.write_bytes(b"beta-data") + bid_b = await service.stage_model_delete( + scanner=ScannerForStage([root]), + target_dir=str(root), + file_name="beta", + main_extension=".safetensors", + original_file_path=str(b1), + cached_entry=None, + ) + assert bid_b is not None + + before_merge = int(time.time()) + merged = await service.merge_batches([bid_a, bid_b]) + assert merged == bid_a + + winner_dir = root / PENDING_DELETE_DIR_NAME / bid_a + loser_dir = root / PENDING_DELETE_DIR_NAME / bid_b + manifest = json.loads((winner_dir / "manifest.json").read_text(encoding="utf-8")) + assert manifest["batch_id"] == bid_a + assert len(manifest["entries"]) == 3 + assert before_merge + PENDING_DELETE_TTL_SECONDS - 2 <= manifest["expires_at"] <= before_merge + PENDING_DELETE_TTL_SECONDS + 2 + + staged_paths = [entry["staged"] for entry in manifest["entries"]] + assert len(staged_paths) == 3 + for staged in staged_paths: + assert str(staged).startswith(str(winner_dir)) + assert os.path.exists(staged) + + # Byte-compare: no file dropped. + assert (winner_dir / "alpha.safetensors").read_bytes() == b"alpha-data" + assert (winner_dir / "alpha.metadata.json").read_bytes() == b"alpha-meta" + assert (winner_dir / "beta.safetensors").read_bytes() == b"beta-data" + + # Loser batch dir removed (after being empty). + assert not loser_dir.exists() + + +# --------------------------------------------------------------------------- +# (k2) MERGE THEN UNDO -> every file restored to its ORIGINAL path +# --------------------------------------------------------------------------- +async def test_k2_merge_then_undo_restores_every_file(tmp_path: Path, monkeypatch) -> None: + root = tmp_path / "loras" + root.mkdir() + _spy_purge_timers(monkeypatch) + + service = await PendingDeleteService.get_instance() + a1 = root / "alpha.safetensors" + a1.write_bytes(b"alpha-data") + a2 = root / "alpha.metadata.json" + a2.write_bytes(b"alpha-meta") + bid_a = await service.stage_model_delete( + scanner=ScannerForStage([root]), + target_dir=str(root), + file_name="alpha", + main_extension=".safetensors", + original_file_path=str(a1), + cached_entry=None, + ) + assert bid_a is not None + b1 = root / "beta.safetensors" + b1.write_bytes(b"beta-data") + bid_b = await service.stage_model_delete( + scanner=ScannerForStage([root]), + target_dir=str(root), + file_name="beta", + main_extension=".safetensors", + original_file_path=str(b1), + cached_entry=None, + ) + assert bid_b is not None + + assert await service.merge_batches([bid_a, bid_b]) == bid_a + await service.undo(bid_a) + + assert a1.read_bytes() == b"alpha-data" + assert a2.read_bytes() == b"alpha-meta" + assert b1.read_bytes() == b"beta-data" + assert not (root / PENDING_DELETE_DIR_NAME / bid_a).exists() + assert not (root / PENDING_DELETE_DIR_NAME / bid_b).exists() + + +# --------------------------------------------------------------------------- +# (k3) MERGE THEN PURGE -> merged batch fully purged +# --------------------------------------------------------------------------- +async def test_k3_merge_then_purge_empties_and_removes_winner_dir(tmp_path: Path, monkeypatch) -> None: + root = tmp_path / "loras" + root.mkdir() + _spy_purge_timers(monkeypatch) + + service = await PendingDeleteService.get_instance() + a1 = root / "alpha.safetensors" + a1.write_bytes(b"alpha-data") + bid_a = await _stage_simple(service, root, "alpha", cached_entry=None) + bid_b = await _stage_simple(service, root, "beta") + assert await service.merge_batches([bid_a, bid_b]) == bid_a + + winner_dir = root / PENDING_DELETE_DIR_NAME / bid_a + manifest_path = winner_dir / "manifest.json" + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + manifest["expires_at"] = int(time.time()) - 10 + manifest_path.write_text(json.dumps(manifest)) + + await _register_model_root(monkeypatch, lora_roots=[root]) + await service.purge_expired() + + assert not winner_dir.exists() + assert not (root / "alpha.safetensors").exists() + assert not (root / "beta.safetensors").exists() + + +# --------------------------------------------------------------------------- +# (l) MERGE MOVE FAILURE -> rollback, all batches intact, sequential undo works +# --------------------------------------------------------------------------- +async def test_l_merge_move_failure_rolls_back(tmp_path: Path, monkeypatch) -> None: + root = tmp_path / "loras" + root.mkdir() + _spy_purge_timers(monkeypatch) + + service = await PendingDeleteService.get_instance() + a1 = root / "alpha.safetensors" + a1.write_bytes(b"alpha-data") + bid_a = await _stage_simple(service, root, "alpha") + # Loser has TWO files so a move fails after the first was already moved. + b1 = root / "beta.safetensors" + b1.write_bytes(b"beta-data") + b2 = root / "beta.metadata.json" + b2.write_bytes(b"beta-meta") + bid_b = await service.stage_model_delete( + scanner=ScannerForStage([root]), + target_dir=str(root), + file_name="beta", + main_extension=".safetensors", + original_file_path=str(b1), + cached_entry=None, + ) + assert bid_b is not None + + real_rename = os.rename + calls = {"n": 0} + fail_next = {"enabled": True} + + def flaky_rename(src: str, dst: str) -> None: + calls["n"] += 1 + if fail_next["enabled"] and calls["n"] == 2: + raise OSError("simulated move failure") + return real_rename(src, dst) + + monkeypatch.setattr("py.services.pending_delete_service.os.rename", flaky_rename) + + result = await service.merge_batches([bid_a, bid_b]) + assert result is None + + # Already-moved file moved back; both batch dirs + manifests + files intact. + winner_dir = root / PENDING_DELETE_DIR_NAME / bid_a + loser_dir = root / PENDING_DELETE_DIR_NAME / bid_b + assert winner_dir.is_dir() + assert loser_dir.is_dir() + assert (winner_dir / "manifest.json").exists() + assert (loser_dir / "manifest.json").exists() + assert (loser_dir / "beta.safetensors").read_bytes() == b"beta-data" + assert (loser_dir / "beta.metadata.json").read_bytes() == b"beta-meta" + assert (winner_dir / "alpha.safetensors").read_bytes() == b"alpha-data" + + # Sequential undo of each constituent batch restores every file. + fail_next["enabled"] = False + await service.undo(bid_a) + await service.undo(bid_b) + assert a1.read_bytes() == b"alpha-data" + assert b1.read_bytes() == b"beta-data" + assert b2.read_bytes() == b"beta-meta" + + +# --------------------------------------------------------------------------- +# (l2) MERGE SAME-BASENAME COLLISION -> abort + rollback, never overwrite +# --------------------------------------------------------------------------- +async def test_l2_merge_basename_collision_aborts_without_dropping_files( + tmp_path: Path, monkeypatch +) -> None: + root = tmp_path / "loras" + root.mkdir() + _spy_purge_timers(monkeypatch) + + service = await PendingDeleteService.get_instance() + sub_a = root / "a" + sub_a.mkdir() + sub_b = root / "b" + sub_b.mkdir() + # Two distinct files that share the same basename after staging. + bid_a = await _stage_simple(service, sub_a, "model") + bid_b = await _stage_simple(service, sub_b, "model") + + result = await service.merge_batches([bid_a, bid_b]) + assert result is None + + # No file dropped: both staged files exist in their own batch dirs. + a_dir = sub_a / PENDING_DELETE_DIR_NAME / bid_a + b_dir = sub_b / PENDING_DELETE_DIR_NAME / bid_b + assert (a_dir / "model.safetensors").read_bytes() == b"model-data" + assert (b_dir / "model.safetensors").read_bytes() == b"model-data" + assert a_dir.is_dir() and b_dir.is_dir() + assert (a_dir / "manifest.json").exists() + assert (b_dir / "manifest.json").exists() + + # Sequential undo of each constituent batch restores every original. + await service.undo(bid_a) + await service.undo(bid_b) + assert (sub_a / "model.safetensors").read_bytes() == b"model-data" + assert (sub_b / "model.safetensors").read_bytes() == b"model-data" + + +# --------------------------------------------------------------------------- +# (m) delete_undo_enabled=false -> stage returns None, nothing created +# --------------------------------------------------------------------------- +async def test_m_undo_disabled_returns_none(tmp_path: Path) -> None: + root = tmp_path / "loras" + root.mkdir() + model = root / "model.safetensors" + model.write_bytes(b"data") + + get_settings_manager().settings["delete_undo_enabled"] = False + + service = await PendingDeleteService.get_instance() + batch_id = await service.stage_model_delete( + scanner=ScannerForStage([root]), + target_dir=str(root), + file_name="model", + main_extension=".safetensors", + original_file_path=str(model), + cached_entry=None, + ) + + assert batch_id is None + assert model.exists() + assert not (root / PENDING_DELETE_DIR_NAME).exists() + + +# --------------------------------------------------------------------------- +# (n) simulated OSError during staging -> rollback, no orphaned batch dir +# --------------------------------------------------------------------------- +async def test_n_staging_oserror_rolls_back(tmp_path: Path, monkeypatch) -> None: + root = tmp_path / "loras" + root.mkdir() + a = root / "model.safetensors" + a.write_bytes(b"a-bytes") + b = root / "model.metadata.json" + b.write_bytes(b"b-bytes") + + service = await PendingDeleteService.get_instance() + + real_rename = os.rename + calls = {"n": 0} + + def flaky_rename(src: str, dst: str) -> None: + calls["n"] += 1 + if calls["n"] == 2: + raise OSError("simulated staging failure") + return real_rename(src, dst) + + monkeypatch.setattr("py.services.pending_delete_service.os.rename", flaky_rename) + + batch_id = await service.stage_model_delete( + scanner=ScannerForStage([root]), + target_dir=str(root), + file_name="model", + main_extension=".safetensors", + original_file_path=str(a), + cached_entry=None, + ) + + assert batch_id is None + # First file renamed back; nothing orphaned. + assert a.read_bytes() == b"a-bytes" + assert b.read_bytes() == b"b-bytes" + staging = root / PENDING_DELETE_DIR_NAME + if staging.exists(): + assert not any(staging.iterdir()) + + +# --------------------------------------------------------------------------- +# (o) DEFAULT_SETTINGS contains delete_undo_enabled=True +# --------------------------------------------------------------------------- +def test_o_default_settings_contains_undo_enabled() -> None: + assert DEFAULT_SETTINGS.get("delete_undo_enabled") is True + + +# --------------------------------------------------------------------------- +# (p) SCANNER EXCLUSION +# --------------------------------------------------------------------------- +class DummyScannerForWalk(ModelScanner): + """Real ModelScanner subclass exercising the real directory walks.""" + + def __init__(self, root: Path) -> None: + super().__init__( + model_type="lora", + model_class=LoraMetadata, + file_extensions={".safetensors"}, + hash_index=ModelHashIndex(), + ) + self._roots = [str(root)] + + def get_model_roots(self) -> List[str]: + return list(self._roots) + + async def _process_model_file( + self, + file_path: str, + root_path: str, + *, + hash_index: Any = None, + excluded_models: Optional[List[str]] = None, + ) -> Optional[Dict[str, Any]]: + rel_path = os.path.relpath(file_path, root_path) + name = os.path.splitext(os.path.basename(file_path))[0] + return { + "file_path": file_path.replace(os.sep, "/"), + "folder": os.path.dirname(rel_path).replace(os.sep, "/"), + "sha256": f"hash-{name}", + "tags": ["alpha"], + "model_name": name, + "size": 1, + "modified": 1.0, + } + + +async def test_p_model_walk_excludes_staging_dir(tmp_path: Path, monkeypatch) -> None: + root = tmp_path / "loras" + root.mkdir() + (root / "normal.safetensors").write_bytes(b"normal") + staging = root / PENDING_DELETE_DIR_NAME / "x" + staging.mkdir(parents=True) + (staging / "model.safetensors").write_bytes(b"ghost") + (staging / "model.metadata.json").write_bytes(b'{"hash_status": "pending"}') + + # Stub the registration side effects the scanner constructor triggers. + from py.services import model_scanner as model_scanner_module + + async def _noop_register(*_args: Any, **_kwargs: Any) -> None: + return None + + monkeypatch.setattr(model_scanner_module.ServiceRegistry, "register_service", _noop_register) + monkeypatch.setenv("LORA_MANAGER_DISABLE_PERSISTENT_CACHE", "1") + + scanner = DummyScannerForWalk(root) + + # Full walk produces NO entry whose path contains the staging dir. + result = await scanner._gather_model_data() + paths = [entry["file_path"] for entry in result.raw_data] + assert any(PENDING_DELETE_DIR_NAME not in p for p in paths) + assert not any(PENDING_DELETE_DIR_NAME in p for p in paths) + + # The file-count walk also excludes it. + assert scanner._count_model_files() == 1 + + +async def test_p_checkpoint_pending_discovery_excludes_staging(tmp_path: Path, monkeypatch) -> None: + root = tmp_path / "checkpoints" + root.mkdir() + (root / "real.safetensors").write_bytes(b"real") + (root / "real.metadata.json").write_text('{"hash_status": "pending"}') + + staging = root / PENDING_DELETE_DIR_NAME / "x" + staging.mkdir(parents=True) + (staging / "model.safetensors").write_bytes(b"ghost") + (staging / "model.metadata.json").write_text('{"hash_status": "pending"}') + + from py.services import checkpoint_scanner as checkpoint_scanner_module + from py.services import model_scanner as model_scanner_module + + async def _noop_register(*_args: Any, **_kwargs: Any) -> None: + return None + + monkeypatch.setattr(model_scanner_module.ServiceRegistry, "register_service", _noop_register) + monkeypatch.setenv("LORA_MANAGER_DISABLE_PERSISTENT_CACHE", "1") + + scanner = checkpoint_scanner_module.CheckpointScanner() + monkeypatch.setattr(scanner, "get_model_roots", lambda: [str(root)]) + + pending = await scanner._find_pending_models_from_filesystem() + paths = [entry["file_path"] for entry in pending] + assert str(root / "real.safetensors") in paths + assert not any(PENDING_DELETE_DIR_NAME in p for p in paths) + + +async def test_p_usage_stats_lookup_excludes_staging(tmp_path: Path) -> None: + root = tmp_path / "checkpoints" + root.mkdir() + (root / "mycheckpoint.safetensors").write_bytes(b"real") + staging = root / PENDING_DELETE_DIR_NAME / "x" + staging.mkdir(parents=True) + (staging / "mycheckpoint.safetensors").write_bytes(b"ghost") + + from py.utils.usage_stats import UsageStats + + stats = object.__new__(UsageStats) # avoid singleton side effects (bg task) + result = await stats._find_checkpoint_file_on_disk( + CheckpointScannerStub(root), "mycheckpoint" + ) + + # Staged file is not matched; only the real one is returned. + assert result == str(root / "mycheckpoint.safetensors") + + +# --------------------------------------------------------------------------- +# (q) MERGE TIMER -> fresh task for winner; fire-time expiry re-read purges +# --------------------------------------------------------------------------- +async def test_q_merge_arms_fresh_timer_and_purges_at_expiry(tmp_path: Path, monkeypatch) -> None: + root = tmp_path / "loras" + root.mkdir() + created = _spy_purge_timers(monkeypatch) + + service = await PendingDeleteService.get_instance() + bid_a = await _stage_simple(service, root, "alpha") + bid_b = await _stage_simple(service, root, "beta") + + merged = await service.merge_batches([bid_a, bid_b]) + assert merged == bid_a + # Staging arms one timer per batch; merge arms a FRESH timer for the winner + # with the re-anchored expiry (the original winner timer would no-op after + # re-reading the later expiry). + assert created == [ + f"pending_delete_purge_{bid_a}", + f"pending_delete_purge_{bid_b}", + f"pending_delete_purge_{bid_a}", + ] + + # Simulate the re-anchored expiry passing, then fire a purge. + winner_dir = root / PENDING_DELETE_DIR_NAME / bid_a + manifest_path = winner_dir / "manifest.json" + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + manifest["expires_at"] = int(time.time()) - 5 + manifest_path.write_text(json.dumps(manifest)) + + await _register_model_root(monkeypatch, lora_roots=[root]) + await service.purge_batch(bid_a) + + assert not winner_dir.exists() + assert not (root / "alpha.safetensors").exists() + assert not (root / "beta.safetensors").exists() + + +# --------------------------------------------------------------------------- +# (r) CROSS-TYPE PURGE ENUMERATION +# --------------------------------------------------------------------------- +async def test_r_purge_expired_enumerates_all_scanner_types_and_recipe_dir( + tmp_path: Path, monkeypatch +) -> None: + lora_root = tmp_path / "loras" + lora_root.mkdir() + ckpt_root = tmp_path / "checkpoints" + ckpt_root.mkdir() + emb_root = tmp_path / "embeddings" + emb_root.mkdir() + + for tag, root in (("lora", lora_root), ("ckpt", ckpt_root), ("emb", emb_root)): + batch_dir = root / PENDING_DELETE_DIR_NAME / f"{tag}-batch" + batch_dir.mkdir(parents=True) + (batch_dir / f"{tag}.safetensors").write_bytes(tag.encode()) + _write_batch_manifest( + batch_dir, + batch_id=f"{tag}-batch", + kind="model", + model_type=f"{tag}s", + expires_at=int(time.time()) - 10, + entries=[ + { + "staged": str(batch_dir / f"{tag}.safetensors"), + "original": str(root / f"{tag}.safetensors"), + "restored": False, + } + ], + ) + + recipe_batch = Path(settings_paths.get_settings_dir()) / PENDING_DELETE_DIR_NAME / "recipe-batch" + recipe_batch.mkdir(parents=True) + (recipe_batch / "recipe.json").write_bytes(b"{}") + _write_batch_manifest( + recipe_batch, + batch_id="recipe-batch", + kind="recipe", + expires_at=int(time.time()) - 10, + entries=[ + { + "staged": str(recipe_batch / "recipe.json"), + "original": str(tmp_path / "recipe.json"), + "restored": False, + } + ], + recipe_snapshot={"id": "r1"}, + ) + + await _register_model_root( + monkeypatch, + lora_roots=[lora_root], + checkpoint_roots=[ckpt_root], + embedding_roots=[emb_root], + ) + + service = await PendingDeleteService.get_instance() + purged = await service.purge_expired() + + assert purged >= 4 + for root in (lora_root, ckpt_root, emb_root): + staging = root / PENDING_DELETE_DIR_NAME + assert not staging.exists() or not any(staging.iterdir()) + assert not recipe_batch.exists() + + +# --------------------------------------------------------------------------- +# (s) PARTIALLY-RESTORED PURGE -> missing staged file treated as already-purged +# --------------------------------------------------------------------------- +async def test_s_partially_restored_purge_removes_remaining(tmp_path: Path, monkeypatch) -> None: + root = tmp_path / "loras" + root.mkdir() + batch_dir = root / PENDING_DELETE_DIR_NAME / "partial" + batch_dir.mkdir(parents=True) + # Entry 1 restored:true and its staged file is absent. + (batch_dir / "entry2.safetensors").write_bytes(b"present") + _write_batch_manifest( + batch_dir, + batch_id="partial", + kind="model", + model_type="loras", + expires_at=int(time.time()) - 10, + entries=[ + { + "staged": str(batch_dir / "entry1.safetensors"), + "original": str(root / "entry1.safetensors"), + "restored": True, + }, + { + "staged": str(batch_dir / "entry2.safetensors"), + "original": str(root / "entry2.safetensors"), + "restored": False, + }, + ], + ) + + await _register_model_root(monkeypatch, lora_roots=[root]) + service = await PendingDeleteService.get_instance() + + await service.purge_batch("partial") # no exception + + assert not batch_dir.exists() + assert not (root / "entry2.safetensors").exists() + + +# --------------------------------------------------------------------------- +# (t) QUARANTINE IS TERMINAL +# --------------------------------------------------------------------------- +async def test_t_quarantine_is_terminal(tmp_path: Path, monkeypatch) -> None: + root = tmp_path / "loras" + root.mkdir() + staging = root / PENDING_DELETE_DIR_NAME + batch_dir = staging / "qbatch" + batch_dir.mkdir(parents=True) + (batch_dir / "model.safetensors").write_bytes(b"data") + + await _register_model_root(monkeypatch, lora_roots=[root]) + service = await PendingDeleteService.get_instance() + + await service.purge_expired() + orphaned = staging / "qbatch.orphaned" + assert orphaned.is_dir() + + # Second sweep must NOT re-rename or delete the quarantined dir. + await service.purge_expired() + assert orphaned.is_dir() + assert (orphaned / "model.safetensors").read_bytes() == b"data" + assert not batch_dir.exists() + + +# --------------------------------------------------------------------------- +# (u) LOCK NO-DEADLOCK: stage/undo interleaved with purge +# --------------------------------------------------------------------------- +async def test_u_lock_no_deadlock_with_concurrent_purge(tmp_path: Path, monkeypatch) -> None: + root = tmp_path / "loras" + root.mkdir() + + expired_batch = root / PENDING_DELETE_DIR_NAME / "expired" + expired_batch.mkdir(parents=True) + (expired_batch / "old.safetensors").write_bytes(b"old") + _write_batch_manifest( + expired_batch, + batch_id="expired", + kind="model", + model_type="loras", + expires_at=int(time.time()) - 10, + entries=[ + { + "staged": str(expired_batch / "old.safetensors"), + "original": str(root / "old.safetensors"), + "restored": False, + } + ], + ) + await _register_model_root(monkeypatch, lora_roots=[root]) + + service = await PendingDeleteService.get_instance() + + new_model = root / "new.safetensors" + new_model.write_bytes(b"new") + + async def do_stage() -> Optional[str]: + return await service.stage_model_delete( + scanner=ScannerForStage([root]), + target_dir=str(root), + file_name="new", + main_extension=".safetensors", + original_file_path=str(new_model), + cached_entry=None, + ) + + purge_task = asyncio.create_task(service.purge_expired()) + stage_task = asyncio.create_task(do_stage()) + results = await asyncio.gather(purge_task, stage_task, return_exceptions=True) + + assert not isinstance(results[0], BaseException) + assert not isinstance(results[1], BaseException) + assert results[1] is not None + + # Expired batch fully purged (never partially), new batch staged intact. + assert not expired_batch.exists() + new_batch_dir = root / PENDING_DELETE_DIR_NAME / results[1] + assert new_batch_dir.is_dir() + assert (new_batch_dir / "new.safetensors").read_bytes() == b"new" + + +# --------------------------------------------------------------------------- +# Extra: recipe staging happy path + missing-image skip + undo +# --------------------------------------------------------------------------- +async def test_recipe_stage_copies_json_and_image_then_undo_restores(tmp_path: Path) -> None: + settings_dir = Path(settings_paths.get_settings_dir()) + recipe_json = tmp_path / "my_recipe.recipe.json" + recipe_json.write_text('{"id": "r1"}') + image = tmp_path / "preview.png" + image.write_bytes(b"img") + + service = await PendingDeleteService.get_instance() + batch_id = await service.stage_recipe_delete( + recipe_json_path=str(recipe_json), + image_path=str(image), + recipe_data={"id": "r1", "name": "Recipe One"}, + ) + assert batch_id is not None + batch_dir = settings_dir / PENDING_DELETE_DIR_NAME / batch_id + assert batch_dir.is_dir() + assert (batch_dir / "my_recipe.recipe.json").read_text() == '{"id": "r1"}' + assert (batch_dir / "preview.png").read_bytes() == b"img" + + manifest = json.loads((batch_dir / "manifest.json").read_text(encoding="utf-8")) + assert manifest["kind"] == "recipe" + assert manifest["model_type"] is None + assert manifest["model_snapshot"] is None + assert manifest["recipe_snapshot"] == {"id": "r1", "name": "Recipe One"} + assert len(manifest["entries"]) == 2 + + # The caller removes the originals after staging (see plan todo 4). + recipe_json.unlink() + image.unlink() + + await service.undo(batch_id) + assert recipe_json.read_text() == '{"id": "r1"}' + assert image.read_bytes() == b"img" + assert not batch_dir.exists() + + +async def test_recipe_stage_skips_missing_image(tmp_path: Path) -> None: + settings_dir = Path(settings_paths.get_settings_dir()) + recipe_json = tmp_path / "r2.recipe.json" + recipe_json.write_text("{}") + + service = await PendingDeleteService.get_instance() + batch_id = await service.stage_recipe_delete( + recipe_json_path=str(recipe_json), + image_path=str(tmp_path / "missing.png"), + recipe_data={"id": "r2"}, + ) + assert batch_id is not None + + batch_dir = settings_dir / PENDING_DELETE_DIR_NAME / batch_id + manifest = json.loads((batch_dir / "manifest.json").read_text(encoding="utf-8")) + assert len(manifest["entries"]) == 1 + assert (batch_dir / "r2.recipe.json").exists() + + +# --------------------------------------------------------------------------- +# Task 6 (a) STAGING ARMS A PURGE TIMER -> task named pending_delete_purge_* +# --------------------------------------------------------------------------- +async def test_t6_stage_model_arms_purge_timer(tmp_path: Path, monkeypatch) -> None: + root = tmp_path / "loras" + root.mkdir() + created = _spy_purge_timers(monkeypatch) + + service = await PendingDeleteService.get_instance() + batch_id = await _stage_simple(service, root, "model") + + # Exactly one timer task armed, with the house naming prefix. + assert created == [f"pending_delete_purge_{batch_id}"] + + +async def test_t6_stage_recipe_arms_purge_timer(tmp_path: Path, monkeypatch) -> None: + created = _spy_purge_timers(monkeypatch) + recipe_json = tmp_path / "r_t6.recipe.json" + recipe_json.write_text("{}") + + service = await PendingDeleteService.get_instance() + batch_id = await service.stage_recipe_delete( + recipe_json_path=str(recipe_json), + image_path=None, + recipe_data={"id": "r_t6"}, + ) + + assert batch_id is not None + assert created == [f"pending_delete_purge_{batch_id}"] + + +# --------------------------------------------------------------------------- +# Task 6 (b) OPPORTUNISTIC PURGE -> awaited at stage/undo entry (lock-free) +# --------------------------------------------------------------------------- +async def test_t6_purge_expired_awaited_at_stage_and_undo_entries( + tmp_path: Path, monkeypatch +) -> None: + root = tmp_path / "loras" + root.mkdir() + service = await PendingDeleteService.get_instance() + + calls: List[str] = [] + + async def counting_purge() -> int: + calls.append("purge") + return 0 + + monkeypatch.setattr(service, "purge_expired", counting_purge) + + batch_id = await _stage_simple(service, root, "model") + assert calls == ["purge"] + + recipe_json = tmp_path / "r_t6b.recipe.json" + recipe_json.write_text("{}") + await service.stage_recipe_delete( + recipe_json_path=str(recipe_json), + image_path=None, + recipe_data=None, + ) + assert calls == ["purge", "purge"] + + await service.undo(batch_id) + assert calls == ["purge", "purge", "purge"] + + +# --------------------------------------------------------------------------- +# Task 6 NON-EXPIRED BATCH SURVIVES THE STARTUP SWEEP +# --------------------------------------------------------------------------- +async def test_t6_non_expired_batch_survives_startup_sweep( + tmp_path: Path, monkeypatch +) -> None: + root = tmp_path / "loras" + root.mkdir() + await _register_model_root(monkeypatch, lora_roots=[root]) + + service = await PendingDeleteService.get_instance() + batch_id = await _stage_simple(service, root, "model") + batch_dir = root / PENDING_DELETE_DIR_NAME / batch_id + assert batch_dir.is_dir() + + # Startup sweep must skip the not-yet-expired batch (undo survives restart). + await service.purge_expired() + + assert batch_dir.is_dir() + assert (batch_dir / "model.safetensors").exists() + assert (batch_dir / "manifest.json").exists() + assert not (root / "model.safetensors").exists() + + +# --------------------------------------------------------------------------- +# F3 EXDEV-1: undo() survives a cross-device staging parent (copy fallback) +# --------------------------------------------------------------------------- +async def test_exdev1_recipe_undo_falls_back_to_copy_across_devices( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Recipe staging copies into ``{settings_dir}/.lm-pending-delete`` which + can live on a DIFFERENT filesystem than the recipes dir. undo() must NOT + die with EXDEV: fall back to copy2+remove so bytes are restored and the + staged copies are gone (no data loss).""" + settings_dir = Path(settings_paths.get_settings_dir()) + recipe_json = tmp_path / "exdev_recipe.recipe.json" + recipe_json.write_text('{"id": "exdev"}') + image = tmp_path / "exdev_preview.png" + image.write_bytes(b"img-bytes") + + service = await PendingDeleteService.get_instance() + batch_id = await service.stage_recipe_delete( + recipe_json_path=str(recipe_json), + image_path=str(image), + recipe_data={"id": "exdev"}, + ) + assert batch_id is not None + batch_dir = settings_dir / PENDING_DELETE_DIR_NAME / batch_id + assert batch_dir.is_dir() + + # The todo-4 caller removes the originals after staging. + recipe_json.unlink() + image.unlink() + + real_rename = os.rename + + def exdev_rename(src: str, dst: str) -> None: + if PENDING_DELETE_DIR_NAME in str(src): + raise OSError(errno.EXDEV, "Invalid cross-device link", str(src), str(dst)) + return real_rename(src, dst) + + monkeypatch.setattr("py.services.pending_delete_service.os.rename", exdev_rename) + + await service.undo(batch_id) + + # Byte-identical content restored; staged copies + batch dir gone. + assert recipe_json.read_text(encoding="utf-8") == '{"id": "exdev"}' + assert image.read_bytes() == b"img-bytes" + assert not batch_dir.exists() + staging = settings_dir / PENDING_DELETE_DIR_NAME + assert not staging.exists() or not any(staging.iterdir()) + + +# --------------------------------------------------------------------------- +# F3 EXDEV-2: partial EXDEV completes within ONE undo (copy fallback inline) +# --------------------------------------------------------------------------- +async def test_exdev2_partial_exdev_completes_within_one_undo( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """When the SECOND restore hits EXDEV, the first entry is restored via + rename, the second via the copy fallback, and the whole batch finishes in a + single undo call (no retry needed): both originals byte-identical, staged + copies gone, batch dir removed.""" + settings_dir = Path(settings_paths.get_settings_dir()) + recipe_json = tmp_path / "exdev2_recipe.recipe.json" + recipe_json.write_text('{"id": "exdev2"}') + image = tmp_path / "exdev2_preview.png" + image.write_bytes(b"img2-bytes") + + service = await PendingDeleteService.get_instance() + batch_id = await service.stage_recipe_delete( + recipe_json_path=str(recipe_json), + image_path=str(image), + recipe_data={"id": "exdev2"}, + ) + assert batch_id is not None + batch_dir = settings_dir / PENDING_DELETE_DIR_NAME / batch_id + assert batch_dir.is_dir() + + recipe_json.unlink() + image.unlink() + + real_rename = os.rename + calls = {"n": 0} + + def exdev_on_second_rename(src: str, dst: str) -> None: + calls["n"] += 1 + if calls["n"] == 2: + raise OSError(errno.EXDEV, "Invalid cross-device link", str(src), str(dst)) + return real_rename(src, dst) + + monkeypatch.setattr( + "py.services.pending_delete_service.os.rename", exdev_on_second_rename + ) + + await service.undo(batch_id) + + assert recipe_json.read_text(encoding="utf-8") == '{"id": "exdev2"}' + assert image.read_bytes() == b"img2-bytes" + assert not batch_dir.exists() + assert calls["n"] >= 2 + + +# --------------------------------------------------------------------------- +# F3 SNAP-1: stage_model_delete attaches the snapshot to the MAIN-file entry +# --------------------------------------------------------------------------- +async def test_snap1_stage_model_writes_snapshot_on_main_file_entry( + tmp_path: Path, +) -> None: + root = tmp_path / "loras" + root.mkdir() + model = root / "model.safetensors" + model.write_bytes(b"data") + metadata = root / "model.metadata.json" + metadata.write_bytes(b"{}") + + cached_entry = {"file_path": str(model), "sha256": "abc", "tags": ["t"]} + + service = await PendingDeleteService.get_instance() + batch_id = await service.stage_model_delete( + scanner=ScannerForStage([root]), + target_dir=str(root), + file_name="model", + main_extension=".safetensors", + original_file_path=str(model), + cached_entry=cached_entry, + ) + assert batch_id is not None + batch_dir = root / PENDING_DELETE_DIR_NAME / batch_id + manifest = json.loads((batch_dir / "manifest.json").read_text(encoding="utf-8")) + + entries = manifest["entries"] + assert len(entries) == 2 + main_entry = next(e for e in entries if e["original"] == str(model)) + meta_entry = next(e for e in entries if e["original"] == str(metadata)) + assert main_entry["snapshot"] == cached_entry + assert "snapshot" not in meta_entry + # Top-level snapshot kept for backward compat / single-delete path. + assert manifest["model_snapshot"] == cached_entry + + +async def test_snap1_none_snapshot_is_fine(tmp_path: Path) -> None: + """cached_entry=None still attaches a (None) snapshot on the main entry.""" + root = tmp_path / "loras" + root.mkdir() + model = root / "model.safetensors" + model.write_bytes(b"data") + + service = await PendingDeleteService.get_instance() + batch_id = await service.stage_model_delete( + scanner=ScannerForStage([root]), + target_dir=str(root), + file_name="model", + main_extension=".safetensors", + original_file_path=str(model), + cached_entry=None, + ) + assert batch_id is not None + batch_dir = root / PENDING_DELETE_DIR_NAME / batch_id + manifest = json.loads((batch_dir / "manifest.json").read_text(encoding="utf-8")) + main_entry = next(iter(manifest["entries"])) + assert "snapshot" in main_entry + assert main_entry["snapshot"] is None + + +# --------------------------------------------------------------------------- +# F3 SNAP-2: merged manifest entries carry BOTH snapshots +# --------------------------------------------------------------------------- +async def test_snap2_merge_keeps_both_snapshots( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + root = tmp_path / "loras" + root.mkdir() + _spy_purge_timers(monkeypatch) + + service = await PendingDeleteService.get_instance() + a1 = root / "alpha.safetensors" + a1.write_bytes(b"alpha-data") + b1 = root / "beta.safetensors" + b1.write_bytes(b"beta-data") + bid_a = await service.stage_model_delete( + scanner=ScannerForStage([root]), + target_dir=str(root), + file_name="alpha", + main_extension=".safetensors", + original_file_path=str(a1), + cached_entry={"file_path": str(a1), "tags": ["alpha"]}, + ) + bid_b = await service.stage_model_delete( + scanner=ScannerForStage([root]), + target_dir=str(root), + file_name="beta", + main_extension=".safetensors", + original_file_path=str(b1), + cached_entry={"file_path": str(b1), "tags": ["beta"]}, + ) + assert bid_a is not None + assert bid_b is not None + + assert await service.merge_batches([bid_a, bid_b]) == bid_a + winner_dir = root / PENDING_DELETE_DIR_NAME / bid_a + manifest = json.loads((winner_dir / "manifest.json").read_text(encoding="utf-8")) + + snap_entries = [e for e in manifest["entries"] if e.get("snapshot")] + assert len(snap_entries) == 2 + assert {e["snapshot"]["file_path"] for e in snap_entries} == {str(a1), str(b1)} diff --git a/tests/services/test_recipe_persistence.py b/tests/services/test_recipe_persistence.py new file mode 100644 index 00000000..82f79b2d --- /dev/null +++ b/tests/services/test_recipe_persistence.py @@ -0,0 +1,338 @@ +"""Tests for recipe delete staging in :mod:`py.services.recipes.persistence_service`. + +Covers the delete-undo wiring (plan todo 4): ``delete_recipe`` and +``bulk_delete`` stage recipe JSON + preview image into the global pending-delete +staging dir when undo is enabled, fall back to the existing hard delete when it +is disabled, and expose the batch field(s) in the result payload. Merge failure +falls back to a ``batch_ids`` array (same no-merge contract as the model bulk +path). + +Deterministic time control: no real sleeps - the re-anchored ``expires_at`` is +compared against a loose before/after window instead. +""" + +from __future__ import annotations + +import json +import logging +import time +from collections.abc import Iterator +from pathlib import Path +from typing import Any, Dict, List, Optional + +import pytest + +from py.services.pending_delete_service import ( + PENDING_DELETE_DIR_NAME, + PENDING_DELETE_TTL_SECONDS, + _reset_pending_delete_service, +) +from py.services.recipes.persistence_service import ( + PersistenceResult, + RecipePersistenceService, +) +from py.services.settings_manager import get_settings_manager +from py.utils import settings_paths + + +class DummyExifUtils: + """Exif double matching the persistence service constructor contract.""" + + def __init__(self) -> None: + self.appended = None + self.optimized_calls = 0 + + def optimize_image(self, image_data, target_width, format, quality, preserve_metadata): + self.optimized_calls += 1 + return image_data, ".webp" + + def append_recipe_metadata(self, image_path, recipe_data): + self.appended = (image_path, recipe_data) + + def extract_image_metadata(self, path): + return {} + + +class RecipeScannerStub: + """Scanner double exposing the persistence methods used by delete flows.""" + + def __init__(self, root: Path) -> None: + self.recipes_dir = str(root) + self.removed: List[str] = [] + self.bulk_removed: List[str] = [] + self._json_paths: Dict[str, str] = {} + + def register_recipe(self, recipe_id: str, json_path: Path) -> None: + self._json_paths[str(recipe_id)] = str(json_path) + + async def get_recipe_json_path(self, recipe_id: str) -> Optional[str]: + return self._json_paths.get(str(recipe_id)) + + async def remove_recipe(self, recipe_id: str) -> bool: + self.removed.append(str(recipe_id)) + return True + + async def bulk_remove(self, recipe_ids) -> int: + self.bulk_removed.extend(str(recipe_id) for recipe_id in recipe_ids) + return len(list(recipe_ids)) + + +@pytest.fixture(autouse=True) +def _reset_service_singleton() -> Iterator[None]: + """Reset the pending-delete singleton before and after each test.""" + _reset_pending_delete_service() + yield + _reset_pending_delete_service() + + +@pytest.fixture(autouse=True) +def _stub_scanner_registry(monkeypatch) -> None: + """Prevent purge enumeration from instantiating real scanner singletons.""" + from py.services.service_registry import ServiceRegistry + + async def _none(*_args: Any, **_kwargs: Any) -> None: + return None + + monkeypatch.setattr(ServiceRegistry, "get_lora_scanner", _none) + monkeypatch.setattr(ServiceRegistry, "get_checkpoint_scanner", _none) + monkeypatch.setattr(ServiceRegistry, "get_embedding_scanner", _none) + + +def _make_service() -> RecipePersistenceService: + return RecipePersistenceService( + exif_utils=DummyExifUtils(), + card_preview_width=512, + logger=logging.getLogger("test"), + ) + + +def _write_recipe(root: Path, recipe_id: str) -> tuple[Path, Path, Dict[str, Any]]: + """Write a recipe JSON + preview image; return (json_path, image_path, data).""" + recipes_dir = root / "recipes" + recipes_dir.mkdir(parents=True, exist_ok=True) + image_path = recipes_dir / f"{recipe_id}.webp" + image_path.write_bytes(f"{recipe_id}-image".encode()) + json_path = recipes_dir / f"{recipe_id}.recipe.json" + recipe_data: Dict[str, Any] = { + "id": recipe_id, + "title": f"Recipe {recipe_id}", + "file_path": str(image_path), + "loras": [], + } + json_path.write_text(json.dumps(recipe_data), encoding="utf-8") + return json_path, image_path, recipe_data + + +def _write_json_only_recipe(root: Path, recipe_id: str) -> tuple[Path, Dict[str, Any]]: + """Write a recipe JSON whose preview image does NOT exist.""" + recipes_dir = root / "recipes" + recipes_dir.mkdir(parents=True, exist_ok=True) + json_path = recipes_dir / f"{recipe_id}.recipe.json" + recipe_data: Dict[str, Any] = { + "id": recipe_id, + "title": f"Recipe {recipe_id}", + "file_path": str(recipes_dir / f"{recipe_id}.missing.webp"), + "loras": [], + } + json_path.write_text(json.dumps(recipe_data), encoding="utf-8") + return json_path, recipe_data + + +def _staging_parent() -> Path: + # Resolve through the module namespace so the conftest settings-dir + # isolation patch takes effect at call time. + return Path(settings_paths.get_settings_dir()) / PENDING_DELETE_DIR_NAME + + +def _batch_dirs() -> List[Path]: + parent = _staging_parent() + if not parent.is_dir(): + return [] + return [p for p in parent.iterdir() if p.is_dir()] + + +# --------------------------------------------------------------------------- +# (1) delete_recipe with undo enabled -> staged JSON + image, originals gone, +# payload batch_id set, manifest recipe_snapshot present +# --------------------------------------------------------------------------- +async def test_delete_recipe_stages_json_and_image(tmp_path: Path) -> None: + scanner = RecipeScannerStub(tmp_path) + json_path, image_path, recipe_data = _write_recipe(tmp_path, "r1") + scanner.register_recipe("r1", json_path) + json_bytes = json_path.read_bytes() + image_bytes = image_path.read_bytes() + + result = await _make_service().delete_recipe( + recipe_scanner=scanner, recipe_id="r1" + ) + + assert isinstance(result, PersistenceResult) + batch_id = result.payload["batch_id"] + assert batch_id is not None + + # JSON + image exist in the global staging dir; originals removed. + batch_dir = _staging_parent() / batch_id + assert batch_dir.is_dir() + assert not json_path.exists() + assert not image_path.exists() + + # QA: staged copies match the original bytes. + assert (batch_dir / json_path.name).read_bytes() == json_bytes + assert (batch_dir / image_path.name).read_bytes() == image_bytes + + manifest = json.loads((batch_dir / "manifest.json").read_text(encoding="utf-8")) + assert manifest["batch_id"] == batch_id + assert manifest["kind"] == "recipe" + assert manifest["model_type"] is None + assert manifest["state"] == "staged" + assert manifest["recipe_snapshot"] == recipe_data + assert manifest["model_snapshot"] is None + assert len(manifest["entries"]) == 2 + originals = {entry["original"] for entry in manifest["entries"]} + assert originals == {str(json_path), str(image_path)} + + # Scanner cache removal still runs. + assert scanner.removed == ["r1"] + + +# --------------------------------------------------------------------------- +# (2) recipe with missing preview image -> only JSON staged, no crash +# --------------------------------------------------------------------------- +async def test_delete_recipe_skips_missing_preview_image(tmp_path: Path) -> None: + scanner = RecipeScannerStub(tmp_path) + json_path, recipe_data = _write_json_only_recipe(tmp_path, "r2") + scanner.register_recipe("r2", json_path) + + result = await _make_service().delete_recipe( + recipe_scanner=scanner, recipe_id="r2" + ) + + batch_id = result.payload["batch_id"] + assert batch_id is not None + batch_dir = _staging_parent() / batch_id + assert batch_dir.is_dir() + assert (batch_dir / "r2.recipe.json").read_text(encoding="utf-8") == json.dumps( + recipe_data + ) + + manifest = json.loads((batch_dir / "manifest.json").read_text(encoding="utf-8")) + assert len(manifest["entries"]) == 1 + assert manifest["recipe_snapshot"] == recipe_data + + assert not json_path.exists() + assert scanner.removed == ["r2"] + + +# --------------------------------------------------------------------------- +# (3) undo disabled -> no staging, payload batch_id None, existing behavior +# --------------------------------------------------------------------------- +async def test_delete_recipe_undo_disabled_no_staging(tmp_path: Path) -> None: + get_settings_manager().settings["delete_undo_enabled"] = False + + scanner = RecipeScannerStub(tmp_path) + json_path, image_path, _recipe_data = _write_recipe(tmp_path, "r3") + scanner.register_recipe("r3", json_path) + + result = await _make_service().delete_recipe( + recipe_scanner=scanner, recipe_id="r3" + ) + + assert result.payload["batch_id"] is None + # No staging leftovers when undo is disabled. + assert not _staging_parent().exists() + # Existing hard delete behavior unchanged. + assert not json_path.exists() + assert not image_path.exists() + assert scanner.removed == ["r3"] + + +# --------------------------------------------------------------------------- +# (4) bulk_delete with 2 ids -> single batch_id, one batch dir with both +# recipes, re-anchored expires_at in the merged manifest +# --------------------------------------------------------------------------- +async def test_bulk_delete_merges_into_single_batch(tmp_path: Path) -> None: + scanner = RecipeScannerStub(tmp_path) + json_a, img_a, data_a = _write_recipe(tmp_path, "ra") + json_b, img_b, data_b = _write_recipe(tmp_path, "rb") + scanner.register_recipe("ra", json_a) + scanner.register_recipe("rb", json_b) + json_a_bytes = json_a.read_bytes() + image_a_bytes = img_a.read_bytes() + json_b_bytes = json_b.read_bytes() + image_b_bytes = img_b.read_bytes() + before = int(time.time()) + + result = await _make_service().bulk_delete( + recipe_scanner=scanner, recipe_ids=["ra", "rb"] + ) + + batch_id = result.payload["batch_id"] + assert batch_id is not None + assert "batch_ids" not in result.payload + assert len(_batch_dirs()) == 1, "loser batch dir must be removed after merge" + + batch_dir = _staging_parent() / batch_id + manifest = json.loads((batch_dir / "manifest.json").read_text(encoding="utf-8")) + assert manifest["batch_id"] == batch_id + assert len(manifest["entries"]) == 4 + + # Re-anchored expires_at: now + TTL at merge time (not the earlier of the + # two staged expiries). Loose window avoids any timing flakiness. + assert ( + before + PENDING_DELETE_TTL_SECONDS - 2 + <= manifest["expires_at"] + <= int(time.time()) + PENDING_DELETE_TTL_SECONDS + 2 + ) + + # Both recipes' files live under ONE batch dir, byte-identical to originals. + assert (batch_dir / "ra.recipe.json").read_bytes() == json_a_bytes + assert (batch_dir / "ra.webp").read_bytes() == image_a_bytes + assert (batch_dir / "rb.recipe.json").read_bytes() == json_b_bytes + assert (batch_dir / "rb.webp").read_bytes() == image_b_bytes + + # Originals removed; both snapshots present. + # Originals removed; merged manifest holds the winner's recipe snapshot. + assert not json_a.exists() + assert not json_b.exists() + assert manifest["recipe_snapshot"] in (data_a, data_b) + assert all(entry["restored"] is False for entry in manifest["entries"]) + + assert scanner.bulk_removed == ["ra", "rb"] + + +# --------------------------------------------------------------------------- +# (5) merge failure fallback -> batch_ids array of length 2, batches intact +# --------------------------------------------------------------------------- +async def test_bulk_delete_merge_failure_falls_back_to_batch_ids( + tmp_path: Path, monkeypatch +) -> None: + scanner = RecipeScannerStub(tmp_path) + json_a, _img_a, _data_a = _write_recipe(tmp_path, "ra") + json_b, _img_b, _data_b = _write_recipe(tmp_path, "rb") + scanner.register_recipe("ra", json_a) + scanner.register_recipe("rb", json_b) + + def failing_rename(src: str, dst: str) -> None: + raise OSError("simulated merge move failure") + + monkeypatch.setattr("py.services.pending_delete_service.os.rename", failing_rename) + + result = await _make_service().bulk_delete( + recipe_scanner=scanner, recipe_ids=["ra", "rb"] + ) + + batch_ids = result.payload["batch_ids"] + assert "batch_id" not in result.payload + assert len(batch_ids) == 2 + assert len(_batch_dirs()) == 2, "both constituent batches stay intact" + + # Each constituent batch is complete and individually undoable. + for batch_id in batch_ids: + batch_dir = _staging_parent() / batch_id + assert batch_dir.is_dir() + assert (batch_dir / "manifest.json").exists() + assert any(entry["original"] == str(json_a) or entry["original"] == str(json_b) for entry in json.loads((batch_dir / "manifest.json").read_text(encoding="utf-8"))["entries"]) + + assert not json_a.exists() + assert not json_b.exists() + assert scanner.bulk_removed == ["ra", "rb"] diff --git a/tests/services/test_settings_manager.py b/tests/services/test_settings_manager.py index 3597b189..dadaab47 100644 --- a/tests/services/test_settings_manager.py +++ b/tests/services/test_settings_manager.py @@ -1178,3 +1178,8 @@ def test_skip_previously_downloaded_model_versions_coerces_string_input(manager) assert manager.get_skip_previously_downloaded_model_versions() is True assert manager.settings["skip_previously_downloaded_model_versions"] is True + + +def test_delete_undo_enabled_defaults_true(manager): + assert settings_manager_module.DEFAULT_SETTINGS.get("delete_undo_enabled") is True + assert manager.get("delete_undo_enabled") is True