fix(ui): reduce reconnect requests

- Add checks to the "recovery" logic for socket connect events to reduce the number of network requests.
- Remove the `isInitialized` state from `systemSlice` and make it a nanostore local to the socketConnected listener. It didn't need to be global state. It's also now more clearly named `isFirstConnection`.
- Export the queue status selector (minor improvement, memoizes it correctly).

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

@@ -1,4 +1,4 @@
-import { queueApi } from 'services/api/endpoints/queue';
+import { queueApi, selectQueueStatus } from 'services/api/endpoints/queue';
 
 import { startAppListening } from '..';
 
@@ -6,7 +6,7 @@ export const addAnyEnqueuedListener = () => {
   startAppListening({
     matcher: queueApi.endpoints.enqueueBatch.matchFulfilled,
     effect: async (_, { dispatch, getState }) => {
-      const { data } = queueApi.endpoints.getQueueStatus.select()(getState());
+      const { data } = selectQueueStatus(getState());
 
       if (!data || data.processor.is_started) {
         return;

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

@@ -1,7 +1,9 @@
 import { logger } from 'app/logging/logger';
-import { isInitializedChanged } from 'features/system/store/systemSlice';
+import { $baseUrl } from 'app/store/nanostores/baseUrl';
 import { size } from 'lodash-es';
+import { atom } from 'nanostores';
 import { api } from 'services/api';
+import { queueApi, selectQueueStatus } from 'services/api/endpoints/queue';
 import { receivedOpenAPISchema } from 'services/api/thunks/schema';
 import { socketConnected } from 'services/events/actions';
 
@@ -9,25 +11,95 @@ import { startAppListening } from '../..';
 
 const log = logger('socketio');
 
+const $isFirstConnection = atom(true);
+
 export const addSocketConnectedEventListener = () => {
   startAppListening({
     actionCreator: socketConnected,
-    effect: (action, { dispatch, getState }) => {
+    effect: async (
+      action,
+      { dispatch, getState, cancelActiveListeners, delay }
+    ) => {
       log.debug('Connected');
 
-      const { nodeTemplates, config, system } = getState();
+      /**
+       * The rest of this listener has recovery logic for when the socket disconnects and reconnects.
+       * If we had generations in progress, we need to reset the API state to re-fetch everything.
+       */
 
-      const { disabledTabs } = config;
+      if ($isFirstConnection.get()) {
-
+        // Bail on the recovery logic if this is the first connection
-      if (!size(nodeTemplates.templates) && !disabledTabs.includes('nodes')) {
+        $isFirstConnection.set(false);
-        dispatch(receivedOpenAPISchema());
+        return;
       }
 
-      if (system.isInitialized) {
+      const { data: prevQueueStatusData } = selectQueueStatus(getState());
-        // only reset the query caches if this connect event is a *reconnect* event
+
+      // Bail if queue was empty before - we have nothing to recover
+      if (
+        !prevQueueStatusData?.queue.in_progress &&
+        !prevQueueStatusData?.queue.pending
+      ) {
+        return;
+      }
+
+      // Else, we need to re-fetch the queue status to see 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
+        cancelActiveListeners();
+        // Add artificial jitter to the debounce
+        await delay(1000 + Math.random() * 1000);
+      }
+
+      try {
+        // Fetch the queue status again
+        const queueStatusRequest = dispatch(
+          await queueApi.endpoints.getQueueStatus.initiate(undefined, {
+            forceRefetch: true,
+          })
+        );
+
+        const nextQueueStatusData = await queueStatusRequest.unwrap();
+        queueStatusRequest.unsubscribe();
+
+        // If we haven't completed more items since the last check, bail
+        if (
+          prevQueueStatusData.queue.completed ===
+          nextQueueStatusData.queue.completed
+        ) {
+          return;
+        }
+
+        /**
+         * Else, we need to reset the API state to update everything.
+         * 
+         * TODO: This is rather inefficient. We don't actually need to re-fetch *all* network requests,
+         * but determining which ones to re-fetch is non-trivial. It's at least the queue related ones
+         * and gallery, but likely others. We'd also need to keep track of which requests need to be
+         * re-fetch in this situation, which opens the door for bugs.
+         * 
+         * Optimize this later.
+         */
         dispatch(api.util.resetApiState());
-      } else {
+      } catch {
-        dispatch(isInitializedChanged(true));
+        // no-op
+        log.debug('Unable to get current queue status on reconnect');
+      }
+    },
+  });
+
+  startAppListening({
+    actionCreator: socketConnected,
+    effect: async (action, { dispatch, getState }) => {
+      const { nodeTemplates, config } = getState();
+      if (
+        !size(nodeTemplates.templates) &&
+        !config.disabledTabs.includes('nodes')
+      ) {
+        // This request is a createAsyncThunk - resetting API state as in the above listener
+        // will not trigger this request, so we need to manually do it.
+        dispatch(receivedOpenAPISchema());
       }
     },
   });

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

@@ -26,7 +26,6 @@ import type { Language, SystemState } from './types';
 
 export const initialSystemState: SystemState = {
   _version: 1,
-  isInitialized: false,
   isConnected: false,
   shouldConfirmOnDelete: true,
   enableImageDebugging: false,
@@ -85,9 +84,6 @@ export const systemSlice = createSlice({
     ) {
       state.shouldEnableInformationalPopovers = action.payload;
     },
-    isInitializedChanged(state, action: PayloadAction<boolean>) {
-      state.isInitialized = action.payload;
-    },
   },
   extraReducers(builder) {
     /**
@@ -206,7 +202,6 @@ export const {
   shouldUseNSFWCheckerChanged,
   shouldUseWatermarkerChanged,
   setShouldEnableInformationalPopovers,
-  isInitializedChanged,
 } = systemSlice.actions;
 
 export default systemSlice.reducer;

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

@@ -44,7 +44,6 @@ export const isLanguage = (v: unknown): v is Language =>
 
 export interface SystemState {
   _version: 1;
-  isInitialized: boolean;
   isConnected: boolean;
   shouldConfirmOnDelete: boolean;
   enableImageDebugging: boolean;

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

@@ -339,6 +339,8 @@ export const {
   useGetBatchStatusQuery,
 } = queueApi;
 
+export const selectQueueStatus = queueApi.endpoints.getQueueStatus.select();
+
 const resetListQueryData = (
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   dispatch: ThunkDispatch<any, any, UnknownAction>