fix(ui): staging area left/right hotkeys

Was running into difficultlies reasoning about the logic and couldn't
write tests because it was all in react.

Moved logic outside react, updated context, make it testable.

.gitignore

@@ -190,3 +190,5 @@ installer/update.bat
 installer/update.sh
 installer/InvokeAI-Installer/
 .aider*
+
+.claude/

invokeai/frontend/web/.prettierignore

@@ -14,3 +14,4 @@ static/
 src/theme/css/overlayscrollbars.css
 src/theme_/css/overlayscrollbars.css
 pnpm-lock.yaml
+.claude

invokeai/frontend/web/package.json

@@ -56,7 +56,7 @@
     "chakra-react-select": "^4.9.2",
     "cmdk": "^1.1.1",
     "compare-versions": "^6.1.1",
-    "dockview": "^4.4.0",
+    "dockview": "^4.4.1",
     "es-toolkit": "^1.39.7",
     "filesize": "^10.1.6",
     "fracturedjsonjs": "^4.1.0",

invokeai/frontend/web/pnpm-lock.yaml

@@ -60,8 +60,8 @@ importers:
         specifier: ^6.1.1
         version: 6.1.1
       dockview:
-        specifier: ^4.4.0
-        version: 4.4.0(react@18.3.1)
+        specifier: ^4.4.1
+        version: 4.4.1(react@18.3.1)
       es-toolkit:
         specifier: ^1.39.7
         version: 1.39.7
@@ -2247,11 +2247,11 @@ packages:
   discontinuous-range@1.0.0:
     resolution: {integrity: sha512-c68LpLbO+7kP/b1Hr1qs8/BJ09F5khZGTxqxZuhzxpmwJKOgRFHJWIb9/KmqnqHhLdO55aOxFH/EGBvUQbL/RQ==}
 
-  dockview-core@4.4.0:
-    resolution: {integrity: sha512-UsBJwS3lfZXM+gaTA+bJs8rAxLd7ZEmNcUf5CbKKhiPeKIPJrNCxXxTLcnQb3IXMJUGkE0aX1ZJ4BDaZGMtzlA==}
+  dockview-core@4.4.1:
+    resolution: {integrity: sha512-pDQPlVfDYDuN3zSebVUMVn2x21bpYPGD1ybGYrKJMI1KDkSQSqy57FJRJXi7yEnkcrmBUF0xEEo4d0Yi3j2vGA==}
 
-  dockview@4.4.0:
-    resolution: {integrity: sha512-cWi5R40R5kDky69vAqsKGznRx5tA0gk3Mdqe5aS2r4ollK951mWNJ/EeMmac+UP/juw4cbl0/APhXTV+EMnAbg==}
+  dockview@4.4.1:
+    resolution: {integrity: sha512-XEAwl+VYVZGkBd3hprF6kRLspWSF/hydbRHuV3KEg3BHev1i5xc+H+Kjp+u5DHTQ97klFAATl5MntNoVXQeg0w==}
     peerDependencies:
       react: ^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0
 
@@ -6601,11 +6601,11 @@ snapshots:
 
   discontinuous-range@1.0.0: {}
 
-  dockview-core@4.4.0: {}
+  dockview-core@4.4.1: {}
 
-  dockview@4.4.0(react@18.3.1):
+  dockview@4.4.1(react@18.3.1):
     dependencies:
-      dockview-core: 4.4.0
+      dockview-core: 4.4.1
       react: 18.3.1
 
   doctrine@2.1.0:

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

@@ -30,16 +30,16 @@ const App = ({ config = DEFAULT_CONFIG, studioInitAction }: Props) => {
   }, [clearStorage]);
 
   return (
-    <ErrorBoundary onReset={handleReset} FallbackComponent={AppErrorBoundaryFallback}>
-      <ThemeLocaleProvider>
+    <ThemeLocaleProvider>
+      <ErrorBoundary onReset={handleReset} FallbackComponent={AppErrorBoundaryFallback}>
         <Box id="invoke-app-wrapper" w="100dvw" h="100dvh" position="relative" overflow="hidden">
           <AppContent />
           {!didStudioInit && <Loading />}
         </Box>
         <GlobalHookIsolator config={config} studioInitAction={studioInitAction} />
         <GlobalModalIsolator />
-      </ThemeLocaleProvider>
-    </ErrorBoundary>
+      </ErrorBoundary>
+    </ThemeLocaleProvider>
   );
 };
 

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

@@ -2,6 +2,7 @@ import { useGlobalModifiersInit } from '@invoke-ai/ui-library';
 import { setupListeners } from '@reduxjs/toolkit/query';
 import type { StudioInitAction } from 'app/hooks/useStudioInitAction';
 import { useStudioInitAction } from 'app/hooks/useStudioInitAction';
+import { useSyncLangDirection } from 'app/hooks/useSyncLangDirection';
 import { useSyncQueueStatus } from 'app/hooks/useSyncQueueStatus';
 import { useLogger } from 'app/logging/useLogger';
 import { useSyncLoggingConfig } from 'app/logging/useSyncLoggingConfig';
