fix(ui): add additional socket event layer to gate handling socket events

Some socket events should not be handled by the slice reducers. For example generation progress should not be handled for a canceled session.

Added another layer of socket actions.

Example:
- `socketGeneratorProgress` is dispatched when the actual socket event is received
- Listener middleware exclusively handles this event and determines if the application should also handle it
- If so, it dispatches `appSocketGeneratorProgress`, which the slices can handle

Needed to fix issues related to canceling invocations.

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/index.ts

@@ -26,15 +26,15 @@ import { addCanvasSavedToGalleryListener } from './listeners/canvasSavedToGaller
 import { addCanvasDownloadedAsImageListener } from './listeners/canvasDownloadedAsImage';
 import { addCanvasCopiedToClipboardListener } from './listeners/canvasCopiedToClipboard';
 import { addCanvasMergedListener } from './listeners/canvasMerged';
-import { addGeneratorProgressListener } from './listeners/socketio/generatorProgress';
-import { addGraphExecutionStateCompleteListener } from './listeners/socketio/graphExecutionStateComplete';
-import { addInvocationCompleteListener } from './listeners/socketio/invocationComplete';
-import { addInvocationErrorListener } from './listeners/socketio/invocationError';
-import { addInvocationStartedListener } from './listeners/socketio/invocationStarted';
-import { addSocketConnectedListener } from './listeners/socketio/socketConnected';
-import { addSocketDisconnectedListener } from './listeners/socketio/socketDisconnected';
-import { addSocketSubscribedListener } from './listeners/socketio/socketSubscribed';
-import { addSocketUnsubscribedListener } from './listeners/socketio/socketUnsubscribed';
+import { addGeneratorProgressEventListener as addGeneratorProgressListener } from './listeners/socketio/socketGeneratorProgress';
+import { addGraphExecutionStateCompleteEventListener as addGraphExecutionStateCompleteListener } from './listeners/socketio/socketGraphExecutionStateComplete';
+import { addInvocationCompleteEventListener as addInvocationCompleteListener } from './listeners/socketio/socketInvocationComplete';
+import { addInvocationErrorEventListener as addInvocationErrorListener } from './listeners/socketio/socketInvocationError';
+import { addInvocationStartedEventListener as addInvocationStartedListener } from './listeners/socketio/socketInvocationStarted';
+import { addSocketConnectedEventListener as addSocketConnectedListener } from './listeners/socketio/socketConnected';
+import { addSocketDisconnectedEventListener as addSocketDisconnectedListener } from './listeners/socketio/socketDisconnected';
+import { addSocketSubscribedEventListener as addSocketSubscribedListener } from './listeners/socketio/socketSubscribed';
+import { addSocketUnsubscribedEventListener as addSocketUnsubscribedListener } from './listeners/socketio/socketUnsubscribed';
 import { addSessionReadyToInvokeListener } from './listeners/sessionReadyToInvoke';
 import {
   addImageMetadataReceivedFulfilledListener,
@@ -126,7 +126,19 @@ addCanvasCopiedToClipboardListener();
 addCanvasMergedListener();
 addStagingAreaImageSavedListener();
 
-// socketio
+/**
+ * Socket.IO Events - these handle SIO events directly and pass on internal application actions.
+ * We don't handle SIO events in slices via `extraReducers` because some of these events shouldn't
+ * actually be handled at all.
+ *
+ * For example, we don't want to respond to progress events for canceled sessions. To avoid
+ * duplicating the logic to determine if an event should be responded to, we handle all of that
+ * "is this session canceled?" logic in these listeners.
+ *
+ * The `socketGeneratorProgress` listener will then only dispatch the `appSocketGeneratorProgress`
+ * action if it should be handled by the rest of the application. It is this `appSocketGeneratorProgress`
+ * action that is handled by reducers in slices.
+ */
 addGeneratorProgressListener();
 addGraphExecutionStateCompleteListener();
 addInvocationCompleteListener();

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

@@ -1,13 +1,13 @@
 import { startAppListening } from '../..';
 import { log } from 'app/logging/useLogger';
-import { socketConnected } from 'services/events/actions';
+import { appSocketConnected, socketConnected } from 'services/events/actions';
 import { receivedPageOfImages } from 'services/thunks/image';
 import { receivedModels } from 'services/thunks/model';
 import { receivedOpenAPISchema } from 'services/thunks/schema';
 
 const moduleLog = log.child({ namespace: 'socketio' });
 
-export const addSocketConnectedListener = () => {
+export const addSocketConnectedEventListener = () => {
   startAppListening({
     actionCreator: socketConnected,
     effect: (action, { dispatch, getState }) => {
@@ -30,6 +30,9 @@ export const addSocketConnectedListener = () => {
       if (!nodes.schema && !disabledTabs.includes('nodes')) {
         dispatch(receivedOpenAPISchema());
       }
+
+      // pass along the socket event as an application action
+      dispatch(appSocketConnected(action.payload));
     },
   });
 };

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

@@ -1,14 +1,19 @@
 import { startAppListening } from '../..';
 import { log } from 'app/logging/useLogger';
-import { socketDisconnected } from 'services/events/actions';
+import {
+  socketDisconnected,
+  appSocketDisconnected,
+} from 'services/events/actions';
 
 const moduleLog = log.child({ namespace: 'socketio' });
 
-export const addSocketDisconnectedListener = () => {
+export const addSocketDisconnectedEventListener = () => {
   startAppListening({
     actionCreator: socketDisconnected,
     effect: (action, { dispatch, getState }) => {
       moduleLog.debug(action.payload, 'Disconnected');
+      // pass along the socket event as an application action
+      dispatch(appSocketDisconnected(action.payload));
     },
   });
 };

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

@@ -1,12 +1,15 @@
 import { startAppListening } from '../..';
 import { log } from 'app/logging/useLogger';
-import { generatorProgress } from 'services/events/actions';
+import {
+  appSocketGeneratorProgress,
+  socketGeneratorProgress,
+} from 'services/events/actions';
 
 const moduleLog = log.child({ namespace: 'socketio' });
 
-export const addGeneratorProgressListener = () => {
+export const addGeneratorProgressEventListener = () => {
   startAppListening({
-    actionCreator: generatorProgress,
+    actionCreator: socketGeneratorProgress,
     effect: (action, { dispatch, getState }) => {
       if (
         getState().system.canceledSession ===
@@ -23,6 +26,9 @@ export const addGeneratorProgressListener = () => {
         action.payload,
         `Generator progress (${action.payload.data.node.type})`
       );
+
+      // pass along the socket event as an application action
+      dispatch(appSocketGeneratorProgress(action.payload));
     },
   });
 };

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

@@ -1,17 +1,22 @@
 import { log } from 'app/logging/useLogger';
-import { graphExecutionStateComplete } from 'services/events/actions';
+import {
+  appSocketGraphExecutionStateComplete,
+  socketGraphExecutionStateComplete,
+} from 'services/events/actions';
 import { startAppListening } from '../..';
 
 const moduleLog = log.child({ namespace: 'socketio' });
 
-export const addGraphExecutionStateCompleteListener = () => {
+export const addGraphExecutionStateCompleteEventListener = () => {
   startAppListening({
-    actionCreator: graphExecutionStateComplete,
+    actionCreator: socketGraphExecutionStateComplete,
     effect: (action, { dispatch, getState }) => {
       moduleLog.debug(
         action.payload,
         `Session invocation complete (${action.payload.data.graph_execution_state_id})`
       );
+      // pass along the socket event as an application action
+      dispatch(appSocketGraphExecutionStateComplete(action.payload));
     },
   });
 };

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

@@ -1,7 +1,10 @@
 import { addImageToStagingArea } from 'features/canvas/store/canvasSlice';
 import { startAppListening } from '../..';
 import { log } from 'app/logging/useLogger';
-import { invocationComplete } from 'services/events/actions';
+import {
+  appSocketInvocationComplete,
+  socketInvocationComplete,
+} from 'services/events/actions';
 import { imageMetadataReceived } from 'services/thunks/image';
 import { sessionCanceled } from 'services/thunks/session';
 import { isImageOutput } from 'services/types/guards';
@@ -10,9 +13,9 @@ import { progressImageSet } from 'features/system/store/systemSlice';
 const moduleLog = log.child({ namespace: 'socketio' });
 const nodeDenylist = ['dataURL_image'];
 
-export const addInvocationCompleteListener = () => {
+export const addInvocationCompleteEventListener = () => {
   startAppListening({
-    actionCreator: invocationComplete,
+    actionCreator: socketInvocationComplete,
     effect: async (action, { dispatch, getState, take }) => {
       moduleLog.debug(
         action.payload,
@@ -57,6 +60,8 @@ export const addInvocationCompleteListener = () => {
 
         dispatch(progressImageSet(null));
       }
+      // pass along the socket event as an application action
+      dispatch(appSocketInvocationComplete(action.payload));
     },
   });
 };

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

@@ -1,17 +1,21 @@
 import { startAppListening } from '../..';
 import { log } from 'app/logging/useLogger';
-import { invocationError } from 'services/events/actions';
+import {
+  appSocketInvocationError,
+  socketInvocationError,
+} from 'services/events/actions';
 
 const moduleLog = log.child({ namespace: 'socketio' });
 
-export const addInvocationErrorListener = () => {
+export const addInvocationErrorEventListener = () => {
   startAppListening({
-    actionCreator: invocationError,
+    actionCreator: socketInvocationError,
     effect: (action, { dispatch, getState }) => {
       moduleLog.error(
         action.payload,
         `Invocation error (${action.payload.data.node.type})`
       );
+      dispatch(appSocketInvocationError(action.payload));
     },
   });
 };

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

@@ -1,12 +1,15 @@
 import { startAppListening } from '../..';
 import { log } from 'app/logging/useLogger';
-import { invocationStarted } from 'services/events/actions';
+import {
+  appSocketInvocationStarted,
+  socketInvocationStarted,
+} from 'services/events/actions';
 
 const moduleLog = log.child({ namespace: 'socketio' });
 
-export const addInvocationStartedListener = () => {
+export const addInvocationStartedEventListener = () => {
   startAppListening({
-    actionCreator: invocationStarted,
+    actionCreator: socketInvocationStarted,
     effect: (action, { dispatch, getState }) => {
       if (
         getState().system.canceledSession ===
@@ -23,6 +26,7 @@ export const addInvocationStartedListener = () => {
         action.payload,
         `Invocation started (${action.payload.data.node.type})`
       );
+      dispatch(appSocketInvocationStarted(action.payload));
     },
   });
 };

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

@@ -1,10 +1,10 @@
 import { startAppListening } from '../..';
 import { log } from 'app/logging/useLogger';
-import { socketSubscribed } from 'services/events/actions';
+import { appSocketSubscribed, socketSubscribed } from 'services/events/actions';
 
 const moduleLog = log.child({ namespace: 'socketio' });
 
-export const addSocketSubscribedListener = () => {
+export const addSocketSubscribedEventListener = () => {
   startAppListening({
     actionCreator: socketSubscribed,
     effect: (action, { dispatch, getState }) => {
@@ -12,6 +12,7 @@ export const addSocketSubscribedListener = () => {
         action.payload,
         `Subscribed (${action.payload.sessionId}))`
       );
+      dispatch(appSocketSubscribed(action.payload));
     },
   });
 };

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

@@ -1,10 +1,13 @@
 import { startAppListening } from '../..';
 import { log } from 'app/logging/useLogger';
-import { socketUnsubscribed } from 'services/events/actions';
+import {
+  appSocketUnsubscribed,
+  socketUnsubscribed,
+} from 'services/events/actions';
 
 const moduleLog = log.child({ namespace: 'socketio' });
 
-export const addSocketUnsubscribedListener = () => {
+export const addSocketUnsubscribedEventListener = () => {
   startAppListening({
     actionCreator: socketUnsubscribed,
     effect: (action, { dispatch, getState }) => {
@@ -12,6 +15,7 @@ export const addSocketUnsubscribedListener = () => {
         action.payload,
         `Unsubscribed (${action.payload.sessionId})`
       );
+      dispatch(appSocketUnsubscribed(action.payload));
     },
   });
 };