fix(ui): viewer progress-image handoff: slow reveal, overlay stuck after socket drop, stale-image flash on quick re-generate

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/autoSwitchedSelection.test.ts

@@ -0,0 +1,73 @@
+import { configureStore, createListenerMiddleware } from '@reduxjs/toolkit';
+import type { AppStartListening } from 'app/store/store';
+import { autoSwitchedImages } from 'features/gallery/store/autoSwitchedImages';
+import {
+  boardIdSelected,
+  gallerySliceConfig,
+  imageSelected,
+  selectionChanged,
+} from 'features/gallery/store/gallerySlice';
+import { beforeEach, describe, expect, it } from 'vitest';
+
+import { addAutoSwitchedSelectionListener } from './autoSwitchedSelection';
+
+// A store with the real gallery reducer and the real listener, so the predicate is exercised
+// against actual selection-writing actions rather than a hand-built state pair.
+const buildStore = () => {
+  const listenerMiddleware = createListenerMiddleware();
+  addAutoSwitchedSelectionListener(listenerMiddleware.startListening as unknown as AppStartListening);
+  return configureStore({
+    reducer: { gallery: gallerySliceConfig.slice.reducer },
+    middleware: (getDefaultMiddleware) => getDefaultMiddleware().prepend(listenerMiddleware.middleware),
+  });
+};
+
+describe('addAutoSwitchedSelectionListener', () => {
+  beforeEach(() => {
+    // The marker is a module singleton; drop anything a previous test left on it.
+    autoSwitchedImages.settle(null);
+  });
+
+  it('keeps the marker when the auto-switch selection lands', () => {
+    const store = buildStore();
+    autoSwitchedImages.record('a.png');
+    store.dispatch(imageSelected('a.png'));
+    expect(autoSwitchedImages.consume('a.png')).toBe(true);
+  });
+
+  it('drops the marker once the user selects something else', () => {
+    // The dead click this exists to prevent: the auto-switch to A never rendered because the user
+    // clicked B first, so their later click on A must still get its reveal.
+    const store = buildStore();
+    autoSwitchedImages.record('a.png');
+    store.dispatch(imageSelected('a.png'));
+    store.dispatch(imageSelected('b.png'));
+    store.dispatch(imageSelected('a.png'));
+    expect(autoSwitchedImages.consume('a.png')).toBe(false);
+  });
+
+  it('settles on every action that writes the selection, not just imageSelected', () => {
+    const store = buildStore();
+
+    autoSwitchedImages.record('a.png');
+    store.dispatch(imageSelected('a.png'));
+    store.dispatch(selectionChanged(['b.png']));
+    expect(autoSwitchedImages.consume('a.png')).toBe(false);
+
+    autoSwitchedImages.record('c.png');
+    store.dispatch(imageSelected('c.png'));
+    store.dispatch(boardIdSelected({ boardId: 'other', select: { selection: ['d.png'], galleryView: 'images' } }));
+    expect(autoSwitchedImages.consume('c.png')).toBe(false);
+  });
+
+  it('leaves the marker alone when an action does not move the selection', () => {
+    const store = buildStore();
+    autoSwitchedImages.record('a.png');
+    store.dispatch(imageSelected('a.png'));
+    // Selecting the same item again, and a board switch that carries no selection, must not
+    // discard a marker whose image has not rendered yet.
+    store.dispatch(imageSelected('a.png'));
+    store.dispatch(boardIdSelected({ boardId: 'other' }));
+    expect(autoSwitchedImages.consume('a.png')).toBe(true);
+  });
+});

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

