diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/anyEnqueued.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/anyEnqueued.ts index fa4ad727c4..e2114769f4 100644 --- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/anyEnqueued.ts +++ b/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; diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketConnected.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketConnected.ts index 589d312a9e..755bebeb24 100644 --- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/socketio/socketConnected.ts +++ b/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 (!size(nodeTemplates.templates) && !disabledTabs.includes('nodes')) { - dispatch(receivedOpenAPISchema()); + if ($isFirstConnection.get()) { + // Bail on the recovery logic if this is the first connection + $isFirstConnection.set(false); + return; } - if (system.isInitialized) { - // only reset the query caches if this connect event is a *reconnect* event + const { data: prevQueueStatusData } = selectQueueStatus(getState()); + + // 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 { - dispatch(isInitializedChanged(true)); + } catch { + // 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()); } }, }); diff --git a/invokeai/frontend/web/src/features/system/store/systemSlice.ts b/invokeai/frontend/web/src/features/system/store/systemSlice.ts index 1f07863947..1307bb00aa 100644 --- a/invokeai/frontend/web/src/features/system/store/systemSlice.ts +++ b/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) { - state.isInitialized = action.payload; - }, }, extraReducers(builder) { /** @@ -206,7 +202,6 @@ export const { shouldUseNSFWCheckerChanged, shouldUseWatermarkerChanged, setShouldEnableInformationalPopovers, - isInitializedChanged, } = systemSlice.actions; export default systemSlice.reducer; diff --git a/invokeai/frontend/web/src/features/system/store/types.ts b/invokeai/frontend/web/src/features/system/store/types.ts index 6ee9b4503c..e8e394b950 100644 --- a/invokeai/frontend/web/src/features/system/store/types.ts +++ b/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; diff --git a/invokeai/frontend/web/src/services/api/endpoints/queue.ts b/invokeai/frontend/web/src/services/api/endpoints/queue.ts index 6c8f2875ae..4e206bb354 100644 --- a/invokeai/frontend/web/src/services/api/endpoints/queue.ts +++ b/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