alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-08-14 01:33:21 -03:00
Code Issues Packages Projects Releases Wiki Activity

feat(delete): stage model and recipe deletes for 30s undo

Browse Source
This commit is contained in:
Will Miao
2026-08-11 14:08:15 +08:00
parent 6a259a14fa
commit 2d6cf545b9
12 changed files with 3444 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/services/checkpoint_scanner.py
+3 -2
View File
@@ -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
py/services/model_lifecycle_service.py
+23 -3
View File
@@ -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]:
py/services/model_scanner.py
+73 -5
View File
@@ -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:
py/services/pending_delete_service.py
+974
View File
@@ -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 ``<batch_id>.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()
py/services/recipes/persistence_service.py
+56 -10
View File
@@ -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,
py/services/settings_manager.py
+1
View File
@@ -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": "",
py/utils/usage_stats.py
+3 -1
View File
@@ -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:
tests/services/test_model_lifecycle_service.py
+169
View File
@@ -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/<batch_id>/`` 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())
tests/services/test_model_scanner.py
+264
View File
@@ -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()
tests/services/test_pending_delete_service.py
+1535
View File
File diff suppressed because it is too large Load Diff
tests/services/test_recipe_persistence.py
+338
View File
@@ -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"]
tests/services/test_settings_manager.py
+5
View File
@@ -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