@@ -0,0 +1,27 @@
+import type { AppStartListening } from 'app/store/store';
+import { autoSwitchedImages } from 'features/gallery/store/autoSwitchedImages';
+import { selectLastSelectedItem } from 'features/gallery/store/gallerySelectors';
+
+/**
+ * Keeps the auto-switch marker scoped to the selection it was recorded for.
+ *
+ * onInvocationComplete records the item it is about to auto-switch to, so the viewer's reveal
+ * effect can tell that handoff apart from a user's gallery click. The marker is only meaningful
+ * while that selection stands: once the selection moves on, the recorded auto-switch will never
+ * render, and leaving the marker behind would make the user's next click on that item read as an
+ * auto-switch and get no reveal.
+ *
+ * Matched by state rather than by action type on purpose — the selection is written by several
+ * reducers (imageSelected, selectionChanged, boardIdSelected, comparedImagesSwapped,
+ * showVirtualBoardsChanged, logout), and a new one added later would silently escape an
+ * action-type list, leaving exactly the stale marker this exists to prevent.
+ */
+export const addAutoSwitchedSelectionListener = (startAppListening: AppStartListening) => {
+  startAppListening({
+    predicate: (_action, currentState, previousState) =>
+      selectLastSelectedItem(currentState) !== selectLastSelectedItem(previousState),
+    effect: (_action, { getState }) => {
+      autoSwitchedImages.settle(selectLastSelectedItem(getState()) ?? null);
+    },
+  });
+};

invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/boardIdSelected.test.ts

@@ -0,0 +1,117 @@
+import { configureStore, createListenerMiddleware } from '@reduxjs/toolkit';
+import type { AppStartListening } from 'app/store/store';
+import {
+  gallerySliceConfig,
+  galleryViewChanged,
+  imageSelected,
+  selectionChanged,
+} from 'features/gallery/store/gallerySlice';
+import { api } from 'services/api';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { addBoardIdSelectedListener } from './boardIdSelected';
+
+// The listener waits for the board's item list before auto-selecting, so the store needs the API
+// slice present (the query is never fulfilled here — the point is what happens meanwhile).
+const buildStore = () => {
+  const listenerMiddleware = createListenerMiddleware();
+  addBoardIdSelectedListener(listenerMiddleware.startListening as unknown as AppStartListening);
+  return configureStore({
+    reducer: {
+      gallery: gallerySliceConfig.slice.reducer,
+      [api.reducerPath]: api.reducer,
+    },
+    middleware: (getDefaultMiddleware) =>
+      getDefaultMiddleware({ serializableCheck: false }).prepend(listenerMiddleware.middleware),
+  });
+};
+
+describe('addBoardIdSelectedListener', () => {
+  beforeEach(() => {
+    vi.useFakeTimers();
+  });
+
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
+  it('does not overwrite a selection made while it was waiting for the board list', async () => {
+    // The gallery's auto-switch dispatches galleryViewChanged immediately before imageSelected.
+    // Without the cancel, the probe this starts wakes up on that very selection and re-selects
+    // from a stale (or empty) list, undoing the auto-switch — and the viewer then reveals the
+    // wrong item over the live preview, the flash the auto-switch marker exists to prevent.
+    const store = buildStore();
+
+    store.dispatch(galleryViewChanged('images'));
+    store.dispatch(imageSelected('new.png'));
+
+    // Past the probe's 5s give-up, which would otherwise clear the selection outright.
+    await vi.advanceTimersByTimeAsync(6000);
+
+    expect(store.getState().gallery.selection).toEqual(['new.png']);
+  });
+
+  it('still clears the selection when a board switch finds nothing to show', async () => {
+    // The auto-select probe itself must keep working: a board change with no items selects
+    // nothing rather than leaving the previous board's item highlighted.
+    const store = buildStore();
+    store.dispatch(imageSelected('from-previous-board.png'));
+
+    store.dispatch(galleryViewChanged('assets'));
+    await vi.advanceTimersByTimeAsync(6000);
+
+    expect(store.getState().gallery.selection).toEqual([]);
+  });
+
+  it('does not overwrite a selection made through the gallery grid either', () => {
+    // Thumbnail clicks and keyboard navigation dispatch selectionChanged, not imageSelected, so
+    // matching on the action type alone leaves the ordinary path exposed.
+    const store = buildStore();
+
+    store.dispatch(galleryViewChanged('images'));
+    store.dispatch(selectionChanged(['picked.png']));
+
+    return vi.advanceTimersByTimeAsync(6000).then(() => {
+      expect(store.getState().gallery.selection).toEqual(['picked.png']);
+    });
+  });
+
+  it('does not overwrite a multi-selection made while the probe was waiting', () => {
+    const store = buildStore();
+
+    store.dispatch(galleryViewChanged('images'));
+    store.dispatch(selectionChanged(['first.png', 'second.png']));
+
+    return vi.advanceTimersByTimeAsync(6000).then(() => {
+      expect(store.getState().gallery.selection).toEqual(['first.png', 'second.png']);
+    });
+  });
+
+  it('does not overwrite a multi-selection narrowed while the probe was waiting', () => {
+    // Removing one of two selected thumbnails leaves the *last* selected item unchanged, so a
+    // predicate watching only the active item never fired and the probe survived to replace the
+    // whole selection when it woke.
+    const store = buildStore();
+    store.dispatch(selectionChanged(['first.png', 'second.png']));
+
+    store.dispatch(galleryViewChanged('images'));
+    store.dispatch(selectionChanged(['second.png']));
+
+    return vi.advanceTimersByTimeAsync(6000).then(() => {
+      expect(store.getState().gallery.selection).toEqual(['second.png']);
+    });
+  });
+
+  it('does not overwrite a re-selection of the item already active', () => {
+    // Same shape: the active item does not change, but the user has just said what they want.
+    const store = buildStore();
+    store.dispatch(imageSelected('a.png'));
+
+    store.dispatch(galleryViewChanged('images'));
+    store.dispatch(imageSelected('a.png'));
+
+    return vi.advanceTimersByTimeAsync(6000).then(() => {
+      expect(store.getState().gallery.selection).toEqual(['a.png']);
+    });
+  });
+});

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

