feat(ui): add debug mode & socketOptions

invokeai/frontend/web/src/app/components/App.tsx

@@ -20,6 +20,7 @@ import AppErrorBoundaryFallback from './AppErrorBoundaryFallback';
 import GlobalHotkeys from './GlobalHotkeys';
 import PreselectedImage from './PreselectedImage';
 import Toaster from './Toaster';
+import { useSocketIO } from '../hooks/useSocketIO';
 
 const DEFAULT_CONFIG = {};
 
@@ -33,10 +34,12 @@ interface Props {
 
 const App = ({ config = DEFAULT_CONFIG, selectedImage }: Props) => {
   const language = useAppSelector(languageSelector);
-
   const logger = useLogger('system');
   const dispatch = useAppDispatch();
 
+  // singleton!
+  useSocketIO();
+
   const handleReset = useCallback(() => {
     localStorage.clear();
     location.reload();

invokeai/frontend/web/src/app/components/InvokeAIUI.tsx

@@ -1,26 +1,27 @@
 import { Middleware } from '@reduxjs/toolkit';
-import { store } from 'app/store/store';
-import { PartialAppConfig } from 'app/types/invokeai';
-import React, {
-  lazy,
-  memo,
-  PropsWithChildren,
-  ReactNode,
-  useEffect,
-} from 'react';
-import { Provider } from 'react-redux';
-import { addMiddleware, resetMiddlewares } from 'redux-dynamic-middlewares';
-import { $authToken, $baseUrl, $projectId } from 'services/api/client';
-import { socketMiddleware } from 'services/events/middleware';
-import Loading from '../../common/components/Loading/Loading';
-import '../../i18n';
-import AppDndContext from '../../features/dnd/components/AppDndContext';
 import { $customStarUI, CustomStarUi } from 'app/store/nanostores/customStarUI';
 import { $headerComponent } from 'app/store/nanostores/headerComponent';
+import { $isDebugging } from 'app/store/nanostores/isDebugging';
+import { store } from 'app/store/store';
+import { PartialAppConfig } from 'app/types/invokeai';
 import {
   $queueId,
   DEFAULT_QUEUE_ID,
 } from 'features/queue/store/queueNanoStore';
+import React, {
+  PropsWithChildren,
+  ReactNode,
+  lazy,
+  memo,
+  useEffect,
+} from 'react';
+import { Provider } from 'react-redux';
+import { $authToken, $baseUrl, $projectId } from 'services/api/client';
+import { ManagerOptions, SocketOptions } from 'socket.io-client';
+import Loading from '../../common/components/Loading/Loading';
+import AppDndContext from '../../features/dnd/components/AppDndContext';
+import '../../i18n';
+import { $socketOptions } from '../hooks/useSocketIO';
 
 const App = lazy(() => import('./App'));
 const ThemeLocaleProvider = lazy(() => import('./ThemeLocaleProvider'));
@@ -38,6 +39,8 @@ interface Props extends PropsWithChildren {
     action: 'sendToImg2Img' | 'sendToCanvas' | 'useAllParameters';
   };
   customStarUi?: CustomStarUi;
+  socketOptions?: Partial<ManagerOptions & SocketOptions>;
+  isDebugging?: boolean;
 }
 
 const InvokeAIUI = ({
@@ -50,6 +53,8 @@ const InvokeAIUI = ({
   queueId,
   selectedImage,
   customStarUi,
+  socketOptions,
+  isDebugging = false,
 }: Props) => {
   useEffect(() => {
     // configure API client token
@@ -72,21 +77,6 @@ const InvokeAIUI = ({
       $queueId.set(queueId);
     }
 
-    // reset dynamically added middlewares
-    resetMiddlewares();
-
-    // TODO: at this point, after resetting the middleware, we really ought to clean up the socket
-    // stuff by calling `dispatch(socketReset())`. but we cannot dispatch from here as we are
-    // outside the provider. it's not needed until there is the possibility that we will change
-    // the `apiUrl`/`token` dynamically.
-
-    // rebuild socket middleware with token and apiUrl
-    if (middleware && middleware.length > 0) {
-      addMiddleware(socketMiddleware(), ...middleware);
-    } else {
-      addMiddleware(socketMiddleware());
-    }
-
     return () => {
       // Reset the API client token and base url on unmount
       $baseUrl.set(undefined);
@@ -116,6 +106,24 @@ const InvokeAIUI = ({
     };
   }, [headerComponent]);
 
+  useEffect(() => {
+    if (socketOptions) {
+      $socketOptions.set(socketOptions);
+    }
+    return () => {
+      $socketOptions.set({});
+    };
+  }, [socketOptions]);
+
+  useEffect(() => {
+    if (isDebugging) {
+      $isDebugging.set(isDebugging);
+    }
+    return () => {
+      $isDebugging.set(false);
+    };
+  }, [isDebugging]);
+
   return (
     <React.StrictMode>
       <Provider store={store}>

invokeai/frontend/web/src/app/hooks/useSocketIO.ts

@@ -0,0 +1,96 @@
+import { useStore } from '@nanostores/react';
+import { $isDebugging } from 'app/store/nanostores/isDebugging';
+import { useAppDispatch } from 'app/store/storeHooks';
+import { MapStore, WritableAtom, atom, map } from 'nanostores';
+import { useEffect } from 'react';
+import { $authToken, $baseUrl } from 'services/api/client';
+import {
+  ClientToServerEvents,
+  ServerToClientEvents,
+} from 'services/events/types';
+import { setEventListeners } from 'services/events/util/setEventListeners';
+import { ManagerOptions, Socket, SocketOptions, io } from 'socket.io-client';
+
+declare global {
+  interface Window {
+    $socketOptions?: MapStore<Partial<ManagerOptions & SocketOptions>>;
+    $socketUrl?: WritableAtom<string>;
+  }
+}
+
+const makeSocketOptions = (): Partial<ManagerOptions & SocketOptions> => {
+  const socketOptions: Parameters<typeof io>[0] = {
+    timeout: 60000,
+    path: '/ws/socket.io',
+    autoConnect: false, // achtung! removing this breaks the dynamic middleware
+    forceNew: true,
+  };
+
+  // if building in package mode, replace socket url with open api base url minus the http protocol
+  if (['nodes', 'package'].includes(import.meta.env.MODE)) {
+    const authToken = $authToken.get();
+    if (authToken) {
+      // TODO: handle providing jwt to socket.io
+      socketOptions.auth = { token: authToken };
+    }
+
+    socketOptions.transports = ['websocket', 'polling'];
+  }
+
+  return socketOptions;
+};
+
+const makeSocketUrl = (): string => {
+  const wsProtocol = window.location.protocol === 'https:' ? 'wss' : 'ws';
+  let socketUrl = `${wsProtocol}://${window.location.host}`;
+  if (['nodes', 'package'].includes(import.meta.env.MODE)) {
+    const baseUrl = $baseUrl.get();
+    if (baseUrl) {
+      //eslint-disable-next-line
+      socketUrl = baseUrl.replace(/^https?\:\/\//i, '');
+    }
+  }
+  return socketUrl;
+};
+
+const makeSocket = (): Socket<ServerToClientEvents, ClientToServerEvents> => {
+  const socketOptions = makeSocketOptions();
+  const socketUrl = $socketUrl.get();
+  const socket: Socket<ServerToClientEvents, ClientToServerEvents> = io(
+    socketUrl,
+    { ...socketOptions, ...$socketOptions.get() }
+  );
+  return socket;
+};
+
+export const $socketOptions = map<Partial<ManagerOptions & SocketOptions>>({});
+export const $socketUrl = atom<string>(makeSocketUrl());
+
+export const useSocketIO = () => {
+  const dispatch = useAppDispatch();
+  const socketOptions = useStore($socketOptions);
+  const socketUrl = useStore($socketUrl);
+  const baseUrl = useStore($baseUrl);
+  const authToken = useStore($authToken);
+
+  useEffect(() => {
+    const socket = makeSocket();
+    setEventListeners({ dispatch, socket });
+    socket.connect();
+
+    if ($isDebugging.get()) {
+      window.$socketOptions = $socketOptions;
+      window.$socketUrl = $socketUrl;
+      console.log('Socket initialized', socket);
+    }
+
+    return () => {
+      if ($isDebugging.get()) {
+        window.$socketOptions = undefined;
+        window.$socketUrl = undefined;
+        console.log('Socket teardown', socket);
+      }
+      socket.disconnect();
+    };
+  }, [dispatch, socketOptions, socketUrl, baseUrl, authToken]);
+};

invokeai/frontend/web/src/app/store/nanostores/isDebugging.ts

@@ -0,0 +1,3 @@
+import { atom } from 'nanostores';
+
+export const $isDebugging = atom<boolean>(false);

invokeai/frontend/web/src/app/store/store.ts

@@ -12,17 +12,16 @@ import deleteImageModalReducer from 'features/deleteImageModal/store/slice';
 import dynamicPromptsReducer from 'features/dynamicPrompts/store/dynamicPromptsSlice';
 import galleryReducer from 'features/gallery/store/gallerySlice';
 import loraReducer from 'features/lora/store/loraSlice';
+import modelmanagerReducer from 'features/modelManager/store/modelManagerSlice';
 import nodesReducer from 'features/nodes/store/nodesSlice';
 import generationReducer from 'features/parameters/store/generationSlice';
 import postprocessingReducer from 'features/parameters/store/postprocessingSlice';
+import queueReducer from 'features/queue/store/queueSlice';
 import sdxlReducer from 'features/sdxl/store/sdxlSlice';
 import configReducer from 'features/system/store/configSlice';
 import systemReducer from 'features/system/store/systemSlice';
-import queueReducer from 'features/queue/store/queueSlice';
-import modelmanagerReducer from 'features/modelManager/store/modelManagerSlice';
 import hotkeysReducer from 'features/ui/store/hotkeysSlice';
 import uiReducer from 'features/ui/store/uiSlice';
-import dynamicMiddlewares from 'redux-dynamic-middlewares';
 import { rememberEnhancer, rememberReducer } from 'redux-remember';
 import { api } from 'services/api';
 import { LOCALSTORAGE_PREFIX } from './constants';
@@ -94,7 +93,6 @@ export const store = configureStore({
       immutableCheck: false,
     })
       .concat(api.middleware)
-      .concat(dynamicMiddlewares)
       .prepend(listenerMiddleware.middleware),
   devTools: {
     actionSanitizer,

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

@@ -1,63 +0,0 @@
-import { Middleware, MiddlewareAPI } from '@reduxjs/toolkit';
-import { AppThunkDispatch, RootState } from 'app/store/store';
-import { $authToken, $baseUrl } from 'services/api/client';
-import {
-  ClientToServerEvents,
-  ServerToClientEvents,
-} from 'services/events/types';
-import { setEventListeners } from 'services/events/util/setEventListeners';
-import { Socket, io } from 'socket.io-client';
-
-export const socketMiddleware = () => {
-  let areListenersSet = false;
-
-  const wsProtocol = window.location.protocol === 'https:' ? 'wss' : 'ws';
-  let socketUrl = `${wsProtocol}://${window.location.host}`;
-
-  const socketOptions: Parameters<typeof io>[0] = {
-    timeout: 60000,
-    path: '/ws/socket.io',
-    autoConnect: false, // achtung! removing this breaks the dynamic middleware
-  };
-
-  // if building in package mode, replace socket url with open api base url minus the http protocol
-  if (['nodes', 'package'].includes(import.meta.env.MODE)) {
-    const baseUrl = $baseUrl.get();
-    if (baseUrl) {
-      //eslint-disable-next-line
-      socketUrl = baseUrl.replace(/^https?\:\/\//i, '');
-    }
-
-    const authToken = $authToken.get();
-    if (authToken) {
-      // TODO: handle providing jwt to socket.io
-      socketOptions.auth = { token: authToken };
-    }
-
-    socketOptions.transports = ['websocket', 'polling'];
-  }
-
-  const socket: Socket<ServerToClientEvents, ClientToServerEvents> = io(
-    socketUrl,
-    socketOptions
-  );
-
-  const middleware: Middleware =
-    (storeApi: MiddlewareAPI<AppThunkDispatch, RootState>) =>
-    (next) =>
-    (action) => {
-      // Set listeners for `connect` and `disconnect` events once
-      // Must happen in middleware to get access to `dispatch`
-      if (!areListenersSet) {
-        setEventListeners({ storeApi, socket });
-
-        areListenersSet = true;
-
-        socket.connect();
-      }
-
-      next(action);
-    };
-
-  return middleware;
-};

invokeai/frontend/web/src/services/events/util/setEventListeners.ts

@@ -1,6 +1,4 @@
-import { MiddlewareAPI } from '@reduxjs/toolkit';
-import { logger } from 'app/logging/logger';
-import { AppDispatch, RootState } from 'app/store/store';
+import { AppDispatch } from 'app/store/store';
 import { $queueId } from 'features/queue/store/queueNanoStore';
 import { addToast } from 'features/system/store/systemSlice';
 import { makeToast } from 'features/system/util/makeToast';
@@ -23,20 +21,16 @@ import { ClientToServerEvents, ServerToClientEvents } from '../types';
 
 type SetEventListenersArg = {
   socket: Socket<ServerToClientEvents, ClientToServerEvents>;
-  storeApi: MiddlewareAPI<AppDispatch, RootState>;
+  dispatch: AppDispatch;
 };
 
 export const setEventListeners = (arg: SetEventListenersArg) => {
-  const { socket, storeApi } = arg;
-  const { dispatch } = storeApi;
+  const { socket, dispatch } = arg;
 
   /**
    * Connect
    */
   socket.on('connect', () => {
-    const log = logger('socketio');
-    log.debug('Connected');
-
     dispatch(socketConnected());
     const queue_id = $queueId.get();
     socket.emit('subscribe_queue', { queue_id });