fix(ui): fix canvas soft-lock if canceled before first generation

The canvas needs to be set to staging mode as soon as a canvas-destined batch is enqueued. If the batch is is fully canceled before an image is generated, we need to remove that batch from the canvas `batchIds` watchlist, else canvas gets stuck in staging mode with no way to exit.

The changes here allow the batch status to be tracked, and if a batch has all its items completed, we can remove it from the `batchIds` watchlist. The `batchIds` watchlist now accurately represents *incomplete* canvas batches, fixing this cause of soft lock.

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketQueueItemStatusChanged.ts

@@ -11,44 +11,66 @@ export const addSocketQueueItemStatusChangedEventListener = () => {
     actionCreator: socketQueueItemStatusChanged,
     effect: async (action, { dispatch }) => {
       const log = logger('socketio');
-      const {
+
-        queue_item_id: item_id,
+      const { queue_item, batch_status, queue_status } = action.payload.data;
-        queue_batch_id,
+
-        status,
-      } = action.payload.data;
       log.debug(
         action.payload,
-        `Queue item ${item_id} status updated: ${status}`
+        `Queue item ${queue_item.item_id} status updated: ${queue_item.status}`
       );
       dispatch(appSocketQueueItemStatusChanged(action.payload));
 
       dispatch(
         queueApi.util.updateQueryData('listQueueItems', undefined, (draft) => {
           queueItemsAdapter.updateOne(draft, {
-            id: item_id,
+            id: queue_item.item_id,
-            changes: action.payload.data,
+            changes: queue_item,
           });
         })
       );
 
+      dispatch(
+        queueApi.util.updateQueryData('getQueueStatus', undefined, (draft) => {
+          Object.assign(draft.queue, queue_status);
+        })
+      );
+      dispatch(
+        queueApi.util.updateQueryData('getQueueStatus', undefined, (draft) => {
+          if (!draft) {
+            return;
+          }
+          Object.assign(draft.queue, queue_status);
+        })
+      );
+
+      dispatch(
+        queueApi.util.updateQueryData(
+          'getBatchStatus',
+          { batch_id: batch_status.batch_id },
+          () => batch_status
+        )
+      );
+
+      dispatch(
+        queueApi.util.updateQueryData(
+          'getQueueItem',
+          queue_item.item_id,
+          (draft) => {
+            if (!draft) {
+              return;
+            }
+            Object.assign(draft, queue_item);
+          }
+        )
+      );
+
       dispatch(
         queueApi.util.invalidateTags([
           'CurrentSessionQueueItem',
           'NextSessionQueueItem',
           'InvocationCacheStatus',
-          { type: 'SessionQueueItem', id: item_id },
-          { type: 'SessionQueueItemDTO', id: item_id },
-          { type: 'BatchStatus', id: queue_batch_id },
         ])
       );
-
-      const req = dispatch(
-        queueApi.endpoints.getQueueStatus.initiate(undefined, {
-          forceRefetch: true,
-        })
-      );
-      await req.unwrap();
-      req.unsubscribe();
     },
   });
 };

invokeai/frontend/web/src/features/canvas/store/canvasSlice.ts

@@ -29,6 +29,7 @@ import {
   isCanvasBaseImage,
   isCanvasMaskLine,
 } from './canvasTypes';