@@ -1,16 +1,44 @@
 import { isAnyOf } from '@reduxjs/toolkit';
 import type { AppStartListening } from 'app/store/store';
-import { selectGalleryItemNamesQueryArgs } from 'features/gallery/store/gallerySelectors';
+import { selectGalleryItemNamesQueryArgs, selectSelection } from 'features/gallery/store/gallerySelectors';
 import { boardIdSelected, galleryViewChanged, imageSelected } from 'features/gallery/store/gallerySlice';
 import { galleryApi } from 'services/api/endpoints/gallery';
 
+/** The actions that ask this listener to pick an item for the user. */
+const startsProbe = isAnyOf(boardIdSelected, galleryViewChanged);
+
 export const addBoardIdSelectedListener = (startAppListening: AppStartListening) => {
   startAppListening({
-    matcher: isAnyOf(boardIdSelected, galleryViewChanged),
+    // Two jobs, so this cannot be a plain action matcher. The probe below is started by a board or
+    // view change — but it must also be *cancelled* by any selection that lands while it waits,
+    // and a selection arrives through several actions: imageSelected from the gallery's auto-switch
+    // and keyboard navigation, selectionChanged from thumbnail clicks, boardIdSelected carrying a
+    // selection. Matching the resulting change of the selection covers all of them, including any
+    // writer added later — an action list would silently miss it.
+    //
+    // The whole selection, not just its active item: removing one of several selected thumbnails,
+    // or re-picking the one already active, leaves the last item unchanged while still being the
+    // user settling what they want. Comparing only that item left the probe running through those,
+    // to overwrite their selection when it woke. The state is immutable, so a new array reference
+    // is exactly "the selection was written", and cancelling a probe more often than strictly
+    // needed costs nothing.
+    predicate: (action, currentState, previousState) =>
+      startsProbe(action) || selectSelection(currentState) !== selectSelection(previousState),
     effect: async (action, { getState, dispatch, condition, cancelActiveListeners }) => {
       // Cancel any in-progress instances of this listener, we don't want to select an item from a previous board
       cancelActiveListeners();
 
+      if (!startsProbe(action)) {
+        // A selection landed. It settles what should be displayed, so a probe still waiting on a
+        // board's items must not overwrite it when it resolves. The gallery's auto-switch dispatches
+        // galleryViewChanged immediately before its selection: without this the probe that view
+        // change starts wakes on the selection that follows it, re-selects the first name in a
+        // possibly stale cached list, and undoes the switch — and the viewer then reveals that
+        // wrong image over the live preview, which is the flash the auto-switch marker exists to
+        // prevent. Cancelling above is the whole effect; there is nothing to auto-select here.
+        return;
+      }
+
       if (boardIdSelected.match(action) && action.payload.select) {
         // This action already has a resource selection - skip the below auto-selection logic
         return;

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

@@ -12,6 +12,7 @@ import { errorHandler } from 'app/store/enhancers/reduxRemember/errors';
 import { addAdHocPostProcessingRequestedListener } from 'app/store/middleware/listenerMiddleware/listeners/addAdHocPostProcessingRequestedListener';
 import { addAnyEnqueuedListener } from 'app/store/middleware/listenerMiddleware/listeners/anyEnqueued';
 import { addAppStartedListener } from 'app/store/middleware/listenerMiddleware/listeners/appStarted';
+import { addAutoSwitchedSelectionListener } from 'app/store/middleware/listenerMiddleware/listeners/autoSwitchedSelection';
 import { addBatchEnqueuedListener } from 'app/store/middleware/listenerMiddleware/listeners/batchEnqueued';
 import { addDeleteBoardAndImagesFulfilledListener } from 'app/store/middleware/listenerMiddleware/listeners/boardAndImagesDeleted';
 import { addBoardIdSelectedListener } from 'app/store/middleware/listenerMiddleware/listeners/boardIdSelected';
@@ -326,6 +327,7 @@ addImageAddedToBoardFulfilledListener(startAppListening);
 addImageRemovedFromBoardFulfilledListener(startAppListening);
 addBoardIdSelectedListener(startAppListening);
 addArchivedOrDeletedBoardListener(startAppListening);
+addAutoSwitchedSelectionListener(startAppListening);
 
 // Node schemas
 addGetOpenAPISchemaListener(startAppListening);

invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentImagePreview.test.ts

@@ -0,0 +1,53 @@
+import { readFileSync } from 'node:fs';
+import { fileURLToPath } from 'node:url';
+
+import { describe, expect, it } from 'vitest';
+
+const read = (file: string) => readFileSync(fileURLToPath(new URL(file, import.meta.url)), 'utf8');
+
+// Wiring checks only — this directory has no DOM test environment, so the component cannot be
+// mounted. The lifecycle behavior behind onLoadImage is covered by real tests in
+// viewerProgressLifecycle.test.ts, and the reveal-suppression registry in autoSwitchedImages.test.ts.
+describe('CurrentImagePreview reveal wiring', () => {
+  const currentImagePreview = read('./CurrentImagePreview.tsx');
+
+  it('gates the viewer reveal on the thumbnail rather than the full-resolution image', () => {
+    // Gating on `/full` holds a stale latent preview on screen for the whole multi-megabyte
+    // download on a slow connection.
+    expect(currentImagePreview).toContain('useMediaUrl(imageDTO?.thumbnail_url)');
+    expect(currentImagePreview).toContain('preloader.src = previewSrc');
+    expect(currentImagePreview).not.toMatch(/preloader\.src\s*=\s*imageDTO\.image_url/);
+  });
+
+  it('clears the progress overlay when the preload settles, including on error', () => {
+    // Chakra reports a failed load as onError, not onLoad, so DndImage's onLoad alone is not
+    // enough to guarantee the overlay is ever cleared.
+    expect(currentImagePreview).toContain('preloader.onerror = onReady');
+    const onReady = currentImagePreview.slice(
+      currentImagePreview.indexOf('const onReady ='),
+      currentImagePreview.indexOf('if (typeof window ===')
+    );
+    expect(onReady).toContain('onLoadImage(imageDTO.session_id ?? null)');
+  });
+
+  it('routes the reveal through the shared decision, passing the auto-switch marker to it', () => {
+    // The decision's branches are unit tested in selectedItemReveal.test.ts; what cannot be seen
+    // from there is whether this component still feeds it the marker, or still writes the atom on
+    // the hide path.
+    expect(currentImagePreview).toContain('getSelectedItemRevealDecision({');
+    expect(currentImagePreview).toContain('wasAutoSwitchedTo,');
+    expect(currentImagePreview).toMatch(
+      /if \(decision === 'hide'\) \{\s+\$isTemporarilyShowingSelectedImage\.set\(false\);\s+return;/
+    );
+  });
+
+  it('consumes the auto-switch marker on every rendered-image change', () => {
+    // The consume call must sit before the reveal gates, so the marker is cleared even when the
+    // auto-switched image renders with no progress showing.
+    const revealEffect = currentImagePreview.slice(
+      currentImagePreview.indexOf('const renderedImageName ='),
+      currentImagePreview.indexOf('$isTemporarilyShowingSelectedImage.set(true)')
+    );
+    expect(revealEffect).toContain('autoSwitchedImages.consume(renderedImageName)');
+  });
+});