@@ -50,6 +51,7 @@ export const GlobalHookIsolator = memo(
     useNavigationApi();
     useDndMonitor();
     useSyncNodeErrors();
+    useSyncLangDirection();
 
     // Persistent subscription to the queue counts query - canvas relies on this to know if there are pending
     // and/or in progress canvas sessions.

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

@@ -3,43 +3,39 @@ import 'overlayscrollbars/overlayscrollbars.css';
 import '@xyflow/react/dist/base.css';
 import 'common/components/OverlayScrollbars/overlayscrollbars.css';
 
-import { ChakraProvider, DarkMode, extendTheme, theme as _theme, TOAST_OPTIONS } from '@invoke-ai/ui-library';
+import { ChakraProvider, DarkMode, extendTheme, theme as baseTheme, TOAST_OPTIONS } from '@invoke-ai/ui-library';
+import { useStore } from '@nanostores/react';
+import { $direction } from 'app/hooks/useSyncLangDirection';
 import type { ReactNode } from 'react';
-import { memo, useEffect, useMemo } from 'react';
-import { useTranslation } from 'react-i18next';
+import { memo, useMemo } from 'react';
 
 type ThemeLocaleProviderProps = {
   children: ReactNode;
 };
 
+const buildTheme = (direction: 'ltr' | 'rtl') => {
+  return extendTheme({
+    ...baseTheme,
+    direction,
+    shadows: {
+      ...baseTheme.shadows,
+      selected:
+        'inset 0px 0px 0px 3px var(--invoke-colors-invokeBlue-500), inset 0px 0px 0px 4px var(--invoke-colors-invokeBlue-800)',
+      hoverSelected:
+        'inset 0px 0px 0px 3px var(--invoke-colors-invokeBlue-400), inset 0px 0px 0px 4px var(--invoke-colors-invokeBlue-800)',
+      hoverUnselected:
+        'inset 0px 0px 0px 2px var(--invoke-colors-invokeBlue-300), inset 0px 0px 0px 3px var(--invoke-colors-invokeBlue-800)',
+      selectedForCompare:
+        'inset 0px 0px 0px 3px var(--invoke-colors-invokeGreen-300), inset 0px 0px 0px 4px var(--invoke-colors-invokeGreen-800)',
+      hoverSelectedForCompare:
+        'inset 0px 0px 0px 3px var(--invoke-colors-invokeGreen-200), inset 0px 0px 0px 4px var(--invoke-colors-invokeGreen-800)',
+    },
+  });
+};
+
 function ThemeLocaleProvider({ children }: ThemeLocaleProviderProps) {
-  const { i18n } = useTranslation();
-
-  const direction = i18n.dir();
-
-  const theme = useMemo(() => {
-    return extendTheme({
-      ..._theme,
-      direction,
-      shadows: {
-        ..._theme.shadows,
-        selected:
-          'inset 0px 0px 0px 3px var(--invoke-colors-invokeBlue-500), inset 0px 0px 0px 4px var(--invoke-colors-invokeBlue-800)',
-        hoverSelected:
-          'inset 0px 0px 0px 3px var(--invoke-colors-invokeBlue-400), inset 0px 0px 0px 4px var(--invoke-colors-invokeBlue-800)',
-        hoverUnselected:
-          'inset 0px 0px 0px 2px var(--invoke-colors-invokeBlue-300), inset 0px 0px 0px 3px var(--invoke-colors-invokeBlue-800)',
-        selectedForCompare:
-          'inset 0px 0px 0px 3px var(--invoke-colors-invokeGreen-300), inset 0px 0px 0px 4px var(--invoke-colors-invokeGreen-800)',
-        hoverSelectedForCompare:
-          'inset 0px 0px 0px 3px var(--invoke-colors-invokeGreen-200), inset 0px 0px 0px 4px var(--invoke-colors-invokeGreen-800)',
-      },
-    });
-  }, [direction]);
-
-  useEffect(() => {
-    document.body.dir = direction;
-  }, [direction]);
+  const direction = useStore($direction);
+  const theme = useMemo(() => buildTheme(direction), [direction]);
 
   return (
     <ChakraProvider theme={theme} toastOptions={TOAST_OPTIONS}>

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

@@ -0,0 +1,36 @@
+import { useAssertSingleton } from 'common/hooks/useAssertSingleton';
+import { atom } from 'nanostores';
+import { useEffect } from 'react';
+import { useTranslation } from 'react-i18next';
+
+/**
+ * Global atom storing the language direction, to be consumed by the Chakra theme.
+ *
+ * Why do we need this? We have a kind of catch-22:
+ * - The Chakra theme needs to know the language direction to apply the correct styles.
+ * - The language direction is determined by i18n and the language selection.
+ * - We want our error boundary to be themed.
+ * - It's possible that i18n can throw if the language selection is invalid or not supported.
+ *
+ * Previously, we had the logic in this file in the theme provider, which wrapped the error boundary. The error
+ * was properly themed. But then, if i18n threw in the theme provider, the error boundary does not catch the
+ * error. The app would crash to a white screen.
+ *
+ * We tried swapping the component hierarchy so that the error boundary wraps the theme provider, but then the
+ * error boundary isn't themed!
+ *
+ * The solution is to move this i18n direction logic out of the theme provider and into a hook that we can use
+ * within the error boundary. The error boundary will be themed, _and_ catch any i18n errors.
+ */
+export const $direction = atom<'ltr' | 'rtl'>('ltr');
+
+export const useSyncLangDirection = () => {
+  useAssertSingleton('useSyncLangDirection');
+  const { i18n, t } = useTranslation();
+
+  useEffect(() => {
+    const direction = i18n.dir();
+    $direction.set(direction);
+    document.body.dir = direction;
+  }, [i18n, t]);
+};

invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/context.tsx

@@ -1,498 +0,0 @@
-import { useStore } from '@nanostores/react';
-import { useAppSelector, useAppStore } from 'app/store/storeHooks';
-import { getOutputImageName } from 'features/controlLayers/components/SimpleSession/shared';
-import { loadImage } from 'features/controlLayers/konva/util';
-import { selectStagingAreaAutoSwitch } from 'features/controlLayers/store/canvasSettingsSlice';
-import {
-  buildSelectCanvasQueueItems,
-  canvasQueueItemDiscarded,
-  canvasSessionReset,
-  selectCanvasSessionId,
-} from 'features/controlLayers/store/canvasStagingAreaSlice';
-import type { ProgressImage } from 'features/nodes/types/common';
-import type { Atom, MapStore, StoreValue, WritableAtom } from 'nanostores';
-import { atom, computed, effect, map, subscribeKeys } from 'nanostores';
-import type { PropsWithChildren } from 'react';
-import { createContext, memo, useCallback, useContext, useEffect, useMemo, useState } from 'react';
-import { getImageDTOSafe } from 'services/api/endpoints/images';
-import { queueApi } from 'services/api/endpoints/queue';
-import type { ImageDTO, S } from 'services/api/types';
-import { $socket } from 'services/events/stores';
-import { assert, objectEntries } from 'tsafe';
-
-export type ProgressData = {
-  itemId: number;
-  progressEvent: S['InvocationProgressEvent'] | null;
-  progressImage: ProgressImage | null;
-  imageDTO: ImageDTO | null;
-  imageLoaded: boolean;
-};
-
-const getInitialProgressData = (itemId: number): ProgressData => ({
-  itemId,
-  progressEvent: null,
-  progressImage: null,
-  imageDTO: null,
-  imageLoaded: false,
-});
-
-export const useProgressData = ($progressData: ProgressDataMap, itemId: number): ProgressData => {
-  const getInitialValue = useCallback(
-    () => $progressData.get()[itemId] ?? getInitialProgressData(itemId),
-    [$progressData, itemId]
-  );
-  const [value, setValue] = useState(getInitialValue);
-  useEffect(() => {
-    const unsub = subscribeKeys($progressData, [itemId], (data) => {
-      const progressData = data[itemId];
-      if (!progressData) {
-        return;
-      }
-      setValue(progressData);
-    });
-    return () => {
-      unsub();
-    };
-  }, [$progressData, itemId]);
-
-  return value;
-};
-
-const setProgress = ($progressData: ProgressDataMap, data: S['InvocationProgressEvent']) => {
-  const progressData = $progressData.get();
-  const current = progressData[data.item_id];
-  if (current) {
-    const next = { ...current };
-    next.progressEvent = data;
-    if (data.image) {
-      next.progressImage = data.image;
-    }
-    $progressData.set({
-      ...progressData,
-      [data.item_id]: next,
-    });
-  } else {
-    $progressData.set({
-      ...progressData,
-      [data.item_id]: {
-        itemId: data.item_id,
-        progressEvent: data,
-        progressImage: data.image ?? null,
-        imageDTO: null,
-        imageLoaded: false,
-      },
-    });
-  }
-};
-
-export type ProgressDataMap = MapStore<Record<number, ProgressData | undefined>>;
-
-type CanvasSessionContextValue = {
-  $items: Atom<S['SessionQueueItem'][]>;
-  $itemCount: Atom<number>;
-  $hasItems: Atom<boolean>;
-  $progressData: ProgressDataMap;
-  $selectedItemId: WritableAtom<number | null>;
-  $selectedItem: Atom<S['SessionQueueItem'] | null>;
-  $selectedItemIndex: Atom<number | null>;
-  $selectedItemOutputImageDTO: Atom<ImageDTO | null>;
-  selectNext: () => void;
-  selectPrev: () => void;
-  selectFirst: () => void;
-  selectLast: () => void;
-  discard: (itemId: number) => void;
-  discardAll: () => void;
-};
-
-const CanvasSessionContext = createContext<CanvasSessionContextValue | null>(null);
-
-export const CanvasSessionContextProvider = memo(({ children }: PropsWithChildren) => {
-  /**
-   * For best performance and interop with the Canvas, which is outside react but needs to interact with the react
-   * app, all canvas session state is packaged as nanostores atoms. The trickiest part is syncing the queue items
-   * with a nanostores atom.
-   */
-
-  /**
-   * App store
-   */
-  const store = useAppStore();
-
-  const sessionId = useAppSelector(selectCanvasSessionId);
-
-  const socket = useStore($socket);
-
-  /**
-   * Track the last completed item. Used to implement autoswitch.
-   */
-  const $lastCompletedItemId = useState(() => atom<number | null>(null))[0];
-
-  /**
-   * Manually-synced atom containing queue items for the current session. This is populated from the RTK Query cache
-   * and kept in sync with it via a redux subscription.
-   */
-  const $items = useState(() => atom<S['SessionQueueItem'][]>([]))[0];
-
-  /**
-   * An ephemeral store of progress events and images for all items in the current session.
-   */
-  const $progressData = useState(() => map<StoreValue<ProgressDataMap>>({}))[0];
-
-  /**
-   * The currently selected queue item's ID, or null if one is not selected.
-   */
-  const $selectedItemId = useState(() => atom<number | null>(null))[0];
-
-  /**
-   * The number of items. Computed from the queue items array.
-   */
-  const $itemCount = useState(() => computed([$items], (items) => items.length))[0];
-
-  /**
-   * Whether there are any items. Computed from the queue items array.
-   */
-  const $hasItems = useState(() => computed([$items], (items) => items.length > 0))[0];
-
-  /**
-   * Whether there are any pending or in-progress items. Computed from the queue items array.
-   */
-  const $isPending = useState(() =>
-    computed([$items], (items) => items.some((item) => item.status === 'pending' || item.status === 'in_progress'))
-  )[0];
-
-  /**
-   * The currently selected queue item, or null if one is not selected.
-   */
-  const $selectedItem = useState(() =>
-    computed([$items, $selectedItemId], (items, selectedItemId) => {
-      if (items.length === 0) {
-        return null;
-      }
-      if (selectedItemId === null) {
-        return null;
-      }
-      return items.find(({ item_id }) => item_id === selectedItemId) ?? null;
-    })
-  )[0];
-
-  /**
-   * The currently selected queue item's index in the list of items, or null if one is not selected.
-   */
-  const $selectedItemIndex = useState(() =>
-    computed([$items, $selectedItemId], (items, selectedItemId) => {
-      if (items.length === 0) {
-        return null;
-      }
-      if (selectedItemId === null) {
-        return null;
-      }
-      return items.findIndex(({ item_id }) => item_id === selectedItemId) ?? null;
-    })
-  )[0];
-
-  /**
-   * The currently selected queue item's output image name, or null if one is not selected or there is no output
-   * image recorded.
-   */
-  const $selectedItemOutputImageDTO = useState(() =>
-    computed([$selectedItemId, $progressData], (selectedItemId, progressData) => {
-      if (selectedItemId === null) {
-        return null;
-      }
-      const datum = progressData[selectedItemId];
-      if (!datum) {
-        return null;
-      }
-      return datum.imageDTO;
-    })
-  )[0];
-
-  /**
-   * A redux selector to select all queue items from the RTK Query cache.
-   */
-  const selectQueueItems = useMemo(() => buildSelectCanvasQueueItems(sessionId), [sessionId]);
-
-  const discard = useCallback(
-    (itemId: number) => {
-      store.dispatch(canvasQueueItemDiscarded({ itemId }));
-    },
-    [store]
-  );
-
-  const discardAll = useCallback(() => {
-    store.dispatch(canvasSessionReset());
-  }, [store]);
-
-  const selectNext = useCallback(() => {
-    const selectedItemId = $selectedItemId.get();
-    if (selectedItemId === null) {
-      return;
-    }
-    const items = $items.get();
-    const currentIndex = items.findIndex((item) => item.item_id === selectedItemId);
-    const nextIndex = (currentIndex + 1) % items.length;
-    const nextItem = items[nextIndex];
-    if (!nextItem) {
-      return;
-    }
-    $selectedItemId.set(nextItem.item_id);
-  }, [$items, $selectedItemId]);
-
-  const selectPrev = useCallback(() => {
-    const selectedItemId = $selectedItemId.get();
-    if (selectedItemId === null) {
-      return;
-    }
-    const items = $items.get();
-    const currentIndex = items.findIndex((item) => item.item_id === selectedItemId);
-    const prevIndex = (currentIndex - 1 + items.length) % items.length;
-    const prevItem = items[prevIndex];
-    if (!prevItem) {
-      return;
-    }
-    $selectedItemId.set(prevItem.item_id);
-  }, [$items, $selectedItemId]);
-
-  const selectFirst = useCallback(() => {
-    const items = $items.get();
-    const first = items.at(0);
-    if (!first) {
-      return;
-    }
-    $selectedItemId.set(first.item_id);
-  }, [$items, $selectedItemId]);
-
-  const selectLast = useCallback(() => {
-    const items = $items.get();
-    const last = items.at(-1);
-    if (!last) {
-      return;
-    }
-    $selectedItemId.set(last.item_id);
-  }, [$items, $selectedItemId]);
-
-  // Set up socket listeners
-  useEffect(() => {
-    if (!socket) {
-      return;
-    }
-
-    const onProgress = (data: S['InvocationProgressEvent']) => {
-      if (data.destination !== sessionId) {
-        return;
-      }
-      setProgress($progressData, data);
-    };
-
-    const onQueueItemStatusChanged = (data: S['QueueItemStatusChangedEvent']) => {
-      if (data.destination !== sessionId) {
-        return;
-      }
-      if (data.status === 'completed') {
-        /**
-         * There is an unpleasant bit of indirection here. When an item is completed, and auto-switch is set to
-         * switch_on_finish, we want to load the image and switch to it. In this socket handler, we don't have
-         * access to the full queue item, which we need to get the output image and load it. We get the full
-         * queue items as part of the list query, so it's rather inefficient to fetch it again here.
-         *
-         * To reduce the number of extra network requests, we instead store this item as the last completed item.
-         * Then in the progress data sync effect, we process the queue item load its image.
-         */
-        $lastCompletedItemId.set(data.item_id);
-      }
-      if (data.status === 'in_progress' && selectStagingAreaAutoSwitch(store.getState()) === 'switch_on_start') {
-        $selectedItemId.set(data.item_id);
-      }
-    };
-
-    socket.on('invocation_progress', onProgress);
-    socket.on('queue_item_status_changed', onQueueItemStatusChanged);
-
-    return () => {
-      socket.off('invocation_progress', onProgress);
-      socket.off('queue_item_status_changed', onQueueItemStatusChanged);
-    };
-  }, [$progressData, $selectedItemId, sessionId, socket, $lastCompletedItemId, store]);
-
-  // Set up state subscriptions and effects
-  useEffect(() => {
-    let _prevItems: readonly S['SessionQueueItem'][] = [];
-    // Seed the $items atom with the initial query cache state
-    $items.set(selectQueueItems(store.getState()));
-
-    // Manually keep the $items atom in sync as the query cache is updated
-    const unsubReduxSyncToItemsAtom = store.subscribe(() => {
-      const prevItems = $items.get();
-      const items = selectQueueItems(store.getState());
-      if (items !== prevItems) {
-        _prevItems = prevItems;
-        $items.set(items);
-      }
-    });
-
-    // Handle cases that could result in a nonexistent queue item being selected.
-    const unsubEnsureSelectedItemIdExists = effect([$items, $selectedItemId], (items, selectedItemId) => {
-      if (items.length === 0) {
-        // If there are no items, cannot have a selected item.
-        $selectedItemId.set(null);
-      } else if (selectedItemId === null && items.length > 0) {
-        // If there is no selected item but there are items, select the first one.
-        $selectedItemId.set(items[0]?.item_id ?? null);
-        return;
-      } else if (selectedItemId !== null && items.findIndex(({ item_id }) => item_id === selectedItemId) === -1) {
-        // If an item is selected and it is not in the list of items, un-set it. This effect will run again and we'll
-        // the above case, selecting the first item if there are any.
-        let prevIndex = _prevItems.findIndex(({ item_id }) => item_id === selectedItemId);
-        if (prevIndex >= items.length) {
-          prevIndex = items.length - 1;
-        }
-        const nextItem = items[prevIndex];
-        $selectedItemId.set(nextItem?.item_id ?? null);
-      }
-
-      if (items !== _prevItems) {
-        _prevItems = items;
-      }
-    });
-
-    // Sync progress data - remove canceled/failed items, update progress data with new images, and load images
-    const unsubSyncProgressData = $items.subscribe(async (items) => {
-      const progressData = $progressData.get();
-
-      const toDelete: number[] = [];
-      const toUpdate: ProgressData[] = [];
-
-      for (const [id, datum] of objectEntries(progressData)) {
-        if (!datum) {
-          toDelete.push(id);
-          continue;
-        }
-        const item = items.find(({ item_id }) => item_id === datum.itemId);
-        if (!item) {
-          toDelete.push(datum.itemId);
-        } else if (item.status === 'canceled' || item.status === 'failed') {
-          toUpdate.push({
-            ...datum,
-            progressEvent: null,
-            progressImage: null,
-            imageDTO: null,
-          });
-        }
-      }
-
-      for (const item of items) {
-        const datum = progressData[item.item_id];
-
-        if (datum?.imageDTO) {
-          continue;
-        }
-        const outputImageName = getOutputImageName(item);
-        if (!outputImageName) {
-          continue;
-        }
-        const imageDTO = await getImageDTOSafe(outputImageName);
-        if (!imageDTO) {
-          continue;
-        }
-
-        // This is the load logic mentioned in the comment in the QueueItemStatusChangedEvent handler above.
-        if (
-          $lastCompletedItemId.get() === item.item_id &&
-          selectStagingAreaAutoSwitch(store.getState()) === 'switch_on_finish'
-        ) {
-          loadImage(imageDTO.image_url, true).then(() => {
-            $selectedItemId.set(item.item_id);
-            $lastCompletedItemId.set(null);
-          });
-        }
-
-        toUpdate.push({
-          ...getInitialProgressData(item.item_id),
-          ...datum,
-          imageDTO,
-        });
-      }
-
-      for (const itemId of toDelete) {
-        $progressData.setKey(itemId, undefined);
-      }
-
-      for (const datum of toUpdate) {
-        $progressData.setKey(datum.itemId, datum);
-      }
-    });
-
-    // Create an RTK Query subscription. Without this, the query cache selector will never return anything bc RTK
-    // doesn't know we care about it.
-    const { unsubscribe: unsubQueueItemsQuery } = store.dispatch(
-      queueApi.endpoints.listAllQueueItems.initiate({ destination: sessionId })
-    );
-
-    // Clean up all subscriptions and top-level (i.e. non-computed/derived state)
-    return () => {
-      unsubQueueItemsQuery();
-      unsubReduxSyncToItemsAtom();
-      unsubEnsureSelectedItemIdExists();
-      unsubSyncProgressData();
-      $items.set([]);
-      $progressData.set({});
-      $selectedItemId.set(null);
-    };
-  }, [$items, $progressData, $selectedItemId, selectQueueItems, sessionId, store, $lastCompletedItemId]);
-
-  const value = useMemo<CanvasSessionContextValue>(
-    () => ({
-      $items,
-      $hasItems,
-      $isPending,
-      $progressData,
-      $selectedItemId,
-      $selectedItem,
-      $selectedItemIndex,
-      $selectedItemOutputImageDTO,
-      $itemCount,
-      selectNext,
-      selectPrev,
-      selectFirst,
-      selectLast,
-      discard,
-      discardAll,
-    }),
-    [
-      $items,
-      $hasItems,
-      $isPending,
-      $progressData,
-      $selectedItem,
-      $selectedItemId,
-      $selectedItemIndex,
-      $selectedItemOutputImageDTO,
-      $itemCount,
-      selectNext,
-      selectPrev,
-      selectFirst,
-      selectLast,
-      discard,
-      discardAll,
-    ]
-  );
-
-  return <CanvasSessionContext.Provider value={value}>{children}</CanvasSessionContext.Provider>;
-});
-CanvasSessionContextProvider.displayName = 'CanvasSessionContextProvider';
-
-export const useCanvasSessionContext = () => {
-  const ctx = useContext(CanvasSessionContext);
-  assert(ctx !== null, "'useCanvasSessionContext' must be used within a CanvasSessionContextProvider");
-  return ctx;
-};
-
-export const useOutputImageDTO = (item: S['SessionQueueItem']) => {
-  const ctx = useCanvasSessionContext();
-  const $imageDTO = useState(() =>
-    computed([ctx.$progressData], (progressData) => progressData[item.item_id]?.imageDTO ?? null)
-  )[0];
-  const imageDTO = useStore($imageDTO);
-
-  return imageDTO;
-};

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/QueueItemCircularProgress.tsx

@@ -1,10 +1,11 @@
 import type { CircularProgressProps, SystemStyleObject } from '@invoke-ai/ui-library';
 import { CircularProgress, Tooltip } from '@invoke-ai/ui-library';
-import { useCanvasSessionContext, useProgressData } from 'features/controlLayers/components/SimpleSession/context';
-import { getProgressMessage } from 'features/controlLayers/components/SimpleSession/shared';
+import { getProgressMessage } from 'features/controlLayers/components/StagingArea/shared';
 import { memo } from 'react';
 import type { S } from 'services/api/types';
 
+import { useProgressDatum } from './context';
+
 const circleStyles: SystemStyleObject = {
   circle: {
     transitionProperty: 'none',
@@ -18,8 +19,7 @@ const circleStyles: SystemStyleObject = {
 type Props = { itemId: number; status: S['SessionQueueItem']['status'] } & CircularProgressProps;
 
 export const QueueItemCircularProgress = memo(({ itemId, status, ...rest }: Props) => {
-  const { $progressData } = useCanvasSessionContext();
-  const { progressEvent } = useProgressData($progressData, itemId);
+  const { progressEvent } = useProgressDatum(itemId);
 
   if (status !== 'in_progress') {
     return null;

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/QueueItemNumber.tsx

@@ -1,8 +1,9 @@
 import type { TextProps } from '@invoke-ai/ui-library';
 import { Text } from '@invoke-ai/ui-library';
-import { DROP_SHADOW } from 'features/controlLayers/components/SimpleSession/shared';
 import { memo } from 'react';
 
+import { DROP_SHADOW } from './shared';
+
 export const QueueItemNumber = memo(({ number, ...rest }: { number: number } & TextProps) => {
   return <Text pointerEvents="none" userSelect="none" filter={DROP_SHADOW} {...rest}>{`#${number}`}</Text>;
 });

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/QueueItemPreviewMini.tsx

@@ -1,25 +1,23 @@
 import type { SystemStyleObject } from '@invoke-ai/ui-library';
 import { Flex } from '@invoke-ai/ui-library';
+import { useStore } from '@nanostores/react';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
-import {
-  useCanvasSessionContext,
-  useOutputImageDTO,
-  useProgressData,
-} from 'features/controlLayers/components/SimpleSession/context';
-import { QueueItemCircularProgress } from 'features/controlLayers/components/SimpleSession/QueueItemCircularProgress';
-import { QueueItemNumber } from 'features/controlLayers/components/SimpleSession/QueueItemNumber';
-import { QueueItemProgressImage } from 'features/controlLayers/components/SimpleSession/QueueItemProgressImage';
-import { QueueItemStatusLabel } from 'features/controlLayers/components/SimpleSession/QueueItemStatusLabel';
-import { getQueueItemElementId } from 'features/controlLayers/components/SimpleSession/shared';
+import { QueueItemCircularProgress } from 'features/controlLayers/components/StagingArea/QueueItemCircularProgress';
+import { QueueItemProgressImage } from 'features/controlLayers/components/StagingArea/QueueItemProgressImage';
+import { QueueItemStatusLabel } from 'features/controlLayers/components/StagingArea/QueueItemStatusLabel';
+import { getQueueItemElementId } from 'features/controlLayers/components/StagingArea/shared';
 import {
   selectStagingAreaAutoSwitch,
   settingsStagingAreaAutoSwitchChanged,
 } from 'features/controlLayers/store/canvasSettingsSlice';
 import { DndImage } from 'features/dnd/DndImage';
 import { toast } from 'features/toast/toast';
-import { memo, useCallback } from 'react';
+import { memo, useCallback, useMemo } from 'react';
 import type { S } from 'services/api/types';
 
+import { useOutputImageDTO, useStagingAreaContext } from './context';
+import { QueueItemNumber } from './QueueItemNumber';
+
 const sx = {
   cursor: 'pointer',
   userSelect: 'none',
@@ -41,19 +39,19 @@ const sx = {
 type Props = {
   item: S['SessionQueueItem'];
   index: number;
-  isSelected: boolean;
 };
 
-export const QueueItemPreviewMini = memo(({ item, isSelected, index }: Props) => {
+export const QueueItemPreviewMini = memo(({ item, index }: Props) => {
+  const ctx = useStagingAreaContext();
   const dispatch = useAppDispatch();
-  const ctx = useCanvasSessionContext();
-  const { imageLoaded } = useProgressData(ctx.$progressData, item.item_id);
-  const imageDTO = useOutputImageDTO(item);
+  const $isSelected = useMemo(() => ctx.buildIsSelectedComputed(item.item_id), [ctx, item.item_id]);
+  const isSelected = useStore($isSelected);
+  const imageDTO = useOutputImageDTO(item.item_id);
   const autoSwitch = useAppSelector(selectStagingAreaAutoSwitch);
 
   const onClick = useCallback(() => {
-    ctx.$selectedItemId.set(item.item_id);
-  }, [ctx.$selectedItemId, item.item_id]);
+    ctx.select(item.item_id);
+  }, [ctx, item.item_id]);
 
   const onDoubleClick = useCallback(() => {
     if (autoSwitch !== 'off') {
@@ -74,7 +72,7 @@ export const QueueItemPreviewMini = memo(({ item, isSelected, index }: Props) =>
     >
       <QueueItemStatusLabel item={item} position="absolute" margin="auto" />
       {imageDTO && <DndImage imageDTO={imageDTO} asThumbnail position="absolute" />}
-      {!imageLoaded && <QueueItemProgressImage itemId={item.item_id} position="absolute" />}
+      <QueueItemProgressImage itemId={item.item_id} position="absolute" />
       <QueueItemNumber number={index + 1} position="absolute" top={0} left={1} />
       <QueueItemCircularProgress itemId={item.item_id} status={item.status} position="absolute" top={1} right={2} />
     </Flex>

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/QueueItemProgressImage.tsx

@@ -1,13 +1,13 @@
 import type { ImageProps } from '@invoke-ai/ui-library';
 import { Image } from '@invoke-ai/ui-library';
-import { useCanvasSessionContext, useProgressData } from 'features/controlLayers/components/SimpleSession/context';
 import { memo } from 'react';
 
+import { useProgressDatum } from './context';
+
 type Props = { itemId: number } & ImageProps;
 
 export const QueueItemProgressImage = memo(({ itemId, ...rest }: Props) => {
-  const ctx = useCanvasSessionContext();
-  const { progressImage } = useProgressData(ctx.$progressData, itemId);
+  const { progressImage } = useProgressDatum(itemId);
 
   if (!progressImage) {
     return null;

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/QueueItemStatusLabel.tsx

@@ -1,16 +1,16 @@
 import type { TextProps } from '@invoke-ai/ui-library';
 import { Text } from '@invoke-ai/ui-library';
-import { useCanvasSessionContext, useProgressData } from 'features/controlLayers/components/SimpleSession/context';
 import { memo } from 'react';
 import type { S } from 'services/api/types';
 
+import { useProgressDatum } from './context';
+
 type Props = { item: S['SessionQueueItem'] } & TextProps;
 
 export const QueueItemStatusLabel = memo(({ item, ...rest }: Props) => {
-  const ctx = useCanvasSessionContext();
-  const { progressImage, imageLoaded } = useProgressData(ctx.$progressData, item.item_id);
+  const { progressImage } = useProgressDatum(item.item_id);
 
-  if (progressImage || imageLoaded) {
+  if (progressImage) {
     return null;
   }
 

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaAutoSwitchButtons.tsx

@@ -1,5 +1,7 @@
 import { IconButton } from '@invoke-ai/ui-library';
+import { useStore } from '@nanostores/react';
 import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
+import { useCanvasManager } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
 import {
   selectStagingAreaAutoSwitch,
   settingsStagingAreaAutoSwitchChanged,
@@ -8,6 +10,9 @@ import { memo, useCallback } from 'react';
 import { PiCaretLineRightBold, PiCaretRightBold, PiMoonBold } from 'react-icons/pi';
 
 export const StagingAreaAutoSwitchButtons = memo(() => {
+  const canvasManager = useCanvasManager();
+  const shouldShowStagedImage = useStore(canvasManager.stagingArea.$shouldShowStagedImage);
+
   const autoSwitch = useAppSelector(selectStagingAreaAutoSwitch);
   const dispatch = useAppDispatch();
 
@@ -29,6 +34,7 @@ export const StagingAreaAutoSwitchButtons = memo(() => {
         icon={<PiMoonBold />}
         colorScheme={autoSwitch === 'off' ? 'invokeBlue' : 'base'}
         onClick={onClickOff}
+        isDisabled={!shouldShowStagedImage}
       />
       <IconButton
         aria-label="Switch on start"
@@ -36,6 +42,7 @@ export const StagingAreaAutoSwitchButtons = memo(() => {
         icon={<PiCaretRightBold />}
         colorScheme={autoSwitch === 'switch_on_start' ? 'invokeBlue' : 'base'}
         onClick={onClickSwitchOnStart}
+        isDisabled={!shouldShowStagedImage}
       />
       <IconButton
         aria-label="Switch on finish"
@@ -43,6 +50,7 @@ export const StagingAreaAutoSwitchButtons = memo(() => {
         icon={<PiCaretLineRightBold />}
         colorScheme={autoSwitch === 'switch_on_finish' ? 'invokeBlue' : 'base'}
         onClick={onClickSwitchOnFinished}
+        isDisabled={!shouldShowStagedImage}
       />
     </>
   );

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaItemsList.tsx

@@ -1,16 +1,16 @@
 import { Box, Flex, forwardRef } from '@invoke-ai/ui-library';
 import { useStore } from '@nanostores/react';
 import { logger } from 'app/logging/logger';
-import { useCanvasSessionContext } from 'features/controlLayers/components/SimpleSession/context';
-import { QueueItemPreviewMini } from 'features/controlLayers/components/SimpleSession/QueueItemPreviewMini';
-import { useCanvasManagerSafe } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
+import { QueueItemPreviewMini } from 'features/controlLayers/components/StagingArea/QueueItemPreviewMini';
+import { useCanvasManager } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
 import { useOverlayScrollbars } from 'overlayscrollbars-react';
 import type { CSSProperties, RefObject } from 'react';
-import { memo, useCallback, useEffect, useMemo, useRef, useState } from 'react';
-import type { Components, ItemContent, ListRange, VirtuosoHandle, VirtuosoProps } from 'react-virtuoso';
+import { memo, useCallback, useEffect, useRef, useState } from 'react';
+import type { Components, ComputeItemKey, ItemContent, ListRange, VirtuosoHandle, VirtuosoProps } from 'react-virtuoso';
 import { Virtuoso } from 'react-virtuoso';
 import type { S } from 'services/api/types';
 
+import { useStagingAreaContext } from './context';
 import { getQueueItemElementId } from './shared';
 
 const log = logger('system');
@@ -20,8 +20,6 @@ const virtuosoStyles = {
   height: '72px',
 } satisfies CSSProperties;
 
-type VirtuosoContext = { selectedItemId: number | null };
-
 /**
  * Scroll the item at the given index into view if it is not currently visible.
  */
@@ -132,28 +130,26 @@ const useScrollableStagingArea = (rootRef: RefObject<HTMLDivElement>) => {
 };
 
 export const StagingAreaItemsList = memo(() => {
-  const canvasManager = useCanvasManagerSafe();
-  const ctx = useCanvasSessionContext();
+  const canvasManager = useCanvasManager();
+
+  const ctx = useStagingAreaContext();
   const virtuosoRef = useRef<VirtuosoHandle>(null);
   const rangeRef = useRef<ListRange>({ startIndex: 0, endIndex: 0 });
   const rootRef = useRef<HTMLDivElement>(null);
 
   const items = useStore(ctx.$items);
-  const selectedItemId = useStore(ctx.$selectedItemId);
 
-  const context = useMemo(() => ({ selectedItemId }), [selectedItemId]);
   const scrollerRef = useScrollableStagingArea(rootRef);
 
   useEffect(() => {
-    if (!canvasManager) {
-      return;
-    }
-
-    return canvasManager.stagingArea.connectToSession(ctx.$items, ctx.$selectedItemId, ctx.$progressData);
-  }, [canvasManager, ctx.$progressData, ctx.$selectedItemId, ctx.$items]);
+    return canvasManager.stagingArea.connectToSession(ctx.$items, ctx.$selectedItem);
+  }, [canvasManager, ctx.$progressData, ctx.$items, ctx.$selectedItem]);
 
   useEffect(() => {
-    return ctx.$selectedItemIndex.listen((index) => {
+    return ctx.$selectedItemIndex.listen((selectedItemIndex) => {
+      if (selectedItemIndex === null) {
+        return;
+      }
       if (!virtuosoRef.current) {
         return;
       }
@@ -162,11 +158,7 @@ export const StagingAreaItemsList = memo(() => {
         return;
       }
 
-      if (index === null) {
-        return;
-      }
-
-      scrollIntoView(index, rootRef.current, virtuosoRef.current, rangeRef.current);
+      scrollIntoView(selectedItemIndex, rootRef.current, virtuosoRef.current, rangeRef.current);
     });
   }, [ctx.$selectedItemIndex]);
 
@@ -176,40 +168,46 @@ export const StagingAreaItemsList = memo(() => {
 
   return (
     <Box data-overlayscrollbars-initialize="" ref={rootRef} position="relative" w="full" h="full">
-      <Virtuoso<S['SessionQueueItem'], VirtuosoContext>
+      <Virtuoso<S['SessionQueueItem']>
         ref={virtuosoRef}
-        context={context}
         data={items}
         horizontalDirection
         style={virtuosoStyles}
+        computeItemKey={computeItemKey}
+        increaseViewportBy={2048}
         itemContent={itemContent}
         components={components}
         rangeChanged={onRangeChanged}
         // Virtuoso expects the ref to be of HTMLElement | null | Window, but overlayscrollbars doesn't allow Window
-        scrollerRef={scrollerRef as VirtuosoProps<S['SessionQueueItem'], VirtuosoContext>['scrollerRef']}
+        scrollerRef={scrollerRef as VirtuosoProps<S['SessionQueueItem'], void>['scrollerRef']}
       />
     </Box>
   );
 });
 StagingAreaItemsList.displayName = 'StagingAreaItemsList';
 
-const itemContent: ItemContent<S['SessionQueueItem'], VirtuosoContext> = (index, item, { selectedItemId }) => (
-  <QueueItemPreviewMini
-    key={`${item.item_id}-mini`}
-    item={item}
-    index={index}
-    isSelected={selectedItemId === item.item_id}
-  />
+const computeItemKey: ComputeItemKey<S['SessionQueueItem'], void> = (_, item: S['SessionQueueItem']) => {
+  return item.item_id;
+};
+
+const itemContent: ItemContent<S['SessionQueueItem'], void> = (index, item) => (
+  <QueueItemPreviewMini key={`${item.item_id}-mini`} item={item} index={index} />
 );
 
 const listSx = {
   '& > * + *': {
     pl: 2,
   },
+  '&[data-disabled="true"]': {
+    filter: 'grayscale(1) opacity(0.5)',
+  },
 };
 
-const components: Components<S['SessionQueueItem'], VirtuosoContext> = {
+const components: Components<S['SessionQueueItem']> = {
   List: forwardRef(({ context: _, ...rest }, ref) => {
-    return <Flex ref={ref} sx={listSx} {...rest} />;
+    const canvasManager = useCanvasManager();
+    const shouldShowStagedImage = useStore(canvasManager.stagingArea.$shouldShowStagedImage);
+
+    return <Flex ref={ref} sx={listSx} data-disabled={!shouldShowStagedImage} {...rest} />;
   }),
 };

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaToolbar.tsx

@@ -1,6 +1,5 @@
 import { ButtonGroup, Flex } from '@invoke-ai/ui-library';
-import { useStore } from '@nanostores/react';
-import { useCanvasSessionContext } from 'features/controlLayers/components/SimpleSession/context';
+import { useStagingAreaContext } from 'features/controlLayers/components/StagingArea/context';
 import { StagingAreaToolbarAcceptButton } from 'features/controlLayers/components/StagingArea/StagingAreaToolbarAcceptButton';
 import { StagingAreaToolbarDiscardAllButton } from 'features/controlLayers/components/StagingArea/StagingAreaToolbarDiscardAllButton';
 import { StagingAreaToolbarDiscardSelectedButton } from 'features/controlLayers/components/StagingArea/StagingAreaToolbarDiscardSelectedButton';
@@ -10,17 +9,13 @@ import { StagingAreaToolbarNextButton } from 'features/controlLayers/components/
 import { StagingAreaToolbarPrevButton } from 'features/controlLayers/components/StagingArea/StagingAreaToolbarPrevButton';
 import { StagingAreaToolbarSaveSelectedToGalleryButton } from 'features/controlLayers/components/StagingArea/StagingAreaToolbarSaveSelectedToGalleryButton';
 import { StagingAreaToolbarToggleShowResultsButton } from 'features/controlLayers/components/StagingArea/StagingAreaToolbarToggleShowResultsButton';
-import { useCanvasManager } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
 import { memo } from 'react';
 import { useHotkeys } from 'react-hotkeys-hook';
 
 import { StagingAreaAutoSwitchButtons } from './StagingAreaAutoSwitchButtons';
 
 export const StagingAreaToolbar = memo(() => {
-  const canvasManager = useCanvasManager();
-  const shouldShowStagedImage = useStore(canvasManager.stagingArea.$shouldShowStagedImage);
-
-  const ctx = useCanvasSessionContext();
+  const ctx = useStagingAreaContext();
 
   useHotkeys('meta+left', ctx.selectFirst, { preventDefault: true });
   useHotkeys('meta+right', ctx.selectLast, { preventDefault: true });
@@ -28,22 +23,22 @@ export const StagingAreaToolbar = memo(() => {
   return (
     <Flex gap={2}>
       <ButtonGroup borderRadius="base" shadow="dark-lg">
-        <StagingAreaToolbarPrevButton isDisabled={!shouldShowStagedImage} />
+        <StagingAreaToolbarPrevButton />
         <StagingAreaToolbarImageCountButton />
-        <StagingAreaToolbarNextButton isDisabled={!shouldShowStagedImage} />
+        <StagingAreaToolbarNextButton />
       </ButtonGroup>
       <ButtonGroup borderRadius="base" shadow="dark-lg">
         <StagingAreaToolbarAcceptButton />
         <StagingAreaToolbarToggleShowResultsButton />
         <StagingAreaToolbarSaveSelectedToGalleryButton />
         <StagingAreaToolbarMenu />
-        <StagingAreaToolbarDiscardSelectedButton isDisabled={!shouldShowStagedImage} />
+        <StagingAreaToolbarDiscardSelectedButton />
       </ButtonGroup>
       <ButtonGroup borderRadius="base" shadow="dark-lg">
         <StagingAreaAutoSwitchButtons />
       </ButtonGroup>
       <ButtonGroup borderRadius="base" shadow="dark-lg">
-        <StagingAreaToolbarDiscardAllButton isDisabled={!shouldShowStagedImage} />
+        <StagingAreaToolbarDiscardAllButton />
       </ButtonGroup>
     </Flex>
   );

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaToolbarAcceptButton.tsx

@@ -1,65 +1,32 @@
 import { IconButton } from '@invoke-ai/ui-library';
 import { useStore } from '@nanostores/react';
-import { useAppDispatch, useAppSelector } from 'app/store/storeHooks';
 import { useIsRegionFocused } from 'common/hooks/focus';
-import { useCanvasSessionContext } from 'features/controlLayers/components/SimpleSession/context';
+import { useStagingAreaContext } from 'features/controlLayers/components/StagingArea/context';
 import { useCanvasManager } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
-import { rasterLayerAdded } from 'features/controlLayers/store/canvasSlice';
-import { canvasSessionReset, selectCanvasSessionId } from 'features/controlLayers/store/canvasStagingAreaSlice';
-import { selectBboxRect, selectSelectedEntityIdentifier } from 'features/controlLayers/store/selectors';
-import type { CanvasRasterLayerState } from 'features/controlLayers/store/types';
-import { imageNameToImageObject } from 'features/controlLayers/store/util';
 import { useCancelQueueItemsByDestination } from 'features/queue/hooks/useCancelQueueItemsByDestination';
-import { memo, useCallback } from 'react';
+import { memo } from 'react';
 import { useHotkeys } from 'react-hotkeys-hook';
 import { useTranslation } from 'react-i18next';
 import { PiCheckBold } from 'react-icons/pi';
 
 export const StagingAreaToolbarAcceptButton = memo(() => {
-  const ctx = useCanvasSessionContext();
-  const dispatch = useAppDispatch();
+  const ctx = useStagingAreaContext();
   const canvasManager = useCanvasManager();
-  const canvasSessionId = useAppSelector(selectCanvasSessionId);
-  const bboxRect = useAppSelector(selectBboxRect);
-  const selectedEntityIdentifier = useAppSelector(selectSelectedEntityIdentifier);
   const shouldShowStagedImage = useStore(canvasManager.stagingArea.$shouldShowStagedImage);
   const isCanvasFocused = useIsRegionFocused('canvas');
-  const selectedItemImageDTO = useStore(ctx.$selectedItemOutputImageDTO);
   const cancelQueueItemsByDestination = useCancelQueueItemsByDestination();
+  const acceptSelectedIsEnabled = useStore(ctx.$acceptSelectedIsEnabled);
 
   const { t } = useTranslation();
 
-  const acceptSelected = useCallback(() => {
-    if (!selectedItemImageDTO) {
-      return;
-    }
-    const { x, y, width, height } = bboxRect;
-    const imageObject = imageNameToImageObject(selectedItemImageDTO.image_name, { width, height });
-    const overrides: Partial<CanvasRasterLayerState> = {
-      position: { x, y },
-      objects: [imageObject],
-    };
-
-    dispatch(rasterLayerAdded({ overrides, isSelected: selectedEntityIdentifier?.type === 'raster_layer' }));
-    dispatch(canvasSessionReset());
-    cancelQueueItemsByDestination.trigger(canvasSessionId, { withToast: false });
-  }, [
-    selectedItemImageDTO,
-    bboxRect,
-    dispatch,
-    selectedEntityIdentifier?.type,
-    cancelQueueItemsByDestination,
-    canvasSessionId,
-  ]);
-
   useHotkeys(
     ['enter'],
-    acceptSelected,
+    ctx.acceptSelected,
     {
       preventDefault: true,
-      enabled: isCanvasFocused && shouldShowStagedImage && selectedItemImageDTO !== null,
+      enabled: isCanvasFocused && shouldShowStagedImage && acceptSelectedIsEnabled,
     },
-    [isCanvasFocused, shouldShowStagedImage, selectedItemImageDTO]
+    [ctx.acceptSelected, isCanvasFocused, shouldShowStagedImage, acceptSelectedIsEnabled]
   );
 
   return (
@@ -67,9 +34,9 @@ export const StagingAreaToolbarAcceptButton = memo(() => {
       tooltip={`${t('common.accept')} (Enter)`}
       aria-label={`${t('common.accept')} (Enter)`}
       icon={<PiCheckBold />}
-      onClick={acceptSelected}
+      onClick={ctx.acceptSelected}
       colorScheme="invokeBlue"
-      isDisabled={!selectedItemImageDTO || !shouldShowStagedImage || cancelQueueItemsByDestination.isDisabled}
+      isDisabled={!acceptSelectedIsEnabled || !shouldShowStagedImage || cancelQueueItemsByDestination.isDisabled}
       isLoading={cancelQueueItemsByDestination.isLoading}
     />
   );

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaToolbarDiscardAllButton.tsx

@@ -1,31 +1,28 @@
 import { IconButton } from '@invoke-ai/ui-library';
-import { useAppSelector } from 'app/store/storeHooks';
-import { useCanvasSessionContext } from 'features/controlLayers/components/SimpleSession/context';
-import { selectCanvasSessionId } from 'features/controlLayers/store/canvasStagingAreaSlice';
+import { useStore } from '@nanostores/react';
+import { useStagingAreaContext } from 'features/controlLayers/components/StagingArea/context';
+import { useCanvasManager } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
 import { useCancelQueueItemsByDestination } from 'features/queue/hooks/useCancelQueueItemsByDestination';
-import { memo, useCallback } from 'react';
+import { memo } from 'react';
 import { useTranslation } from 'react-i18next';
 import { PiTrashSimpleBold } from 'react-icons/pi';
 
-export const StagingAreaToolbarDiscardAllButton = memo(({ isDisabled }: { isDisabled?: boolean }) => {
-  const ctx = useCanvasSessionContext();
+export const StagingAreaToolbarDiscardAllButton = memo(() => {
+  const canvasManager = useCanvasManager();
+  const shouldShowStagedImage = useStore(canvasManager.stagingArea.$shouldShowStagedImage);
+
+  const ctx = useStagingAreaContext();
   const { t } = useTranslation();
   const cancelQueueItemsByDestination = useCancelQueueItemsByDestination();
-  const canvasSessionId = useAppSelector(selectCanvasSessionId);
-
-  const discardAll = useCallback(() => {
-    ctx.discardAll();
-    cancelQueueItemsByDestination.trigger(canvasSessionId, { withToast: false });
-  }, [cancelQueueItemsByDestination, ctx, canvasSessionId]);
 
   return (
     <IconButton
       tooltip={`${t('controlLayers.stagingArea.discardAll')} (Esc)`}
       aria-label={t('controlLayers.stagingArea.discardAll')}
       icon={<PiTrashSimpleBold />}
-      onClick={discardAll}
+      onClick={ctx.discardAll}
       colorScheme="error"
-      isDisabled={isDisabled || cancelQueueItemsByDestination.isDisabled}
+      isDisabled={cancelQueueItemsByDestination.isDisabled || !shouldShowStagedImage}
       isLoading={cancelQueueItemsByDestination.isLoading}
     />
   );

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaToolbarDiscardSelectedButton.tsx

@@ -1,34 +1,30 @@
 import { IconButton } from '@invoke-ai/ui-library';
 import { useStore } from '@nanostores/react';
-import { useCanvasSessionContext } from 'features/controlLayers/components/SimpleSession/context';
+import { useStagingAreaContext } from 'features/controlLayers/components/StagingArea/context';
+import { useCanvasManager } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
 import { useCancelQueueItem } from 'features/queue/hooks/useCancelQueueItem';
-import { memo, useCallback } from 'react';
+import { memo } from 'react';
 import { useTranslation } from 'react-i18next';
 import { PiXBold } from 'react-icons/pi';
 
-export const StagingAreaToolbarDiscardSelectedButton = memo(({ isDisabled }: { isDisabled?: boolean }) => {
-  const ctx = useCanvasSessionContext();
+export const StagingAreaToolbarDiscardSelectedButton = memo(() => {
+  const canvasManager = useCanvasManager();
+  const shouldShowStagedImage = useStore(canvasManager.stagingArea.$shouldShowStagedImage);
+
+  const ctx = useStagingAreaContext();
   const cancelQueueItem = useCancelQueueItem();
-  const selectedItemId = useStore(ctx.$selectedItemId);
+  const discardSelectedIsEnabled = useStore(ctx.$discardSelectedIsEnabled);
 
   const { t } = useTranslation();
 
-  const discardSelected = useCallback(async () => {
-    if (selectedItemId === null) {
-      return;
-    }
-    ctx.discard(selectedItemId);
-    await cancelQueueItem.trigger(selectedItemId, { withToast: false });
-  }, [selectedItemId, ctx, cancelQueueItem]);
-
   return (
     <IconButton
       tooltip={t('controlLayers.stagingArea.discard')}
       aria-label={t('controlLayers.stagingArea.discard')}
       icon={<PiXBold />}
-      onClick={discardSelected}
+      onClick={ctx.discardSelected}
       colorScheme="invokeBlue"
-      isDisabled={selectedItemId === null || cancelQueueItem.isDisabled || isDisabled}
+      isDisabled={!discardSelectedIsEnabled || cancelQueueItem.isDisabled || !shouldShowStagedImage}
       isLoading={cancelQueueItem.isLoading}
     />
   );

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaToolbarImageCountButton.tsx

@@ -1,23 +1,27 @@
 import { Button } from '@invoke-ai/ui-library';
 import { useStore } from '@nanostores/react';
-import { useCanvasSessionContext } from 'features/controlLayers/components/SimpleSession/context';
+import { useStagingAreaContext } from 'features/controlLayers/components/StagingArea/context';
+import { useCanvasManager } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
 import { memo, useMemo } from 'react';
 
 export const StagingAreaToolbarImageCountButton = memo(() => {
-  const ctx = useCanvasSessionContext();
-  const selectItemIndex = useStore(ctx.$selectedItemIndex);
+  const canvasManager = useCanvasManager();
+  const shouldShowStagedImage = useStore(canvasManager.stagingArea.$shouldShowStagedImage);
+
+  const ctx = useStagingAreaContext();
+  const selectedItem = useStore(ctx.$selectedItem);
   const itemCount = useStore(ctx.$itemCount);
 
   const counterText = useMemo(() => {
-    if (itemCount > 0 && selectItemIndex !== null) {
-      return `${selectItemIndex + 1} of ${itemCount}`;
+    if (itemCount > 0 && selectedItem !== null) {
+      return `${selectedItem.index + 1} of ${itemCount}`;
     } else {
       return `0 of 0`;
     }
-  }, [itemCount, selectItemIndex]);
+  }, [itemCount, selectedItem]);
 
   return (
-    <Button colorScheme="base" pointerEvents="none" minW={28}>
+    <Button colorScheme="base" pointerEvents="none" minW={28} isDisabled={!shouldShowStagedImage}>
       {counterText}
     </Button>
   );

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaToolbarMenu.tsx

@@ -1,12 +1,23 @@
 import { IconButton, Menu, MenuButton, MenuList } from '@invoke-ai/ui-library';
+import { useStore } from '@nanostores/react';
 import { StagingAreaToolbarNewLayerFromImageMenuItems } from 'features/controlLayers/components/StagingArea/StagingAreaToolbarMenuNewLayerFromImage';
+import { useCanvasManager } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
 import { memo } from 'react';
 import { PiDotsThreeVerticalBold } from 'react-icons/pi';
 
 export const StagingAreaToolbarMenu = memo(() => {
+  const canvasManager = useCanvasManager();
+  const shouldShowStagedImage = useStore(canvasManager.stagingArea.$shouldShowStagedImage);
+
   return (
     <Menu>
-      <MenuButton as={IconButton} icon={<PiDotsThreeVerticalBold />} colorScheme="invokeBlue" />
+      <MenuButton
+        tooltip="Image Actions"
+        as={IconButton}
+        icon={<PiDotsThreeVerticalBold />}
+        colorScheme="invokeBlue"
+        isDisabled={!shouldShowStagedImage}
+      />
       <MenuList>
         <StagingAreaToolbarNewLayerFromImageMenuItems />
       </MenuList>

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaToolbarMenuNewLayerFromImage.tsx

@@ -2,7 +2,7 @@ import { MenuGroup, MenuItem } from '@invoke-ai/ui-library';
 import { useStore } from '@nanostores/react';
 import { useAppStore } from 'app/store/storeHooks';
 import { NewLayerIcon } from 'features/controlLayers/components/common/icons';
-import { useCanvasSessionContext } from 'features/controlLayers/components/SimpleSession/context';
+import { useStagingAreaContext } from 'features/controlLayers/components/StagingArea/context';
 import { useCanvasManager } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
 import { createNewCanvasEntityFromImage } from 'features/imageActions/actions';
 import { toast } from 'features/toast/toast';
@@ -15,8 +15,8 @@ const uploadImageArg = { image_category: 'general', is_intermediate: true, silen
 export const StagingAreaToolbarNewLayerFromImageMenuItems = memo(() => {
   const canvasManager = useCanvasManager();
   const { t } = useTranslation();
-  const ctx = useCanvasSessionContext();
-  const selectedItemOutputImageDTO = useStore(ctx.$selectedItemOutputImageDTO);
+  const ctx = useStagingAreaContext();
+  const selectedItemImageDTO = useStore(ctx.$selectedItemImageDTO);
   const store = useAppStore();
   const shouldShowStagedImage = useStore(canvasManager.stagingArea.$shouldShowStagedImage);
 
@@ -29,11 +29,11 @@ export const StagingAreaToolbarNewLayerFromImageMenuItems = memo(() => {
   }, [t]);
 
   const onClickNewRasterLayerFromImage = useCallback(async () => {
-    if (!selectedItemOutputImageDTO) {
+    if (!selectedItemImageDTO) {
       return;
     }
     const { dispatch, getState } = store;
-    const imageDTO = await copyImage(selectedItemOutputImageDTO.image_name, uploadImageArg);
+    const imageDTO = await copyImage(selectedItemImageDTO.image_name, uploadImageArg);
     createNewCanvasEntityFromImage({
       imageDTO,
       type: 'raster_layer',
@@ -42,14 +42,14 @@ export const StagingAreaToolbarNewLayerFromImageMenuItems = memo(() => {
       overrides: { isEnabled: false }, // We are adding the layer while staging, it should be disabled by default
     });
     toastSentToCanvas();
-  }, [selectedItemOutputImageDTO, store, toastSentToCanvas]);
+  }, [selectedItemImageDTO, store, toastSentToCanvas]);
 
   const onClickNewControlLayerFromImage = useCallback(async () => {
-    if (!selectedItemOutputImageDTO) {
+    if (!selectedItemImageDTO) {
       return;
     }
     const { dispatch, getState } = store;
-    const imageDTO = await copyImage(selectedItemOutputImageDTO.image_name, uploadImageArg);
+    const imageDTO = await copyImage(selectedItemImageDTO.image_name, uploadImageArg);
 
     createNewCanvasEntityFromImage({
       imageDTO,
@@ -59,14 +59,14 @@ export const StagingAreaToolbarNewLayerFromImageMenuItems = memo(() => {
       overrides: { isEnabled: false }, // We are adding the layer while staging, it should be disabled by default
     });
     toastSentToCanvas();
-  }, [selectedItemOutputImageDTO, store, toastSentToCanvas]);
+  }, [selectedItemImageDTO, store, toastSentToCanvas]);
 
   const onClickNewInpaintMaskFromImage = useCallback(async () => {
-    if (!selectedItemOutputImageDTO) {
+    if (!selectedItemImageDTO) {
       return;
     }
     const { dispatch, getState } = store;
-    const imageDTO = await copyImage(selectedItemOutputImageDTO.image_name, uploadImageArg);
+    const imageDTO = await copyImage(selectedItemImageDTO.image_name, uploadImageArg);
 
     createNewCanvasEntityFromImage({
       imageDTO,
@@ -76,14 +76,14 @@ export const StagingAreaToolbarNewLayerFromImageMenuItems = memo(() => {
       overrides: { isEnabled: false }, // We are adding the layer while staging, it should be disabled by default
     });
     toastSentToCanvas();
-  }, [selectedItemOutputImageDTO, store, toastSentToCanvas]);
+  }, [selectedItemImageDTO, store, toastSentToCanvas]);
 
   const onClickNewRegionalGuidanceFromImage = useCallback(async () => {
-    if (!selectedItemOutputImageDTO) {
+    if (!selectedItemImageDTO) {
       return;
     }
     const { dispatch, getState } = store;
-    const imageDTO = await copyImage(selectedItemOutputImageDTO.image_name, uploadImageArg);
+    const imageDTO = await copyImage(selectedItemImageDTO.image_name, uploadImageArg);
 
     createNewCanvasEntityFromImage({
       imageDTO,
@@ -93,35 +93,35 @@ export const StagingAreaToolbarNewLayerFromImageMenuItems = memo(() => {
       overrides: { isEnabled: false }, // We are adding the layer while staging, it should be disabled by default
     });
     toastSentToCanvas();
-  }, [selectedItemOutputImageDTO, store, toastSentToCanvas]);
+  }, [selectedItemImageDTO, store, toastSentToCanvas]);
 
   return (
     <MenuGroup title="New Layer From Image">
       <MenuItem
         icon={<NewLayerIcon />}
         onClickCapture={onClickNewInpaintMaskFromImage}
-        isDisabled={!selectedItemOutputImageDTO || !shouldShowStagedImage}
+        isDisabled={!selectedItemImageDTO || !shouldShowStagedImage}
       >
         {t('controlLayers.inpaintMask')}
       </MenuItem>
       <MenuItem
         icon={<NewLayerIcon />}
         onClickCapture={onClickNewRegionalGuidanceFromImage}
-        isDisabled={!selectedItemOutputImageDTO || !shouldShowStagedImage}
+        isDisabled={!selectedItemImageDTO || !shouldShowStagedImage}
       >
         {t('controlLayers.regionalGuidance')}
       </MenuItem>
       <MenuItem
         icon={<NewLayerIcon />}
         onClickCapture={onClickNewControlLayerFromImage}
-        isDisabled={!selectedItemOutputImageDTO || !shouldShowStagedImage}
+        isDisabled={!selectedItemImageDTO || !shouldShowStagedImage}
       >
         {t('controlLayers.controlLayer')}
       </MenuItem>
       <MenuItem
         icon={<NewLayerIcon />}
         onClickCapture={onClickNewRasterLayerFromImage}
-        isDisabled={!selectedItemOutputImageDTO || !shouldShowStagedImage}
+        isDisabled={!selectedItemImageDTO || !shouldShowStagedImage}
       >
         {t('controlLayers.rasterLayer')}
       </MenuItem>

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaToolbarNextButton.tsx

@@ -1,14 +1,18 @@
 import { IconButton } from '@invoke-ai/ui-library';
 import { useStore } from '@nanostores/react';
 import { useIsRegionFocused } from 'common/hooks/focus';
-import { useCanvasSessionContext } from 'features/controlLayers/components/SimpleSession/context';
+import { useStagingAreaContext } from 'features/controlLayers/components/StagingArea/context';
+import { useCanvasManager } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
 import { memo, useCallback } from 'react';
 import { useHotkeys } from 'react-hotkeys-hook';
 import { useTranslation } from 'react-i18next';
 import { PiArrowRightBold } from 'react-icons/pi';
 
-export const StagingAreaToolbarNextButton = memo(({ isDisabled }: { isDisabled?: boolean }) => {
-  const ctx = useCanvasSessionContext();
+export const StagingAreaToolbarNextButton = memo(() => {
+  const canvasManager = useCanvasManager();
+  const shouldShowStagedImage = useStore(canvasManager.stagingArea.$shouldShowStagedImage);
+
+  const ctx = useStagingAreaContext();
   const itemCount = useStore(ctx.$itemCount);
   const isCanvasFocused = useIsRegionFocused('canvas');
 
@@ -23,9 +27,9 @@ export const StagingAreaToolbarNextButton = memo(({ isDisabled }: { isDisabled?:
     ctx.selectNext,
     {
       preventDefault: true,
-      enabled: isCanvasFocused && !isDisabled && itemCount > 1,
+      enabled: isCanvasFocused && shouldShowStagedImage && itemCount > 1,
     },
-    [isCanvasFocused, isDisabled, itemCount, ctx.selectNext]
+    [isCanvasFocused, shouldShowStagedImage, itemCount, ctx.selectNext]
   );
 
   return (
@@ -35,7 +39,7 @@ export const StagingAreaToolbarNextButton = memo(({ isDisabled }: { isDisabled?:
       icon={<PiArrowRightBold />}
       onClick={selectNext}
       colorScheme="invokeBlue"
-      isDisabled={itemCount <= 1 || isDisabled}
+      isDisabled={itemCount <= 1 || !shouldShowStagedImage}
     />
   );
 });

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaToolbarPrevButton.tsx

@@ -1,14 +1,17 @@
 import { IconButton } from '@invoke-ai/ui-library';
 import { useStore } from '@nanostores/react';
 import { useIsRegionFocused } from 'common/hooks/focus';
-import { useCanvasSessionContext } from 'features/controlLayers/components/SimpleSession/context';
+import { useStagingAreaContext } from 'features/controlLayers/components/StagingArea/context';
+import { useCanvasManager } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
 import { memo, useCallback } from 'react';
 import { useHotkeys } from 'react-hotkeys-hook';
 import { useTranslation } from 'react-i18next';
 import { PiArrowLeftBold } from 'react-icons/pi';
 
-export const StagingAreaToolbarPrevButton = memo(({ isDisabled }: { isDisabled?: boolean }) => {
-  const ctx = useCanvasSessionContext();
+export const StagingAreaToolbarPrevButton = memo(() => {
+  const canvasManager = useCanvasManager();
+  const shouldShowStagedImage = useStore(canvasManager.stagingArea.$shouldShowStagedImage);
+  const ctx = useStagingAreaContext();
   const itemCount = useStore(ctx.$itemCount);
   const isCanvasFocused = useIsRegionFocused('canvas');
 
@@ -23,9 +26,9 @@ export const StagingAreaToolbarPrevButton = memo(({ isDisabled }: { isDisabled?:
     ctx.selectPrev,
     {
       preventDefault: true,
-      enabled: isCanvasFocused && !isDisabled && itemCount > 1,
+      enabled: isCanvasFocused && shouldShowStagedImage && itemCount > 1,
     },
-    [isCanvasFocused, isDisabled, itemCount, ctx.selectPrev]
+    [isCanvasFocused, shouldShowStagedImage, itemCount, ctx.selectPrev]
   );
 
   return (
@@ -35,7 +38,7 @@ export const StagingAreaToolbarPrevButton = memo(({ isDisabled }: { isDisabled?:
       icon={<PiArrowLeftBold />}
       onClick={selectPrev}
       colorScheme="invokeBlue"
-      isDisabled={itemCount <= 1 || isDisabled}
+      isDisabled={itemCount <= 1 || !shouldShowStagedImage}
     />
   );
 });

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaToolbarSaveSelectedToGalleryButton.tsx

@@ -2,7 +2,7 @@ import { IconButton } from '@invoke-ai/ui-library';
 import { useStore } from '@nanostores/react';
 import { useAppSelector } from 'app/store/storeHooks';
 import { withResultAsync } from 'common/util/result';
-import { useCanvasSessionContext } from 'features/controlLayers/components/SimpleSession/context';
+import { useStagingAreaContext } from 'features/controlLayers/components/StagingArea/context';
 import { useCanvasManager } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
 import { selectAutoAddBoardId } from 'features/gallery/store/gallerySelectors';
 import { toast } from 'features/toast/toast';
@@ -16,14 +16,14 @@ const TOAST_ID = 'SAVE_STAGING_AREA_IMAGE_TO_GALLERY';
 export const StagingAreaToolbarSaveSelectedToGalleryButton = memo(() => {
   const canvasManager = useCanvasManager();
   const autoAddBoardId = useAppSelector(selectAutoAddBoardId);
-  const ctx = useCanvasSessionContext();
-  const selectedItemOutputImageDTO = useStore(ctx.$selectedItemOutputImageDTO);
+  const ctx = useStagingAreaContext();
+  const selectedItemImageDTO = useStore(ctx.$selectedItemImageDTO);
   const shouldShowStagedImage = useStore(canvasManager.stagingArea.$shouldShowStagedImage);
 
   const { t } = useTranslation();
 
   const saveSelectedImageToGallery = useCallback(async () => {
-    if (!selectedItemOutputImageDTO) {
+    if (!selectedItemImageDTO) {
       return;
     }
 
@@ -31,7 +31,7 @@ export const StagingAreaToolbarSaveSelectedToGalleryButton = memo(() => {
     // the gallery without borking the canvas, which may need this image to exist.
     const result = await withResultAsync(async () => {
       // Create a new file with the same name, which we will upload
-      await copyImage(selectedItemOutputImageDTO.image_name, {
+      await copyImage(selectedItemImageDTO.image_name, {
         // Image should show up in the Images tab
         image_category: 'general',
         is_intermediate: false,
@@ -55,7 +55,7 @@ export const StagingAreaToolbarSaveSelectedToGalleryButton = memo(() => {
         status: 'error',
       });
     }
-  }, [autoAddBoardId, selectedItemOutputImageDTO, t]);
+  }, [autoAddBoardId, selectedItemImageDTO, t]);
 
   return (
     <IconButton
@@ -64,7 +64,7 @@ export const StagingAreaToolbarSaveSelectedToGalleryButton = memo(() => {
       icon={<PiFloppyDiskBold />}
       onClick={saveSelectedImageToGallery}
       colorScheme="invokeBlue"
-      isDisabled={!selectedItemOutputImageDTO || !shouldShowStagedImage}
+      isDisabled={!selectedItemImageDTO || !shouldShowStagedImage}
     />
   );
 });

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/__mocks__/mockStagingAreaApp.ts

@@ -0,0 +1,193 @@
+import { merge } from 'es-toolkit';
+import type { StagingAreaAppApi } from 'features/controlLayers/components/StagingArea/state';
+import type { AutoSwitchMode } from 'features/controlLayers/store/canvasSettingsSlice';
+import type { ImageDTO, S } from 'services/api/types';
+import type { PartialDeep } from 'type-fest';
+import { vi } from 'vitest';
+
+export const createMockStagingAreaApp = (): StagingAreaAppApi & {
+  // Additional methods for testing
+  _triggerItemsChanged: (items: S['SessionQueueItem'][]) => void;
+  _triggerQueueItemStatusChanged: (data: S['QueueItemStatusChangedEvent']) => void;
+  _triggerInvocationProgress: (data: S['InvocationProgressEvent']) => void;
+  _setAutoSwitchMode: (mode: AutoSwitchMode) => void;
+  _setImageDTO: (imageName: string, imageDTO: ImageDTO | null) => void;
+  _setLoadImageDelay: (delay: number) => void;
+} => {
+  const itemsChangedHandlers = new Set<(items: S['SessionQueueItem'][]) => void>();
+  const queueItemStatusChangedHandlers = new Set<(data: S['QueueItemStatusChangedEvent']) => void>();
+  const invocationProgressHandlers = new Set<(data: S['InvocationProgressEvent']) => void>();
+
+  let autoSwitchMode: AutoSwitchMode = 'switch_on_start';
+  const imageDTOs = new Map<string, ImageDTO | null>();
+  let loadImageDelay = 0;
+
+  return {
+    onDiscard: vi.fn(),
+    onDiscardAll: vi.fn(),
+    onAccept: vi.fn(),
+    onSelect: vi.fn(),
+    onSelectPrev: vi.fn(),
+    onSelectNext: vi.fn(),
+    onSelectFirst: vi.fn(),
+    onSelectLast: vi.fn(),
+    getAutoSwitch: vi.fn(() => autoSwitchMode),
+    onAutoSwitchChange: vi.fn(),
+    getImageDTO: vi.fn((imageName: string) => {
+      return Promise.resolve(imageDTOs.get(imageName) || null);
+    }),
+    loadImage: vi.fn(async (imageName: string) => {
+      if (loadImageDelay > 0) {
+        await new Promise((resolve) => {
+          setTimeout(resolve, loadImageDelay);
+        });
+      }
+      // Mock HTMLImageElement for testing environment
+      const mockImage = {
+        src: imageName,
+        width: 512,
+        height: 512,
+        onload: null,
+        onerror: null,
+      } as HTMLImageElement;
+      return mockImage;
+    }),
+    onItemsChanged: vi.fn((handler) => {
+      itemsChangedHandlers.add(handler);
+      return () => itemsChangedHandlers.delete(handler);
+    }),
+    onQueueItemStatusChanged: vi.fn((handler) => {
+      queueItemStatusChangedHandlers.add(handler);
+      return () => queueItemStatusChangedHandlers.delete(handler);
+    }),
+    onInvocationProgress: vi.fn((handler) => {
+      invocationProgressHandlers.add(handler);
+      return () => invocationProgressHandlers.delete(handler);
+    }),
+
+    // Testing helper methods
+    _triggerItemsChanged: (items: S['SessionQueueItem'][]) => {
+      itemsChangedHandlers.forEach((handler) => handler(items));
+    },
+    _triggerQueueItemStatusChanged: (data: S['QueueItemStatusChangedEvent']) => {
+      queueItemStatusChangedHandlers.forEach((handler) => handler(data));
+    },
+    _triggerInvocationProgress: (data: S['InvocationProgressEvent']) => {
+      invocationProgressHandlers.forEach((handler) => handler(data));
+    },
+    _setAutoSwitchMode: (mode: AutoSwitchMode) => {
+      autoSwitchMode = mode;
+    },
+    _setImageDTO: (imageName: string, imageDTO: ImageDTO | null) => {
+      imageDTOs.set(imageName, imageDTO);
+    },
+    _setLoadImageDelay: (delay: number) => {
+      loadImageDelay = delay;
+    },
+  };
+};
+
+export const createMockQueueItem = (overrides: PartialDeep<S['SessionQueueItem']> = {}): S['SessionQueueItem'] =>
+  merge(
+    {
+      item_id: 1,
+      batch_id: 'test-batch-id',
+      session_id: 'test-session',
+      queue_id: 'test-queue-id',
+      status: 'pending',
+      priority: 0,
+      origin: null,
+      destination: 'test-session',
+      error_type: null,
+      error_message: null,
+      error_traceback: null,
+      created_at: '2024-01-01T00:00:00Z',
+      updated_at: '2024-01-01T00:00:00Z',
+      started_at: null,
+      completed_at: null,
+      field_values: null,
+      retried_from_item_id: null,
+      is_api_validation_run: false,
+      published_workflow_id: null,
+      session: {
+        id: 'test-session',
+        graph: {},
+        execution_graph: {},
+        executed: [],
+        executed_history: [],
+        results: {
+          'test-node-id': {
+            image: {
+              image_name: 'test-image.png',
+            },
+          },
+        },
+        errors: {},
+        prepared_source_mapping: {},
+        source_prepared_mapping: {
+          canvas_output: ['test-node-id'],
+        },
+      },
+      workflow: null,
+    },
+    overrides
+  ) as S['SessionQueueItem'];
+
+export const createMockImageDTO = (overrides: Partial<ImageDTO> = {}): ImageDTO => ({
+  image_name: 'test-image.png',
+  image_url: 'http://test.com/test-image.png',
+  thumbnail_url: 'http://test.com/test-image-thumb.png',
+  image_origin: 'internal',
+  image_category: 'general',
+  width: 512,
+  height: 512,
+  created_at: '2024-01-01T00:00:00Z',
+  updated_at: '2024-01-01T00:00:00Z',
+  deleted_at: null,
+  is_intermediate: false,
+  starred: false,
+  has_workflow: false,
+  session_id: 'test-session',
+  node_id: 'test-node-id',
+  board_id: null,
+  ...overrides,
+});
+
+export const createMockProgressEvent = (
+  overrides: PartialDeep<S['InvocationProgressEvent']> = {}
+): S['InvocationProgressEvent'] =>
+  merge(
+    {
+      timestamp: Date.now(),
+      queue_id: 'test-queue-id',
+      item_id: 1,
+      batch_id: 'test-batch-id',
+      session_id: 'test-session',
+      origin: null,
+      destination: 'test-session',
+      invocation: {},
+      invocation_source_id: 'test-invocation-source-id',
+      message: 'Processing...',
+      percentage: 50,
+      image: null,
+    } as S['InvocationProgressEvent'],
+    overrides
+  );
+
+export const createMockQueueItemStatusChangedEvent = (
+  overrides: PartialDeep<S['QueueItemStatusChangedEvent']> = {}
+): S['QueueItemStatusChangedEvent'] =>
+  merge(
+    {
+      timestamp: Date.now(),
+      queue_id: 'test-queue-id',
+      item_id: 1,
+      batch_id: 'test-batch-id',
+      origin: null,
+      destination: 'test-session',
+      status: 'completed',
+      error_type: null,
+      error_message: null,
+    } as S['QueueItemStatusChangedEvent'],
+    overrides
+  );

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/context.tsx

@@ -0,0 +1,134 @@
+import { useStore } from '@nanostores/react';
+import { useAppStore } from 'app/store/storeHooks';
+import { loadImage } from 'features/controlLayers/konva/util';
+import {
+  selectStagingAreaAutoSwitch,
+  settingsStagingAreaAutoSwitchChanged,
+} from 'features/controlLayers/store/canvasSettingsSlice';
+import { rasterLayerAdded } from 'features/controlLayers/store/canvasSlice';
+import {
+  buildSelectCanvasQueueItems,
+  canvasQueueItemDiscarded,
+  canvasSessionReset,
+} from 'features/controlLayers/store/canvasStagingAreaSlice';
+import { selectBboxRect, selectSelectedEntityIdentifier } from 'features/controlLayers/store/selectors';
+import type { CanvasRasterLayerState } from 'features/controlLayers/store/types';
+import { imageNameToImageObject } from 'features/controlLayers/store/util';
+import type { PropsWithChildren } from 'react';
+import { createContext, memo, useContext, useEffect, useMemo, useState } from 'react';
+import { getImageDTOSafe } from 'services/api/endpoints/images';
+import { queueApi } from 'services/api/endpoints/queue';
+import type { S } from 'services/api/types';
+import { $socket } from 'services/events/stores';
+import { assert } from 'tsafe';
+
+import type { ProgressData, StagingAreaAppApi } from './state';
+import { getInitialProgressData, StagingAreaApi } from './state';
+
+const StagingAreaContext = createContext<StagingAreaApi | null>(null);
+
+export const StagingAreaContextProvider = memo(({ children, sessionId }: PropsWithChildren<{ sessionId: string }>) => {
+  const store = useAppStore();
+  const socket = useStore($socket);
+  const stagingAreaAppApi = useMemo<StagingAreaAppApi>(() => {
+    const selectQueueItems = buildSelectCanvasQueueItems(sessionId);
+
+    const _stagingAreaAppApi: StagingAreaAppApi = {
+      getAutoSwitch: () => selectStagingAreaAutoSwitch(store.getState()),
+      getImageDTO: (imageName: string) => getImageDTOSafe(imageName),
+      loadImage: (imageUrl: string) => loadImage(imageUrl, true),
+      onInvocationProgress: (handler) => {
+        socket?.on('invocation_progress', handler);
+        return () => {
+          socket?.off('invocation_progress', handler);
+        };
+      },
+      onQueueItemStatusChanged: (handler) => {
+        socket?.on('queue_item_status_changed', handler);
+        return () => {
+          socket?.off('queue_item_status_changed', handler);
+        };
+      },
+      onItemsChanged: (handler) => {
+        let prev: S['SessionQueueItem'][] = [];
+        return store.subscribe(() => {
+          const next = selectQueueItems(store.getState());
+          if (prev !== next) {
+            prev = next;
+            handler(next);
+          }
+        });
+      },
+      onDiscard: ({ item_id, status }) => {
+        store.dispatch(canvasQueueItemDiscarded({ itemId: item_id }));
+        if (status === 'in_progress' || status === 'pending') {
+          store.dispatch(queueApi.endpoints.cancelQueueItem.initiate({ item_id }, { track: false }));
+        }
+      },
+      onDiscardAll: () => {
+        store.dispatch(canvasSessionReset());
+        store.dispatch(
+          queueApi.endpoints.cancelQueueItemsByDestination.initiate({ destination: sessionId }, { track: false })
+        );
+      },
+      onAccept: (item, imageDTO) => {
+        const bboxRect = selectBboxRect(store.getState());
+        const { x, y, width, height } = bboxRect;
+        const imageObject = imageNameToImageObject(imageDTO.image_name, { width, height });
+        const selectedEntityIdentifier = selectSelectedEntityIdentifier(store.getState());
+        const overrides: Partial<CanvasRasterLayerState> = {
+          position: { x, y },
+          objects: [imageObject],
+        };
+
+        store.dispatch(rasterLayerAdded({ overrides, isSelected: selectedEntityIdentifier?.type === 'raster_layer' }));
+        store.dispatch(canvasSessionReset());
+        store.dispatch(
+          queueApi.endpoints.cancelQueueItemsByDestination.initiate({ destination: sessionId }, { track: false })
+        );
+      },
+      onAutoSwitchChange: (mode) => {
+        store.dispatch(settingsStagingAreaAutoSwitchChanged(mode));
+      },
+    };
+
+    return _stagingAreaAppApi;
+  }, [sessionId, socket, store]);
+
+  const [stagingAreaApi] = useState(() => new StagingAreaApi());
+
+  useEffect(() => {
+    stagingAreaApi.connectToApp(sessionId, stagingAreaAppApi);
+
+    // We need to subscribe to the queue items query manually to ensure the staging area actually gets the items
+    const { unsubscribe: unsubQueueItemsQuery } = store.dispatch(
+      queueApi.endpoints.listAllQueueItems.initiate({ destination: sessionId })
+    );
+
+    return () => {
+      stagingAreaApi.cleanup();
+      unsubQueueItemsQuery();
+    };
+  }, [sessionId, stagingAreaApi, stagingAreaAppApi, store]);
+
+  return <StagingAreaContext.Provider value={stagingAreaApi}>{children}</StagingAreaContext.Provider>;
+});
+StagingAreaContextProvider.displayName = 'StagingAreaContextProvider';
+
+export const useStagingAreaContext = () => {
+  const ctx = useContext(StagingAreaContext);
+  assert(ctx !== null, "'useStagingAreaContext' must be used within a StagingAreaContextProvider");
+  return ctx;
+};
+
+export const useOutputImageDTO = (itemId: number) => {
+  const ctx = useStagingAreaContext();
+  const allProgressData = useStore(ctx.$progressData, { keys: [itemId] });
+  return allProgressData[itemId]?.imageDTO ?? null;
+};
+
+export const useProgressDatum = (itemId: number): ProgressData => {
+  const ctx = useStagingAreaContext();
+  const allProgressData = useStore(ctx.$progressData, { keys: [itemId] });
+  return allProgressData[itemId] ?? getInitialProgressData(itemId);
+};

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/shared.test.ts

@@ -0,0 +1,205 @@
+import type { S } from 'services/api/types';
+import { describe, expect, it } from 'vitest';
+
+import { getOutputImageName, getProgressMessage, getQueueItemElementId } from './shared';
+
+describe('StagingAreaApi Utility Functions', () => {
+  describe('getProgressMessage', () => {
+    it('should return default message when no data provided', () => {
+      expect(getProgressMessage()).toBe('Generating');
+      expect(getProgressMessage(null)).toBe('Generating');
+    });
+
+    it('should format progress message when data is provided', () => {
+      const progressEvent: S['InvocationProgressEvent'] = {
+        item_id: 1,
+        destination: 'test-session',
+        node_id: 'test-node',
+        source_node_id: 'test-source-node',
+        progress: 0.5,
+        message: 'Processing image...',
+        image: null,
+      } as unknown as S['InvocationProgressEvent'];
+
+      const result = getProgressMessage(progressEvent);
+      expect(result).toBe('Processing image...');
+    });
+  });
+
+  describe('getQueueItemElementId', () => {
+    it('should generate correct element ID for queue item', () => {
+      expect(getQueueItemElementId(0)).toBe('queue-item-preview-0');
+      expect(getQueueItemElementId(5)).toBe('queue-item-preview-5');
+      expect(getQueueItemElementId(99)).toBe('queue-item-preview-99');
+    });
+  });
+
+  describe('getOutputImageName', () => {
+    it('should extract image name from completed queue item', () => {
+      const queueItem: S['SessionQueueItem'] = {
+        item_id: 1,
+        status: 'completed',
+        priority: 0,
+        destination: 'test-session',
+        created_at: '2024-01-01T00:00:00Z',
+        updated_at: '2024-01-01T00:00:00Z',
+        started_at: '2024-01-01T00:00:01Z',
+        completed_at: '2024-01-01T00:01:00Z',
+        error: null,
+        session: {
+          id: 'test-session',
+          source_prepared_mapping: {
+            canvas_output: ['output-node-id'],
+          },
+          results: {
+            'output-node-id': {
+              image: {
+                image_name: 'test-output.png',
+              },
+            },
+          },
+        },
+      } as unknown as S['SessionQueueItem'];
+
+      expect(getOutputImageName(queueItem)).toBe('test-output.png');
+    });
+
+    it('should return null when no canvas output node found', () => {
+      const queueItem = {
+        item_id: 1,
+        status: 'completed',
+        priority: 0,
+        destination: 'test-session',
+        created_at: '2024-01-01T00:00:00Z',
+        updated_at: '2024-01-01T00:00:00Z',
+        started_at: '2024-01-01T00:00:01Z',
+        completed_at: '2024-01-01T00:01:00Z',
+        error: null,
+        session: {
+          id: 'test-session',
+          source_prepared_mapping: {
+            some_other_node: ['other-node-id'],
+          },
+          results: {
+            'other-node-id': {
+              type: 'image_output',
+              image: {
+                image_name: 'test-output.png',
+              },
+              width: 512,
+              height: 512,
+            },
+          },
+        },
+      } as unknown as S['SessionQueueItem'];
+
+      expect(getOutputImageName(queueItem)).toBe(null);
+    });
+
+    it('should return null when output node has no results', () => {
+      const queueItem: S['SessionQueueItem'] = {
+        item_id: 1,
+        status: 'completed',
+        priority: 0,
+        destination: 'test-session',
+        created_at: '2024-01-01T00:00:00Z',
+        updated_at: '2024-01-01T00:00:00Z',
+        started_at: '2024-01-01T00:00:01Z',
+        completed_at: '2024-01-01T00:01:00Z',
+        error: null,
+        session: {
+          id: 'test-session',
+          source_prepared_mapping: {
+            canvas_output: ['output-node-id'],
+          },
+          results: {},
+        },
+      } as unknown as S['SessionQueueItem'];
+
+      expect(getOutputImageName(queueItem)).toBe(null);
+    });
+
+    it('should return null when results contain no image fields', () => {
+      const queueItem: S['SessionQueueItem'] = {
+        item_id: 1,
+        status: 'completed',
+        priority: 0,
+        destination: 'test-session',
+        created_at: '2024-01-01T00:00:00Z',
+        updated_at: '2024-01-01T00:00:00Z',
+        started_at: '2024-01-01T00:00:01Z',
+        completed_at: '2024-01-01T00:01:00Z',
+        error: null,
+        session: {
+          id: 'test-session',
+          source_prepared_mapping: {
+            canvas_output: ['output-node-id'],
+          },
+          results: {
+            'output-node-id': {
+              text: 'some text output',
+              number: 42,
+            },
+          },
+        },
+      } as unknown as S['SessionQueueItem'];
+
+      expect(getOutputImageName(queueItem)).toBe(null);
+    });
+
+    it('should handle multiple outputs and return first image', () => {
+      const queueItem: S['SessionQueueItem'] = {
+        item_id: 1,
+        status: 'completed',
+        priority: 0,
+        destination: 'test-session',
+        created_at: '2024-01-01T00:00:00Z',
+        updated_at: '2024-01-01T00:00:00Z',
+        started_at: '2024-01-01T00:00:01Z',
+        completed_at: '2024-01-01T00:01:00Z',
+        error: null,
+        session: {
+          id: 'test-session',
+          source_prepared_mapping: {
+            canvas_output: ['output-node-id'],
+          },
+          results: {
+            'output-node-id': {
+              text: 'some text',
+              first_image: {
+                image_name: 'first-image.png',
+              },
+              second_image: {
+                image_name: 'second-image.png',
+              },
+            },
+          },
+        },
+      } as unknown as S['SessionQueueItem'];
+
+      const result = getOutputImageName(queueItem);
+      expect(result).toBe('first-image.png');
+    });
+
+    it('should handle empty session mapping', () => {
+      const queueItem: S['SessionQueueItem'] = {
+        item_id: 1,
+        status: 'completed',
+        priority: 0,
+        destination: 'test-session',
+        created_at: '2024-01-01T00:00:00Z',
+        updated_at: '2024-01-01T00:00:00Z',
+        started_at: '2024-01-01T00:00:01Z',
+        completed_at: '2024-01-01T00:01:00Z',
+        error: null,
+        session: {
+          id: 'test-session',
+          source_prepared_mapping: {},
+          results: {},
+        },
+      } as unknown as S['SessionQueueItem'];
+
+      expect(getOutputImageName(queueItem)).toBe(null);
+    });
+  });
+});

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/state.test.ts

@@ -0,0 +1,784 @@
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
+
+import {
+  createMockImageDTO,
+  createMockProgressEvent,
+  createMockQueueItem,
+  createMockQueueItemStatusChangedEvent,
+  createMockStagingAreaApp,
+} from './__mocks__/mockStagingAreaApp';
+import { StagingAreaApi } from './state';
+
+describe('StagingAreaApi', () => {
+  let api: StagingAreaApi;
+  let mockApp: ReturnType<typeof createMockStagingAreaApp>;
+  const sessionId = 'test-session';
+
+  beforeEach(() => {
+    mockApp = createMockStagingAreaApp();
+    api = new StagingAreaApi();
+    api.connectToApp(sessionId, mockApp);
+  });
+
+  afterEach(() => {
+    api.cleanup();
+  });
+
+  describe('Constructor and Setup', () => {
+    it('should initialize with correct session ID', () => {
+      expect(api._sessionId).toBe(sessionId);
+    });
+
+    it('should set up event subscriptions', () => {
+      expect(mockApp.onItemsChanged).toHaveBeenCalledWith(expect.any(Function));
+      expect(mockApp.onQueueItemStatusChanged).toHaveBeenCalledWith(expect.any(Function));
+      expect(mockApp.onInvocationProgress).toHaveBeenCalledWith(expect.any(Function));
+    });
+
+    it('should initialize atoms with default values', () => {
+      expect(api.$lastStartedItemId.get()).toBe(null);
+      expect(api.$lastCompletedItemId.get()).toBe(null);
+      expect(api.$items.get()).toEqual([]);
+      expect(api.$progressData.get()).toEqual({});
+      expect(api.$selectedItemId.get()).toBe(null);
+    });
+  });
+
+  describe('Computed Values', () => {
+    it('should compute item count correctly', () => {
+      expect(api.$itemCount.get()).toBe(0);
+
+      const items = [createMockQueueItem({ item_id: 1 })];
+      api.$items.set(items);
+      expect(api.$itemCount.get()).toBe(1);
+    });
+
+    it('should compute hasItems correctly', () => {
+      expect(api.$hasItems.get()).toBe(false);
+
+      const items = [createMockQueueItem({ item_id: 1 })];
+      api.$items.set(items);
+      expect(api.$hasItems.get()).toBe(true);
+    });
+
+    it('should compute isPending correctly', () => {
+      expect(api.$isPending.get()).toBe(false);
+
+      const items = [
+        createMockQueueItem({ item_id: 1, status: 'pending' }),
+        createMockQueueItem({ item_id: 2, status: 'completed' }),
+      ];
+      api.$items.set(items);
+      expect(api.$isPending.get()).toBe(true);
+    });
+
+    it('should compute selectedItem correctly', () => {
+      expect(api.$selectedItem.get()).toBe(null);
+
+      const items = [createMockQueueItem({ item_id: 1 }), createMockQueueItem({ item_id: 2 })];
+      api.$items.set(items);
+      api.$selectedItemId.set(1);
+
+      const selectedItem = api.$selectedItem.get();
+      expect(selectedItem).not.toBe(null);
+      expect(selectedItem?.item.item_id).toBe(1);
+      expect(selectedItem?.index).toBe(0);
+    });
+
+    it('should compute selectedItemImageDTO correctly', () => {
+      const items = [createMockQueueItem({ item_id: 1 })];
+      const imageDTO = createMockImageDTO();
+
+      api.$items.set(items);
+      api.$selectedItemId.set(1);
+      api.$progressData.setKey(1, {
+        itemId: 1,
+        progressEvent: null,
+        progressImage: null,
+        imageDTO,
+      });
+
+      expect(api.$selectedItemImageDTO.get()).toBe(imageDTO);
+    });
+
+    it('should compute selectedItemIndex correctly', () => {
+      const items = [createMockQueueItem({ item_id: 1 }), createMockQueueItem({ item_id: 2 })];
+      api.$items.set(items);
+      api.$selectedItemId.set(2);
+
+      expect(api.$selectedItemIndex.get()).toBe(1);
+    });
+  });
+
+  describe('Selection Methods', () => {
+    beforeEach(() => {
+      const items = [
+        createMockQueueItem({ item_id: 1 }),
+        createMockQueueItem({ item_id: 2 }),
+        createMockQueueItem({ item_id: 3 }),
+      ];
+      api.$items.set(items);
+    });
+
+    it('should select item by ID', () => {
+      api.select(2);
+      expect(api.$selectedItemId.get()).toBe(2);
+      expect(mockApp.onSelect).toHaveBeenCalledWith(2);
+    });
+
+    it('should select next item', () => {
+      api.$selectedItemId.set(1);
+      api.selectNext();
+      expect(api.$selectedItemId.get()).toBe(2);
+      expect(mockApp.onSelectNext).toHaveBeenCalled();
+    });
+
+    it('should wrap to first item when selecting next from last', () => {
+      api.$selectedItemId.set(3);
+      api.selectNext();
+      expect(api.$selectedItemId.get()).toBe(1);
+    });
+
+    it('should select previous item', () => {
+      api.$selectedItemId.set(2);
+      api.selectPrev();
+      expect(api.$selectedItemId.get()).toBe(1);
+      expect(mockApp.onSelectPrev).toHaveBeenCalled();
+    });
+
+    it('should wrap to last item when selecting previous from first', () => {
+      api.$selectedItemId.set(1);
+      api.selectPrev();
+      expect(api.$selectedItemId.get()).toBe(3);
+    });
+
+    it('should select first item', () => {
+      api.selectFirst();
+      expect(api.$selectedItemId.get()).toBe(1);
+      expect(mockApp.onSelectFirst).toHaveBeenCalled();
+    });
+
+    it('should select last item', () => {
+      api.selectLast();
+      expect(api.$selectedItemId.get()).toBe(3);
+      expect(mockApp.onSelectLast).toHaveBeenCalled();
+    });
+
+    it('should do nothing when no items exist', () => {
+      api.$items.set([]);
+      api.selectNext();
+      api.selectPrev();
+      api.selectFirst();
+      api.selectLast();
+
+      expect(api.$selectedItemId.get()).toBe(null);
+    });
+
+    it('should do nothing when no item is selected', () => {
+      api.$selectedItemId.set(null);
+      api.selectNext();
+      api.selectPrev();
+
+      expect(api.$selectedItemId.get()).toBe(null);
+    });
+  });
+
+  describe('Discard Methods', () => {
+    beforeEach(() => {
+      const items = [
+        createMockQueueItem({ item_id: 1 }),
+        createMockQueueItem({ item_id: 2 }),
+        createMockQueueItem({ item_id: 3 }),
+      ];
+      api.$items.set(items);
+    });
+
+    it('should discard selected item and select next', () => {
+      api.$selectedItemId.set(2);
+      const selectedItem = api.$selectedItem.get();
+
+      api.discardSelected();
+
+      expect(api.$selectedItemId.get()).toBe(3);
+      expect(mockApp.onDiscard).toHaveBeenCalledWith(selectedItem?.item);
+    });
+
+    it('should discard selected item and clamp to last valid index', () => {
+      api.$selectedItemId.set(3);
+      const selectedItem = api.$selectedItem.get();
+
+      api.discardSelected();
+
+      // The logic clamps to the next index, so when discarding last item (index 2),
+      // it tries to select index 3 which clamps to index 2 (item 3)
+      expect(api.$selectedItemId.get()).toBe(3);
+      expect(mockApp.onDiscard).toHaveBeenCalledWith(selectedItem?.item);
+    });
+
+    it('should set selectedItemId to null when discarding last item', () => {
+      api.$items.set([createMockQueueItem({ item_id: 1 })]);
+      api.$selectedItemId.set(1);
+
+      api.discardSelected();
+
+      // When there's only one item, after clamping we get the same item, so it stays selected
+      expect(api.$selectedItemId.get()).toBe(1);
+    });
+
+    it('should do nothing when no item is selected', () => {
+      api.$selectedItemId.set(null);
+      api.discardSelected();
+
+      expect(mockApp.onDiscard).not.toHaveBeenCalled();
+    });
+
+    it('should discard all items', () => {
+      api.$selectedItemId.set(2);
+      api.discardAll();
+
+      expect(api.$selectedItemId.get()).toBe(null);
+      expect(mockApp.onDiscardAll).toHaveBeenCalled();
+    });
+
+    it('should compute discardSelectedIsEnabled correctly', () => {
+      expect(api.$discardSelectedIsEnabled.get()).toBe(false);
+
+      api.$selectedItemId.set(1);
+      expect(api.$discardSelectedIsEnabled.get()).toBe(true);
+    });
+  });
+
+  describe('Accept Methods', () => {
+    beforeEach(() => {
+      const items = [createMockQueueItem({ item_id: 1 })];
+      api.$items.set(items);
+      api.$selectedItemId.set(1);
+    });
+
+    it('should accept selected item when image is available', () => {
+      const imageDTO = createMockImageDTO();
+      api.$progressData.setKey(1, {
+        itemId: 1,
+        progressEvent: null,
+        progressImage: null,
+        imageDTO,
+      });
+
+      const selectedItem = api.$selectedItem.get();
+      api.acceptSelected();
+
+      expect(mockApp.onAccept).toHaveBeenCalledWith(selectedItem?.item, imageDTO);
+    });
+
+    it('should do nothing when no image is available', () => {
+      api.$progressData.setKey(1, {
+        itemId: 1,
+        progressEvent: null,
+        progressImage: null,
+        imageDTO: null,
+      });
+
+      api.acceptSelected();
+
+      expect(mockApp.onAccept).not.toHaveBeenCalled();
+    });
+
+    it('should do nothing when no item is selected', () => {
+      api.$selectedItemId.set(null);
+      api.acceptSelected();
+
+      expect(mockApp.onAccept).not.toHaveBeenCalled();
+    });
+
+    it('should compute acceptSelectedIsEnabled correctly', () => {
+      expect(api.$acceptSelectedIsEnabled.get()).toBe(false);
+
+      const imageDTO = createMockImageDTO();
+      api.$progressData.setKey(1, {
+        itemId: 1,
+        progressEvent: null,
+        progressImage: null,
+        imageDTO,
+      });
+
+      expect(api.$acceptSelectedIsEnabled.get()).toBe(true);
+    });
+  });
+
+  describe('Progress Event Handling', () => {
+    it('should handle invocation progress events', () => {
+      const progressEvent = createMockProgressEvent({
+        item_id: 1,
+        destination: sessionId,
+      });
+
+      api.onInvocationProgressEvent(progressEvent);
+
+      const progressData = api.$progressData.get();
+      expect(progressData[1]?.progressEvent).toBe(progressEvent);
+    });
+
+    it('should ignore events for different sessions', () => {
+      const progressEvent = createMockProgressEvent({
+        item_id: 1,
+        destination: 'different-session',
+      });
+
+      api.onInvocationProgressEvent(progressEvent);
+
+      const progressData = api.$progressData.get();
+      expect(progressData[1]).toBeUndefined();
+    });
+
+    it('should update existing progress data', () => {
+      api.$progressData.setKey(1, {
+        itemId: 1,
+        progressEvent: null,
+        progressImage: null,
+        imageDTO: createMockImageDTO(),
+      });
+
+      const progressEvent = createMockProgressEvent({
+        item_id: 1,
+        destination: sessionId,
+      });
+
+      api.onInvocationProgressEvent(progressEvent);
+
+      const progressData = api.$progressData.get();
+      expect(progressData[1]?.progressEvent).toBe(progressEvent);
+      expect(progressData[1]?.imageDTO).toBeTruthy();
+    });
+  });
+
+  describe('Queue Item Status Change Handling', () => {
+    it('should handle completed status and set last completed item', () => {
+      const statusEvent = createMockQueueItemStatusChangedEvent({
+        item_id: 1,
+        destination: sessionId,
+        status: 'completed',
+      });
+
+      api.onQueueItemStatusChangedEvent(statusEvent);
+
+      expect(api.$lastCompletedItemId.get()).toBe(1);
+    });
+
+    it('should handle in_progress status with switch_on_start', () => {
+      mockApp._setAutoSwitchMode('switch_on_start');
+
+      const statusEvent = createMockQueueItemStatusChangedEvent({
+        item_id: 1,
+        destination: sessionId,
+        status: 'in_progress',
+      });
+
+      api.onQueueItemStatusChangedEvent(statusEvent);
+
+      expect(api.$lastStartedItemId.get()).toBe(1);
+    });
+
+    it('should ignore events for different sessions', () => {
+      const statusEvent = createMockQueueItemStatusChangedEvent({
+        item_id: 1,
+        destination: 'different-session',
+        status: 'completed',
+      });
+
+      api.onQueueItemStatusChangedEvent(statusEvent);
+
+      expect(api.$lastCompletedItemId.get()).toBe(null);
+    });
+  });
+
+  describe('Items Changed Event Handling', () => {
+    it('should update items and auto-select first item', async () => {
+      const items = [createMockQueueItem({ item_id: 1 }), createMockQueueItem({ item_id: 2 })];
+
+      await api.onItemsChangedEvent(items);
+
+      expect(api.$items.get()).toBe(items);
+      expect(api.$selectedItemId.get()).toBe(1);
+    });
+
+    it('should clear selection when no items', async () => {
+      api.$selectedItemId.set(1);
+
+      await api.onItemsChangedEvent([]);
+
+      expect(api.$selectedItemId.get()).toBe(null);
+    });
+
+    it('should not change selection if item already selected', async () => {
+      api.$selectedItemId.set(2);
+
+      const items = [createMockQueueItem({ item_id: 1 }), createMockQueueItem({ item_id: 2 })];
+
+      await api.onItemsChangedEvent(items);
+
+      expect(api.$selectedItemId.get()).toBe(2);
+    });
+
+    it('should load images for completed items', async () => {
+      const imageDTO = createMockImageDTO({ image_name: 'test-image.png' });
+      mockApp._setImageDTO('test-image.png', imageDTO);
+
+      const items = [
+        createMockQueueItem({
+          item_id: 1,
+          status: 'completed',
+        }),
+      ];
+
+      await api.onItemsChangedEvent(items);
+
+      const progressData = api.$progressData.get();
+      expect(progressData[1]?.imageDTO).toBe(imageDTO);
+    });
+
+    it('should handle auto-switch on completion', async () => {
+      mockApp._setAutoSwitchMode('switch_on_finish');
+      api.$lastCompletedItemId.set(1);
+
+      const imageDTO = createMockImageDTO({ image_name: 'test-image.png' });
+      mockApp._setImageDTO('test-image.png', imageDTO);
+
+      const items = [
+        createMockQueueItem({
+          item_id: 1,
+          status: 'completed',
+        }),
+      ];
+
+      await api.onItemsChangedEvent(items);
+
+      // Wait for async image loading - the loadImage promise needs to complete
+      await new Promise((resolve) => {
+        setTimeout(resolve, 50);
+      });
+
+      expect(api.$selectedItemId.get()).toBe(1);
+      // The lastCompletedItemId should be reset after the loadImage promise resolves
+      expect(api.$lastCompletedItemId.get()).toBe(null);
+    });
+
+    it('should clean up progress data for removed items', async () => {
+      api.$progressData.setKey(999, {
+        itemId: 999,
+        progressEvent: null,
+        progressImage: null,
+        imageDTO: null,
+      });
+
+      const items = [createMockQueueItem({ item_id: 1 })];
+
+      await api.onItemsChangedEvent(items);
+
+      const progressData = api.$progressData.get();
+      expect(progressData[999]).toBeUndefined();
+    });
+
+    it('should clear progress data for canceled/failed items', async () => {
+      api.$progressData.setKey(1, {
+        itemId: 1,
+        progressEvent: createMockProgressEvent({ item_id: 1 }),
+        progressImage: null,
+        imageDTO: createMockImageDTO(),
+      });
+
+      const items = [createMockQueueItem({ item_id: 1, status: 'canceled' })];
+
+      await api.onItemsChangedEvent(items);
+
+      const progressData = api.$progressData.get();
+      expect(progressData[1]?.progressEvent).toBe(null);
+      expect(progressData[1]?.progressImage).toBe(null);
+      expect(progressData[1]?.imageDTO).toBe(null);
+    });
+  });
+
+  describe('Auto Switch', () => {
+    it('should set auto switch mode', () => {
+      api.setAutoSwitch('switch_on_finish');
+      expect(mockApp.onAutoSwitchChange).toHaveBeenCalledWith('switch_on_finish');
+    });
+  });
+
+  describe('Utility Methods', () => {
+    it('should build isSelected computed correctly', () => {
+      const isSelected = api.buildIsSelectedComputed(1);
+      expect(isSelected.get()).toBe(false);
+
+      api.$selectedItemId.set(1);
+      expect(isSelected.get()).toBe(true);
+    });
+  });
+
+  describe('Cleanup', () => {
+    it('should reset all state on cleanup', () => {
+      api.$selectedItemId.set(1);
+      api.$items.set([createMockQueueItem({ item_id: 1 })]);
+      api.$lastStartedItemId.set(1);
+      api.$lastCompletedItemId.set(1);
+      api.$progressData.setKey(1, {
+        itemId: 1,
+        progressEvent: null,
+        progressImage: null,
+        imageDTO: null,
+      });
+
+      api.cleanup();
+
+      expect(api.$selectedItemId.get()).toBe(null);
+      expect(api.$items.get()).toEqual([]);
+      expect(api.$lastStartedItemId.get()).toBe(null);
+      expect(api.$lastCompletedItemId.get()).toBe(null);
+      expect(api.$progressData.get()).toEqual({});
+    });
+  });
+
+  describe('Edge Cases and Error Handling', () => {
+    describe('Selection with Empty or Single Item Lists', () => {
+      it('should handle selection operations with single item', () => {
+        const items = [createMockQueueItem({ item_id: 1 })];
+        api.$items.set(items);
+        api.$selectedItemId.set(1);
+
+        // Navigation should wrap around to the same item
+        api.selectNext();
+        expect(api.$selectedItemId.get()).toBe(1);
+
+        api.selectPrev();
+        expect(api.$selectedItemId.get()).toBe(1);
+      });
+
+      it('should handle selection operations with empty list', () => {
+        api.$items.set([]);
+
+        api.selectFirst();
+        api.selectLast();
+        api.selectNext();
+        api.selectPrev();
+
+        expect(api.$selectedItemId.get()).toBe(null);
+      });
+    });
+
+    describe('Progress Data Edge Cases', () => {
+      it('should handle progress updates with image data', () => {
+        const progressEvent = createMockProgressEvent({
+          item_id: 1,
+          destination: sessionId,
+          image: { width: 512, height: 512, dataURL: 'foo' },
+        });
+
+        api.onInvocationProgressEvent(progressEvent);
+
+        const progressData = api.$progressData.get();
+        expect(progressData[1]?.progressImage).toBe(progressEvent.image);
+        expect(progressData[1]?.progressEvent).toBe(progressEvent);
+      });
+
+      it('should preserve imageDTO when updating progress', () => {
+        const imageDTO = createMockImageDTO();
+        api.$progressData.setKey(1, {
+          itemId: 1,
+          progressEvent: null,
+          progressImage: null,
+          imageDTO,
+        });
+
+        const progressEvent = createMockProgressEvent({
+          item_id: 1,
+          destination: sessionId,
+        });
+
+        api.onInvocationProgressEvent(progressEvent);
+
+        const progressData = api.$progressData.get();
+        expect(progressData[1]?.imageDTO).toBe(imageDTO);
+        expect(progressData[1]?.progressEvent).toBe(progressEvent);
+      });
+    });
+
+    describe('Auto-Switch Edge Cases', () => {
+      it('should handle auto-switch when item is not in current items list', async () => {
+        mockApp._setAutoSwitchMode('switch_on_start');
+        api.$lastStartedItemId.set(999); // Non-existent item
+
+        const items = [createMockQueueItem({ item_id: 1 })];
+        await api.onItemsChangedEvent(items);
+
+        // Should not switch to non-existent item
+        expect(api.$selectedItemId.get()).toBe(1);
+        expect(api.$lastStartedItemId.get()).toBe(999);
+      });
+
+      it('should handle auto-switch on finish when image loading fails', async () => {
+        mockApp._setAutoSwitchMode('switch_on_finish');
+        api.$lastCompletedItemId.set(1);
+
+        // Mock image loading failure
+        mockApp._setImageDTO('test-image.png', null);
+
+        const items = [
+          createMockQueueItem({
+            item_id: 1,
+            status: 'completed',
+            session: {
+              id: sessionId,
+              source_prepared_mapping: { canvas_output: ['test-node-id'] },
+              results: {
+                'test-node-id': {
+                  type: 'image_output',
+                  image: { image_name: 'test-image.png' },
+                  width: 512,
+                  height: 512,
+                },
+              },
+            },
+          }),
+        ];
+
+        await api.onItemsChangedEvent(items);
+
+        // Should not switch when image loading fails
+        expect(api.$selectedItemId.get()).toBe(1);
+        expect(api.$lastCompletedItemId.get()).toBe(1);
+      });
+
+      it('should handle auto-switch on finish with slow image loading', async () => {
+        mockApp._setAutoSwitchMode('switch_on_finish');
+        api.$lastCompletedItemId.set(1);
+
+        const imageDTO = createMockImageDTO({ image_name: 'test-image.png' });
+        mockApp._setImageDTO('test-image.png', imageDTO);
+        mockApp._setLoadImageDelay(50); // Add delay to image loading
+
+        const items = [
+          createMockQueueItem({
+            item_id: 1,
+            status: 'completed',
+            session: {
+              id: sessionId,
+              source_prepared_mapping: { canvas_output: ['test-node-id'] },
+              results: { 'test-node-id': { image: { image_name: 'test-image.png' } } },
+            },
+          }),
+        ];
+
+        await api.onItemsChangedEvent(items);
+
+        // Should switch after image loads - wait for both the delay and promise resolution
+        await new Promise((resolve) => {
+          setTimeout(resolve, 150);
+        });
+
+        expect(api.$selectedItemId.get()).toBe(1);
+        // The lastCompletedItemId should be reset after the loadImage promise resolves
+        expect(api.$lastCompletedItemId.get()).toBe(null);
+      });
+    });
+
+    describe('Concurrent Operations', () => {
+      it('should handle rapid item changes', async () => {
+        const items1 = [createMockQueueItem({ item_id: 1 })];
+        const items2 = [createMockQueueItem({ item_id: 2 })];
+
+        // Fire multiple events rapidly
+        const promise1 = api.onItemsChangedEvent(items1);
+        const promise2 = api.onItemsChangedEvent(items2);
+
+        await Promise.all([promise1, promise2]);
+
+        // Should end up with the last set of items
+        expect(api.$items.get()).toBe(items2);
+        // The selectedItemId retains the old value (1) but $selectedItem will be null
+        // because item 1 is no longer in the items list
+        expect(api.$selectedItemId.get()).toBe(1);
+        expect(api.$selectedItem.get()).toBe(null);
+      });
+
+      it('should handle multiple progress events for same item', () => {
+        const event1 = createMockProgressEvent({
+          item_id: 1,
+          destination: sessionId,
+          percentage: 0.3,
+        });
+        const event2 = createMockProgressEvent({
+          item_id: 1,
+          destination: sessionId,
+          percentage: 0.7,
+        });
+
+        api.onInvocationProgressEvent(event1);
+        api.onInvocationProgressEvent(event2);
+
+        const progressData = api.$progressData.get();
+        expect(progressData[1]?.progressEvent).toBe(event2);
+      });
+    });
+
+    describe('Memory Management', () => {
+      it('should clean up progress data for large number of items', async () => {
+        // Create progress data for many items
+        for (let i = 1; i <= 1000; i++) {
+          api.$progressData.setKey(i, {
+            itemId: i,
+            progressEvent: null,
+            progressImage: null,
+            imageDTO: null,
+          });
+        }
+
+        // Only keep a few items
+        const items = [createMockQueueItem({ item_id: 1 }), createMockQueueItem({ item_id: 2 })];
+
+        await api.onItemsChangedEvent(items);
+
+        const progressData = api.$progressData.get();
+        const progressDataKeys = Object.keys(progressData);
+
+        // Should only have progress data for current items
+        expect(progressDataKeys.length).toBeLessThanOrEqual(2);
+        expect(progressData[1]).toBeDefined();
+        expect(progressData[2]).toBeDefined();
+      });
+    });
+
+    describe('Event Subscription Management', () => {
+      it('should handle multiple subscriptions and unsubscriptions', () => {
+        const api2 = new StagingAreaApi();
+        api2.connectToApp(sessionId, mockApp);
+        const api3 = new StagingAreaApi();
+        api3.connectToApp(sessionId, mockApp);
+
+        // All should be subscribed
+        expect(mockApp.onItemsChanged).toHaveBeenCalledTimes(3);
+
+        api2.cleanup();
+        api3.cleanup();
+
+        // Should not affect original api
+        expect(api.$items.get()).toBeDefined();
+      });
+
+      it('should handle events after cleanup', () => {
+        api.cleanup();
+
+        // These should not crash
+        const progressEvent = createMockProgressEvent({
+          item_id: 1,
+          destination: sessionId,
+        });
+
+        api.onInvocationProgressEvent(progressEvent);
+
+        // State should remain clean - but the event handler still works
+        // so it will add progress data even after cleanup
+        const progressData = api.$progressData.get();
+        expect(progressData[1]).toBeDefined();
+      });
+    });
+  });
+});

invokeai/frontend/web/src/features/controlLayers/components/StagingArea/state.ts

@@ -0,0 +1,426 @@
+import { clamp } from 'es-toolkit';
+import type { AutoSwitchMode } from 'features/controlLayers/store/canvasSettingsSlice';
+import type { ProgressImage } from 'features/nodes/types/common';
+import type { MapStore } from 'nanostores';
+import { atom, computed, map } from 'nanostores';
+import type { ImageDTO, S } from 'services/api/types';
+import { objectEntries } from 'tsafe';
+
+import { getOutputImageName } from './shared';
+
+/**
+ * Interface for the app-level API that the StagingAreaApi depends on.
+ * This provides the connection between the staging area and the rest of the application.
+ */
+export type StagingAreaAppApi = {
+  onDiscard?: (item: S['SessionQueueItem']) => void;
+  onDiscardAll?: () => void;
+  onAccept?: (item: S['SessionQueueItem'], imageDTO: ImageDTO) => void;
+  onSelect?: (itemId: number) => void;
+  onSelectPrev?: () => void;
+  onSelectNext?: () => void;
+  onSelectFirst?: () => void;
+  onSelectLast?: () => void;
+  getAutoSwitch: () => AutoSwitchMode;
+  onAutoSwitchChange?: (mode: AutoSwitchMode) => void;
+  getImageDTO: (imageName: string) => Promise<ImageDTO | null>;
+  loadImage: (imageName: string) => Promise<HTMLImageElement>;
+  onItemsChanged: (handler: (data: S['SessionQueueItem'][]) => Promise<void> | void) => () => void;
+  onQueueItemStatusChanged: (handler: (data: S['QueueItemStatusChangedEvent']) => Promise<void> | void) => () => void;
+  onInvocationProgress: (handler: (data: S['InvocationProgressEvent']) => Promise<void> | void) => () => void;
+};
+
+/** Progress data for a single queue item */
+export type ProgressData = {
+  itemId: number;
+  progressEvent: S['InvocationProgressEvent'] | null;
+  progressImage: ProgressImage | null;
+  imageDTO: ImageDTO | null;
+};
+
+/** Combined data for the currently selected item */
+export type SelectedItemData = {
+  item: S['SessionQueueItem'];
+  index: number;
+  progressData: ProgressData;
+};
+
+/** Creates initial progress data for a queue item */
+export const getInitialProgressData = (itemId: number): ProgressData => ({
+  itemId,
+  progressEvent: null,
+  progressImage: null,
+  imageDTO: null,
+});
+type ProgressDataMap = Record<number, ProgressData | undefined>;
+
+/**
+ * API for managing the Canvas Staging Area - a view of the image generation queue.
+ * Provides reactive state management for pending, in-progress, and completed images.
+ * Users can accept images to place on canvas, discard them, navigate between items,
+ * and configure auto-switching behavior.
+ */
+export class StagingAreaApi {
+  /** The current session ID. */
+  _sessionId: string | null = null;
+
+  /** The app API */
+  _app: StagingAreaAppApi | null = null;
+
+  /** A set of subscriptions to be cleaned up when we are finished with a session */
+  _subscriptions = new Set<() => void>();
+
+  /** Item ID of the last started item. Used for auto-switch on start. */
+  $lastStartedItemId = atom<number | null>(null);
+
+  /** Item ID of the last completed item. Used for auto-switch on completion. */
+  $lastCompletedItemId = atom<number | null>(null);
+
+  /** All queue items for the current session. */
+  $items = atom<S['SessionQueueItem'][]>([]);
+
+  /** Progress data for all items including events, images, and ImageDTOs. */
+  $progressData = map<ProgressDataMap>({});
+
+  /** ID of the currently selected queue item, or null if none selected. */
+  $selectedItemId = atom<number | null>(null);
+
+  /** Total number of items in the queue. */
+  $itemCount = computed([this.$items], (items) => items.length);
+
+  /** Whether there are any items in the queue. */
+  $hasItems = computed([this.$items], (items) => items.length > 0);
+
+  /** Whether there are any pending or in-progress items. */
+  $isPending = computed([this.$items], (items) =>
+    items.some((item) => item.status === 'pending' || item.status === 'in_progress')
+  );
+
+  /** The currently selected queue item with its index and progress data, or null if none selected. */
+  $selectedItem = computed(
+    [this.$items, this.$selectedItemId, this.$progressData],
+    (items, selectedItemId, progressData) => {
+      if (items.length === 0) {
+        return null;
+      }
+      if (selectedItemId === null) {
+        return null;
+      }
+      const item = items.find(({ item_id }) => item_id === selectedItemId);
+      if (!item) {
+        return null;
+      }
+
+      return {
+        item,
+        index: items.findIndex(({ item_id }) => item_id === selectedItemId),
+        progressData: progressData[selectedItemId] || getInitialProgressData(selectedItemId),
+      };
+    }
+  );
+
+  /** The ImageDTO of the currently selected item, or null if none available. */
+  $selectedItemImageDTO = computed([this.$selectedItem], (selectedItem) => {
+    return selectedItem?.progressData.imageDTO ?? null;
+  });
+
+  /** The index of the currently selected item, or null if none selected. */
+  $selectedItemIndex = computed([this.$selectedItem], (selectedItem) => {
+    return selectedItem?.index ?? null;
+  });
+
+  /** Selects a queue item by ID. */
+  select = (itemId: number) => {
+    this.$selectedItemId.set(itemId);
+    this._app?.onSelect?.(itemId);
+  };
+
+  /** Selects the next item in the queue, wrapping to the first item if at the end. */
+  selectNext = () => {
+    const selectedItem = this.$selectedItem.get();
+    if (selectedItem === null) {
+      return;
+    }
+    const items = this.$items.get();
+    const nextIndex = (selectedItem.index + 1) % items.length;
+    const nextItem = items[nextIndex];
+    if (!nextItem) {
+      return;
+    }
+    this.$selectedItemId.set(nextItem.item_id);
+    this._app?.onSelectNext?.();
+  };
+
+  /** Selects the previous item in the queue, wrapping to the last item if at the beginning. */
+  selectPrev = () => {
+    const selectedItem = this.$selectedItem.get();
+    if (selectedItem === null) {
+      return;
+    }
+    const items = this.$items.get();
+    const prevIndex = (selectedItem.index - 1 + items.length) % items.length;
+    const prevItem = items[prevIndex];
+    if (!prevItem) {
+      return;
+    }
+    this.$selectedItemId.set(prevItem.item_id);
+    this._app?.onSelectPrev?.();
+  };
+
+  /** Selects the first item in the queue. */
+  selectFirst = () => {
+    const items = this.$items.get();
+    const first = items.at(0);
+    if (!first) {
+      return;
+    }
+    this.$selectedItemId.set(first.item_id);
+    this._app?.onSelectFirst?.();
+  };
+
+  /** Selects the last item in the queue. */
+  selectLast = () => {
+    const items = this.$items.get();
+    const last = items.at(-1);
+    if (!last) {
+      return;
+    }
+    this.$selectedItemId.set(last.item_id);
+    this._app?.onSelectLast?.();
+  };
+
+  /** Discards the currently selected item and selects the next available item. */
+  discardSelected = () => {
+    const selectedItem = this.$selectedItem.get();
+    if (selectedItem === null) {
+      return;
+    }
+    const items = this.$items.get();
+    const nextIndex = clamp(selectedItem.index + 1, 0, items.length - 1);
+    const nextItem = items[nextIndex];
+    if (nextItem) {
+      this.$selectedItemId.set(nextItem.item_id);
+    } else {
+      this.$selectedItemId.set(null);
+    }
+    this._app?.onDiscard?.(selectedItem.item);
+  };
+
+  /** Whether the discard selected action is enabled. */
+  $discardSelectedIsEnabled = computed([this.$selectedItem], (selectedItem) => {
+    if (selectedItem === null) {
+      return false;
+    }
+    return true;
+  });
+
+  /** Connects to the app, registering listeners and such */
+  connectToApp = (sessionId: string, app: StagingAreaAppApi) => {
+    if (this._sessionId !== sessionId) {
+      this.cleanup();
+      this._sessionId = sessionId;
+    }
+    this._app = app;
+
+    this._subscriptions.add(this._app.onItemsChanged(this.onItemsChangedEvent));
+    this._subscriptions.add(this._app.onQueueItemStatusChanged(this.onQueueItemStatusChangedEvent));
+    this._subscriptions.add(this._app.onInvocationProgress(this.onInvocationProgressEvent));
+  };
+
+  /** Discards all items in the queue. */
+  discardAll = () => {
+    this.$selectedItemId.set(null);
+    this._app?.onDiscardAll?.();
+  };
+
+  /** Accepts the currently selected item if an image is available. */
+  acceptSelected = () => {
+    const selectedItem = this.$selectedItem.get();
+    if (selectedItem === null) {
+      return;
+    }
+    const progressData = this.$progressData.get();
+    const datum = progressData[selectedItem.item.item_id];
+    if (!datum || !datum.imageDTO) {
+      return;
+    }
+    this._app?.onAccept?.(selectedItem.item, datum.imageDTO);
+  };
+
+  /** Whether the accept selected action is enabled. */
+  $acceptSelectedIsEnabled = computed([this.$selectedItem, this.$progressData], (selectedItem, progressData) => {
+    if (selectedItem === null) {
+      return false;
+    }
+    const datum = progressData[selectedItem.item.item_id];
+    return !!datum && !!datum.imageDTO;
+  });
+
+  /** Sets the auto-switch mode. */
+  setAutoSwitch = (mode: AutoSwitchMode) => {
+    this._app?.onAutoSwitchChange?.(mode);
+  };
+
+  /** Handles invocation progress events from the WebSocket. */
+  onInvocationProgressEvent = (data: S['InvocationProgressEvent']) => {
+    if (data.destination !== this._sessionId) {
+      return;
+    }
+    setProgress(this.$progressData, data);
+  };
+
+  /** Handles queue item status change events from the WebSocket. */
+  onQueueItemStatusChangedEvent = (data: S['QueueItemStatusChangedEvent']) => {
+    if (data.destination !== this._sessionId) {
+      return;
+    }
+    if (data.status === 'completed') {
+      /**
+       * There is an unpleasant bit of indirection here. When an item is completed, and auto-switch is set to
+       * switch_on_finish, we want to load the image and switch to it. In this socket handler, we don't have
+       * access to the full queue item, which we need to get the output image and load it. We get the full
+       * queue items as part of the list query, so it's rather inefficient to fetch it again here.
+       *
+       * To reduce the number of extra network requests, we instead store this item as the last completed item.
+       * Then in the progress data sync effect, we process the queue item load its image.
+       */
+      this.$lastCompletedItemId.set(data.item_id);
+    }
+    if (data.status === 'in_progress' && this._app?.getAutoSwitch() === 'switch_on_start') {
+      this.$lastStartedItemId.set(data.item_id);
+    }
+  };
+
+  /**
+   * Handles queue items changed events. Updates items, manages progress data,
+   * handles auto-selection, and implements auto-switch behavior.
+   */
+  onItemsChangedEvent = async (items: S['SessionQueueItem'][]) => {
+    const oldItems = this.$items.get();
+
+    if (items === oldItems) {
+      return;
+    }
+
+    if (items.length === 0) {
+      // If there are no items, cannot have a selected item.
+      this.$selectedItemId.set(null);
+    } else if (this.$selectedItemId.get() === null && items.length > 0) {
+      // If there is no selected item but there are items, select the first one.
+      this.$selectedItemId.set(items[0]?.item_id ?? null);
+    }
+
+    const progressData = this.$progressData.get();
+
+    const toDelete: number[] = [];
+    const toUpdate: ProgressData[] = [];
+
+    for (const [id, datum] of objectEntries(progressData)) {
+      if (!datum) {
+        toDelete.push(id);
+        continue;
+      }
+      const item = items.find(({ item_id }) => item_id === datum.itemId);
+      if (!item) {
+        toDelete.push(datum.itemId);
+      } else if (item.status === 'canceled' || item.status === 'failed') {
+        toUpdate.push({
+          ...datum,
+          progressEvent: null,
+          progressImage: null,
+          imageDTO: null,
+        });
+      }
+    }
+
+    for (const item of items) {
+      const datum = progressData[item.item_id];
+
+      if (this.$lastStartedItemId.get() === item.item_id && this._app?.getAutoSwitch() === 'switch_on_start') {
+        this.$selectedItemId.set(item.item_id);
+        this.$lastStartedItemId.set(null);
+      }
+
+      if (datum?.imageDTO) {
+        continue;
+      }
+      const outputImageName = getOutputImageName(item);
+      if (!outputImageName) {
+        continue;
+      }
+      const imageDTO = await this._app?.getImageDTO(outputImageName);
+      if (!imageDTO) {
+        continue;
+      }
+
+      // This is the load logic mentioned in the comment in the QueueItemStatusChangedEvent handler above.
+      if (this.$lastCompletedItemId.get() === item.item_id && this._app?.getAutoSwitch() === 'switch_on_finish') {
+        this._app.loadImage(imageDTO.image_url).then(() => {
+          this.$selectedItemId.set(item.item_id);
+          this.$lastCompletedItemId.set(null);
+        });
+      }
+
+      toUpdate.push({
+        ...getInitialProgressData(item.item_id),
+        ...datum,
+        imageDTO,
+      });
+    }
+
+    for (const itemId of toDelete) {
+      this.$progressData.setKey(itemId, undefined);
+    }
+
+    for (const datum of toUpdate) {
+      this.$progressData.setKey(datum.itemId, datum);
+    }
+
+    this.$items.set(items);
+  };
+
+  /** Creates a computed value that returns true if the given item ID is selected. */
+  buildIsSelectedComputed = (itemId: number) => {
+    return computed([this.$selectedItemId], (selectedItemId) => {
+      return selectedItemId === itemId;
+    });
+  };
+
+  /** Cleans up all state and unsubscribes from all events. */
+  cleanup = () => {
+    this.$lastStartedItemId.set(null);
+    this.$lastCompletedItemId.set(null);
+    this.$items.set([]);
+    this.$progressData.set({});
+    this.$selectedItemId.set(null);
+    this._subscriptions.forEach((unsubscribe) => unsubscribe());
+    this._subscriptions.clear();
+  };
+}
+
+/** Updates progress data for a queue item with the latest progress event. */
+const setProgress = ($progressData: MapStore<ProgressDataMap>, data: S['InvocationProgressEvent']) => {
+  const progressData = $progressData.get();
+  const current = progressData[data.item_id];
+  if (current) {
+    const next = { ...current };
+    next.progressEvent = data;
+    if (data.image) {
+      next.progressImage = data.image;
+    }
+    $progressData.set({
+      ...progressData,
+      [data.item_id]: next,
+    });
+  } else {
+    $progressData.set({
+      ...progressData,
+      [data.item_id]: {
+        itemId: data.item_id,
+        progressEvent: data,
+        progressImage: data.image ?? null,
+        imageDTO: null,
+      },
+    });
+  }
+};

invokeai/frontend/web/src/features/controlLayers/konva/CanvasStagingAreaModule.ts

@@ -1,5 +1,5 @@
 import { Mutex } from 'async-mutex';
-import type { ProgressData, ProgressDataMap } from 'features/controlLayers/components/SimpleSession/context';
+import type { SelectedItemData } from 'features/controlLayers/components/StagingArea/state';
 import type { CanvasManager } from 'features/controlLayers/konva/CanvasManager';
 import { CanvasModuleBase } from 'features/controlLayers/konva/CanvasModuleBase';
 import { CanvasObjectImage } from 'features/controlLayers/konva/CanvasObject/CanvasObjectImage';
@@ -149,33 +149,24 @@ export class CanvasStagingAreaModule extends CanvasModuleBase {
     this.render();
   };
 
-  connectToSession = (
-    $items: Atom<S['SessionQueueItem'][]>,
-    $selectedItemId: Atom<number | null>,
-    $progressData: ProgressDataMap
-  ) => {
-    const imageSrcListener = (
-      selectedItemId: number | null,
-      progressData: Record<number, ProgressData | undefined>
-    ) => {
-      if (!selectedItemId) {
+  connectToSession = ($items: Atom<S['SessionQueueItem'][]>, $selectedItem: Atom<SelectedItemData | null>) => {
+    const imageSrcListener = (selectedItem: SelectedItemData | null) => {
+      if (!selectedItem) {
         this.$imageSrc.set(null);
         return;
       }
 
-      const datum = progressData[selectedItemId];
-
-      if (datum?.imageDTO) {
-        this.$imageSrc.set({ type: 'imageName', data: datum.imageDTO.image_name });
+      if (selectedItem.progressData.imageDTO) {
+        this.$imageSrc.set({ type: 'imageName', data: selectedItem.progressData.imageDTO.image_name });
         return;
-      } else if (datum?.progressImage) {
-        this.$imageSrc.set({ type: 'dataURL', data: datum.progressImage.dataURL });
+      } else if (selectedItem.progressData?.progressImage) {
+        this.$imageSrc.set({ type: 'dataURL', data: selectedItem.progressData.progressImage.dataURL });
         return;
       } else {
         this.$imageSrc.set(null);
       }
     };
-    const unsubImageSrc = effect([$selectedItemId, $progressData], imageSrcListener);
+    const unsubImageSrc = effect([$selectedItem], imageSrcListener);
 
     const isPendingListener = (items: S['SessionQueueItem'][]) => {
       this.$isPending.set(items.some((item) => item.status === 'pending' || item.status === 'in_progress'));
@@ -190,7 +181,7 @@ export class CanvasStagingAreaModule extends CanvasModuleBase {
     // Run the effects & forcibly render once to initialize
     isStagingListener($items.get());
     isPendingListener($items.get());
-    imageSrcListener($selectedItemId.get(), $progressData.get());
+    imageSrcListener($selectedItem.get());
     this.render();
 
     return () => {

invokeai/frontend/web/src/features/controlLayers/store/canvasSettingsSlice.ts

@@ -5,6 +5,7 @@ import { zRgbaColor } from 'features/controlLayers/store/types';
 import { z } from 'zod';
 
 const zAutoSwitchMode = z.enum(['off', 'switch_on_start', 'switch_on_finish']);
+export type AutoSwitchMode = z.infer<typeof zAutoSwitchMode>;
 
 const zCanvasSettingsState = z.object({
   /**

invokeai/frontend/web/src/features/nodes/components/sidePanel/inspector/InspectorDetailsTab.tsx

@@ -2,6 +2,7 @@ import { Box, Flex, FormControl, FormLabel, HStack, Text } from '@invoke-ai/ui-l
 import { useAppSelector } from 'app/store/storeHooks';
 import { IAINoContentFallback } from 'common/components/IAIImageFallback';
 import ScrollableContent from 'common/components/OverlayScrollbars/ScrollableContent';
+import { InvocationNodeContextProvider } from 'features/nodes/components/flow/nodes/Invocation/context';
 import { InvocationNodeNotesTextarea } from 'features/nodes/components/flow/nodes/Invocation/InvocationNodeNotesTextarea';
 import { TemplateGate } from 'features/nodes/components/sidePanel/inspector/NodeTemplateGate';
 import { useNodeNeedsUpdate } from 'features/nodes/hooks/useNodeNeedsUpdate';
@@ -22,12 +23,14 @@ const InspectorDetailsTab = () => {
   }
 
   return (
-    <TemplateGate
-      nodeId={lastSelectedNodeId}
-      fallback={<IAINoContentFallback label={t('nodes.noNodeSelected')} icon={null} />}
-    >
-      <Content nodeId={lastSelectedNodeId} />
-    </TemplateGate>
+    <InvocationNodeContextProvider nodeId={lastSelectedNodeId}>
+      <TemplateGate
+        nodeId={lastSelectedNodeId}
+        fallback={<IAINoContentFallback label={t('nodes.noNodeSelected')} icon={null} />}
+      >
+        <Content nodeId={lastSelectedNodeId} />
+      </TemplateGate>
+    </InvocationNodeContextProvider>
   );
 };
 

invokeai/frontend/web/src/features/nodes/components/sidePanel/inspector/InspectorOutputsTab.tsx

@@ -3,6 +3,7 @@ import { useAppSelector } from 'app/store/storeHooks';
 import { IAINoContentFallback } from 'common/components/IAIImageFallback';
 import ScrollableContent from 'common/components/OverlayScrollbars/ScrollableContent';
 import DataViewer from 'features/gallery/components/ImageMetadataViewer/DataViewer';
+import { InvocationNodeContextProvider } from 'features/nodes/components/flow/nodes/Invocation/context';
 import { TemplateGate } from 'features/nodes/components/sidePanel/inspector/NodeTemplateGate';
 import { useNodeExecutionState } from 'features/nodes/hooks/useNodeExecutionState';
 import { useNodeTemplateOrThrow } from 'features/nodes/hooks/useNodeTemplateOrThrow';
@@ -22,12 +23,14 @@ const InspectorOutputsTab = () => {
   }
 
   return (
-    <TemplateGate
-      nodeId={lastSelectedNodeId}
-      fallback={<IAINoContentFallback label={t('nodes.noNodeSelected')} icon={null} />}
-    >
-      <Content nodeId={lastSelectedNodeId} />
-    </TemplateGate>
+    <InvocationNodeContextProvider nodeId={lastSelectedNodeId}>
+      <TemplateGate
+        nodeId={lastSelectedNodeId}
+        fallback={<IAINoContentFallback label={t('nodes.noNodeSelected')} icon={null} />}
+      >
+        <Content nodeId={lastSelectedNodeId} />
+      </TemplateGate>
+    </InvocationNodeContextProvider>
   );
 };
 

invokeai/frontend/web/src/features/nodes/components/sidePanel/inspector/InspectorTemplateTab.tsx

@@ -1,6 +1,7 @@
 import { useAppSelector } from 'app/store/storeHooks';
 import { IAINoContentFallback } from 'common/components/IAIImageFallback';
 import DataViewer from 'features/gallery/components/ImageMetadataViewer/DataViewer';
+import { InvocationNodeContextProvider } from 'features/nodes/components/flow/nodes/Invocation/context';
 import { TemplateGate } from 'features/nodes/components/sidePanel/inspector/NodeTemplateGate';
 import { useNodeTemplateOrThrow } from 'features/nodes/hooks/useNodeTemplateOrThrow';
 import { selectLastSelectedNodeId } from 'features/nodes/store/selectors';
@@ -16,12 +17,14 @@ const NodeTemplateInspector = () => {
   }
 
   return (
-    <TemplateGate
-      nodeId={lastSelectedNodeId}
-      fallback={<IAINoContentFallback label={t('nodes.noNodeSelected')} icon={null} />}
-    >
-      <Content nodeId={lastSelectedNodeId} />
-    </TemplateGate>
+    <InvocationNodeContextProvider nodeId={lastSelectedNodeId}>
+      <TemplateGate
+        nodeId={lastSelectedNodeId}
+        fallback={<IAINoContentFallback label={t('nodes.noNodeSelected')} icon={null} />}
+      >
+        <Content nodeId={lastSelectedNodeId} />
+      </TemplateGate>
+    </InvocationNodeContextProvider>
   );
 };
 

invokeai/frontend/web/src/features/nodes/components/sidePanel/viewMode/EmptyState.tsx

@@ -1,9 +1,9 @@
 import { Flex, Heading, Icon, Link, Text } from '@invoke-ai/ui-library';
 import { useAppDispatch } from 'app/store/storeHooks';
-import { LaunchpadButton } from 'features/controlLayers/components/SimpleSession/LaunchpadButton';
 import { useIsWorkflowUntouched } from 'features/nodes/components/sidePanel/workflow/IsolatedWorkflowBuilderWatcher';
 import { useWorkflowLibraryModal } from 'features/nodes/store/workflowLibraryModal';
 import { workflowModeChanged } from 'features/nodes/store/workflowLibrarySlice';
+import { LaunchpadButton } from 'features/ui/layouts/LaunchpadButton';
 import { useCallback } from 'react';
 import { Trans, useTranslation } from 'react-i18next';
 import { PiFolderOpenBold, PiPlusBold } from 'react-icons/pi';

invokeai/frontend/web/src/features/ui/layouts/CanvasWorkspacePanel.tsx

@@ -12,11 +12,12 @@ import { Filter } from 'features/controlLayers/components/Filters/Filter';
 import { CanvasHUD } from 'features/controlLayers/components/HUD/CanvasHUD';
 import { InvokeCanvasComponent } from 'features/controlLayers/components/InvokeCanvasComponent';
 import { SelectObject } from 'features/controlLayers/components/SelectObject/SelectObject';
-import { CanvasSessionContextProvider } from 'features/controlLayers/components/SimpleSession/context';
+import { StagingAreaContextProvider } from 'features/controlLayers/components/StagingArea/context';
 import { CanvasToolbar } from 'features/controlLayers/components/Toolbar/CanvasToolbar';
 import { Transform } from 'features/controlLayers/components/Transform/Transform';
 import { CanvasManagerProviderGate } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
 import { selectDynamicGrid, selectShowHUD } from 'features/controlLayers/store/canvasSettingsSlice';
+import { selectCanvasSessionId } from 'features/controlLayers/store/canvasStagingAreaSlice';
 import { memo, useCallback } from 'react';
 import { PiDotsThreeOutlineVerticalFill } from 'react-icons/pi';
 
@@ -49,74 +50,75 @@ const canvasBgSx = {
 export const CanvasWorkspacePanel = memo(() => {
   const dynamicGrid = useAppSelector(selectDynamicGrid);
   const showHUD = useAppSelector(selectShowHUD);
+  const sessionId = useAppSelector(selectCanvasSessionId);
 
   const renderMenu = useCallback(() => {
     return <MenuContent />;
   }, []);
 
   return (
-    <Flex
-      borderRadius="base"
-      position="relative"
-      flexDirection="column"
-      height="full"
-      width="full"
-      gap={2}
-      alignItems="center"
-      justifyContent="center"
-      overflow="hidden"
-    >
-      <CanvasManagerProviderGate>
-        <CanvasToolbar />
-      </CanvasManagerProviderGate>
-      <Divider />
-      <ContextMenu<HTMLDivElement> renderMenu={renderMenu} withLongPress={false}>
-        {(ref) => (
-          <Flex ref={ref} sx={canvasBgSx} data-dynamic-grid={dynamicGrid}>
-            <InvokeCanvasComponent />
-            <CanvasManagerProviderGate>
-              <Flex
-                position="absolute"
-                flexDir="column"
-                top={1}
-                insetInlineStart={1}
-                pointerEvents="none"
-                gap={2}
-                alignItems="flex-start"
-              >
-                {showHUD && <CanvasHUD />}
-                <CanvasAlertsSaveAllImagesToGallery />
-                <CanvasAlertsSelectedEntityStatus />
-                <CanvasAlertsPreserveMask />
-                <CanvasAlertsInvocationProgress />
-              </Flex>
-              <Flex position="absolute" top={1} insetInlineEnd={1}>
-                <Menu>
-                  <MenuButton as={IconButton} icon={<PiDotsThreeOutlineVerticalFill />} colorScheme="base" />
-                  <MenuContent />
-                </Menu>
-              </Flex>
-              <CanvasBusySpinner position="absolute" insetInlineEnd={2} bottom={2} />
-            </CanvasManagerProviderGate>
-          </Flex>
-        )}
-      </ContextMenu>
-      <CanvasManagerProviderGate>
-        <CanvasSessionContextProvider>
-          <StagingArea />
-        </CanvasSessionContextProvider>
-      </CanvasManagerProviderGate>
-      <Flex position="absolute" bottom={4}>
+    <StagingAreaContextProvider sessionId={sessionId}>
+      <Flex
+        borderRadius="base"
+        position="relative"
+        flexDirection="column"
+        height="full"
+        width="full"
+        gap={2}
+        alignItems="center"
+        justifyContent="center"
+        overflow="hidden"
+      >
         <CanvasManagerProviderGate>
-          <Filter />
-          <Transform />
-          <SelectObject />
+          <CanvasToolbar />
+        </CanvasManagerProviderGate>
+        <Divider />
+        <ContextMenu<HTMLDivElement> renderMenu={renderMenu} withLongPress={false}>
+          {(ref) => (
+            <Flex ref={ref} sx={canvasBgSx} data-dynamic-grid={dynamicGrid}>
+              <InvokeCanvasComponent />
+              <CanvasManagerProviderGate>
+                <Flex
+                  position="absolute"
+                  flexDir="column"
+                  top={1}
+                  insetInlineStart={1}
+                  pointerEvents="none"
+                  gap={2}
+                  alignItems="flex-start"
+                >
+                  {showHUD && <CanvasHUD />}
+                  <CanvasAlertsSaveAllImagesToGallery />
+                  <CanvasAlertsSelectedEntityStatus />
+                  <CanvasAlertsPreserveMask />
+                  <CanvasAlertsInvocationProgress />
+                </Flex>
+                <Flex position="absolute" top={1} insetInlineEnd={1}>
+                  <Menu>
+                    <MenuButton as={IconButton} icon={<PiDotsThreeOutlineVerticalFill />} colorScheme="base" />
+                    <MenuContent />
+                  </Menu>
+                </Flex>
+                <CanvasBusySpinner position="absolute" insetInlineEnd={2} bottom={2} />
+              </CanvasManagerProviderGate>
+            </Flex>
+          )}
+        </ContextMenu>
+        <CanvasManagerProviderGate>
+          <StagingArea />
+        </CanvasManagerProviderGate>
+        <Flex position="absolute" bottom={4}>
+          <CanvasManagerProviderGate>
+            <Filter />
+            <Transform />
+            <SelectObject />
+          </CanvasManagerProviderGate>
+        </Flex>
+        <CanvasManagerProviderGate>
+          <CanvasDropArea />
         </CanvasManagerProviderGate>
       </Flex>
-      <CanvasManagerProviderGate>
-        <CanvasDropArea />
-      </CanvasManagerProviderGate>
-    </Flex>
+    </StagingAreaContextProvider>
   );
 });
 CanvasWorkspacePanel.displayName = 'CanvasWorkspacePanel';

invokeai/frontend/web/src/features/ui/layouts/DockviewTab.tsx

@@ -5,7 +5,6 @@ import type { IDockviewPanelHeaderProps } from 'dockview';
 import { memo, useCallback, useRef } from 'react';
 
 import type { PanelParameters } from './auto-layout-context';
-import { useHackOutDvTabDraggable } from './use-hack-out-dv-tab-draggable';
 
 export const DockviewTab = memo((props: IDockviewPanelHeaderProps<PanelParameters>) => {
   const ref = useRef<HTMLDivElement>(null);
@@ -21,8 +20,6 @@ export const DockviewTab = memo((props: IDockviewPanelHeaderProps<PanelParameter
     setFocusedRegion(props.params.focusRegion);
   }, [props.params.focusRegion]);
 
-  useHackOutDvTabDraggable(ref);
-
   return (
     <Flex ref={ref} alignItems="center" h="full" onPointerDown={onPointerDown}>
       <Text userSelect="none" px={4}>

invokeai/frontend/web/src/features/ui/layouts/DockviewTabCanvasViewer.tsx

@@ -8,7 +8,6 @@ import { memo, useCallback, useRef } from 'react';
 import { useIsGenerationInProgress } from 'services/api/endpoints/queue';
 
 import type { PanelParameters } from './auto-layout-context';
-import { useHackOutDvTabDraggable } from './use-hack-out-dv-tab-draggable';
 
 export const DockviewTabCanvasViewer = memo((props: IDockviewPanelHeaderProps<PanelParameters>) => {
   const isGenerationInProgress = useIsGenerationInProgress();
@@ -27,8 +26,6 @@ export const DockviewTabCanvasViewer = memo((props: IDockviewPanelHeaderProps<Pa
     setFocusedRegion(props.params.focusRegion);
   }, [props.params.focusRegion]);
 
-  useHackOutDvTabDraggable(ref);
-
   return (
     <Flex ref={ref} position="relative" alignItems="center" h="full" onPointerDown={onPointerDown}>
       <Text userSelect="none" px={4}>

invokeai/frontend/web/src/features/ui/layouts/DockviewTabCanvasWorkspace.tsx

@@ -10,7 +10,6 @@ import { memo, useCallback, useRef } from 'react';
 import { useIsGenerationInProgress } from 'services/api/endpoints/queue';
 
 import type { PanelParameters } from './auto-layout-context';
-import { useHackOutDvTabDraggable } from './use-hack-out-dv-tab-draggable';
 
 export const DockviewTabCanvasWorkspace = memo((props: IDockviewPanelHeaderProps<PanelParameters>) => {
   const isGenerationInProgress = useIsGenerationInProgress();
@@ -30,8 +29,6 @@ export const DockviewTabCanvasWorkspace = memo((props: IDockviewPanelHeaderProps
     setFocusedRegion(props.params.focusRegion);
   }, [props.params.focusRegion]);
 
-  useHackOutDvTabDraggable(ref);
-
   return (
     <Flex ref={ref} position="relative" alignItems="center" h="full" onPointerDown={onPointerDown}>
       <Text userSelect="none" px={4}>

invokeai/frontend/web/src/features/ui/layouts/DockviewTabLaunchpad.tsx

@@ -16,8 +16,6 @@ import {
   PiTextAaBold,
 } from 'react-icons/pi';
 
-import { useHackOutDvTabDraggable } from './use-hack-out-dv-tab-draggable';
-
 const TAB_ICONS: Record<TabName, IconType> = {
   generate: PiTextAaBold,
   canvas: PiBoundingBoxBold,
@@ -43,8 +41,6 @@ export const DockviewTabLaunchpad = memo((props: IDockviewPanelHeaderProps) => {
     setFocusedRegion(props.params.focusRegion);
   }, [props.params.focusRegion]);
 
-  useHackOutDvTabDraggable(ref);
-
   return (
     <Flex ref={ref} alignItems="center" h="full" px={4} gap={3} onPointerDown={onPointerDown}>
       <Icon as={TAB_ICONS[activeTab]} color="invokeYellow.300" boxSize={5} />

invokeai/frontend/web/src/features/ui/layouts/DockviewTabProgress.tsx

@@ -7,7 +7,6 @@ import { memo, useCallback, useRef } from 'react';
 import { useIsGenerationInProgress } from 'services/api/endpoints/queue';
 
 import type { PanelParameters } from './auto-layout-context';
-import { useHackOutDvTabDraggable } from './use-hack-out-dv-tab-draggable';
 
 export const DockviewTabProgress = memo((props: IDockviewPanelHeaderProps<PanelParameters>) => {
   const isGenerationInProgress = useIsGenerationInProgress();
@@ -25,8 +24,6 @@ export const DockviewTabProgress = memo((props: IDockviewPanelHeaderProps<PanelP
     setFocusedRegion(props.params.focusRegion);
   }, [props.params.focusRegion]);
 
-  useHackOutDvTabDraggable(ref);
-
   return (
     <Flex ref={ref} position="relative" alignItems="center" h="full" onPointerDown={onPointerDown}>
       <Text userSelect="none" px={4}>

invokeai/frontend/web/src/features/ui/layouts/GenerateLaunchpadPanel.tsx

@@ -1,9 +1,9 @@
 import { Alert, Button, Flex, Grid, Text } from '@invoke-ai/ui-library';
-import { InitialStateMainModelPicker } from 'features/controlLayers/components/SimpleSession/InitialStateMainModelPicker';
-import { LaunchpadAddStyleReference } from 'features/controlLayers/components/SimpleSession/LaunchpadAddStyleReference';
 import { navigationApi } from 'features/ui/layouts/navigation-api';
 import { memo, useCallback } from 'react';
 
+import { InitialStateMainModelPicker } from './InitialStateMainModelPicker';
+import { LaunchpadAddStyleReference } from './LaunchpadAddStyleReference';
 import { LaunchpadContainer } from './LaunchpadContainer';
 import { LaunchpadGenerateFromTextButton } from './LaunchpadGenerateFromTextButton';
 

invokeai/frontend/web/src/features/ui/layouts/LaunchpadAddStyleReference.tsx

@@ -1,12 +1,12 @@
 import { Flex, Heading, Icon, Text } from '@invoke-ai/ui-library';
 import { useAppStore } from 'app/store/storeHooks';
 import { useImageUploadButton } from 'common/hooks/useImageUploadButton';
-import { LaunchpadButton } from 'features/controlLayers/components/SimpleSession/LaunchpadButton';
 import { getDefaultRefImageConfig } from 'features/controlLayers/hooks/addLayerHooks';
 import { refImageAdded } from 'features/controlLayers/store/refImagesSlice';
 import { imageDTOToImageWithDims } from 'features/controlLayers/store/util';
 import { addGlobalReferenceImageDndTarget, newCanvasFromImageDndTarget } from 'features/dnd/dnd';
 import { DndDropTarget } from 'features/dnd/DndDropTarget';
+import { LaunchpadButton } from 'features/ui/layouts/LaunchpadButton';
 import { memo, useMemo } from 'react';
 import { PiUploadBold, PiUserCircleGearBold } from 'react-icons/pi';
 import type { ImageDTO } from 'services/api/types';

invokeai/frontend/web/src/features/ui/layouts/LaunchpadEditImageButton.tsx

@@ -1,10 +1,10 @@
 import { Flex, Heading, Icon, Text } from '@invoke-ai/ui-library';
 import { useAppStore } from 'app/store/storeHooks';
 import { useImageUploadButton } from 'common/hooks/useImageUploadButton';
-import { LaunchpadButton } from 'features/controlLayers/components/SimpleSession/LaunchpadButton';
 import { newCanvasFromImageDndTarget } from 'features/dnd/dnd';
 import { DndDropTarget } from 'features/dnd/DndDropTarget';
 import { newCanvasFromImage } from 'features/imageActions/actions';
+import { LaunchpadButton } from 'features/ui/layouts/LaunchpadButton';
 import { memo, useCallback } from 'react';
 import { PiPencilBold, PiUploadBold } from 'react-icons/pi';
 import type { ImageDTO } from 'services/api/types';

invokeai/frontend/web/src/features/ui/layouts/LaunchpadGenerateFromTextButton.tsx

@@ -1,5 +1,5 @@
 import { Flex, Heading, Icon, Text } from '@invoke-ai/ui-library';
-import { LaunchpadButton } from 'features/controlLayers/components/SimpleSession/LaunchpadButton';
+import { LaunchpadButton } from 'features/ui/layouts/LaunchpadButton';
 import { memo, useCallback } from 'react';
 import { PiCursorTextBold, PiTextAaBold } from 'react-icons/pi';
 

invokeai/frontend/web/src/features/ui/layouts/StagingArea.tsx

@@ -1,5 +1,5 @@
 import { Flex } from '@invoke-ai/ui-library';
-import { StagingAreaItemsList } from 'features/controlLayers/components/SimpleSession/StagingAreaItemsList';
+import { StagingAreaItemsList } from 'features/controlLayers/components/StagingArea/StagingAreaItemsList';
 import { StagingAreaToolbar } from 'features/controlLayers/components/StagingArea/StagingAreaToolbar';
 import { useCanvasIsStaging } from 'features/controlLayers/store/canvasStagingAreaSlice';
 import { memo } from 'react';

invokeai/frontend/web/src/features/ui/layouts/canvas-tab-auto-layout.tsx

@@ -1,7 +1,6 @@
 import type { DockviewApi, GridviewApi, IDockviewReactProps, IGridviewReactProps } from 'dockview';
 import { DockviewReact, GridviewReact, LayoutPriority, Orientation } from 'dockview';
 import { CanvasLayersPanel } from 'features/controlLayers/components/CanvasLayersPanelContent';
-import { CanvasLaunchpadPanel } from 'features/controlLayers/components/SimpleSession/CanvasLaunchpadPanel';
 import { BoardsPanel } from 'features/gallery/components/BoardsListPanelContent';
 import { GalleryPanel } from 'features/gallery/components/Gallery';
 import { GenerationProgressPanel } from 'features/gallery/components/ImageViewer/GenerationProgressPanel';
@@ -15,6 +14,7 @@ import type {
   RootLayoutGridviewComponents,
 } from 'features/ui/layouts/auto-layout-context';
 import { AutoLayoutProvider, useAutoLayoutContext, withPanelContainer } from 'features/ui/layouts/auto-layout-context';
+import { CanvasLaunchpadPanel } from 'features/ui/layouts/CanvasLaunchpadPanel';
 import type { TabName } from 'features/ui/store/uiTypes';
 import { dockviewTheme } from 'features/ui/styles/theme';
 import { memo, useCallback, useEffect } from 'react';

invokeai/frontend/web/src/features/ui/layouts/generate-tab-auto-layout.tsx

@@ -1,6 +1,5 @@
 import type { DockviewApi, GridviewApi, IDockviewReactProps, IGridviewReactProps } from 'dockview';
 import { DockviewReact, GridviewReact, LayoutPriority, Orientation } from 'dockview';
-import { GenerateLaunchpadPanel } from 'features/controlLayers/components/SimpleSession/GenerateLaunchpadPanel';
 import { BoardsPanel } from 'features/gallery/components/BoardsListPanelContent';
 import { GalleryPanel } from 'features/gallery/components/Gallery';
 import { GenerationProgressPanel } from 'features/gallery/components/ImageViewer/GenerationProgressPanel';
@@ -21,6 +20,7 @@ import { memo, useCallback, useEffect } from 'react';
 import { DockviewTab } from './DockviewTab';
 import { DockviewTabLaunchpad } from './DockviewTabLaunchpad';
 import { DockviewTabProgress } from './DockviewTabProgress';
+import { GenerateLaunchpadPanel } from './GenerateLaunchpadPanel';
 import { GenerateTabLeftPanel } from './GenerateTabLeftPanel';
 import { navigationApi } from './navigation-api';
 import { PanelHotkeysLogical } from './PanelHotkeysLogical';

invokeai/frontend/web/src/features/ui/layouts/upscaling-tab-auto-layout.tsx

@@ -1,6 +1,5 @@
 import type { DockviewApi, GridviewApi, IDockviewReactProps, IGridviewReactProps } from 'dockview';
 import { DockviewReact, GridviewReact, LayoutPriority, Orientation } from 'dockview';
-import { UpscalingLaunchpadPanel } from 'features/controlLayers/components/SimpleSession/UpscalingLaunchpadPanel';
 import { BoardsPanel } from 'features/gallery/components/BoardsListPanelContent';
 import { GalleryPanel } from 'features/gallery/components/Gallery';
 import { GenerationProgressPanel } from 'features/gallery/components/ImageViewer/GenerationProgressPanel';
@@ -43,6 +42,7 @@ import {
   SETTINGS_PANEL_ID,
   VIEWER_PANEL_ID,
 } from './shared';
+import { UpscalingLaunchpadPanel } from './UpscalingLaunchpadPanel';
 import { UpscalingTabLeftPanel } from './UpscalingTabLeftPanel';
 
 const tabComponents = {

invokeai/frontend/web/src/features/ui/layouts/use-hack-out-dv-tab-draggable.ts

@@ -1,34 +0,0 @@
-import type { RefObject } from 'react';
-import { useEffect } from 'react';
-
-/**
- * Prevent undesired dnd behavior in Dockview tabs.
- *
- * Dockview always sets the draggable flag on its tab elements, even when dnd is disabled. This hook traverses
- * up from the provided ref to find the closest tab element and sets its `draggable` attribute to `false`.
- *
- * TODO: Remove this when https://github.com/mathuo/dockview/pull/961 is shipped.
- */
-export const useHackOutDvTabDraggable = (ref: RefObject<HTMLElement>) => {
-  useEffect(() => {
-    const el = ref.current;
-    if (!el) {
-      return;
-    }
-    const parentTab = el.closest('.dv-tab');
-    if (!parentTab) {
-      return;
-    }
-    parentTab.setAttribute('draggable', 'false');
-
-    const tabContainer = parentTab.closest('.dv-tabs-and-actions-container');
-    if (!tabContainer) {
-      return;
-    }
-    const voidContainer = tabContainer.querySelector('.dv-void-container');
-    if (!voidContainer) {
-      return;
-    }
-    voidContainer.setAttribute('draggable', 'false');
-  }, [ref]);
-};

invokeai/frontend/web/src/features/ui/layouts/workflows-tab-auto-layout.tsx

@@ -1,6 +1,5 @@
 import type { DockviewApi, GridviewApi, IDockviewReactProps, IGridviewReactProps } from 'dockview';
 import { DockviewReact, GridviewReact, LayoutPriority, Orientation } from 'dockview';
-import { WorkflowsLaunchpadPanel } from 'features/controlLayers/components/SimpleSession/WorkflowsLaunchpadPanel';
 import { BoardsPanel } from 'features/gallery/components/BoardsListPanelContent';
 import { GalleryPanel } from 'features/gallery/components/Gallery';
 import { GenerationProgressPanel } from 'features/gallery/components/ImageViewer/GenerationProgressPanel';
@@ -46,6 +45,7 @@ import {
   VIEWER_PANEL_ID,
   WORKSPACE_PANEL_ID,
 } from './shared';
+import { WorkflowsLaunchpadPanel } from './WorkflowsLaunchpadPanel';
 
 const tabComponents = {
   [DOCKVIEW_TAB_ID]: DockviewTab,

invokeai/version/invokeai_version.py

@@ -1 +1 @@
-__version__ = "6.1.0rc1"
+__version__ = "6.1.0rc2"