From bc3550f23820540eee8c949a7fa637660d589ae0 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Thu, 5 Jun 2025 20:03:47 +1000
Subject: [PATCH] feat(ui): finish generation when discarding last item

---
 .../StagingAreaToolbarDiscardSelectedButton.tsx          | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaToolbarDiscardSelectedButton.tsx b/invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaToolbarDiscardSelectedButton.tsx
index 7006a30aa3..e1c0738fe9 100644
--- a/invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaToolbarDiscardSelectedButton.tsx
+++ b/invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaToolbarDiscardSelectedButton.tsx
@@ -1,12 +1,15 @@
 import { IconButton } from '@invoke-ai/ui-library';
 import { useStore } from '@nanostores/react';
+import { useAppDispatch } from 'app/store/storeHooks';
 import { useCanvasSessionContext } from 'features/controlLayers/components/SimpleSession/context';
+import { canvasSessionGenerationFinished } from 'features/controlLayers/store/canvasStagingAreaSlice';
 import { memo, useCallback } from 'react';
 import { useTranslation } from 'react-i18next';
 import { PiXBold } from 'react-icons/pi';
 import { useDeleteQueueItemMutation } from 'services/api/endpoints/queue';
 
 export const StagingAreaToolbarDiscardSelectedButton = memo(({ isDisabled }: { isDisabled?: boolean }) => {
+  const dispatch = useAppDispatch();
   const ctx = useCanvasSessionContext();
   const [deleteQueueItem] = useDeleteQueueItemMutation();
   const selectedItemId = useStore(ctx.$selectedItemId);
@@ -17,8 +20,12 @@ export const StagingAreaToolbarDiscardSelectedButton = memo(({ isDisabled }: { i
     if (selectedItemId === null) {
       return;
     }
+    const itemCount = ctx.$itemCount.get();
     deleteQueueItem({ item_id: selectedItemId });
-  }, [selectedItemId, deleteQueueItem]);
+    if (itemCount <= 1) {
+      dispatch(canvasSessionGenerationFinished());
+    }
+  }, [selectedItemId, ctx.$itemCount, deleteQueueItem, dispatch]);
 
   return (
     <IconButton