refactor(ui): track discarded items instead of using delete method

Author: psychedelicious

invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/context.tsx

@@ -4,6 +4,7 @@ import { EMPTY_ARRAY } from 'app/store/constants';
 import { useAppStore } from 'app/store/storeHooks';
 import { buildZodTypeGuard } from 'common/util/zodUtils';
 import { getOutputImageName } from 'features/controlLayers/components/SimpleSession/shared';
+import { canvasQueueItemDiscarded, selectDiscardedItems } from 'features/controlLayers/store/canvasStagingAreaSlice';
 import type { ProgressImage } from 'features/nodes/types/common';
 import type { Atom, MapStore, StoreValue, WritableAtom } from 'nanostores';
 import { atom, computed, effect, map, subscribeKeys } from 'nanostores';
@@ -104,6 +105,7 @@ type CanvasSessionContextValue = {
   selectFirst: () => void;
   selectLast: () => void;
   onImageLoad: (itemId: number) => void;
+  discard: (itemId: number) => void;
 };
 
 const CanvasSessionContext = createContext<CanvasSessionContextValue | null>(null);
@@ -138,14 +140,7 @@ export const CanvasSessionContextProvider = memo(
      * Manually-synced atom containing queue items for the current session. This is populated from the RTK Query cache
      * and kept in sync with it via a redux subscription.
      */
-    const $allItems = useState(() => atom<S['SessionQueueItem'][]>([]))[0];
-
-    /**
-     * All queue items for the current session, excluding canceled and failed items.
-     */
-    const $items = useState(() =>
-      computed([$allItems], (allItems) => allItems.filter(({ status }) => status !== 'canceled' && status !== 'failed'))
-    )[0];
+    const $items = useState(() => atom<S['SessionQueueItem'][]>([]))[0];
 
     /**
      * Whether auto-switch is enabled.
@@ -239,12 +234,27 @@ export const CanvasSessionContextProvider = memo(
      */
     const selectQueueItems = useMemo(
       () =>
-        createSelector(queueApi.endpoints.listAllQueueItems.select({ destination: session.id }), ({ data }) => {
-          return data ?? EMPTY_ARRAY;
-        }),
+        createSelector(
+          [queueApi.endpoints.listAllQueueItems.select({ destination: session.id }), selectDiscardedItems],
+          ({ data }, discardedItems) => {
+            if (!data) {
+              return EMPTY_ARRAY;
+            }
+            return data.filter(
+              ({ status, item_id }) => status !== 'canceled' && status !== 'failed' && !discardedItems.includes(item_id)
+            );
+          }
+        ),
       [session.id]
     );
 
+    const discard = useCallback(
+      (itemId: number) => {
+        store.dispatch(canvasQueueItemDiscarded({ itemId }));
+      },
+      [store]
+    );
+
     const selectNext = useCallback(() => {
       const selectedItemId = $selectedItemId.get();
       if (selectedItemId === null) {
@@ -350,22 +360,20 @@ export const CanvasSessionContextProvider = memo(
 
     // Set up state subscriptions and effects
     useEffect(() => {
-      let _prevAllItems: readonly S['SessionQueueItem'][] = [];
-      // Seed the $allItems atom with the initial query cache state
-      $allItems.set(selectQueueItems(store.getState()));
+      let _prevItems: readonly S['SessionQueueItem'][] = [];
+      // Seed the $items atom with the initial query cache state
+      $items.set(selectQueueItems(store.getState()));
 
-      // Manually keep the $allItems atom in sync as the query cache is updated
+      // Manually keep the $items atom in sync as the query cache is updated
       const unsubReduxSyncToItemsAtom = store.subscribe(() => {
-        const prevItems = $allItems.get();
+        const prevItems = $items.get();
         const items = selectQueueItems(store.getState());
         if (items !== prevItems) {
-          _prevAllItems = prevItems;
-          $allItems.set(items);
+          _prevItems = prevItems;
+          $items.set(items);
         }
       });
 
-      let _prevItems: readonly S['SessionQueueItem'][] = [];
-
       // Handle cases that could result in a nonexistent queue item being selected.
       const unsubEnsureSelectedItemIdExists = effect(
         [$items, $selectedItemId, $lastStartedItemId],
@@ -500,14 +508,13 @@ export const CanvasSessionContextProvider = memo(
         unsubReduxSyncToItemsAtom();
         unsubEnsureSelectedItemIdExists();
         unsubCleanUpProgressData();
-        $allItems.set([]);
+        $items.set([]);
         $progressData.set({});
         $selectedItemId.set(null);
       };
     }, [
-      $allItems,
-      $autoSwitch,
       $items,
+      $autoSwitch,
       $lastLoadedItemId,
       $lastStartedItemId,
       $progressData,
@@ -535,6 +542,7 @@ export const CanvasSessionContextProvider = memo(
         selectFirst,
         selectLast,
         onImageLoad,
+        discard,
       }),
       [
         $autoSwitch,
@@ -553,6 +561,7 @@ export const CanvasSessionContextProvider = memo(
         selectFirst,
         selectLast,
         onImageLoad,
+        discard,
       ]
     );
 

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaToolbarDiscardSelectedButton.tsx

@@ -3,15 +3,15 @@ import { useStore } from '@nanostores/react';
 import { useAppDispatch } from 'app/store/storeHooks';
 import { useCanvasSessionContext } from 'features/controlLayers/components/SimpleSession/context';
 import { canvasSessionReset, generateSessionReset } from 'features/controlLayers/store/canvasStagingAreaSlice';
-import { useDeleteQueueItem } from 'features/queue/hooks/useDeleteQueueItem';
+import { useCancelQueueItem } from 'features/queue/hooks/useCancelQueueItem';
 import { memo, useCallback } from 'react';
 import { useTranslation } from 'react-i18next';
 import { PiXBold } from 'react-icons/pi';
 
 export const StagingAreaToolbarDiscardSelectedButton = memo(({ isDisabled }: { isDisabled?: boolean }) => {
   const dispatch = useAppDispatch();
   const ctx = useCanvasSessionContext();
-  const deleteQueueItem = useDeleteQueueItem();
+  const cancelQueueItem = useCancelQueueItem();
   const selectedItemId = useStore(ctx.$selectedItemId);
 
   const { t } = useTranslation();
@@ -20,7 +20,8 @@ export const StagingAreaToolbarDiscardSelectedButton = memo(({ isDisabled }: { i
     if (selectedItemId === null) {
       return;
     }
-    await deleteQueueItem.trigger(selectedItemId, { withToast: false });
+    ctx.discard(selectedItemId);
+    await cancelQueueItem.trigger(selectedItemId, { withToast: false });
     const itemCount = ctx.$itemCount.get();
     if (itemCount <= 1) {
       if (ctx.session.type === 'advanced') {
@@ -30,7 +31,7 @@ export const StagingAreaToolbarDiscardSelectedButton = memo(({ isDisabled }: { i
         dispatch(generateSessionReset());
       }
     }
-  }, [selectedItemId, deleteQueueItem, ctx.$itemCount, ctx.session.type, dispatch]);
+  }, [selectedItemId, ctx, cancelQueueItem, dispatch]);
 
   return (
     <IconButton
@@ -40,8 +41,8 @@ export const StagingAreaToolbarDiscardSelectedButton = memo(({ isDisabled }: { i
       onClick={discardSelected}
       colorScheme="invokeBlue"
       fontSize={16}
-      isDisabled={selectedItemId === null || deleteQueueItem.isDisabled || isDisabled}
-      isLoading={deleteQueueItem.isLoading}
+      isDisabled={selectedItemId === null || cancelQueueItem.isDisabled || isDisabled}
+      isLoading={cancelQueueItem.isLoading}
     />
   );
 });

invokeai/frontend/web/src/features/controlLayers/store/canvasStagingAreaSlice.ts

@@ -6,11 +6,13 @@ import { canvasReset } from 'features/controlLayers/store/actions';
 type CanvasStagingAreaState = {
   generateSessionId: string | null;
   canvasSessionId: string | null;
+  canvasDiscardedQueueItems: number[];
 };
 
 const INITIAL_STATE: CanvasStagingAreaState = {
   generateSessionId: null,
   canvasSessionId: null,
+  canvasDiscardedQueueItems: [],
 };
 
 const getInitialState = (): CanvasStagingAreaState => deepClone(INITIAL_STATE);
@@ -26,12 +28,20 @@ export const canvasSessionSlice = createSlice({
     generateSessionReset: (state) => {
       state.generateSessionId = null;
     },
+    canvasQueueItemDiscarded: (state, action: PayloadAction<{ itemId: number }>) => {
+      const { itemId } = action.payload;
+      if (!state.canvasDiscardedQueueItems.includes(itemId)) {
+        state.canvasDiscardedQueueItems.push(itemId);
+      }
+    },
     canvasSessionIdChanged: (state, action: PayloadAction<{ id: string }>) => {
       const { id } = action.payload;
       state.canvasSessionId = id;
+      state.canvasDiscardedQueueItems = [];
     },
     canvasSessionReset: (state) => {
       state.canvasSessionId = null;
+      state.canvasDiscardedQueueItems = [];
     },
   },
   extraReducers(builder) {
@@ -41,8 +51,13 @@ export const canvasSessionSlice = createSlice({
   },
 });
 
-export const { generateSessionIdChanged, generateSessionReset, canvasSessionIdChanged, canvasSessionReset } =
-  canvasSessionSlice.actions;
+export const {
+  generateSessionIdChanged,
+  generateSessionReset,
+  canvasSessionIdChanged,
+  canvasSessionReset,
+  canvasQueueItemDiscarded,
+} = canvasSessionSlice.actions;
 
 /* eslint-disable-next-line @typescript-eslint/no-explicit-any */
 const migrate = (state: any): any => {
@@ -64,3 +79,7 @@ export const selectGenerateSessionId = createSelector(
   ({ generateSessionId }) => generateSessionId
 );
 export const selectIsStaging = createSelector(selectCanvasSessionId, (canvasSessionId) => canvasSessionId !== null);
+export const selectDiscardedItems = createSelector(
+  selectCanvasSessionSlice,
+  ({ canvasDiscardedQueueItems }) => canvasDiscardedQueueItems
+);