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.

Author: psychedelicious

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'],
 };

invokeai/frontend/web/src/features/system/store/types.ts

@@ -55,4 +55,5 @@ export interface SystemState {
   shouldUseWatermarker: boolean;
   status: SystemStatus;
   shouldEnableInformationalPopovers: boolean;
+  cancellations: string[]
 }