From 2e0824a799bff3c5166f5297d0430febf0b8fcce Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Mon, 16 Jun 2025 17:41:47 +1000
Subject: [PATCH] feat(ui): make autoswitch on/off

When the invocation cache is used, we might skip all progress images. This can prevent auto-switch-on-first-progress from working, as we don't get any of those events.

It's much easier to only support auto-switch on complete.
---
 .../components/SimpleSession/context.tsx           | 14 +++-----------
 .../StagingAreaToolbarMenuAutoSwitch.tsx           | 14 +++++---------
 2 files changed, 8 insertions(+), 20 deletions(-)

diff --git a/invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/context.tsx b/invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/context.tsx
index 94357c5e0f..88f6f4f2d0 100644
--- a/invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/context.tsx
+++ b/invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/context.tsx
@@ -13,10 +13,6 @@ import { queueApi } from 'services/api/endpoints/queue';
 import type { ImageDTO, S } from 'services/api/types';
 import { $socket } from 'services/events/stores';
 import { assert } from 'tsafe';
-import { z } from 'zod';
-
-export const zAutoSwitchMode = z.enum(['off', 'first_progress', 'completed']);
-export type AutoSwitchMode = z.infer<typeof zAutoSwitchMode>;
 
 export type ProgressData = {
   itemId: number;
@@ -91,7 +87,7 @@ type CanvasSessionContextValue = {
   $selectedItem: Atom<S['SessionQueueItem'] | null>;
   $selectedItemIndex: Atom<number | null>;
   $selectedItemOutputImageDTO: Atom<ImageDTO | null>;
-  $autoSwitch: WritableAtom<AutoSwitchMode>;
+  $autoSwitch: WritableAtom<boolean>;
   $lastLoadedItemId: WritableAtom<number | null>;
   selectNext: () => void;
   selectPrev: () => void;
@@ -126,7 +122,7 @@ export const CanvasSessionContextProvider = memo(
     /**
      * Whether auto-switch is enabled.
      */
-    const $autoSwitch = useState(() => atom<AutoSwitchMode>('first_progress'))[0];
+    const $autoSwitch = useState(() => atom(true))[0];
 
     /**
      * An internal flag used to work around race conditions with auto-switch switching to queue items before their
@@ -273,11 +269,7 @@ export const CanvasSessionContextProvider = memo(
         if (data.destination !== session.id) {
           return;
         }
-        const isFirstProgressImage = !$progressData.get()[data.item_id]?.progressImage && !!data.image;
         setProgress($progressData, data);
-        if ($autoSwitch.get() === 'first_progress' && isFirstProgressImage) {
-          $selectedItemId.set(data.item_id);
-        }
       };
 
       socket.on('invocation_progress', onProgress);
@@ -413,7 +405,7 @@ export const CanvasSessionContextProvider = memo(
         if (lastLoadedItemId === null) {
           return;
         }
-        if ($autoSwitch.get() === 'completed') {
+        if ($autoSwitch.get()) {
           $selectedItemId.set(lastLoadedItemId);
         }
         $lastLoadedItemId.set(null);
diff --git a/invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaToolbarMenuAutoSwitch.tsx b/invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaToolbarMenuAutoSwitch.tsx
index 567e6beb5e..ff8a83ff87 100644
--- a/invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaToolbarMenuAutoSwitch.tsx
+++ b/invokeai/frontend/web/src/features/controlLayers/components/StagingArea/StagingAreaToolbarMenuAutoSwitch.tsx
@@ -1,7 +1,7 @@
 /* eslint-disable i18next/no-literal-string */
 import { MenuItemOption, MenuOptionGroup } from '@invoke-ai/ui-library';
 import { useStore } from '@nanostores/react';
-import { useCanvasSessionContext, zAutoSwitchMode } from 'features/controlLayers/components/SimpleSession/context';
+import { useCanvasSessionContext } from 'features/controlLayers/components/SimpleSession/context';
 import { memo, useCallback } from 'react';
 
 export const StagingAreaToolbarMenuAutoSwitch = memo(() => {
@@ -10,22 +10,18 @@ export const StagingAreaToolbarMenuAutoSwitch = memo(() => {
 
   const onChange = useCallback(
     (val: string | string[]) => {
-      const newAutoSwitch = zAutoSwitchMode.parse(val);
-      ctx.$autoSwitch.set(newAutoSwitch);
+      ctx.$autoSwitch.set(val === 'on');
     },
     [ctx.$autoSwitch]
   );
 
   return (
-    <MenuOptionGroup value={autoSwitch} onChange={onChange} title="Auto Switch" type="radio">
+    <MenuOptionGroup value={autoSwitch ? 'on' : 'off'} onChange={onChange} title="Auto Switch" type="radio">
       <MenuItemOption value="off" closeOnSelect={false}>
         Off
       </MenuItemOption>
-      <MenuItemOption value="first_progress" closeOnSelect={false}>
-        First Progress
-      </MenuItemOption>
-      <MenuItemOption value="completed" closeOnSelect={false}>
-        Completed
+      <MenuItemOption value="on" closeOnSelect={false}>
+        On
       </MenuItemOption>
     </MenuOptionGroup>
   );