fix(images): make single-image and intermediate deletion transactional

invokeai/app/api/routers/images.py

@@ -27,6 +27,7 @@
     ImageCategory,
     ImageNamesResult,
     ImageRecordChanges,
+    ImageRecordNotFoundException,
     ResourceOrigin,
 )
 from invokeai.app.services.images.images_common import (
@@ -207,22 +208,28 @@ async def delete_image(
     _assert_image_owner(image_name, current_user)
     assert_image_move_maintenance_inactive()
 
-    deleted_images: set[str] = set()
-    affected_boards: set[str] = set()
-
+    # Let service-level failures surface as errors rather than swallowing them and returning
+    # a success-shaped response. A previous version of this handler caught everything and
+    # returned an empty ``deleted_images`` list with HTTP 200; the frontend treated that as
+    # success and dropped the item from its cache even though the record was still live.
     try:
         image_dto = ApiDependencies.invoker.services.images.get_dto(image_name)
-        board_id = image_dto.board_id or "none"
+    except ImageRecordNotFoundException:
+        raise HTTPException(status_code=404, detail="Image not found")
+    except Exception:
+        # A record/URL/board lookup failure for an image that does exist is a server fault, not a
+        # missing image — reporting it as 404 would tell the frontend to drop a live item.
+        raise HTTPException(status_code=500, detail="Failed to delete image")
+
+    board_id = image_dto.board_id or "none"
+    try:
         ApiDependencies.invoker.services.images.delete(image_name)
-        deleted_images.add(image_name)
-        affected_boards.add(board_id)
     except Exception:
-        # TODO: Does this need any exception handling at all?
-        pass
+        raise HTTPException(status_code=500, detail="Failed to delete image")
 
     return DeleteImagesResult(
-        deleted_images=list(deleted_images),
-        affected_boards=list(affected_boards),
+        deleted_images=[image_name],
+        affected_boards=[board_id],
     )
 
 

invokeai/app/services/image_records/image_records_base.py

@@ -70,8 +70,18 @@ def delete_many(self, image_names: list[str]) -> None:
         pass
 
     @abstractmethod
-    def delete_intermediates(self) -> list[tuple[str, str]]:
-        """Deletes all intermediate image records, returning a list of (image_name, image_subfolder) tuples."""
+    def get_intermediates(self) -> list[tuple[str, str]]:
+        """Gets all intermediate image records as (image_name, image_subfolder) tuples, without deleting them."""
+        pass
+
+    @abstractmethod
+    def delete_intermediates_by_names(self, image_names: list[str]) -> tuple[list[str], list[str]]:
+        """Deletes the named image records, skipping any that are no longer intermediates.
+
+        Returns ``(deleted, retained)``: the names whose records were removed, and the names whose
+        records are still present because they are no longer intermediates. Names whose records were
+        already gone appear in neither list, so a caller holding their files must not restore them.
+        """
         pass
 
     @abstractmethod

invokeai/app/services/image_records/image_records_sqlite.py

@@ -23,24 +23,27 @@
 
 
 class SqliteImageRecordStorage(ImageRecordStorageBase):
+    # Conservative bound on bound parameters per statement. SQLITE_MAX_VARIABLE_NUMBER defaults to
+    # 999 on SQLite builds older than 3.32, and an image library can hold far more intermediates.
+    _MAX_SQL_VARIABLES = 500
+
     def __init__(self, db: SqliteDatabase) -> None:
         super().__init__()
         self._db = db
 
     def get(self, image_name: str) -> ImageRecord:
+        # A query failure means the database is unavailable, not that the image is missing. Reporting
+        # it as "not found" makes callers (and the routes above them) delete live images from view.
         with self._db.transaction() as cursor:
-            try:
-                cursor.execute(
-                    f"""--sql
-                    SELECT {IMAGE_DTO_COLS} FROM images
-                    WHERE image_name = ?;
-                    """,
-                    (image_name,),
-                )
+            cursor.execute(
+                f"""--sql
+                SELECT {IMAGE_DTO_COLS} FROM images
+                WHERE image_name = ?;
+                """,
+                (image_name,),
+            )
 
-                result = cast(Optional[sqlite3.Row], cursor.fetchone())
-            except sqlite3.Error as e:
-                raise ImageRecordNotFoundException from e
+            result = cast(Optional[sqlite3.Row], cursor.fetchone())
 
         if not result:
             raise ImageRecordNotFoundException
@@ -62,20 +65,17 @@ def get_user_id(self, image_name: str) -> Optional[str]:
             return cast(Optional[str], dict(result).get("user_id"))
 
     def get_metadata(self, image_name: str) -> Optional[MetadataField]:
+        # As in get(): a query failure is a database fault, not a missing record.
         with self._db.transaction() as cursor:
-            try:
-                cursor.execute(
-                    """--sql
-                    SELECT metadata FROM images
-                    WHERE image_name = ?;
-                    """,
-                    (image_name,),
-                )
-
-                result = cast(Optional[sqlite3.Row], cursor.fetchone())
+            cursor.execute(
+                """--sql
+                SELECT metadata FROM images
+                WHERE image_name = ?;
+                """,
+                (image_name,),
+            )
 
-            except sqlite3.Error as e:
-                raise ImageRecordNotFoundException from e
+            result = cast(Optional[sqlite3.Row], cursor.fetchone())
 
             if not result:
                 raise ImageRecordNotFoundException
@@ -302,30 +302,59 @@ def get_intermediates_count(self, user_id: Optional[str] = None) -> int:
             count = cast(int, cursor.fetchone()[0])
         return count
 
-    def delete_intermediates(self) -> list[tuple[str, str]]:
-        """Deletes all intermediate image records.
+    def get_intermediates(self) -> list[tuple[str, str]]:
+        """Gets all intermediate image records without deleting them.
 
-        Returns a list of (image_name, image_subfolder) tuples for file cleanup.
+        Returns a list of (image_name, image_subfolder) tuples for staged file deletion.
         """
         with self._db.transaction() as cursor:
-            try:
-                cursor.execute(
-                    """--sql
-                    SELECT image_name, image_subfolder FROM images
-                    WHERE is_intermediate = TRUE;
-                    """
-                )
-                result = cast(list[sqlite3.Row], cursor.fetchall())
-                image_name_subfolder_pairs = [(r[0], r[1]) for r in result]
-                cursor.execute(
-                    """--sql
-                    DELETE FROM images
-                    WHERE is_intermediate = TRUE;
-                    """
-                )
-            except sqlite3.Error as e:
-                raise ImageRecordDeleteException from e
-        return image_name_subfolder_pairs
+            cursor.execute(
+                """--sql
+                SELECT image_name, image_subfolder FROM images
+                WHERE is_intermediate = TRUE;
+                """
+            )
+            result = cast(list[sqlite3.Row], cursor.fetchall())
+        return [(r[0], r[1]) for r in result]
+
+    def delete_intermediates_by_names(self, image_names: list[str]) -> tuple[list[str], list[str]]:
+        """Deletes the named image records, skipping any that are no longer intermediates.
+
+        The ``is_intermediate`` predicate rides on the DELETE itself rather than on a preceding
+        SELECT, so an image promoted out of intermediate status keeps its record however the
+        promotion interleaves with this call. (Python's legacy sqlite3 transaction control opens a
+        transaction only before a write, so a SELECT here holds no read lock to rely on.)
+
+        Returns ``(deleted, retained)``: the names whose records this call removed, and the names
+        whose records are still present because they are no longer intermediates. Names whose
+        records were already gone appear in neither list — the caller must not restore their files.
+        """
+        deleted: list[str] = []
+        retained: list[str] = []
+        try:
+            with self._db.transaction() as cursor:
+                # Chunked to stay under SQLITE_MAX_VARIABLE_NUMBER; every chunk runs inside the one
+                # transaction above.
+                for start in range(0, len(image_names), self._MAX_SQL_VARIABLES):
+                    chunk = image_names[start : start + self._MAX_SQL_VARIABLES]
+                    placeholders = ",".join("?" for _ in chunk)
+                    select_query = f"SELECT image_name FROM images WHERE image_name IN ({placeholders})"
+
+                    cursor.execute(select_query, chunk)
+                    present_before = {cast(str, r[0]) for r in cursor.fetchall()}
+                    cursor.execute(
+                        f"DELETE FROM images WHERE image_name IN ({placeholders}) AND is_intermediate = TRUE",
+                        chunk,
+                    )
+                    cursor.execute(select_query, chunk)
+                    present_after = {cast(str, r[0]) for r in cursor.fetchall()}
+
+                    deleted.extend(name for name in chunk if name in present_before and name not in present_after)
+                    retained.extend(name for name in chunk if name in present_after)
+        except sqlite3.Error as e:
+            # The try wraps the context manager so a failure in its commit is reported too.
+            raise ImageRecordDeleteException from e
+        return deleted, retained
 
     def save(
         self,

invokeai/app/services/images/images_default.py

@@ -276,18 +276,43 @@ def get_many(
             raise e
 
     def delete(self, image_name: str):
+        # Stage the file deletion first so a database failure can be rolled back by
+        # restoring the files, keeping the record and files consistent either way.
+        token: object | None = None
+        record_deleted = False
         try:
             record = self.__invoker.services.image_records.get(image_name)
-            self.__invoker.services.image_files.delete(image_name, image_subfolder=record.image_subfolder)
+            token = self.__invoker.services.image_files.stage_delete(image_name, image_subfolder=record.image_subfolder)
             self.__invoker.services.image_records.delete(image_name)
+            record_deleted = True
+            try:
+                self.__invoker.services.image_files.commit_delete(token)
+            except Exception as cleanup_error:
+                # The record is gone; a failed purge only leaves a staging directory
+                # behind, which startup recovery will clean up. Not a delete failure.
+                self.__invoker.services.logger.error(f"Failed to purge staged image files: {cleanup_error}")
             self._on_deleted(image_name)
         except ImageRecordDeleteException:
+            if token is not None:
+                try:
+                    self.__invoker.services.image_files.rollback_delete(token)
+                except Exception as rollback_error:
+                    self.__invoker.services.logger.error(
+                        f"Failed to restore staged image files for {image_name}: {rollback_error}"
+                    )
             self.__invoker.services.logger.error("Failed to delete image record")
             raise
         except ImageFileDeleteException:
             self.__invoker.services.logger.error("Failed to delete image file")
             raise
         except Exception as e:
+            if token is not None and not record_deleted:
+                try:
+                    self.__invoker.services.image_files.rollback_delete(token)
+                except Exception as rollback_error:
+                    self.__invoker.services.logger.error(
+                        f"Failed to restore staged image files for {image_name}: {rollback_error}"
+                    )
             self.__invoker.services.logger.error("Problem deleting image record and file")
             raise e
 
@@ -346,14 +371,77 @@ def delete_images_on_board(self, board_id: str, user_id: Optional[str] = None) -
             self.__invoker.services.logger.error(f"Problem deleting image records and files: {str(e)}")
             raise e
 
+    def _record_still_exists(self, image_name: str) -> bool:
+        """Whether an image record is still present, erring towards "yes".
+
+        Used to decide whether staged files must be restored. A restore is only correct while the
+        record is live; if we cannot tell, restoring is the safer error, because leftover files can
+        be cleaned up later but files purged against a live record are gone.
+        """
+        try:
+            self.__invoker.services.image_records.get(image_name)
+            return True
+        except ImageRecordNotFoundException:
+            return False
+        except Exception as e:
+            self.__invoker.services.logger.error(f"Could not confirm whether {image_name} still exists: {e}")
+            return True
+
     def delete_intermediates(self) -> int:
+        # All-or-nothing transaction: stage every file first, then delete the records in one
+        # operation, then purge the stages. Any staging or database failure rolls back every staged
+        # file, so a live record's files are never destroyed — though a failed rollback leaves them
+        # in the staging dir until startup recovery restores them.
         try:
-            image_name_subfolder_pairs = self.__invoker.services.image_records.delete_intermediates()
-            count = len(image_name_subfolder_pairs)
-            for image_name, image_subfolder in image_name_subfolder_pairs:
-                self.__invoker.services.image_files.delete(image_name, image_subfolder=image_subfolder)
+            image_name_subfolder_pairs = self.__invoker.services.image_records.get_intermediates()
+            staged_deletes: list[tuple[str, object]] = []
+            try:
+                for image_name, image_subfolder in image_name_subfolder_pairs:
+                    token = self.__invoker.services.image_files.stage_delete(
+                        image_name, image_subfolder=image_subfolder
+                    )
+                    staged_deletes.append((image_name, token))
+                # Deletion is conditional on the row still being an intermediate. An image can be
+                # promoted out of intermediate status between the snapshot above and this call, and
+                # such an image must keep both its record and its files.
+                deleted, retained = self.__invoker.services.image_records.delete_intermediates_by_names(
+                    [name for name, _ in staged_deletes]
+                )
+                deleted_names = set(deleted)
+                retained_names = set(retained)
+            except Exception:
+                for image_name, token in staged_deletes:
+                    try:
+                        self.__invoker.services.image_files.rollback_delete(token)
+                    except Exception as rollback_error:
+                        self.__invoker.services.logger.error(
+                            f"Failed to restore staged image files for {image_name}: {rollback_error}"
+                        )
+                raise
+            deleted_image_names: list[str] = []
+            for image_name, token in staged_deletes:
+                # Only a record that is still there earns a restore. A name in neither list had its
+                # record removed by someone else while we held its files, and a retained record can
+                # still be deleted while this loop works through the other names — restoring either
+                # would strand the files on disk with no record and no staging dir to recover from.
+                # Re-checking here shrinks that window from the whole loop to a single lookup.
+                if image_name in retained_names and self._record_still_exists(image_name):
+                    try:
+                        self.__invoker.services.image_files.rollback_delete(token)
+                    except Exception as rollback_error:
+                        self.__invoker.services.logger.error(
+                            f"Failed to restore staged image files for {image_name}: {rollback_error}"
+                        )
+                    continue
+                try:
+                    self.__invoker.services.image_files.commit_delete(token)
+                except Exception as cleanup_error:
+                    self.__invoker.services.logger.error(f"Failed to purge staged image files: {cleanup_error}")
+                if image_name in deleted_names:
+                    deleted_image_names.append(image_name)
+            for image_name in deleted_image_names:
                 self._on_deleted(image_name)
-            return count
+            return len(deleted_image_names)
         except ImageRecordDeleteException:
             self.__invoker.services.logger.error("Failed to delete image records")
             raise