+import { appSocketQueueItemStatusChanged } from 'services/events/actions';
 
 export const initialLayerState: CanvasLayerState = {
   objects: [],
@@ -786,6 +787,18 @@ export const canvasSlice = createSlice({
     },
   },
   extraReducers: (builder) => {
+    builder.addCase(appSocketQueueItemStatusChanged, (state, action) => {
+      const batch_status = action.payload.data.batch_status;
+      if (!state.batchIds.includes(batch_status.batch_id)) {
+        return;
+      }
+
+      if (batch_status.in_progress === 0 && batch_status.pending === 0) {
+        state.batchIds = state.batchIds.filter(
+          (id) => id !== batch_status.batch_id
+        );
+      }
+    });
     builder.addCase(setAspectRatio, (state, action) => {
       const ratio = action.payload;
       if (ratio) {

invokeai/frontend/web/src/features/nodes/store/nodesSlice.ts

@@ -984,7 +984,7 @@ const nodesSlice = createSlice({
       }
     });
     builder.addCase(appSocketQueueItemStatusChanged, (state, action) => {
-      if (['in_progress'].includes(action.payload.data.status)) {
+      if (['in_progress'].includes(action.payload.data.queue_item.status)) {
         forEach(state.nodeExecutionStates, (nes) => {
           nes.status = NodeStatus.PENDING;
           nes.error = null;

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

@@ -164,7 +164,9 @@ export const systemSlice = createSlice({
 
     builder.addCase(appSocketQueueItemStatusChanged, (state, action) => {
       if (
-        ['completed', 'canceled', 'failed'].includes(action.payload.data.status)
+        ['completed', 'canceled', 'failed'].includes(
+          action.payload.data.queue_item.status
+        )
       ) {
         state.status = 'CONNECTED';
         state.denoiseProgress = null;

invokeai/frontend/web/src/services/api/endpoints/queue.ts

@@ -135,12 +135,7 @@ export const queueApi = api.injectEndpoints({
         url: `queue/${$queueId.get()}/prune`,
         method: 'PUT',
       }),
-      invalidatesTags: [
+      invalidatesTags: ['SessionQueueStatus', 'BatchStatus'],
-        'SessionQueueStatus',
-        'BatchStatus',
-        'SessionQueueItem',
-        'SessionQueueItemDTO',
-      ],
       onQueryStarted: async (arg, api) => {
         const { dispatch, queryFulfilled } = api;
         try {
@@ -165,8 +160,6 @@ export const queueApi = api.injectEndpoints({
         'BatchStatus',
         'CurrentSessionQueueItem',
         'NextSessionQueueItem',
-        'SessionQueueItem',
-        'SessionQueueItemDTO',
       ],
       onQueryStarted: async (arg, api) => {
         const { dispatch, queryFulfilled } = api;
@@ -218,7 +211,6 @@ export const queueApi = api.injectEndpoints({
         url: `queue/${$queueId.get()}/status`,
         method: 'GET',
       }),
-
       providesTags: ['SessionQueueStatus'],
     }),
     getBatchStatus: build.query<
@@ -269,7 +261,11 @@ export const queueApi = api.injectEndpoints({
               (draft) => {
                 queueItemsAdapter.updateOne(draft, {
                   id: item_id,
-                  changes: { status: data.status },
+                  changes: {
+                    status: data.status,
+                    completed_at: data.completed_at,
+                    updated_at: data.updated_at,
+                  },
                 });
               }
             )
@@ -284,7 +280,6 @@ export const queueApi = api.injectEndpoints({
         }
         return [
           { type: 'SessionQueueItem', id: result.item_id },
-          { type: 'SessionQueueItemDTO', id: result.item_id },
           { type: 'BatchStatus', id: result.batch_id },
         ];
       },
@@ -307,11 +302,7 @@ export const queueApi = api.injectEndpoints({
           // no-op
         }
       },
-      invalidatesTags: [
+      invalidatesTags: ['SessionQueueStatus', 'BatchStatus'],
-        'SessionQueueItem',
-        'SessionQueueItemDTO',
-        'BatchStatus',
-      ],
     }),
     listQueueItems: build.query<
       EntityState<components['schemas']['SessionQueueItemDTO']> & {

invokeai/frontend/web/src/services/api/index.ts

@@ -21,8 +21,6 @@ export const tagTypes = [
   'ImageMetadataFromFile',
   'IntermediatesCount',
   'SessionQueueItem',
-  'SessionQueueItemDTO',
-  'SessionQueueItemDTOList',
   'SessionQueueStatus',
   'SessionProcessorStatus',
   'CurrentSessionQueueItem',

invokeai/frontend/web/src/services/api/schema.d.ts

@@ -9701,11 +9701,23 @@ export type components = {
       ui_order?: number;
     };
     /**
-     * StableDiffusionOnnxModelFormat
+     * T2IAdapterModelFormat
      * @description An enumeration.
      * @enum {string}
      */
-    StableDiffusionOnnxModelFormat: "olive" | "onnx";
+    T2IAdapterModelFormat: "diffusers";
+    /**
+     * ControlNetModelFormat
+     * @description An enumeration.
+     * @enum {string}
+     */
+    ControlNetModelFormat: "checkpoint" | "diffusers";
+    /**
+     * StableDiffusion2ModelFormat
+     * @description An enumeration.
+     * @enum {string}
+     */
+    StableDiffusion2ModelFormat: "checkpoint" | "diffusers";
     /**
      * StableDiffusionXLModelFormat
      * @description An enumeration.
@@ -9713,11 +9725,11 @@ export type components = {
      */
     StableDiffusionXLModelFormat: "checkpoint" | "diffusers";
     /**
-     * StableDiffusion2ModelFormat
+     * StableDiffusionOnnxModelFormat
      * @description An enumeration.
      * @enum {string}
      */
-    StableDiffusion2ModelFormat: "checkpoint" | "diffusers";
+    StableDiffusionOnnxModelFormat: "olive" | "onnx";
     /**
      * CLIPVisionModelFormat
      * @description An enumeration.
@@ -9736,18 +9748,6 @@ export type components = {
      * @enum {string}
      */
     StableDiffusion1ModelFormat: "checkpoint" | "diffusers";
-    /**
-     * T2IAdapterModelFormat
-     * @description An enumeration.
-     * @enum {string}
-     */
-    T2IAdapterModelFormat: "diffusers";
-    /**
-     * ControlNetModelFormat
-     * @description An enumeration.
-     * @enum {string}
-     */
-    ControlNetModelFormat: "checkpoint" | "diffusers";
   };
   responses: never;
   parameters: never;

invokeai/frontend/web/src/services/events/types.ts

@@ -170,10 +170,11 @@ export type InvocationRetrievalErrorEvent = {
  */
 export type QueueItemStatusChangedEvent = {
   queue_id: string;
-  queue_item_id: number;
+  queue_item: {
-  queue_batch_id: string;
+    queue_id: string;
+    item_id: number;
+    batch_id: string;
     session_id: string;
-  graph_execution_state_id: string;
     status: components['schemas']['SessionQueueItemDTO']['status'];
     error: string | undefined;
     created_at: string;
@@ -181,6 +182,29 @@ export type QueueItemStatusChangedEvent = {
     started_at: string | undefined;
     completed_at: string | undefined;
   };
+  batch_status: {
+    queue_id: string;
+    batch_id: string;
+    pending: number;
+    in_progress: number;
+    completed: number;
+    failed: number;
+    canceled: number;
+    total: number;
+  };
+  queue_status: {
+    queue_id: string;
+    item_id?: number;
+    batch_id?: string;
+    session_id?: string;
+    pending: number;
+    in_progress: number;
+    completed: number;
+    failed: number;
+    canceled: number;
+    total: number;
+  };
+};
 
 export type ClientEmitSubscribeQueue = {
   queue_id: string;