diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/autoSwitchedSelection.test.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/autoSwitchedSelection.test.ts new file mode 100644 index 00000000000..54c1ff2bde2 --- /dev/null +++ b/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); + }); +}); diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/autoSwitchedSelection.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/autoSwitchedSelection.ts new file mode 100644 index 00000000000..0b14616f0b3 --- /dev/null +++ b/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); + }, + }); +}; diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/boardIdSelected.test.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/boardIdSelected.test.ts new file mode 100644 index 00000000000..f457fac2280 --- /dev/null +++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/boardIdSelected.test.ts @@ -0,0 +1,60 @@ +import { configureStore, createListenerMiddleware } from '@reduxjs/toolkit'; +import type { AppStartListening } from 'app/store/store'; +import { gallerySliceConfig, galleryViewChanged, imageSelected } 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([]); + }); +}); diff --git a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/boardIdSelected.ts b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/boardIdSelected.ts index 05dd8e9f208..588b39881a8 100644 --- a/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/boardIdSelected.ts +++ b/invokeai/frontend/web/src/app/store/middleware/listenerMiddleware/listeners/boardIdSelected.ts @@ -6,11 +6,22 @@ import { galleryApi } from 'services/api/endpoints/gallery'; export const addBoardIdSelectedListener = (startAppListening: AppStartListening) => { startAppListening({ - matcher: isAnyOf(boardIdSelected, galleryViewChanged), + matcher: isAnyOf(boardIdSelected, galleryViewChanged, imageSelected), 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 (imageSelected.match(action)) { + // An explicit selection 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 imageSelected: 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; diff --git a/invokeai/frontend/web/src/app/store/store.ts b/invokeai/frontend/web/src/app/store/store.ts index 59551225314..7ea4ffbf45c 100644 --- a/invokeai/frontend/web/src/app/store/store.ts +++ b/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); diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentImagePreview.test.ts b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentImagePreview.test.ts new file mode 100644 index 00000000000..73c81cab943 --- /dev/null +++ b/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)'); + }); +}); diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentImagePreview.tsx b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentImagePreview.tsx index ba76ca299bb..55688c037bc 100644 --- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentImagePreview.tsx +++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentImagePreview.tsx @@ -1,11 +1,13 @@ import { Box, Flex } from '@invoke-ai/ui-library'; import { useStore } from '@nanostores/react'; import { useAppSelector } from 'app/store/storeHooks'; +import { useMediaUrl } from 'features/auth/store/mediaCookieRefresh'; import { CanvasAlertsInvocationProgress } from 'features/controlLayers/components/CanvasAlerts/CanvasAlertsInvocationProgress'; import { DndImage } from 'features/dnd/DndImage'; import ImageMetadataViewer from 'features/gallery/components/ImageMetadataViewer/ImageMetadataViewer'; import NextPrevItemButtons from 'features/gallery/components/NextPrevItemButtons'; import { useNextPrevItemNavigation } from 'features/gallery/components/useNextPrevItemNavigation'; +import { autoSwitchedImages } from 'features/gallery/store/autoSwitchedImages'; import { selectLastSelectedItem } from 'features/gallery/store/gallerySelectors'; import { useRegisteredHotkeys } from 'features/system/components/HotkeysModal/useHotkeyData'; import { navigationApi } from 'features/ui/layouts/navigation-api'; @@ -24,6 +26,7 @@ import { NoContentForViewer } from './NoContentForViewer'; import { ProgressImage } from './ProgressImage2'; import { ProgressImageTiles } from './ProgressImageTiles'; import { ProgressIndicator } from './ProgressIndicator2'; +import { getSelectedItemRevealDecision } from './selectedItemReveal'; export const CurrentImagePreview = memo(({ imageDTO }: { imageDTO: ImageDTO | null }) => { const activeTab = useAppSelector(selectActiveTab); @@ -48,6 +51,20 @@ export const CurrentImagePreview = memo(({ imageDTO }: { imageDTO: ImageDTO | nu const previousRenderedImageNameRef = useRef(null); const selectedImageRevealTimeoutId = useRef(0); + // The reveal gate below deliberately preloads the *thumbnail*, not the full-resolution image. The + // progress overlay covers this element until onLoadImage fires, so gating on the multi-megabyte + // `/full` response would hold a stale latent preview on screen for that entire download on a slow + // connection. The 256px thumbnail is roughly 100x smaller and is typically higher resolution than + // the preview it replaces; DndImage renders it via Chakra's `fallbackSrc` and swaps the full image + // in, in place, once that finishes loading. + // + // The URL must go through useMediaUrl so it is byte-identical to the one DndImage requests. The + // media cookie version is a query parameter, so a mismatch is a different key and the bytes are + // fetched twice (measured: 2 requests mismatched vs 1 matched). Note the reuse here is the + // document's list of available images, which is keyed by URL and is not the HTTP cache — it still + // holds in multiuser mode, where images are served `Cache-Control: private, no-store`. + const previewSrc = useMediaUrl(imageDTO?.thumbnail_url); + useEffect(() => { if (!selectedImageName) { setImageToRender(null); @@ -65,9 +82,15 @@ export const CurrentImagePreview = memo(({ imageDTO }: { imageDTO: ImageDTO | nu return; } setImageToRender(imageDTO); + // Resolve the progress overlay as soon as the thumbnail settles — on success *or* error. + // Relying on DndImage's onLoad alone leaves the overlay stuck whenever the image fails to + // load, because Chakra reports that as onError instead. The session id lets the lifecycle + // attribute the load, so a late-settling thumbnail from an earlier session cannot cut a + // different session's resolve illusion short. + onLoadImage(imageDTO.session_id ?? null); }; - if (typeof window === 'undefined') { + if (typeof window === 'undefined' || !previewSrc) { onReady(); return; } @@ -76,7 +99,7 @@ export const CurrentImagePreview = memo(({ imageDTO }: { imageDTO: ImageDTO | nu preloader.onload = onReady; preloader.onerror = onReady; - preloader.src = imageDTO.image_url; + preloader.src = previewSrc; if (preloader.complete) { onReady(); @@ -87,7 +110,7 @@ export const CurrentImagePreview = memo(({ imageDTO }: { imageDTO: ImageDTO | nu preloader.onload = null; preloader.onerror = null; }; - }, [imageDTO, imageToRender?.image_name, selectedImageName]); + }, [imageDTO, imageToRender?.image_name, onLoadImage, previewSrc, selectedImageName]); const hasProgressImage = progressImage !== null; @@ -96,20 +119,33 @@ export const CurrentImagePreview = memo(({ imageDTO }: { imageDTO: ImageDTO | nu const previousRenderedImageName = previousRenderedImageNameRef.current; previousRenderedImageNameRef.current = renderedImageName; + // Consume on every change of the rendered image, not only when the reveal conditions below + // hold — in the common case the auto-switched image renders with no progress showing, and the + // marker must not outlive the render it was recorded for. The marker is also dropped whenever + // the selection moves on before rendering (see addAutoSwitchedSelectionListener), so it can + // never suppress a genuine user selection of the same image later. + const wasAutoSwitchedTo = + renderedImageName !== null && + renderedImageName !== previousRenderedImageName && + autoSwitchedImages.consume(renderedImageName); + + // Clearing the timer before deciding is why getSelectedItemRevealDecision has no third + // "leave it alone" outcome: a path that returned without writing the atom would leave a raised + // reveal with nothing left to lower it, hiding the live preview for the rest of the render. window.clearTimeout(selectedImageRevealTimeoutId.current); - if ( - !shouldShowProgressInViewer || - !hasProgressImage || - isProgressImageResolving || - !renderedImageName || - renderedImageName !== selectedImageName - ) { - $isTemporarilyShowingSelectedImage.set(false); - return; - } + const decision = getSelectedItemRevealDecision({ + shouldShowProgressInViewer, + hasProgressImage, + isProgressImageResolving, + renderedItemName: renderedImageName, + selectedItemName: selectedImageName ?? null, + previousRenderedItemName: previousRenderedImageName, + wasAutoSwitchedTo, + }); - if (previousRenderedImageName === null || previousRenderedImageName === renderedImageName) { + if (decision === 'hide') { + $isTemporarilyShowingSelectedImage.set(false); return; } diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/selectedItemReveal.test.ts b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/selectedItemReveal.test.ts new file mode 100644 index 00000000000..b9228abe693 --- /dev/null +++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/selectedItemReveal.test.ts @@ -0,0 +1,69 @@ +import { describe, expect, it } from 'vitest'; + +import type { SelectedItemRevealInputs } from './selectedItemReveal'; +import { getSelectedItemRevealDecision } from './selectedItemReveal'; + +// A mid-generation user click on b.png, with the previous item still remembered: the case the +// reveal exists for. Each test below perturbs exactly one input away from it. +const userClickMidGeneration: SelectedItemRevealInputs = { + shouldShowProgressInViewer: true, + hasProgressImage: true, + isProgressImageResolving: false, + renderedItemName: 'b.png', + selectedItemName: 'b.png', + previousRenderedItemName: 'a.png', + wasAutoSwitchedTo: false, +}; + +describe('getSelectedItemRevealDecision', () => { + it('reveals a mid-generation user click so it visibly lands', () => { + expect(getSelectedItemRevealDecision(userClickMidGeneration)).toBe('reveal'); + }); + + it('does not reveal an auto-switch to a just-finished item', () => { + // Without this the finished image flashes over the next generation's live preview for 2s. + expect(getSelectedItemRevealDecision({ ...userClickMidGeneration, wasAutoSwitchedTo: true })).toBe('hide'); + }); + + it('hides when no progress preview is covering the viewer', () => { + expect(getSelectedItemRevealDecision({ ...userClickMidGeneration, hasProgressImage: false })).toBe('hide'); + expect(getSelectedItemRevealDecision({ ...userClickMidGeneration, shouldShowProgressInViewer: false })).toBe( + 'hide' + ); + }); + + it('hides while a finished generation is resolving into its final image', () => { + expect(getSelectedItemRevealDecision({ ...userClickMidGeneration, isProgressImageResolving: true })).toBe('hide'); + }); + + it('hides while the render still lags the selection', () => { + // The preload has not settled, so the item on screen is not the one that was clicked. + expect(getSelectedItemRevealDecision({ ...userClickMidGeneration, renderedItemName: 'a.png' })).toBe('hide'); + expect(getSelectedItemRevealDecision({ ...userClickMidGeneration, renderedItemName: null })).toBe('hide'); + }); + + it('hides when the displayed item did not change', () => { + expect(getSelectedItemRevealDecision({ ...userClickMidGeneration, previousRenderedItemName: 'b.png' })).toBe( + 'hide' + ); + }); + + it('hides on the first render, which is not a click', () => { + expect(getSelectedItemRevealDecision({ ...userClickMidGeneration, previousRenderedItemName: null })).toBe('hide'); + }); + + it('never answers anything but reveal or hide', () => { + // The caller has already cleared the running reveal's timer by the time it asks, so a third + // "leave it alone" outcome would wedge the overlay off for the rest of the render. + const inputs: SelectedItemRevealInputs[] = [ + userClickMidGeneration, + { ...userClickMidGeneration, wasAutoSwitchedTo: true }, + { ...userClickMidGeneration, hasProgressImage: false }, + { ...userClickMidGeneration, renderedItemName: null }, + { ...userClickMidGeneration, previousRenderedItemName: null }, + ]; + for (const input of inputs) { + expect(['reveal', 'hide']).toContain(getSelectedItemRevealDecision(input)); + } + }); +}); diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/selectedItemReveal.ts b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/selectedItemReveal.ts new file mode 100644 index 00000000000..58bb94e4393 --- /dev/null +++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/selectedItemReveal.ts @@ -0,0 +1,74 @@ +/** + * Whether the viewer should lift its progress overlay to show the selected item. + * + * A generation covers the viewer with an opaque progress overlay, which would otherwise swallow + * every gallery click for the whole render: the selection changes underneath, but nothing visibly + * happens. The reveal makes that click land — the clicked item is shown for + * SELECTED_ITEM_REVEAL_DURATION_MS, then the live preview returns. + * + * Only two outcomes, deliberately. The caller clears the running reveal's timer before asking, so a + * "leave it as it is" answer would strand the reveal on for the rest of the render with no timer + * left to end it — the decision must always name the state the overlay ends up in. + */ +type SelectedItemRevealDecision = 'reveal' | 'hide'; + +export type SelectedItemRevealInputs = { + /** The user's "show progress in viewer" setting. */ + shouldShowProgressInViewer: boolean; + /** Whether a live progress preview exists to be covering the viewer at all. */ + hasProgressImage: boolean; + /** Whether a finished generation's preview is mid-handoff to its final image. */ + isProgressImageResolving: boolean; + /** The item the viewer is rendering right now, or null when it has nothing to show. */ + renderedItemName: string | null; + /** The item the gallery selection points at. */ + selectedItemName: string | null; + /** The item the viewer rendered before this change — null on the first render. */ + previousRenderedItemName: string | null; + /** Whether this render is the gallery auto-switching to a just-finished item, rather than the + * user picking one. See features/gallery/store/autoSwitchedImages. */ + wasAutoSwitchedTo: boolean; +}; + +export const getSelectedItemRevealDecision = ({ + shouldShowProgressInViewer, + hasProgressImage, + isProgressImageResolving, + renderedItemName, + selectedItemName, + previousRenderedItemName, + wasAutoSwitchedTo, +}: SelectedItemRevealInputs): SelectedItemRevealDecision => { + // Nothing is covering the viewer, so there is nothing to reveal from under. + if (!shouldShowProgressInViewer || !hasProgressImage) { + return 'hide'; + } + + // The preview is already resolving into a finished generation's final image; a reveal here would + // fight that handoff. + if (isProgressImageResolving) { + return 'hide'; + } + + // Render lagging the selection (the preload has not settled yet): whatever is on screen is not + // what the user picked, so showing it would not make their click land. + if (renderedItemName === null || renderedItemName !== selectedItemName) { + return 'hide'; + } + + // Not a change of displayed item — nothing happened that needs to be made visible. The first + // render after the viewer opens is not a click either. + if (previousRenderedItemName === null || previousRenderedItemName === renderedItemName) { + return 'hide'; + } + + // The reveal exists for *user* selections. An auto-switch to a just-finished item lands + // asynchronously and can reach the viewer after the next generation's first progress event, so + // timing cannot tell it apart from a click — treating it as one flashes the finished item over + // the new generation's live preview for two seconds. + if (wasAutoSwitchedTo) { + return 'hide'; + } + + return 'reveal'; +}; diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.test.ts b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.test.ts index a3181382a7f..a807d9ba7d3 100644 --- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.test.ts +++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.test.ts @@ -130,6 +130,34 @@ describe('viewerProgressLifecycle', () => { expect(stores.$progressEvent.get()?.item_id).toBe(1); }); + it('keeps promoting through a chain of terminations while any session is still generating', () => { + // Three sessions, each promotion making the next owner: the shared preview must never be + // torn down while something is still generating, no matter how many owners terminate. Each + // terminal event here arrives with no further progress event behind it, so a promotion that + // left the new owner unprotected would show up as a blank viewer on the next termination. + const eventB = buildProgressEvent({ item_id: 2, session_id: 'session-2', image: buildProgressImage(2) }); + const eventC = buildProgressEvent({ item_id: 3, session_id: 'session-3', image: buildProgressImage(3) }); + const eventA = buildProgressEvent({ item_id: 1, session_id: 'session-1', image: buildProgressImage(1) }); + lifecycle.recordProgress(eventC); + lifecycle.recordProgress(eventB); + lifecycle.recordProgress(eventA); + + // A owns the preview and completes: B reported most recently of the two left, so B takes it. + lifecycle.onTerminal(buildTerminalEvent({ item_id: 1, status: 'completed' }), true); + expect(stores.$progressEvent.get()).toBe(eventB); + + // B is then canceled while C is still generating — C must take the preview, not lose it. + lifecycle.onTerminal(buildTerminalEvent({ item_id: 2, status: 'canceled' }), true); + expect(stores.$progressEvent.get()).toBe(eventC); + expect(stores.$progressImage.get()).toBe(eventC.image); + expect(stores.$isProgressImageResolving.get()).toBe(false); + + // Only when the last session ends does the preview come down. + lifecycle.onTerminal(buildTerminalEvent({ item_id: 3, status: 'canceled' }), true); + expect(stores.$progressEvent.get()).toBeNull(); + expect(stores.$progressImage.get()).toBeNull(); + }); + it('leaves the shared preview alone when a non-owner terminates', () => { const { eventA } = startTwoSessions(); lifecycle.onTerminal(buildTerminalEvent({ item_id: 2, status: 'canceled' }), true); diff --git a/invokeai/frontend/web/src/features/gallery/store/autoSwitchedImages.test.ts b/invokeai/frontend/web/src/features/gallery/store/autoSwitchedImages.test.ts new file mode 100644 index 00000000000..899397eeaac --- /dev/null +++ b/invokeai/frontend/web/src/features/gallery/store/autoSwitchedImages.test.ts @@ -0,0 +1,79 @@ +import { describe, expect, it } from 'vitest'; + +import { createAutoSwitchedSelectionMarker } from './autoSwitchedImages'; + +describe('createAutoSwitchedSelectionMarker', () => { + it('reports the auto-switched item on its first render, once', () => { + const marker = createAutoSwitchedSelectionMarker(); + marker.record('a.png'); + marker.settle('a.png'); + expect(marker.consume('a.png')).toBe(true); + expect(marker.consume('a.png')).toBe(false); + }); + + it('reports nothing for an item that was never auto-switched to', () => { + const marker = createAutoSwitchedSelectionMarker(); + expect(marker.consume('a.png')).toBe(false); + }); + + it('drops a marker whose selection was superseded before it rendered', () => { + // The auto-switch to A is dispatched, then the user clicks B before A's preload settles. A can + // never render as that auto-switch, and the user's later click on A must reveal it. + const marker = createAutoSwitchedSelectionMarker(); + marker.record('a.png'); + marker.settle('a.png'); + marker.settle('b.png'); + marker.settle('a.png'); + expect(marker.consume('a.png')).toBe(false); + }); + + it('keeps only the last of several auto-switches settled in one batch', () => { + // Two sessions completing together record two names, but only the last selection stands. + const marker = createAutoSwitchedSelectionMarker(); + marker.record('a.png'); + marker.settle('a.png'); + marker.record('b.png'); + marker.settle('b.png'); + expect(marker.consume('b.png')).toBe(true); + expect(marker.consume('a.png')).toBe(false); + }); + + it('survives settles that do not move the selection', () => { + // Re-selecting the same item (or any other action that leaves the selection alone) must not + // discard a marker whose item has not rendered yet. + const marker = createAutoSwitchedSelectionMarker(); + marker.record('a.png'); + marker.settle('a.png'); + marker.settle('a.png'); + expect(marker.consume('a.png')).toBe(true); + }); + + it('drops the marker when the selection is cleared', () => { + const marker = createAutoSwitchedSelectionMarker(); + marker.record('a.png'); + marker.settle('a.png'); + marker.settle(null); + expect(marker.consume('a.png')).toBe(false); + }); + + it('does not report a second render of an item after its marker was consumed', () => { + // The user navigates away and back: that second render is a user selection and must reveal. + const marker = createAutoSwitchedSelectionMarker(); + marker.record('a.png'); + marker.settle('a.png'); + expect(marker.consume('a.png')).toBe(true); + marker.settle('b.png'); + marker.settle('a.png'); + expect(marker.consume('a.png')).toBe(false); + }); + + it('does not clear a pending marker when a different item is consumed', () => { + // The rendered item can lag a just-auto-switched selection: the viewer consumes for whatever + // is still on screen, and that must not spend the marker the incoming item is owed. + const marker = createAutoSwitchedSelectionMarker(); + marker.record('a.png'); + marker.settle('a.png'); + expect(marker.consume('b.png')).toBe(false); + expect(marker.consume('a.png'), 'the marker survives an unrelated consume').toBe(true); + }); +}); diff --git a/invokeai/frontend/web/src/features/gallery/store/autoSwitchedImages.ts b/invokeai/frontend/web/src/features/gallery/store/autoSwitchedImages.ts new file mode 100644 index 00000000000..d7ab9003458 --- /dev/null +++ b/invokeai/frontend/web/src/features/gallery/store/autoSwitchedImages.ts @@ -0,0 +1,61 @@ +/** + * Marks the gallery selection the auto-switch made, until the viewer renders it. + * + * The viewer briefly reveals a newly selected item over the progress overlay so that a + * mid-generation gallery click is not invisible (see the reveal effect in CurrentImagePreview). + * Auto-switch selections must not trigger that reveal — but they land asynchronously: the switch is + * dispatched only after onInvocationComplete's DTO fetch resolves, and the viewer renders it only + * after the thumbnail preload settles. When the next generation is started quickly, its first + * invocation_progress event slots into that window and resets $isProgressImageResolving, so by the + * time the auto-switch selection reaches the viewer it is indistinguishable from a user click and + * the reveal flashes the previous result over the live preview for 2 seconds. + * + * The marker is scoped to the selection it was recorded for, not to the item name: it survives only + * as long as that selection stands. A name-keyed marker cannot tell "this render is the auto-switch + * landing" from "the user picked that same item later", so an auto-switch that never rendered — + * because the user clicked elsewhere first — would swallow their later click on it, the very dead + * click the reveal exists to prevent. Settling on every selection change closes that: once the + * selection moves on, the recorded auto-switch can never render, and the marker goes with it. + * + * That scoping is also why no expiry or bound is needed. At most one marker exists, and only while + * its selection is the current one. + */ +type AutoSwitchedSelectionMarker = { + /** Records that the gallery is auto-switching the selection to this item. Call immediately + * before dispatching that selection, so the settle it triggers sees the marker. */ + record: (itemName: string) => void; + /** + * Points the marker at the selection that now stands, dropping it unless it is still the one + * recorded. Call on every selection change (see addAutoSwitchedSelectionListener). + */ + settle: (selectedItemName: string | null) => void; + /** + * Returns whether the item now rendering is the one the auto-switch selected, clearing the + * marker. Call on every change of the rendered item. + */ + consume: (itemName: string) => boolean; +}; + +export const createAutoSwitchedSelectionMarker = (): AutoSwitchedSelectionMarker => { + let pendingItemName: string | null = null; + + return { + record: (itemName) => { + pendingItemName = itemName; + }, + settle: (selectedItemName) => { + if (pendingItemName !== selectedItemName) { + pendingItemName = null; + } + }, + consume: (itemName) => { + if (pendingItemName !== itemName) { + return false; + } + pendingItemName = null; + return true; + }, + }; +}; + +export const autoSwitchedImages = createAutoSwitchedSelectionMarker(); diff --git a/invokeai/frontend/web/src/services/api/endpoints/videos.ts b/invokeai/frontend/web/src/services/api/endpoints/videos.ts index 49043d763a0..3f7327ca49c 100644 --- a/invokeai/frontend/web/src/services/api/endpoints/videos.ts +++ b/invokeai/frontend/web/src/services/api/endpoints/videos.ts @@ -352,14 +352,13 @@ export const { useDeleteUncategorizedVideosMutation, } = videosApi; +export const { useGetVideoMetadataQuery, useGetVideoWorkflowQuery, useLazyGetVideoWorkflowQuery } = videosApi; + /** @knipignore Reserved for follow-up phases (bulk delete / intermediate toggle / video-only views). * useDeleteVideoMutation is here because the only call site uses videosApi.endpoints.deleteVideo.initiate * via the delete-video modal, but a future bulk/multi-select flow may want the React hook form. */ export const { useListVideosQuery, - useGetVideoMetadataQuery, - useGetVideoWorkflowQuery, - useLazyGetVideoWorkflowQuery, useGetVideoNamesQuery, useDeleteVideoMutation, useDeleteVideosMutation, diff --git a/invokeai/frontend/web/src/services/events/onInvocationComplete.test.ts b/invokeai/frontend/web/src/services/events/onInvocationComplete.test.ts index 0a811360ae0..35a2e5c11bb 100644 --- a/invokeai/frontend/web/src/services/events/onInvocationComplete.test.ts +++ b/invokeai/frontend/web/src/services/events/onInvocationComplete.test.ts @@ -117,8 +117,11 @@ vi.mock('services/events/stores', () => ({ // Import AFTER the mocks above are declared (vi.mock is hoisted; explicit ordering here // is for the human reader). +import { autoSwitchedImages } from 'features/gallery/store/autoSwitchedImages'; +import { selectAutoSwitch } from 'features/gallery/store/gallerySelectors'; import { getImageDTOSafe } from 'services/api/endpoints/images'; import { getVideoDTOSafe } from 'services/api/endpoints/videos'; +import { $lastProgressEvent } from 'services/events/stores'; import { buildOnForeignInvocationComplete, @@ -157,6 +160,10 @@ const buildImageCompleteEvent = (): S['InvocationCompleteEvent'] => describe('onInvocationComplete polymorphic gallery cache', () => { beforeEach(() => { vi.clearAllMocks(); + // clearAllMocks resets calls, not implementations — restore the factory default so a test that + // turns auto-switch on cannot leak it into the next one. + vi.mocked(selectAutoSwitch).mockReturnValue(false); + autoSwitchedImages.settle(null); }); it('invalidates GalleryItemNameList + GalleryItemList when an image output completes', async () => { @@ -288,6 +295,407 @@ describe('onInvocationComplete polymorphic gallery cache', () => { expect(galleryInvalidation?.payload).toContainEqual({ type: 'BoardVideosTotal', id: 'board-123' }); expect(galleryInvalidation?.payload).toContain('VirtualBoards'); }); + + it('processes each completion event exactly once — a duplicate delivery does no gallery work', async () => { + // Re-running the gallery handling on a duplicate double-counts the optimistic board totals and + // re-records the auto-switch marker after it was consumed, which suppresses a later genuine + // gallery click on that image. + const dispatch = vi.fn(() => ({ unwrap: () => Promise.resolve(undefined) })); + const getState = vi.fn(() => ({})); + + const handler = buildOnInvocationComplete( + // eslint-disable-next-line @typescript-eslint/no-explicit-any + getState as any, + // eslint-disable-next-line @typescript-eslint/no-explicit-any + dispatch as any, + new Map() + ); + + await handler(buildImageCompleteEvent()); + const dispatchCountAfterFirst = dispatch.mock.calls.length; + expect(getImageDTOSafe).toHaveBeenCalledTimes(1); + + await handler(buildImageCompleteEvent()); + expect(getImageDTOSafe).toHaveBeenCalledTimes(1); + expect(dispatch.mock.calls.length).toBe(dispatchCountAfterFirst); + }); + + it('rejects a duplicate that arrives while the first delivery is still awaiting its DTO fetch', async () => { + // The gallery work awaits a DTO fetch, so a duplicate can land mid-flight. The handler must + // mark the event as processed before the first await, not after. + const dispatch = vi.fn(() => ({ unwrap: () => Promise.resolve(undefined) })); + const getState = vi.fn(() => ({})); + + const handler = buildOnInvocationComplete( + // eslint-disable-next-line @typescript-eslint/no-explicit-any + getState as any, + // eslint-disable-next-line @typescript-eslint/no-explicit-any + dispatch as any, + new Map() + ); + + await Promise.all([handler(buildImageCompleteEvent()), handler(buildImageCompleteEvent())]); + expect(getImageDTOSafe).toHaveBeenCalledTimes(1); + }); + + it('lets a re-delivery retry when the first attempt lost its DTO to a transient failure', async () => { + // getImageDTOSafe swallows fetch errors and returns null, so a transient failure silently + // drops the image from the gallery. Keeping the dedupe key in that case would make the loss + // permanent — the redelivery that could fix it is turned away as a duplicate. + vi.mocked(getImageDTOSafe).mockResolvedValueOnce(null); + + const dispatched: unknown[] = []; + const dispatch = vi.fn((action: unknown) => { + dispatched.push(action); + return { unwrap: () => Promise.resolve(undefined) }; + }); + const getState = vi.fn(() => ({})); + + const handler = buildOnInvocationComplete( + // eslint-disable-next-line @typescript-eslint/no-explicit-any + getState as any, + // eslint-disable-next-line @typescript-eslint/no-explicit-any + dispatch as any, + new Map() + ); + + await handler(buildImageCompleteEvent()); + expect( + dispatched.some((action) => { + const payload = (action as { payload?: unknown }).payload; + return Array.isArray(payload) && payload.includes('GalleryItemNameList'); + }), + 'a failed lookup produces no gallery work' + ).toBe(false); + + await handler(buildImageCompleteEvent()); + expect(getImageDTOSafe).toHaveBeenCalledTimes(2); + expect( + dispatched.some((action) => { + const payload = (action as { payload?: unknown }).payload; + return Array.isArray(payload) && payload.includes('GalleryItemNameList'); + }), + 'the retry must do the gallery work the first delivery lost' + ).toBe(true); + }); + + it('retries only the output whose lookup failed, leaving the ones that landed alone', async () => { + // One of two image lookups fails. The successful one's board totals and optimistic insert have + // already gone out, so a re-delivery must not touch it — but the lost one is only recoverable + // here, since nothing re-emits the event on its own. + vi.mocked(getImageDTOSafe).mockResolvedValueOnce(null); + + const dispatch = vi.fn(() => ({ unwrap: () => Promise.resolve(undefined) })); + const getState = vi.fn(() => ({})); + + const handler = buildOnInvocationComplete( + // eslint-disable-next-line @typescript-eslint/no-explicit-any + getState as any, + // eslint-disable-next-line @typescript-eslint/no-explicit-any + dispatch as any, + new Map() + ); + + const twoImages = buildImageCompleteEvent(); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (twoImages as any).result = { + image_1: { image_name: 'first.png' }, + image_2: { image_name: 'second.png' }, + }; + + await handler(twoImages); + expect(getImageDTOSafe).toHaveBeenCalledTimes(2); + + await handler(twoImages); + // Exactly one more lookup, and it is the one that failed. + expect(getImageDTOSafe).toHaveBeenCalledTimes(3); + expect(vi.mocked(getImageDTOSafe).mock.calls.at(-1)?.[0]).toBe('first.png'); + + // Once everything has landed, the event is closed again. + await handler(twoImages); + expect(getImageDTOSafe).toHaveBeenCalledTimes(3); + }); + + it('retries a failed lookup inside an image collection', async () => { + // Collections are where partial failure is actually plausible: the user sees most of a batch + // and silently misses one image, its auto-switch, and its board count. + vi.mocked(getImageDTOSafe).mockResolvedValueOnce(null); + + const dispatch = vi.fn(() => ({ unwrap: () => Promise.resolve(undefined) })); + const getState = vi.fn(() => ({})); + + const handler = buildOnInvocationComplete( + // eslint-disable-next-line @typescript-eslint/no-explicit-any + getState as any, + // eslint-disable-next-line @typescript-eslint/no-explicit-any + dispatch as any, + new Map() + ); + + const collection = buildImageCompleteEvent(); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (collection as any).result = { + collection: [{ image_name: 'batch-1.png' }, { image_name: 'batch-2.png' }, { image_name: 'batch-3.png' }], + }; + + await handler(collection); + expect(getImageDTOSafe).toHaveBeenCalledTimes(3); + + await handler(collection); + expect(getImageDTOSafe).toHaveBeenCalledTimes(4); + expect(vi.mocked(getImageDTOSafe).mock.calls.at(-1)?.[0]).toBe('batch-1.png'); + }); + + it('lets a duplicate that overlapped a lost delivery become the retry', async () => { + // The duplicate arrives while the first delivery is still fetching, so it cannot be told yet + // that the fetch will fail. Rejecting it outright strands the output: there is no third event. + let resolveFirstLookup: (value: null) => void = () => {}; + vi.mocked(getImageDTOSafe).mockImplementationOnce( + () => + new Promise((resolve) => { + resolveFirstLookup = resolve; + }) + ); + + const dispatch = vi.fn(() => ({ unwrap: () => Promise.resolve(undefined) })); + const getState = vi.fn(() => ({})); + + const handler = buildOnInvocationComplete( + // eslint-disable-next-line @typescript-eslint/no-explicit-any + getState as any, + // eslint-disable-next-line @typescript-eslint/no-explicit-any + dispatch as any, + new Map() + ); + + const event = buildImageCompleteEvent(); + const first = handler(event); + const duplicate = handler(event); + + // The first delivery's only lookup fails, losing the image. + resolveFirstLookup(null); + await Promise.all([first, duplicate]); + + // The duplicate picked the work back up rather than being discarded. + expect(getImageDTOSafe).toHaveBeenCalledTimes(2); + expect(vi.mocked(getImageDTOSafe).mock.calls.at(-1)?.[0]).toBe('fresh-image.png'); + }); + + it('marks the selection it auto-switches to, so the viewer does not reveal it as a user click', async () => { + // Without the marker the finished image flashes over the next generation's live preview for + // two seconds — the regression this PR exists to fix. The marker is the only signal that can + // tell the handoff from a click, since it lands after the next generation's first progress + // event has already reset the timing-based guard. + vi.mocked(selectAutoSwitch).mockReturnValue(true); + autoSwitchedImages.settle(null); + + const dispatch = vi.fn(() => ({ unwrap: () => Promise.resolve(undefined) })); + const getState = vi.fn(() => ({})); + + const handler = buildOnInvocationComplete( + // eslint-disable-next-line @typescript-eslint/no-explicit-any + getState as any, + // eslint-disable-next-line @typescript-eslint/no-explicit-any + dispatch as any, + new Map() + ); + + await handler(buildImageCompleteEvent()); + + expect(autoSwitchedImages.consume('fresh-image.png')).toBe(true); + }); + + it('does not mark anything when auto-switch is off — nothing is selected to be revealed', async () => { + vi.mocked(selectAutoSwitch).mockReturnValue(false); + autoSwitchedImages.settle(null); + + const dispatch = vi.fn(() => ({ unwrap: () => Promise.resolve(undefined) })); + const getState = vi.fn(() => ({})); + + const handler = buildOnInvocationComplete( + // eslint-disable-next-line @typescript-eslint/no-explicit-any + getState as any, + // eslint-disable-next-line @typescript-eslint/no-explicit-any + dispatch as any, + new Map() + ); + + await handler(buildImageCompleteEvent()); + + expect(autoSwitchedImages.consume('fresh-image.png')).toBe(false); + }); + + it('retries only the image half of a mixed result, not the video that landed', async () => { + // The video lookup succeeded, so its board invalidation and auto-switch already went out; a + // re-delivery re-running them would invalidate twice and move the selection a second time. + vi.mocked(selectAutoSwitch).mockReturnValue(true); + vi.mocked(getImageDTOSafe).mockResolvedValueOnce(null); + vi.mocked(getVideoDTOSafe).mockResolvedValueOnce({ + video_name: 'fresh-video.mp4', + video_url: 'mock://fresh-video.mp4', + thumbnail_url: 'mock://thumb/fresh-video.mp4', + is_intermediate: false, + is_starred: false, + board_id: 'board-123', + created_at: '2026-01-01', + updated_at: '2026-01-01', + session_id: 'test-session', + node_id: 'test-node', + // eslint-disable-next-line @typescript-eslint/no-explicit-any + } as any); + + const dispatch = vi.fn(() => ({ unwrap: () => Promise.resolve(undefined) })); + const getState = vi.fn(() => ({})); + + const handler = buildOnInvocationComplete( + // eslint-disable-next-line @typescript-eslint/no-explicit-any + getState as any, + // eslint-disable-next-line @typescript-eslint/no-explicit-any + dispatch as any, + new Map() + ); + + const mixed = buildImageCompleteEvent(); + mixed.invocation.type = 'wan_l2v'; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (mixed as any).result = { + image: { image_name: 'lost-image.png' }, + video: { video_name: 'fresh-video.mp4' }, + }; + + await handler(mixed); + expect(getVideoDTOSafe).toHaveBeenCalledTimes(1); + expect(getImageDTOSafe).toHaveBeenCalledTimes(1); + + // The re-delivery refetches the lost image and leaves the video alone. + await handler(mixed); + expect(getImageDTOSafe).toHaveBeenCalledTimes(2); + expect(getVideoDTOSafe).toHaveBeenCalledTimes(1); + }); + + it('redoes only the gallery work on a retry, not the global side effects', async () => { + // The canvas processing flag and $lastProgressEvent are global: by the time a re-delivery + // arrives the user may have started another run, and clearing them again would stop that run's + // spinner and blank its progress. + vi.mocked(getImageDTOSafe).mockResolvedValueOnce(null); + + const dispatched: unknown[] = []; + const dispatch = vi.fn((action: unknown) => { + dispatched.push(action); + return { unwrap: () => Promise.resolve(undefined) }; + }); + const getState = vi.fn(() => ({})); + + const handler = buildOnInvocationComplete( + // eslint-disable-next-line @typescript-eslint/no-explicit-any + getState as any, + // eslint-disable-next-line @typescript-eslint/no-explicit-any + dispatch as any, + new Map() + ); + + const canvasEvent = buildImageCompleteEvent(); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (canvasEvent as any).origin = 'canvas_workflow_integration'; + + await handler(canvasEvent); + const canvasClears = () => dispatched.filter((a) => (a as { type?: string }).type === 'mock/canvasComplete').length; + expect(canvasClears()).toBe(1); + expect($lastProgressEvent.set).toHaveBeenCalledTimes(1); + + // The retry must redo the lost gallery fetch... + await handler(canvasEvent); + expect(getImageDTOSafe).toHaveBeenCalledTimes(2); + // ...and nothing else. + expect(canvasClears()).toBe(1); + expect($lastProgressEvent.set).toHaveBeenCalledTimes(1); + }); + + it('retries the lost output alongside an intermediate one, without redoing the intermediate', async () => { + // An intermediate image never reaches the gallery, so there is nothing to redo for it — but the + // lookup that failed beside it must still be recoverable. + vi.mocked(getImageDTOSafe) + .mockResolvedValueOnce(null) + // eslint-disable-next-line @typescript-eslint/no-explicit-any + .mockResolvedValueOnce({ image_name: 'intermediate.png', is_intermediate: true, board_id: null } as any); + + const dispatch = vi.fn(() => ({ unwrap: () => Promise.resolve(undefined) })); + const getState = vi.fn(() => ({})); + + const handler = buildOnInvocationComplete( + // eslint-disable-next-line @typescript-eslint/no-explicit-any + getState as any, + // eslint-disable-next-line @typescript-eslint/no-explicit-any + dispatch as any, + new Map() + ); + + const twoImages = buildImageCompleteEvent(); + // eslint-disable-next-line @typescript-eslint/no-explicit-any + (twoImages as any).result = { + image_1: { image_name: 'lost.png' }, + image_2: { image_name: 'intermediate.png' }, + }; + + await handler(twoImages); + expect(getImageDTOSafe).toHaveBeenCalledTimes(2); + + // Only the lost output is refetched — the intermediate one never had gallery work to redo. + await handler(twoImages); + expect(getImageDTOSafe).toHaveBeenCalledTimes(3); + expect(vi.mocked(getImageDTOSafe).mock.calls.at(-1)?.[0]).toBe('lost.png'); + }); + + it('serializes several duplicates waiting on one lost delivery into a single retry', async () => { + // All three duplicates are parked on the same in-flight delivery. When it fails they all wake + // up; only one may pick the work back up, or the retried output lands two or three times over. + let resolveFirstLookup: (value: null) => void = () => {}; + vi.mocked(getImageDTOSafe).mockImplementationOnce( + () => + new Promise((resolve) => { + resolveFirstLookup = resolve; + }) + ); + + const dispatch = vi.fn(() => ({ unwrap: () => Promise.resolve(undefined) })); + const getState = vi.fn(() => ({})); + + const handler = buildOnInvocationComplete( + // eslint-disable-next-line @typescript-eslint/no-explicit-any + getState as any, + // eslint-disable-next-line @typescript-eslint/no-explicit-any + dispatch as any, + new Map() + ); + + const event = buildImageCompleteEvent(); + const deliveries = [handler(event), handler(event), handler(event), handler(event)]; + + resolveFirstLookup(null); + await Promise.all(deliveries); + + // One failed lookup plus exactly one retry. + expect(getImageDTOSafe).toHaveBeenCalledTimes(2); + }); + + it('still processes distinct invocations of the same queue item', async () => { + const dispatch = vi.fn(() => ({ unwrap: () => Promise.resolve(undefined) })); + const getState = vi.fn(() => ({})); + + const handler = buildOnInvocationComplete( + // eslint-disable-next-line @typescript-eslint/no-explicit-any + getState as any, + // eslint-disable-next-line @typescript-eslint/no-explicit-any + dispatch as any, + new Map() + ); + + await handler(buildImageCompleteEvent()); + const secondNode = buildImageCompleteEvent(); + secondNode.invocation.id = 'prepared-node-2'; + await handler(secondNode); + expect(getImageDTOSafe).toHaveBeenCalledTimes(2); + }); }); describe('buildOnForeignInvocationComplete', () => { diff --git a/invokeai/frontend/web/src/services/events/onInvocationComplete.tsx b/invokeai/frontend/web/src/services/events/onInvocationComplete.tsx index 5318043749d..d5fe150d645 100644 --- a/invokeai/frontend/web/src/services/events/onInvocationComplete.tsx +++ b/invokeai/frontend/web/src/services/events/onInvocationComplete.tsx @@ -1,6 +1,7 @@ import { logger } from 'app/logging/logger'; import type { AppDispatch, AppGetState } from 'app/store/store'; import { canvasWorkflowIntegrationProcessingCompleted } from 'features/controlLayers/store/canvasWorkflowIntegrationSlice'; +import { autoSwitchedImages } from 'features/gallery/store/autoSwitchedImages'; import { selectAutoSwitch, selectGalleryView, @@ -11,6 +12,7 @@ import { import { boardIdSelected, galleryViewChanged, imageSelected } from 'features/gallery/store/gallerySlice'; import { $nodeExecutionStates, upsertExecutionState } from 'features/nodes/hooks/useNodeExecutionState'; import { isImageField, isImageFieldCollection, isVideoField } from 'features/nodes/types/common'; +import { LRUCache } from 'lru-cache'; import type { ApiTagDescription } from 'services/api'; import { api, LIST_ALL_TAG, LIST_TAG } from 'services/api'; import { boardsApi } from 'services/api/endpoints/boards'; @@ -31,6 +33,17 @@ import type { JsonObject } from 'type-fest'; const log = logger('events'); +/** + * What a completion event still has outstanding. + * + * `done` rejects any re-delivery. `retryable` names the outputs whose DTO lookup failed — the ones + * that never reached the gallery — so a re-delivery can fetch exactly those again. Tracking the + * missing outputs rather than the event as a whole is what makes a partial failure recoverable: an + * event-wide flag would either re-run the outputs that did land (double-counting their board totals + * and optimistic inserts) or abandon the one that did not. + */ +type ProcessedInvocationState = { status: 'done' } | { status: 'retryable'; missingNames: ReadonlySet }; + // These nodes are passthrough nodes. They do not add images/videos to the gallery — their // outputs reference an existing asset — so we must skip the gallery handling for them. // Without 'video' here, a Video Primitive completing mid-run would invalidate the gallery @@ -53,13 +66,40 @@ export const buildOnInvocationComplete = ( dispatch: AppDispatch, completedInvocationKeysByItemId: Map> ) => { - const addImagesToGallery = async (data: S['InvocationCompleteEvent']) => { + // A duplicate delivery of a completion event must not repeat the work below: re-running the + // gallery handling double-counts the optimistic board totals and re-dispatches the auto-switch + // selection. The shared completedInvocationKeysByItemId map cannot detect this — the + // workflow coordinator pre-marks first-delivery events for non-active workflow items before this + // handler runs — so the handler tracks what it has processed itself. The invocation id half of + // the key is the prepared node's per-execution UUID, so keys cannot collide across distinct + // executions even where item ids restart (in-memory DB); the LRU bounds memory. + const processedInvocations = new LRUCache({ max: 1000 }); + // Deliveries currently fetching, so a duplicate can wait for one instead of being discarded. + // Entries live only for the duration of a pass; the LRU above is what bounds long-term memory. + const inFlightDeliveries = new Map>(); + + // `retryNames`, when set, restricts this pass to the outputs a previous delivery lost to a failed + // lookup — everything else already landed and must not be dispatched twice. + const addImagesToGallery = async ( + data: S['InvocationCompleteEvent'], + onLookupFailure: (imageName: string) => void, + retryNames: ReadonlySet | null + ) => { + // A retry exists to land an output the gallery lost, not to re-run the handoff around it. By + // the time a re-delivery arrives the user has had time to select something else, and pulling + // the selection back would be a worse failure than the one being repaired. + const isRetry = retryNames !== null; if (nodeTypeDenylist.includes(data.invocation.type)) { log.trace(`Skipping denylisted node type (${data.invocation.type})`); return; } - const imageDTOs = await getResultImageDTOs(data); + const fetchedImageDTOs = await getResultImageDTOs(data, onLookupFailure, retryNames); + // Intermediates never reach the gallery. Dropping them here rather than bailing out of the + // whole pass on the first one (which is what this used to do) matters now that a delivery's + // outputs are tracked individually: a sibling abandoned that way is in nobody's missing set, + // so no re-delivery could ever recover it. Mirrors addVideosToGallery. + const imageDTOs = fetchedImageDTOs.filter((imageDTO) => !imageDTO.is_intermediate); if (imageDTOs.length === 0) { return; } @@ -70,10 +110,6 @@ export const buildOnInvocationComplete = ( const getImageNamesArg = selectGetImageNamesQueryArgs(getState()); for (const imageDTO of imageDTOs) { - if (imageDTO.is_intermediate) { - return; - } - const board_id = imageDTO.board_id ?? 'none'; // update the total images for the board boardTotalAdditions[board_id] = (boardTotalAdditions[board_id] || 0) + 1; @@ -182,7 +218,7 @@ export const buildOnInvocationComplete = ( const autoSwitch = selectAutoSwitch(getState()); - if (!autoSwitch) { + if (!autoSwitch || isRetry) { return; } @@ -196,6 +232,12 @@ export const buildOnInvocationComplete = ( const { image_name } = lastImageDTO; const board_id = lastImageDTO.board_id ?? 'none'; + // Both branches below auto-switch the selection to this image. Record that so the viewer's + // reveal effect can tell the handoff apart from a user's gallery click — this dispatch happens + // after an async DTO fetch, so it can land after the next generation's first progress event has + // already reset $isProgressImageResolving, and timing alone cannot distinguish the two. + autoSwitchedImages.record(image_name); + // With optimistic updates, we can immediately switch to the new image const selectedBoardId = selectSelectedBoardId(getState()); @@ -222,35 +264,70 @@ export const buildOnInvocationComplete = ( } }; - const getResultImageDTOs = async (data: S['InvocationCompleteEvent']): Promise => { + // getImageDTOSafe swallows fetch errors and returns null, which downstream is indistinguishable + // from "this node produced no image". onLookupFailure separates the two, naming the output that + // was lost so a re-delivery can fetch just that one again. + const getResultImageDTOs = async ( + data: S['InvocationCompleteEvent'], + onLookupFailure: (imageName: string) => void, + retryNames: ReadonlySet | null + ): Promise => { const { result } = data; const imageDTOs: ImageDTO[] = []; + const fetched = new Set(); + const fetch = async (imageName: string) => { + if (retryNames !== null && !retryNames.has(imageName)) { + return; + } + // A result can name the same image twice (an image collection concatenates its inputs + // without deduping). It is still one image: fetching it once keeps its board total counted + // once, and keeps the retry set — which is keyed by name — from re-admitting an occurrence + // that already landed. + if (fetched.has(imageName)) { + return; + } + fetched.add(imageName); + const imageDTO = await getImageDTOSafe(imageName); + if (imageDTO) { + imageDTOs.push(imageDTO); + } else { + onLookupFailure(imageName); + } + }; for (const [_name, value] of objectEntries(result)) { if (isImageField(value)) { - const imageDTO = await getImageDTOSafe(value.image_name); - if (imageDTO) { - imageDTOs.push(imageDTO); - } + await fetch(value.image_name); } else if (isImageFieldCollection(value)) { for (const imageField of value) { - const imageDTO = await getImageDTOSafe(imageField.image_name); - if (imageDTO) { - imageDTOs.push(imageDTO); - } + await fetch(imageField.image_name); } } } return imageDTOs; }; - const getResultVideoDTOs = async (data: S['InvocationCompleteEvent']): Promise => { + const getResultVideoDTOs = async ( + data: S['InvocationCompleteEvent'], + onLookupFailure: (videoName: string) => void, + retryNames: ReadonlySet | null + ): Promise => { const { result } = data; const videoDTOs: VideoDTO[] = []; + const fetched = new Set(); for (const [_name, value] of objectEntries(result)) { if (isVideoField(value)) { + if (retryNames !== null && !retryNames.has(value.video_name)) { + continue; + } + if (fetched.has(value.video_name)) { + continue; + } + fetched.add(value.video_name); const videoDTO = await getVideoDTOSafe(value.video_name); if (videoDTO) { videoDTOs.push(videoDTO); + } else { + onLookupFailure(value.video_name); } } } @@ -266,12 +343,18 @@ export const buildOnInvocationComplete = ( // (DndImage onLoad) to clear them. When auto-switching to a video, the viewer swaps // CurrentImagePreview for CurrentVideoPreview, which unmounts the stale progress overlay // so the stuck "Saving video" spinner goes away on its own. - const addVideosToGallery = async (data: S['InvocationCompleteEvent']) => { + const addVideosToGallery = async ( + data: S['InvocationCompleteEvent'], + onLookupFailure: (videoName: string) => void, + retryNames: ReadonlySet | null + ) => { + // See addImagesToGallery: a retry does not redo the auto-switch. + const isRetry = retryNames !== null; if (nodeTypeDenylist.includes(data.invocation.type)) { return; } - const videoDTOs = await getResultVideoDTOs(data); + const videoDTOs = await getResultVideoDTOs(data, onLookupFailure, retryNames); if (videoDTOs.length === 0) { return; } @@ -296,7 +379,7 @@ export const buildOnInvocationComplete = ( dispatch(galleryApi.util.invalidateTags(getTagsToInvalidateForBoardAffectingMutation(affectedBoards))); const autoSwitch = selectAutoSwitch(getState()); - if (!autoSwitch) { + if (!autoSwitch || isRetry) { return; } @@ -350,35 +433,122 @@ export const buildOnInvocationComplete = ( } }; + /** + * One pass over a completion event. `retryNames` is null for a first delivery and otherwise names + * the outputs an earlier delivery lost, restricting this pass to those. + * + * Returns the outputs whose lookup failed this time. Bookkeeping is done here, before the + * returned promise settles, so a duplicate waiting on it always observes the final state. + */ + const deliver = async ( + data: S['InvocationCompleteEvent'], + invocationKey: string, + retryNames: ReadonlySet | null + ): Promise => { + const isRetry = retryNames !== null; + const missingNames = new Set(); + const onLookupFailure = (itemName: string) => { + missingNames.add(itemName); + }; + + // A retry redoes the lost gallery work and nothing else: the rest of this handler is not + // idempotent against a *later* generation, because the canvas processing flag and + // $lastProgressEvent are global. Re-running them would end the spinner and blank the progress + // of whatever is running by then. + if (!isRetry) { + const nodeExecutionState = $nodeExecutionStates.get()[data.invocation_source_id]; + const updatedNodeExecutionState = getUpdatedNodeExecutionStateOnInvocationComplete( + nodeExecutionState, + data, + completedInvocationKeysByItemId + ); + + if (nodeExecutionState && !updatedNodeExecutionState) { + log.trace( + { data } as JsonObject, + `Ignoring duplicate invocation complete (${data.invocation.type}, ${data.invocation_source_id})` + ); + } + + if (updatedNodeExecutionState) { + upsertExecutionState(updatedNodeExecutionState.nodeId, updatedNodeExecutionState); + } + + // Clear canvas workflow integration processing state if needed + clearCanvasWorkflowIntegrationProcessing(data); + } + + try { + // Add images to gallery (canvas workflow integration results go to staging area automatically) + await addImagesToGallery(data, onLookupFailure, retryNames); + await addVideosToGallery(data, onLookupFailure, retryNames); + + if (!isRetry) { + $lastProgressEvent.set(null); + } + } catch (error) { + // Both call sites discard this handler's promise, so a throw here would surface as an + // unhandled rejection and nothing else. Log it and let the bookkeeping below run: the + // outputs whose lookups failed are still worth recording as retryable. + log.error({ data, error } as JsonObject, `Error handling invocation complete: ${String(error)}`); + } finally { + if (missingNames.size > 0) { + processedInvocations.set(invocationKey, { status: 'retryable', missingNames }); + } + } + }; + return async (data: S['InvocationCompleteEvent']) => { log.debug({ data } as JsonObject, `Invocation complete (${data.invocation.type}, ${data.invocation_source_id})`); - const nodeExecutionState = $nodeExecutionStates.get()[data.invocation_source_id]; - const updatedNodeExecutionState = getUpdatedNodeExecutionStateOnInvocationComplete( - nodeExecutionState, - data, - completedInvocationKeysByItemId - ); - - if (nodeExecutionState && !updatedNodeExecutionState) { + const invocationKey = `${data.item_id}:${data.invocation.id}`; + const logDuplicate = () => { log.trace( { data } as JsonObject, `Ignoring duplicate invocation complete (${data.invocation.type}, ${data.invocation_source_id})` ); - } + }; - if (updatedNodeExecutionState) { - upsertExecutionState(updatedNodeExecutionState.nodeId, updatedNodeExecutionState); + // A duplicate that lands while the first delivery is still fetching waits for it rather than + // being discarded: if that delivery lost outputs to failed lookups, this duplicate is the only + // thing that can recover them — nothing re-emits the event on its own. This check must come + // before the 'done' one, which the in-flight delivery has already written by now. + // + // Several duplicates can be waiting here at once. They resume one at a time, and the first to + // find work marks the event 'done' again before it suspends, so the rest fall through to the + // duplicate branch rather than starting parallel retries. + const inFlight = inFlightDeliveries.get(invocationKey); + if (inFlight) { + await inFlight; + if (processedInvocations.get(invocationKey)?.status !== 'retryable') { + logDuplicate(); + return; + } + } else if (processedInvocations.get(invocationKey)?.status === 'done') { + logDuplicate(); + return; } - // Clear canvas workflow integration processing state if needed - clearCanvasWorkflowIntegrationProcessing(data); - - // Add images to gallery (canvas workflow integration results go to staging area automatically) - await addImagesToGallery(data); - await addVideosToGallery(data); - - $lastProgressEvent.set(null); + const state = processedInvocations.get(invocationKey); + const retryNames = state?.status === 'retryable' ? state.missingNames : null; + // Mark before the awaits below so a duplicate cannot start a second pass over the same outputs + // while this one is in flight; the handshake above is what lets it retry afterwards instead. + processedInvocations.set(invocationKey, { status: 'done' }); + + const delivery = deliver(data, invocationKey, retryNames); + // Waiters only need to know when the pass finished, not whether it threw. + const settled = delivery.then( + () => undefined, + () => undefined + ); + inFlightDeliveries.set(invocationKey, settled); + try { + await delivery; + } finally { + if (inFlightDeliveries.get(invocationKey) === settled) { + inFlightDeliveries.delete(invocationKey); + } + } }; };