feat(ui): optimize reconnect queries

Add `FetchOnReconnect` tag, tagging relevant queries with it. This tag is invalidated in the socketConnected listener, when it is determined that the queue changed.

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

@@ -25,8 +25,11 @@ export const addSocketConnectedEventListener = () => {
       /**
        * The rest of this listener has recovery logic for when the socket disconnects and reconnects.
        *
-       * If the queue totals have changed while we were disconnected, re-fetch everything, updating
-       * the gallery and queue to recover.
+       * We need to re-fetch if something has changed while we were disconnected. In practice, the only
+       * thing that could change while disconnected is a queue item finishes processing.
+       *
+       * The queue status is a proxy for this - if the queue status has changed, we need to re-fetch
+       * the queries that may have changed while we were disconnected.
        */
 
       // Bail on the recovery logic if this is the first connection - we don't need to recover anything
@@ -35,10 +38,8 @@ export const addSocketConnectedEventListener = () => {
         return;
       }
 
-      /**
-       * Else, we need to compare the last-known queue status with the current queue status, re-fetching
-       * everything if it has changed.
-       */
+      // Else, we need to compare the last-known queue status with the current queue status, re-fetching
+      // everything if it has changed.
 
       if ($baseUrl.get()) {
         // If we have a baseUrl (e.g. not localhost), we need to debounce the re-fetch to not hammer server
@@ -59,32 +60,14 @@ export const addSocketConnectedEventListener = () => {
         const nextQueueStatusData = await queueStatusRequest.unwrap();
         queueStatusRequest.unsubscribe();
 
-        // If the queue hasn't changed, we don't need to recover
+        // If the queue hasn't changed, we don't need to do anything.
         if (isEqual(prevQueueStatusData?.queue, nextQueueStatusData.queue)) {
           return;
         }
 
-        /**
-         * The queue has changed. We need to reset the API state to update everything and recover
-         * from the disconnect.
-         *
-         * TODO: This is rather inefficient. We don't actually need to re-fetch *all* queries, but
-         * determining which queries to re-fetch and how to re-initialize them is non-trivial:
-         *
-         * - We need to keep track of which queries might have different data in this scenario. This
-         * could be handled via tags, but it feels risky - if we miss tagging a critical query, we
-         * could end up with a de-sync'd UI.
-         *
-         * - We need to re-initialize the queries with *the right query args*. This is very tricky,
-         * because the query args are not stored in the API state, but rather in the component state.
-         *
-         * By totally resetting the API state, we also re-fetch things like model lists, which is
-         * probably a good idea anyways.
-         *
-         * PS: RTKQ provides a related abstraction for recovery:
-         * https://redux-toolkit.js.org/rtk-query/api/setupListeners
-         */
-        dispatch(api.util.resetApiState());
+        //The queue has changed. We need to re-fetch everything that may have changed while we were
+        // disconnected.
+        dispatch(api.util.invalidateTags(['FetchOnReconnect']));
       } catch {
         // no-op
         log.debug('Unable to get current queue status on reconnect');

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

@@ -29,7 +29,7 @@ export const appInfoApi = api.injectEndpoints({
         url: `app/invocation_cache/status`,
         method: 'GET',
       }),
-      providesTags: ['InvocationCacheStatus'],
+      providesTags: ['InvocationCacheStatus', 'FetchOnReconnect'],
     }),
     clearInvocationCache: build.mutation<void, void>({
       query: () => ({

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

@@ -23,7 +23,10 @@ export const boardsApi = api.injectEndpoints({
       query: (arg) => ({ url: 'boards/', params: arg }),
       providesTags: (result) => {
         // any list of boards
-        const tags: ApiTagDescription[] = [{ type: 'Board', id: LIST_TAG }];
+        const tags: ApiTagDescription[] = [
+          { type: 'Board', id: LIST_TAG },
+          'FetchOnReconnect',
+        ];
 
         if (result) {
           // and individual tags for each board
@@ -46,7 +49,10 @@ export const boardsApi = api.injectEndpoints({
       }),
       providesTags: (result) => {
         // any list of boards
-        const tags: ApiTagDescription[] = [{ type: 'Board', id: LIST_TAG }];
+        const tags: ApiTagDescription[] = [
+          { type: 'Board', id: LIST_TAG },
+          'FetchOnReconnect',
+        ];
 
         if (result) {
           // and individual tags for each board
@@ -68,6 +74,7 @@ export const boardsApi = api.injectEndpoints({
       }),
       providesTags: (result, error, arg) => [
         { type: 'ImageNameList', id: arg },
+        'FetchOnReconnect',
       ],
       keepUnusedDataFor: 0,
     }),
@@ -85,6 +92,7 @@ export const boardsApi = api.injectEndpoints({
       }),
       providesTags: (result, error, arg) => [
         { type: 'BoardImagesTotal', id: arg ?? 'none' },
+        'FetchOnReconnect',
       ],
       transformResponse: (response: OffsetPaginatedResults_ImageDTO_) => {
         return { total: response.total };
@@ -104,6 +112,7 @@ export const boardsApi = api.injectEndpoints({
       }),
       providesTags: (result, error, arg) => [
         { type: 'BoardAssetsTotal', id: arg ?? 'none' },
+        'FetchOnReconnect',
       ],
       transformResponse: (response: OffsetPaginatedResults_ImageDTO_) => {
         return { total: response.total };

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

@@ -47,6 +47,7 @@ export const imagesApi = api.injectEndpoints({
       providesTags: (result, error, { board_id, categories }) => [
         // Make the tags the same as the cache key
         { type: 'ImageList', id: getListImagesUrl({ board_id, categories }) },
+        'FetchOnReconnect',
       ],
       serializeQueryArgs: ({ queryArgs }) => {
         // Create cache & key based on board_id and categories - skip the other args.
@@ -100,7 +101,7 @@ export const imagesApi = api.injectEndpoints({
     }),
     getIntermediatesCount: build.query<number, void>({
       query: () => ({ url: 'images/intermediates' }),
-      providesTags: ['IntermediatesCount'],
+      providesTags: ['IntermediatesCount', 'FetchOnReconnect'],
     }),
     clearIntermediates: build.mutation<number, void>({
       query: () => ({ url: `images/intermediates`, method: 'DELETE' }),

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

@@ -164,7 +164,10 @@ export const queueApi = api.injectEndpoints({
         method: 'GET',
       }),
       providesTags: (result) => {
-        const tags: ApiTagDescription[] = ['CurrentSessionQueueItem'];
+        const tags: ApiTagDescription[] = [
+          'CurrentSessionQueueItem',
+          'FetchOnReconnect',
+        ];
         if (result) {
           tags.push({ type: 'SessionQueueItem', id: result.item_id });
         }
@@ -180,7 +183,10 @@ export const queueApi = api.injectEndpoints({
         method: 'GET',
       }),
       providesTags: (result) => {
-        const tags: ApiTagDescription[] = ['NextSessionQueueItem'];
+        const tags: ApiTagDescription[] = [
+          'NextSessionQueueItem',
+          'FetchOnReconnect',
+        ];
         if (result) {
           tags.push({ type: 'SessionQueueItem', id: result.item_id });
         }
@@ -195,7 +201,7 @@ export const queueApi = api.injectEndpoints({
         url: `queue/${$queueId.get()}/status`,
         method: 'GET',
       }),
-      providesTags: ['SessionQueueStatus'],
+      providesTags: ['SessionQueueStatus', 'FetchOnReconnect'],
     }),
     getBatchStatus: build.query<
       paths['/api/v1/queue/{queue_id}/b/{batch_id}/status']['get']['responses']['200']['content']['application/json'],
@@ -206,10 +212,11 @@ export const queueApi = api.injectEndpoints({
         method: 'GET',
       }),
       providesTags: (result) => {
-        if (!result) {
-          return [];
+        const tags: ApiTagDescription[] = ['FetchOnReconnect'];
+        if (result) {
+          tags.push({ type: 'BatchStatus', id: result.batch_id });
         }
-        return [{ type: 'BatchStatus', id: result.batch_id }];
+        return tags;
       },
     }),
     getQueueItem: build.query<
@@ -221,10 +228,11 @@ export const queueApi = api.injectEndpoints({
         method: 'GET',
       }),
       providesTags: (result) => {
-        if (!result) {
-          return [];
+        const tags: ApiTagDescription[] = ['FetchOnReconnect'];
+        if (result) {
+          tags.push({ type: 'SessionQueueItem', id: result.item_id });
         }
-        return [{ type: 'SessionQueueItem', id: result.item_id }];
+        return tags;
       },
     }),
     cancelQueueItem: build.mutation<
@@ -319,6 +327,7 @@ export const queueApi = api.injectEndpoints({
       },
       forceRefetch: ({ currentArg, previousArg }) => currentArg !== previousArg,
       keepUnusedDataFor: 60 * 5, // 5 minutes
+      providesTags: ['FetchOnReconnect'],
     }),
   }),
 });

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

@@ -14,6 +14,9 @@ export const utilitiesApi = api.injectEndpoints({
         method: 'POST',
       }),
       keepUnusedDataFor: 86400, // 24 hours
+      // We need to fetch this on reconnect bc the user may have changed the text field while
+      // disconnected.
+      providesTags: ['FetchOnReconnect'],
     }),
   }),
 });

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

@@ -11,6 +11,7 @@ export const workflowsApi = api.injectEndpoints({
       query: (workflow_id) => `workflows/i/${workflow_id}`,
       providesTags: (result, error, workflow_id) => [
         { type: 'Workflow', id: workflow_id },
+        'FetchOnReconnect',
       ],
       onQueryStarted: async (arg, api) => {
         const { dispatch, queryFulfilled } = api;
@@ -74,7 +75,7 @@ export const workflowsApi = api.injectEndpoints({
         url: 'workflows/',
         params,
       }),
-      providesTags: [{ type: 'Workflow', id: LIST_TAG }],
+      providesTags: ['FetchOnReconnect', { type: 'Workflow', id: LIST_TAG }],
     }),
   }),
 });

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

@@ -40,6 +40,9 @@ export const tagTypes = [
   'SDXLRefinerModel',
   'Workflow',
   'WorkflowsRecent',
+  // This is invalidated on reconnect. It should be used for queries that have changing data,
+  // especially related to the queue and generation.
+  'FetchOnReconnect',
 ] as const;
 export type ApiTagDescription = TagDescription<(typeof tagTypes)[number]>;
 export const LIST_TAG = 'LIST';