mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-08-14 09:43:22 -03:00
feat(delete): stage model and recipe deletes for 30s undo
This commit is contained in:
@@ -13,7 +13,7 @@ from ..utils.models import CheckpointMetadata
|
|||||||
from ..utils.file_utils import find_preview_file, normalize_path, calculate_autov3
|
from ..utils.file_utils import find_preview_file, normalize_path, calculate_autov3
|
||||||
from ..utils.metadata_manager import MetadataManager
|
from ..utils.metadata_manager import MetadataManager
|
||||||
from ..config import config
|
from ..config import config
|
||||||
from .model_scanner import ModelScanner
|
from .model_scanner import ModelScanner, _is_excluded_dir
|
||||||
from .model_hash_index import ModelHashIndex
|
from .model_hash_index import ModelHashIndex
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
@@ -328,7 +328,8 @@ class CheckpointScanner(ModelScanner):
|
|||||||
if not os.path.exists(root_path):
|
if not os.path.exists(root_path):
|
||||||
continue
|
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:
|
for filename in filenames:
|
||||||
if not filename.endswith(".metadata.json"):
|
if not filename.endswith(".metadata.json"):
|
||||||
continue
|
continue
|
||||||
|
|||||||
@@ -7,6 +7,7 @@ import os
|
|||||||
from typing import Any, Awaitable, Callable, Dict, Iterable, List, Mapping, Optional, TYPE_CHECKING, cast
|
from typing import Any, Awaitable, Callable, Dict, Iterable, List, Mapping, Optional, TYPE_CHECKING, cast
|
||||||
|
|
||||||
from ..services.service_registry import ServiceRegistry
|
from ..services.service_registry import ServiceRegistry
|
||||||
|
from ..services.pending_delete_service import get_pending_delete_service
|
||||||
from ..utils.constants import PREVIEW_EXTENSIONS
|
from ..utils.constants import PREVIEW_EXTENSIONS
|
||||||
from ..utils.metadata_manager import MetadataManager
|
from ..utils.metadata_manager import MetadataManager
|
||||||
|
|
||||||
@@ -129,9 +130,24 @@ class ModelLifecycleService:
|
|||||||
target_dir = os.path.dirname(file_path)
|
target_dir = os.path.dirname(file_path)
|
||||||
base_name = os.path.basename(file_path)
|
base_name = os.path.basename(file_path)
|
||||||
file_name, main_extension = os.path.splitext(base_name)
|
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:
|
if cache:
|
||||||
cache.raw_data = [
|
cache.raw_data = [
|
||||||
@@ -151,7 +167,11 @@ class ModelLifecycleService:
|
|||||||
if callable(persist_current_cache):
|
if callable(persist_current_cache):
|
||||||
await cast(Awaitable[Any], 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
|
@staticmethod
|
||||||
def _extract_model_id_from_payload(payload: Any) -> Optional[int]:
|
def _extract_model_id_from_payload(payload: Any) -> Optional[int]:
|
||||||
|
|||||||
@@ -19,12 +19,28 @@ from .service_registry import ServiceRegistry
|
|||||||
from .websocket_manager import ws_manager
|
from .websocket_manager import ws_manager
|
||||||
from .persistent_model_cache import get_persistent_cache
|
from .persistent_model_cache import get_persistent_cache
|
||||||
from .settings_manager import get_settings_manager
|
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_entry_validator import CacheEntryValidator
|
||||||
from .cache_health_monitor import CacheHealthMonitor, CacheHealthStatus
|
from .cache_health_monitor import CacheHealthMonitor, CacheHealthStatus
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
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
|
@dataclass
|
||||||
class CacheBuildResult:
|
class CacheBuildResult:
|
||||||
"""Represents the outcome of scanning model files for cache building."""
|
"""Represents the outcome of scanning model files for cache building."""
|
||||||
@@ -711,6 +727,8 @@ class ModelScanner:
|
|||||||
if ext in self.file_extensions:
|
if ext in self.file_extensions:
|
||||||
total_files += 1
|
total_files += 1
|
||||||
elif entry.is_dir(follow_symlinks=True):
|
elif entry.is_dir(follow_symlinks=True):
|
||||||
|
if _is_excluded_dir(entry.name):
|
||||||
|
continue
|
||||||
count_recursive(entry.path)
|
count_recursive(entry.path)
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Error counting files in entry {entry.path}: {e}")
|
logger.error(f"Error counting files in entry {entry.path}: {e}")
|
||||||
@@ -864,7 +882,8 @@ class ModelScanner:
|
|||||||
continue
|
continue
|
||||||
|
|
||||||
# Recursively scan directory
|
# 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)
|
real_root = os.path.realpath(root)
|
||||||
if real_root in visited_real_paths:
|
if real_root in visited_real_paths:
|
||||||
continue
|
continue
|
||||||
@@ -1137,6 +1156,11 @@ class ModelScanner:
|
|||||||
hash_index = hash_index or self._hash_index
|
hash_index = hash_index or self._hash_index
|
||||||
excluded_models = excluded_models if excluded_models is not None else self._excluded_models
|
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)
|
metadata, should_skip = await MetadataManager.load_metadata(file_path, self.model_class)
|
||||||
|
|
||||||
if should_skip:
|
if should_skip:
|
||||||
@@ -1456,6 +1480,8 @@ class ModelScanner:
|
|||||||
if self.is_cancelled():
|
if self.is_cancelled():
|
||||||
return
|
return
|
||||||
elif entry.is_dir(follow_symlinks=True):
|
elif entry.is_dir(follow_symlinks=True):
|
||||||
|
if _is_excluded_dir(entry.name):
|
||||||
|
continue
|
||||||
await scan_recursive(entry.path, root_path, visited_paths)
|
await scan_recursive(entry.path, root_path, visited_paths)
|
||||||
except Exception as entry_error:
|
except Exception as entry_error:
|
||||||
logger.error(f"Error processing entry {entry.path}: {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
|
# Track deleted models to update cache once
|
||||||
deleted_models = []
|
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:
|
for file_path in file_paths:
|
||||||
if self.is_cancelled():
|
if self.is_cancelled():
|
||||||
logger.info(f"{self.model_type.capitalize()} Scanner: Bulk delete cancelled by user")
|
logger.info(f"{self.model_type.capitalize()} Scanner: Bulk delete cancelled by user")
|
||||||
@@ -2218,12 +2249,36 @@ class ModelScanner:
|
|||||||
base_name = os.path.basename(file_path)
|
base_name = os.path.basename(file_path)
|
||||||
file_name, main_extension = os.path.splitext(base_name)
|
file_name, main_extension = os.path.splitext(base_name)
|
||||||
|
|
||||||
deleted_files = await delete_model_artifacts(
|
# Snapshot the cache entry BEFORE the cache mutation that
|
||||||
target_dir,
|
# runs after the loop - the manifest needs it for undo.
|
||||||
file_name,
|
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,
|
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:
|
if deleted_files:
|
||||||
deleted_models.append(file_path)
|
deleted_models.append(file_path)
|
||||||
results.append({
|
results.append({
|
||||||
@@ -2246,6 +2301,18 @@ class ModelScanner:
|
|||||||
'error': str(e)
|
'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
|
# Batch update cache if any models were deleted
|
||||||
if deleted_models:
|
if deleted_models:
|
||||||
# Update the cache in a batch operation
|
# Update the cache in a batch operation
|
||||||
@@ -2257,7 +2324,8 @@ class ModelScanner:
|
|||||||
'total_deleted': total_deleted,
|
'total_deleted': total_deleted,
|
||||||
'total_attempted': len(file_paths),
|
'total_attempted': len(file_paths),
|
||||||
'cache_updated': cache_updated,
|
'cache_updated': cache_updated,
|
||||||
'results': results
|
'results': results,
|
||||||
|
**batch_field
|
||||||
}
|
}
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
|||||||
@@ -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()
|
||||||
@@ -14,6 +14,7 @@ from typing import Any, Awaitable, Dict, Iterable, Optional, cast
|
|||||||
from ...config import config
|
from ...config import config
|
||||||
from ...recipes.constants import GEN_PARAM_KEYS
|
from ...recipes.constants import GEN_PARAM_KEYS
|
||||||
from ...utils.utils import calculate_recipe_fingerprint
|
from ...utils.utils import calculate_recipe_fingerprint
|
||||||
|
from ..pending_delete_service import get_pending_delete_service
|
||||||
from .errors import RecipeNotFoundError, RecipeValidationError
|
from .errors import RecipeNotFoundError, RecipeValidationError
|
||||||
|
|
||||||
|
|
||||||
@@ -201,12 +202,31 @@ class RecipePersistenceService:
|
|||||||
recipe_data = json.load(file_obj)
|
recipe_data = json.load(file_obj)
|
||||||
|
|
||||||
image_path = recipe_data.get("file_path")
|
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)
|
os.remove(recipe_json_path)
|
||||||
if image_path and os.path.exists(image_path):
|
if image_path and os.path.exists(image_path):
|
||||||
os.remove(image_path)
|
os.remove(image_path)
|
||||||
|
|
||||||
await recipe_scanner.remove_recipe(recipe_id)
|
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:
|
async def update_recipe(self, *, recipe_scanner, recipe_id: str, updates: dict[str, Any]) -> PersistenceResult:
|
||||||
"""Update persisted metadata for a recipe."""
|
"""Update persisted metadata for a recipe."""
|
||||||
@@ -450,6 +470,9 @@ class RecipePersistenceService:
|
|||||||
|
|
||||||
deleted_recipes: list[str] = []
|
deleted_recipes: list[str] = []
|
||||||
failed_recipes: list[dict[str, Any]] = []
|
failed_recipes: list[dict[str, Any]] = []
|
||||||
|
batch_ids: list[str] = []
|
||||||
|
|
||||||
|
pending_delete_service = await get_pending_delete_service()
|
||||||
|
|
||||||
for recipe_id in recipe_ids:
|
for recipe_id in recipe_ids:
|
||||||
recipe_json_path = await recipe_scanner.get_recipe_json_path(recipe_id)
|
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:
|
with open(recipe_json_path, "r", encoding="utf-8") as file_obj:
|
||||||
recipe_data = json.load(file_obj)
|
recipe_data = json.load(file_obj)
|
||||||
image_path = recipe_data.get("file_path")
|
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)
|
os.remove(recipe_json_path)
|
||||||
if image_path and os.path.exists(image_path):
|
if image_path and os.path.exists(image_path):
|
||||||
os.remove(image_path)
|
os.remove(image_path)
|
||||||
@@ -471,15 +505,27 @@ class RecipePersistenceService:
|
|||||||
if deleted_recipes:
|
if deleted_recipes:
|
||||||
await recipe_scanner.bulk_remove(deleted_recipes)
|
await recipe_scanner.bulk_remove(deleted_recipes)
|
||||||
|
|
||||||
return PersistenceResult(
|
payload: dict[str, Any] = {
|
||||||
{
|
"success": True,
|
||||||
"success": True,
|
"deleted": deleted_recipes,
|
||||||
"deleted": deleted_recipes,
|
"failed": failed_recipes,
|
||||||
"failed": failed_recipes,
|
"total_deleted": len(deleted_recipes),
|
||||||
"total_deleted": len(deleted_recipes),
|
"total_failed": len(failed_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(
|
async def save_recipe_from_widget(
|
||||||
self,
|
self,
|
||||||
|
|||||||
@@ -111,6 +111,7 @@ DEFAULT_SETTINGS: Dict[str, Any] = {
|
|||||||
"backup_retention_count": 5,
|
"backup_retention_count": 5,
|
||||||
"use_new_license_icons": True,
|
"use_new_license_icons": True,
|
||||||
"group_by_model": False,
|
"group_by_model": False,
|
||||||
|
"delete_undo_enabled": True,
|
||||||
# AI / LLM provider configuration (BYOK)
|
# AI / LLM provider configuration (BYOK)
|
||||||
"llm_provider": "openai", # "openai" | "ollama" | "custom"
|
"llm_provider": "openai", # "openai" | "ollama" | "custom"
|
||||||
"llm_api_key": "",
|
"llm_api_key": "",
|
||||||
|
|||||||
@@ -10,6 +10,7 @@ from typing import Any, Awaitable, Dict, Set, cast
|
|||||||
|
|
||||||
from ..config import config
|
from ..config import config
|
||||||
from ..services.service_registry import ServiceRegistry
|
from ..services.service_registry import ServiceRegistry
|
||||||
|
from ..services.model_scanner import _is_excluded_dir
|
||||||
from ..utils.settings_paths import get_settings_dir
|
from ..utils.settings_paths import get_settings_dir
|
||||||
|
|
||||||
# Check if running in standalone mode
|
# Check if running in standalone mode
|
||||||
@@ -421,7 +422,8 @@ class UsageStats:
|
|||||||
if not os.path.exists(root_path):
|
if not os.path.exists(root_path):
|
||||||
continue
|
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:
|
for filename in filenames:
|
||||||
extension = os.path.splitext(filename)[1].lower()
|
extension = os.path.splitext(filename)[1].lower()
|
||||||
if extension not in supported_extensions:
|
if extension not in supported_extensions:
|
||||||
|
|||||||
@@ -1,11 +1,16 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
|
from collections.abc import Iterator
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any, Dict, cast
|
from typing import Any, Dict, cast
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from py.services.model_lifecycle_service import ModelLifecycleService, _require_path_in_library_roots
|
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.metadata_manager import MetadataManager
|
||||||
from py.utils.models import LoraMetadata
|
from py.utils.models import LoraMetadata
|
||||||
|
|
||||||
@@ -901,3 +906,167 @@ async def test_extract_model_id_handles_string_values():
|
|||||||
|
|
||||||
payload = {"civitai": {"modelId": "54321"}}
|
payload = {"civitai": {"modelId": "54321"}}
|
||||||
assert service._extract_model_id_from_payload(payload) == 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())
|
||||||
|
|||||||
@@ -1,6 +1,11 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
|
import json
|
||||||
import os
|
import os
|
||||||
import sqlite3
|
import sqlite3
|
||||||
|
import time
|
||||||
|
from collections.abc import Iterator
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any, Dict, List, Optional
|
from typing import Any, Dict, List, Optional
|
||||||
from types import MethodType
|
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_cache import ModelCache
|
||||||
from py.services.model_hash_index import ModelHashIndex
|
from py.services.model_hash_index import ModelHashIndex
|
||||||
from py.services.model_scanner import CacheBuildResult, ModelScanner
|
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.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.civitai_utils import build_license_flags
|
||||||
from py.utils.models import BaseModelMetadata
|
from py.utils.models import BaseModelMetadata
|
||||||
|
|
||||||
@@ -104,6 +115,27 @@ def stub_register_service(monkeypatch):
|
|||||||
monkeypatch.setattr(model_scanner.ServiceRegistry, "register_service", noop)
|
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]:
|
def _create_files(root: Path) -> tuple[Path, Path, Path]:
|
||||||
first = root / "one.txt"
|
first = root / "one.txt"
|
||||||
first.write_text("one", encoding="utf-8")
|
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 changed is True
|
||||||
assert resort_calls == 1
|
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()
|
||||||
|
|||||||
File diff suppressed because it is too large
Load Diff
@@ -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"]
|
||||||
@@ -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.get_skip_previously_downloaded_model_versions() is True
|
||||||
assert manager.settings["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
|
||||||
|
|||||||
Reference in New Issue
Block a user