From a7cebbd970f821d72371cfb036841d848f4bfe65 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Mon, 29 May 2023 19:06:37 +1000
Subject: [PATCH] feat(ui): cancel session when staging image accepted

---
 .../middleware/listenerMiddleware/index.ts    |  2 +
 .../addCommitStagingAreaImageListener.ts      | 37 +++++++++++++++++++
 .../IAICanvasStagingAreaToolbar.tsx           | 21 +++++++++--
 .../src/features/canvas/store/canvasSlice.ts  |  5 ++-
 4 files changed, 60 insertions(+), 5 deletions(-)
 create mode 100644 invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/addCommitStagingAreaImageListener.ts

diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts
index cf4544e4ea..26f32252d1 100644
--- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts
+++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts
@@ -68,6 +68,7 @@ import {
   addReceivedPageOfImagesRejectedListener,
 } from './listeners/receivedPageOfImages';
 import { addStagingAreaImageSavedListener } from './listeners/stagingAreaImageSaved';
+import { addCommitStagingAreaImageListener } from './listeners/addCommitStagingAreaImageListener';
 
 export const listenerMiddleware = createListenerMiddleware();
 
@@ -125,6 +126,7 @@ addCanvasDownloadedAsImageListener();
 addCanvasCopiedToClipboardListener();
 addCanvasMergedListener();
 addStagingAreaImageSavedListener();
+addCommitStagingAreaImageListener();
 
 /**
  * Socket.IO Events - these handle SIO events directly and pass on internal application actions.
diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/addCommitStagingAreaImageListener.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/addCommitStagingAreaImageListener.ts
new file mode 100644
index 0000000000..428ecf9c62
--- /dev/null
+++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/addCommitStagingAreaImageListener.ts
@@ -0,0 +1,37 @@
+import { startAppListening } from '..';
+import { log } from 'app/logging/useLogger';
+import { commitStagingAreaImage } from 'features/canvas/store/canvasSlice';
+import { sessionCanceled } from 'services/thunks/session';
+
+const moduleLog = log.child({ namespace: 'canvas' });
+
+export const addCommitStagingAreaImageListener = () => {
+  startAppListening({
+    actionCreator: commitStagingAreaImage,
+    effect: async (action, { dispatch, getState }) => {
+      const state = getState();
+      const { sessionId } = state.system;
+      const canvasSessionId = action.payload;
+
+      if (!canvasSessionId) {
+        moduleLog.debug('No canvas session, skipping cancel');
+        return;
+      }
+
+      if (canvasSessionId !== sessionId) {
+        moduleLog.debug(
+          {
+            data: {
+              canvasSessionId,
+              sessionId,
+            },
+          },
+          'Canvas session does not match global session, skipping cancel'
+        );
+        return;
+      }
+
+      dispatch(sessionCanceled({ sessionId }));
+    },
+  });
+};
diff --git a/invokeai/frontend/web/src/features/canvas/components/IAICanvasStagingAreaToolbar.tsx b/invokeai/frontend/web/src/features/canvas/components/IAICanvasStagingAreaToolbar.tsx
index 68bc15bbaa..76ffdcf082 100644
--- a/invokeai/frontend/web/src/features/canvas/components/IAICanvasStagingAreaToolbar.tsx
+++ b/invokeai/frontend/web/src/features/canvas/components/IAICanvasStagingAreaToolbar.tsx
@@ -32,7 +32,7 @@ const selector = createSelector(
   (canvas) => {
     const {
       layerState: {
-        stagingArea: { images, selectedImageIndex },
+        stagingArea: { images, selectedImageIndex, sessionId },
       },
       shouldShowStagingOutline,
       shouldShowStagingImage,
@@ -45,6 +45,7 @@ const selector = createSelector(
       isOnLastImage: selectedImageIndex === images.length - 1,
       shouldShowStagingImage,
       shouldShowStagingOutline,
+      sessionId,
     };
   },
   {
@@ -61,6 +62,7 @@ const IAICanvasStagingAreaToolbar = () => {
     isOnLastImage,
     currentStagingAreaImage,
     shouldShowStagingImage,
+    sessionId,
   } = useAppSelector(selector);
 
   const { t } = useTranslation();
@@ -106,9 +108,20 @@ const IAICanvasStagingAreaToolbar = () => {
     }
   );
 
-  const handlePrevImage = () => dispatch(prevStagingAreaImage());
-  const handleNextImage = () => dispatch(nextStagingAreaImage());
-  const handleAccept = () => dispatch(commitStagingAreaImage());
+  const handlePrevImage = useCallback(
+    () => dispatch(prevStagingAreaImage()),
+    [dispatch]
+  );
+
+  const handleNextImage = useCallback(
+    () => dispatch(nextStagingAreaImage()),
+    [dispatch]
+  );
+
+  const handleAccept = useCallback(
+    () => dispatch(commitStagingAreaImage(sessionId)),
+    [dispatch, sessionId]
+  );
 
   if (!currentStagingAreaImage) return null;
 
diff --git a/invokeai/frontend/web/src/features/canvas/store/canvasSlice.ts b/invokeai/frontend/web/src/features/canvas/store/canvasSlice.ts
index 0ebe5b264c..ad0581e42f 100644
--- a/invokeai/frontend/web/src/features/canvas/store/canvasSlice.ts
+++ b/invokeai/frontend/web/src/features/canvas/store/canvasSlice.ts
@@ -696,7 +696,10 @@ export const canvasSlice = createSlice({
         0
       );
     },
-    commitStagingAreaImage: (state) => {
+    commitStagingAreaImage: (
+      state,
+      action: PayloadAction<string | undefined>
+    ) => {
       if (!state.layerState.stagingArea.images.length) {
         return;
       }