refactor(ui): move staging area logic out side react

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.

invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/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 { memo } from 'react';
 import type { S } from 'services/api/types';
 
+import { useProgressDatum } from './context2';
+
 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/SimpleSession/QueueItemPreviewMini.tsx

@@ -1,11 +1,7 @@
 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';
@@ -17,9 +13,11 @@ import {
 } 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 './context2';
+
 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/SimpleSession/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 './context2';
+
 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/SimpleSession/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 './context2';
+
 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/SimpleSession/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 { 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 './context2';
 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/SimpleSession/context2.tsx

@@ -0,0 +1,119 @@
+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, useMemo } 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 value = useMemo(() => new StagingAreaApi(sessionId, stagingAreaAppApi), [sessionId, stagingAreaAppApi]);
+
+  return <StagingAreaContext.Provider value={value}>{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/SimpleSession/state.ts

@@ -0,0 +1,399 @@
+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';
+
+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;
+};
+
+export type ProgressData = {
+  itemId: number;
+  progressEvent: S['InvocationProgressEvent'] | null;
+  progressImage: ProgressImage | null;
+  imageDTO: ImageDTO | null;
+};
+
+export type SelectedItemData = {
+  item: S['SessionQueueItem'];
+  index: number;
+  progressData: ProgressData;
+};
+
+export const getInitialProgressData = (itemId: number): ProgressData => ({
+  itemId,
+  progressEvent: null,
+  progressImage: null,
+  imageDTO: null,
+});
+export type ProgressDataMap = Record<number, ProgressData | undefined>;
+
+export class StagingAreaApi {
+  sessionId: string;
+  _app: StagingAreaAppApi;
+  _subscriptions = new Set<() => void>();
+
+  constructor(sessionId: string, app: StagingAreaAppApi) {
+    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));
+  }
+
+  /**
+   * Track the last started item. Used to implement autoswitch.
+   */
+  $lastStartedItemId = atom<number | null>(null);
+
+  /**
+   * Track the last completed item. Used to implement autoswitch.
+   */
+  $lastCompletedItemId = atom<number | null>(null);
+
+  /**
+   * 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.
+   */
+  $items = atom<S['SessionQueueItem'][]>([]);
+
+  /**
+   * An ephemeral store of progress events and images for all items in the current session.
+   */
+  $progressData = map<ProgressDataMap>({});
+
+  /**
+   * The currently selected queue item's ID, or null if one is not selected.
+   */
+  $selectedItemId = atom<number | null>(null);
+
+  /**
+   * The number of items. Computed from the queue items array.
+   */
+  $itemCount = computed([this.$items], (items) => items.length);
+
+  /**
+   * Whether there are any items. Computed from the queue items array.
+   */
+  $hasItems = computed([this.$items], (items) => items.length > 0);
+
+  /**
+   * Whether there are any pending or in-progress items. Computed from the queue items array.
+   */
+  $isPending = computed([this.$items], (items) =>
+    items.some((item) => item.status === 'pending' || item.status === 'in_progress')
+  );
+
+  /**
+   * The currently selected queue item, its index and progress data - or null, if one is not 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),
+      };
+    }
+  );
+
+  $selectedItemImageDTO = computed([this.$selectedItem], (selectedItem) => {
+    return selectedItem?.progressData.imageDTO ?? null;
+  });
+
+  $selectedItemIndex = computed([this.$selectedItem], (selectedItem) => {
+    return selectedItem?.index ?? null;
+  });
+
+  select = (itemId: number) => {
+    this.$selectedItemId.set(itemId);
+    this._app.onSelect?.(itemId);
+  };
+
+  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?.();
+  };
+
+  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?.();
+  };
+
+  selectFirst = () => {
+    const items = this.$items.get();
+    const first = items.at(0);
+    if (!first) {
+      return;
+    }
+    this.$selectedItemId.set(first.item_id);
+    this._app.onSelectFirst?.();
+  };
+
+  selectLast = () => {
+    const items = this.$items.get();
+    const last = items.at(-1);
+    if (!last) {
+      return;
+    }
+    this.$selectedItemId.set(last.item_id);
+    this._app.onSelectLast?.();
+  };
+
+  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);
+  };
+  $discardSelectedIsEnabled = computed([this.$selectedItem], (selectedItem) => {
+    if (selectedItem === null) {
+      return false;
+    }
+    return true;
+  });
+
+  discardAll = () => {
+    this.$selectedItemId.set(null);
+    this._app.onDiscardAll?.();
+  };
+
+  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);
+  };
+  $acceptSelectedIsEnabled = computed([this.$selectedItem, this.$progressData], (selectedItem, progressData) => {
+    if (selectedItem === null) {
+      return false;
+    }
+    const datum = progressData[selectedItem.item.item_id];
+    return !!datum && !!datum.imageDTO;
+  });
+
+  setAutoSwitch = (mode: AutoSwitchMode) => {
+    this._app.onAutoSwitchChange?.(mode);
+  };
+
+  onInvocationProgressEvent = (data: S['InvocationProgressEvent']) => {
+    if (data.destination !== this.sessionId) {
+      return;
+    }
+    setProgress(this.$progressData, data);
+  };
+
+  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);
+    }
+  };
+
+  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);
+  };
+
+  buildIsSelectedComputed = (itemId: number) => {
+    return computed([this.$selectedItemId], (selectedItemId) => {
+      return selectedItemId === itemId;
+    });
+  };
+
+  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();
+  };
+}
+
+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/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/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/SimpleSession/context2';
 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/SimpleSession/context2';
 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/SimpleSession/context2';
+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/SimpleSession/context2';
+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/SimpleSession/context2';
+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/SimpleSession/context2';
 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/SimpleSession/context2';
+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/SimpleSession/context2';
+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/SimpleSession/context2';
 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/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/SimpleSession/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/ui/layouts/CanvasWorkspacePanel.tsx

@@ -13,10 +13,12 @@ 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/SimpleSession/context2';
 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 +51,77 @@ 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>
+          <CanvasSessionContextProvider>
+            <StagingArea />
+          </CanvasSessionContextProvider>
+        </CanvasManagerProviderGate>
+        <Flex position="absolute" bottom={4}>
+          <CanvasManagerProviderGate>
+            <Filter />
+            <Transform />
+            <SelectObject />
+          </CanvasManagerProviderGate>
+        </Flex>
+        <CanvasManagerProviderGate>
+          <CanvasDropArea />
         </CanvasManagerProviderGate>
       </Flex>
-      <CanvasManagerProviderGate>
-        <CanvasDropArea />
-      </CanvasManagerProviderGate>
-    </Flex>
+    </StagingAreaContextProvider>
   );
 });
 CanvasWorkspacePanel.displayName = 'CanvasWorkspacePanel';