From 0758e9cb9b3290e5af3831596bb512cea476d010 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Fri, 24 May 2024 12:01:02 +1000
Subject: [PATCH] fix(ui): race condition with progress

There's a race condition where a canceled session may emit a progress event or two after it's been canceled, and the progress image isn't cleared out.

To resolve this, the system slice tracks canceled session ids. When a progress event comes in, we check the cancellations and skip setting the progress if canceled.
---
 .../web/src/features/system/store/systemSlice.ts      | 11 ++++++++++-
 .../frontend/web/src/features/system/store/types.ts   |  1 +
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/invokeai/frontend/web/src/features/system/store/systemSlice.ts b/invokeai/frontend/web/src/features/system/store/systemSlice.ts
index 65903460ed..4d87b2c4ec 100644
--- a/invokeai/frontend/web/src/features/system/store/systemSlice.ts
+++ b/invokeai/frontend/web/src/features/system/store/systemSlice.ts
@@ -31,6 +31,7 @@ const initialSystemState: SystemState = {
   shouldUseWatermarker: false,
   shouldEnableInformationalPopovers: false,
   status: 'DISCONNECTED',
+  cancellations: [],
 };
 
 export const systemSlice = createSlice({
@@ -88,6 +89,7 @@ export const systemSlice = createSlice({
      * Invocation Started
      */
     builder.addCase(socketInvocationStarted, (state) => {
+      state.cancellations = [];
       state.denoiseProgress = null;
       state.status = 'PROCESSING';
     });
@@ -105,6 +107,12 @@ export const systemSlice = createSlice({
         queue_batch_id: batch_id,
       } = action.payload.data;
 
+      if (state.cancellations.includes(session_id)) {
+        // Do not update the progress if this session has been cancelled. This prevents a race condition where we get a
+        // progress update after the session has been cancelled.
+        return;
+      }
+
       state.denoiseProgress = {
         step,
         total_steps,
@@ -146,6 +154,7 @@ export const systemSlice = createSlice({
       if (['completed', 'canceled', 'failed'].includes(action.payload.data.queue_item.status)) {
         state.status = 'CONNECTED';
         state.denoiseProgress = null;
+        state.cancellations.push(action.payload.data.queue_item.session_id);
       }
     });
   },
@@ -177,5 +186,5 @@ export const systemPersistConfig: PersistConfig<SystemState> = {
   name: systemSlice.name,
   initialState: initialSystemState,
   migrate: migrateSystemState,
-  persistDenylist: ['isConnected', 'denoiseProgress', 'status'],
+  persistDenylist: ['isConnected', 'denoiseProgress', 'status', 'cancellations'],
 };
diff --git a/invokeai/frontend/web/src/features/system/store/types.ts b/invokeai/frontend/web/src/features/system/store/types.ts
index d8bc8cd08a..d32e347870 100644
--- a/invokeai/frontend/web/src/features/system/store/types.ts
+++ b/invokeai/frontend/web/src/features/system/store/types.ts
@@ -55,4 +55,5 @@ export interface SystemState {
   shouldUseWatermarker: boolean;
   status: SystemStatus;
   shouldEnableInformationalPopovers: boolean;
+  cancellations: string[]
 }