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 1978a7fc1ab..ba76ca299bb 100644
--- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentImagePreview.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentImagePreview.tsx
@@ -193,6 +193,12 @@ export const CurrentImagePreview = memo(({ imageDTO }: { imageDTO: ImageDTO | nu
dependencies: [onHotkeyNextImage],
});
+ // The loaded image identifies its session so the viewer can tell a late load from an earlier
+ // session apart from the one whose preview is currently retained (see onLoadImage).
+ const onLoadRenderedImage = useCallback(() => {
+ onLoadImage(imageToRender?.session_id ?? null);
+ }, [imageToRender?.session_id, onLoadImage]);
+
const withProgress = shouldShowProgressInViewer && hasProgressImage && !isTemporarilyShowingSelectedImage;
// When more than one session is generating concurrently (multi-GPU), tile their previews instead of
// showing only the most recent one.
@@ -210,7 +216,7 @@ export const CurrentImagePreview = memo(({ imageDTO }: { imageDTO: ImageDTO | nu
>
{imageToRender && (
-
+
)}
{!imageToRender && }
diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentVideoPreview.tsx b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentVideoPreview.tsx
index 9329a7173c7..19aacb1944b 100644
--- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentVideoPreview.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentVideoPreview.tsx
@@ -302,7 +302,7 @@ export const CurrentVideoPreview = memo(({ videoDTO }: Props) => {
// video frame until playback or a seek — the element just shows its black background.
// Setting currentTime to 0.0001 nudges the decoder to paint without measurably advancing.
const handleLoadedMetadata = useCallback(() => {
- onLoadImage();
+ onLoadImage(videoDTO?.session_id ?? null);
const el = videoRef.current;
if (el && !isPlaying && el.currentTime === 0) {
try {
@@ -311,7 +311,7 @@ export const CurrentVideoPreview = memo(({ videoDTO }: Props) => {
// Some browsers throw if metadata isn't fully ready yet; harmless.
}
}
- }, [isPlaying, onLoadImage]);
+ }, [isPlaying, onLoadImage, videoDTO?.session_id]);
if (!videoDTO) {
return ;
diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/context.tsx b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/context.tsx
index 74fb418761f..19530ffebb9 100644
--- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/context.tsx
+++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/context.tsx
@@ -1,27 +1,25 @@
import { useStore } from '@nanostores/react';
import { logger } from 'app/logging/logger';
import { useAppSelector, useAppStore } from 'app/store/storeHooks';
+import { selectCurrentUser } from 'features/auth/store/authSlice';
+import type {
+ ViewerProgressDataMap,
+ ViewerProgressDatum,
+} from 'features/gallery/components/ImageViewer/viewerProgressLifecycle';
+import { createViewerProgressLifecycle } from 'features/gallery/components/ImageViewer/viewerProgressLifecycle';
import { selectAutoSwitch } from 'features/gallery/store/gallerySelectors';
import type { ProgressImage as ProgressImageType } from 'features/nodes/types/common';
import { LRUCache } from 'lru-cache';
import { type Atom, atom, computed, map, type MapStore, type WritableAtom } from 'nanostores';
import type { PropsWithChildren } from 'react';
-import { createContext, memo, useCallback, useContext, useEffect, useMemo, useRef, useState } from 'react';
+import { createContext, memo, useCallback, useContext, useEffect, useMemo, useState } from 'react';
import type { S } from 'services/api/types';
import { getEventScope } from 'services/events/eventScope';
import { $socket } from 'services/events/stores';
import { assert } from 'tsafe';
import type { JsonObject } from 'type-fest';
-/** Live progress for a single in-flight session (queue item). Used to tile the viewer when several
- * sessions run concurrently (multi-GPU). Only items that have produced a preview image are tracked. */
-export type ViewerProgressDatum = {
- itemId: number;
- progressEvent: S['InvocationProgressEvent'];
- progressImage: ProgressImageType;
-};
-
-type ViewerProgressDataMap = Record;
+export type { ViewerProgressDatum } from 'features/gallery/components/ImageViewer/viewerProgressLifecycle';
type ImageViewerContextValue = {
$progressEvent: Atom;
@@ -33,7 +31,11 @@ type ImageViewerContextValue = {
$activeProgressData: Atom;
$isProgressImageResolving: Atom;
$isTemporarilyShowingSelectedImage: WritableAtom;
- onLoadImage: () => void;
+ /**
+ * The viewer finished loading the final image/video for the given session (its DTO's
+ * `session_id`, or null when it has none). Ends the completed session's "resolve" illusion.
+ */
+ onLoadImage: (sessionId: string | null) => void;
};
const ImageViewerContext = createContext(null);
@@ -58,10 +60,24 @@ export const ImageViewerContextProvider = memo((props: PropsWithChildren) => {
)[0];
const $isProgressImageResolving = useState(() => atom(false))[0];
const $isTemporarilyShowingSelectedImage = useState(() => atom(false))[0];
- const shouldClearProgressImageOnLoadRef = useRef(false);
// We can have race conditions where we receive a progress event for a queue item that has already finished. Easiest
// way to handle this is to keep track of finished queue items in a cache and ignore progress events for those.
const [finishedQueueItemIds] = useState(() => new LRUCache({ max: 200 }));
+ // Session id -> queue item id, learned from progress events. Outlives the item's terminal event
+ // so a late final-image load can be attributed to the session that produced it.
+ const [itemIdBySessionId] = useState(() => new LRUCache({ max: 200 }));
+ // All store mutations live in the lifecycle (extracted for unit testing); the effects below own
+ // the socket subscriptions and the ownership/scope checks on incoming events.
+ const lifecycle = useState(() =>
+ createViewerProgressLifecycle({
+ $progressEvent,
+ $progressImage,
+ $progressData,
+ $isProgressImageResolving,
+ finishedQueueItemIds,
+ itemIdBySessionId,
+ })
+ )[0];
useEffect(() => {
if (!socket) {
@@ -74,24 +90,11 @@ export const ImageViewerContextProvider = memo((props: PropsWithChildren) => {
if (getEventScope(store.getState, data) !== 'own') {
return;
}
- if (finishedQueueItemIds.has(data.item_id)) {
+ if (!lifecycle.recordProgress(data)) {
log.trace(
{ data } as JsonObject,
`Received InvocationProgressEvent event for already-finished queue item ${data.item_id}`
);
- return;
- }
- shouldClearProgressImageOnLoadRef.current = false;
- $isProgressImageResolving.set(false);
- $progressEvent.set(data);
- if (data.image) {
- $progressImage.set(data.image);
- // Track per-session so the viewer can tile concurrent sessions (multi-GPU).
- $progressData.setKey(data.item_id, {
- itemId: data.item_id,
- progressEvent: data,
- progressImage: data.image,
- });
}
};
@@ -100,7 +103,7 @@ export const ImageViewerContextProvider = memo((props: PropsWithChildren) => {
return () => {
socket.off('invocation_progress', onInvocationProgress);
};
- }, [$isProgressImageResolving, $progressData, $progressEvent, $progressImage, finishedQueueItemIds, socket, store]);
+ }, [lifecycle, socket, store]);
useEffect(() => {
if (!socket) {
@@ -116,51 +119,22 @@ export const ImageViewerContextProvider = memo((props: PropsWithChildren) => {
if (getEventScope(store.getState, data) !== 'own') {
return;
}
- if (finishedQueueItemIds.has(data.item_id)) {
+ if (data.status === 'in_progress') {
+ // Track the claim itself, not just the progress events that follow it: a queue clear
+ // cancels the items already running before it deletes the rows, so a worker that claims an
+ // item in between never gets a terminal event and its first progress event lands after the
+ // clear (see the lifecycle's onItemStarted).
+ lifecycle.onItemStarted(data.item_id);
+ return;
+ }
+ if (data.status !== 'completed' && data.status !== 'canceled' && data.status !== 'failed') {
+ return;
+ }
+ if (!lifecycle.onTerminal(data, autoSwitch)) {
log.trace(
{ data } as JsonObject,
`Received QueueItemStatusChangedEvent event for already-finished queue item ${data.item_id}`
);
- return;
- }
- if (data.status === 'completed' || data.status === 'canceled' || data.status === 'failed') {
- finishedQueueItemIds.set(data.item_id, true);
- // Remove this session's tile from the multi-session preview as soon as it reaches a terminal
- // state. The single-image "resolve" illusion below is handled separately via onLoadImage.
- $progressData.setKey(data.item_id, undefined);
- // The shared $progressEvent/$progressImage globals may currently hold a DIFFERENT session's
- // latest preview (multi-GPU). Only the item that owns them may clear them — otherwise
- // canceling item A would blank item B's still-running preview until B's next image event.
- const globalProgressEvent = $progressEvent.get();
- if (globalProgressEvent !== null && globalProgressEvent.item_id !== data.item_id) {
- return;
- }
- // Completed queue items have the progress event cleared by the onLoadImage callback. This allows the viewer to
- // create the illusion of the progress image "resolving" into the final image. If we cleared the progress image
- // now, there would be a flicker where the progress image disappears before the final image appears, and the
- // last-selected gallery image should be shown for a brief moment.
- //
- // When gallery auto-switch is disabled, we do not need to create this illusion, because we are not going to
- // switch to the final image automatically. In this case, we clear the progress image immediately.
- //
- // We also clear the progress image if the queue item is canceled or failed, as there is no final image to show.
- if (
- data.status === 'canceled' ||
- data.status === 'failed' ||
- !autoSwitch ||
- // When the origin is 'canvas' and destination is 'canvas' (without a ':' suffix), that means the
- // image is going to be added to the staging area. In this case, we need to clear the progress image else it
- // will be stuck on the viewer.
- (data.origin === 'canvas' && data.destination !== 'canvas')
- ) {
- shouldClearProgressImageOnLoadRef.current = false;
- $isProgressImageResolving.set(false);
- $progressEvent.set(null);
- $progressImage.set(null);
- } else {
- shouldClearProgressImageOnLoadRef.current = true;
- $isProgressImageResolving.set(true);
- }
}
};
@@ -169,27 +143,55 @@ export const ImageViewerContextProvider = memo((props: PropsWithChildren) => {
return () => {
socket.off('queue_item_status_changed', onQueueItemStatusChanged);
};
- }, [
- $isProgressImageResolving,
- $progressData,
- $progressEvent,
- $progressImage,
- autoSwitch,
- finishedQueueItemIds,
- socket,
- store,
- ]);
-
- const onLoadImage = useCallback(() => {
- if (!shouldClearProgressImageOnLoadRef.current) {
+ }, [autoSwitch, lifecycle, socket, store]);
+
+ useEffect(() => {
+ if (!socket) {
+ return;
+ }
+
+ const onQueueCleared = (data: S['QueueClearedEvent']) => {
+ // Scope is decided inside the lifecycle: it needs the current user id to tell whether the
+ // clear could have deleted this client's items (see onQueueCleared's docstring).
+ const currentUserId = selectCurrentUser(store.getState())?.user_id ?? null;
+ lifecycle.onQueueCleared(data, currentUserId);
+ };
+
+ socket.on('queue_cleared', onQueueCleared);
+
+ return () => {
+ socket.off('queue_cleared', onQueueCleared);
+ };
+ }, [lifecycle, socket, store]);
+
+ useEffect(() => {
+ if (!socket) {
return;
}
- shouldClearProgressImageOnLoadRef.current = false;
- $isProgressImageResolving.set(false);
- $progressEvent.set(null);
- $progressImage.set(null);
- }, [$isProgressImageResolving, $progressEvent, $progressImage]);
+ const onDisconnect = () => {
+ // Mirrors the app-wide progress stores (see setEventListeners): a disconnected socket may
+ // miss terminal events, so previews from before the gap cannot be trusted to ever resolve.
+ lifecycle.reset();
+ };
+
+ socket.on('disconnect', onDisconnect);
+
+ return () => {
+ socket.off('disconnect', onDisconnect);
+ // The socket is being replaced (e.g. an auth-token change swapped $socket for a different
+ // user's connection) or the viewer is unmounting: any tracked sessions belong to the old
+ // connection and will never emit another terminal event here.
+ lifecycle.reset();
+ };
+ }, [lifecycle, socket]);
+
+ const onLoadImage = useCallback(
+ (sessionId: string | null) => {
+ lifecycle.onFinalImageLoaded(sessionId);
+ },
+ [lifecycle]
+ );
const value = useMemo(
() => ({
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
new file mode 100644
index 00000000000..a3181382a7f
--- /dev/null
+++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.test.ts
@@ -0,0 +1,357 @@
+import type { ProgressImage as ProgressImageType } from 'features/nodes/types/common';
+import { atom, map } from 'nanostores';
+import type { S } from 'services/api/types';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import type { ViewerProgressDataMap, ViewerProgressStores } from './viewerProgressLifecycle';
+import { createViewerProgressLifecycle, RESOLVE_TIMEOUT_MS } from './viewerProgressLifecycle';
+
+const buildProgressImage = (itemId: number): ProgressImageType =>
+ ({
+ dataURL: `data:image/png;base64,item-${itemId}`,
+ width: 512,
+ height: 512,
+ }) as ProgressImageType;
+
+const buildProgressEvent = (overrides: Partial = {}): S['InvocationProgressEvent'] =>
+ ({
+ queue_id: 'default',
+ item_id: 1,
+ batch_id: 'batch-1',
+ origin: null,
+ destination: null,
+ user_id: 'user-1',
+ session_id: 'session-1',
+ invocation_source_id: 'node-1',
+ invocation: { id: 'node-1', type: 'test_node' },
+ message: 'denoising',
+ percentage: 0.5,
+ image: null,
+ ...overrides,
+ }) as S['InvocationProgressEvent'];
+
+const buildTerminalEvent = (
+ overrides: Partial = {}
+): S['QueueItemStatusChangedEvent'] =>
+ ({
+ queue_id: 'default',
+ item_id: 1,
+ batch_id: 'batch-1',
+ origin: null,
+ destination: null,
+ user_id: 'user-1',
+ status: 'completed',
+ ...overrides,
+ }) as S['QueueItemStatusChangedEvent'];
+
+const buildQueueClearedEvent = (userId: string | null): S['QueueClearedEvent'] =>
+ ({ queue_id: 'default', user_id: userId }) as S['QueueClearedEvent'];
+
+describe('viewerProgressLifecycle', () => {
+ let stores: ViewerProgressStores;
+ let lifecycle: ReturnType;
+
+ beforeEach(() => {
+ vi.useFakeTimers();
+ stores = {
+ $progressEvent: atom(null),
+ $progressImage: atom(null),
+ $progressData: map({}),
+ $isProgressImageResolving: atom(false),
+ finishedQueueItemIds: new Map(),
+ itemIdBySessionId: new Map(),
+ };
+ lifecycle = createViewerProgressLifecycle(stores);
+ });
+
+ afterEach(() => {
+ vi.useRealTimers();
+ });
+
+ const startTwoSessions = () => {
+ // B posts a preview first, then A — A owns the shared single-image preview.
+ const eventB = buildProgressEvent({ item_id: 2, session_id: 'session-2', image: buildProgressImage(2) });
+ const eventA = buildProgressEvent({ item_id: 1, session_id: 'session-1', image: buildProgressImage(1) });
+ lifecycle.recordProgress(eventB);
+ lifecycle.recordProgress(eventA);
+ return { eventA, eventB };
+ };
+
+ describe('recordProgress', () => {
+ it('tracks per-session data and hands the shared preview to the latest reporter', () => {
+ const { eventA } = startTwoSessions();
+ expect(stores.$progressEvent.get()).toBe(eventA);
+ expect(stores.$progressImage.get()).toBe(eventA.image);
+ expect(stores.$progressData.get()[1]?.itemId).toBe(1);
+ expect(stores.$progressData.get()[2]?.itemId).toBe(2);
+ });
+
+ it('ignores events for finished items so trailing progress cannot repopulate the preview', () => {
+ stores.finishedQueueItemIds.set(1, true);
+ const handled = lifecycle.recordProgress(buildProgressEvent({ item_id: 1, image: buildProgressImage(1) }));
+ expect(handled).toBe(false);
+ expect(stores.$progressEvent.get()).toBeNull();
+ expect(stores.$progressData.get()[1]).toBeUndefined();
+ });
+ });
+
+ describe('onTerminal', () => {
+ it.each(['completed', 'canceled', 'failed'] as const)(
+ 'hands the shared preview to the remaining session when its owner reaches %s',
+ (status) => {
+ const { eventB } = startTwoSessions();
+ lifecycle.onTerminal(buildTerminalEvent({ item_id: 1, status }), true);
+ // B immediately becomes the single visible progress image and indicator.
+ expect(stores.$progressEvent.get()).toBe(eventB);
+ expect(stores.$progressImage.get()).toBe(eventB.image);
+ expect(stores.$progressData.get()[1]).toBeUndefined();
+ expect(stores.$progressData.get()[2]?.itemId).toBe(2);
+ // No resolve illusion may be pending — it would swap B's live preview for A's final image.
+ expect(stores.$isProgressImageResolving.get()).toBe(false);
+ lifecycle.onFinalImageLoaded('session-1');
+ expect(stores.$progressEvent.get()).toBe(eventB);
+ expect(stores.$progressImage.get()).toBe(eventB.image);
+ }
+ );
+
+ it('hands the shared preview to the remaining session even when auto-switch is off', () => {
+ const { eventB } = startTwoSessions();
+ lifecycle.onTerminal(buildTerminalEvent({ item_id: 1, status: 'canceled' }), false);
+ expect(stores.$progressEvent.get()).toBe(eventB);
+ expect(stores.$progressImage.get()).toBe(eventB.image);
+ });
+
+ it('promotes the most recently updated remaining session when several remain', () => {
+ startTwoSessions();
+ const eventC = buildProgressEvent({ item_id: 3, session_id: 'session-3', image: buildProgressImage(3) });
+ lifecycle.recordProgress(eventC);
+ lifecycle.onTerminal(buildTerminalEvent({ item_id: 3, status: 'canceled' }), true);
+ // C owned the preview; A reported more recently than B, so A takes over.
+ expect(stores.$progressEvent.get()?.item_id).toBe(1);
+ });
+
+ it('leaves the shared preview alone when a non-owner terminates', () => {
+ const { eventA } = startTwoSessions();
+ lifecycle.onTerminal(buildTerminalEvent({ item_id: 2, status: 'canceled' }), true);
+ expect(stores.$progressEvent.get()).toBe(eventA);
+ expect(stores.$progressImage.get()).toBe(eventA.image);
+ expect(stores.$progressData.get()[2]).toBeUndefined();
+ });
+
+ it('clears immediately when the last session is canceled', () => {
+ const eventA = buildProgressEvent({ item_id: 1, image: buildProgressImage(1) });
+ lifecycle.recordProgress(eventA);
+ lifecycle.onTerminal(buildTerminalEvent({ item_id: 1, status: 'canceled' }), true);
+ expect(stores.$progressEvent.get()).toBeNull();
+ expect(stores.$progressImage.get()).toBeNull();
+ expect(stores.$isProgressImageResolving.get()).toBe(false);
+ });
+
+ it('runs the resolve illusion when the last session completes with auto-switch on', () => {
+ const eventA = buildProgressEvent({ item_id: 1, image: buildProgressImage(1) });
+ lifecycle.recordProgress(eventA);
+ lifecycle.onTerminal(buildTerminalEvent({ item_id: 1, status: 'completed' }), true);
+ // The preview is retained until the final image loads, "resolving" into it.
+ expect(stores.$progressEvent.get()).toBe(eventA);
+ expect(stores.$isProgressImageResolving.get()).toBe(true);
+ lifecycle.onFinalImageLoaded('session-1');
+ expect(stores.$progressEvent.get()).toBeNull();
+ expect(stores.$progressImage.get()).toBeNull();
+ expect(stores.$isProgressImageResolving.get()).toBe(false);
+ });
+
+ it('runs no resolve illusion when the completed item never reported progress', () => {
+ // Nothing is retained, so arming the illusion would only leave the resolving flag stuck on.
+ lifecycle.onTerminal(buildTerminalEvent({ item_id: 1, status: 'completed' }), true);
+ expect(stores.$isProgressImageResolving.get()).toBe(false);
+ expect(stores.$progressEvent.get()).toBeNull();
+ });
+
+ it('ignores repeat terminal events for an already-finished item', () => {
+ const { eventA } = startTwoSessions();
+ lifecycle.onTerminal(buildTerminalEvent({ item_id: 2, status: 'canceled' }), true);
+ expect(lifecycle.onTerminal(buildTerminalEvent({ item_id: 2, status: 'canceled' }), true)).toBe(false);
+ expect(stores.$progressEvent.get()).toBe(eventA);
+ });
+ });
+
+ describe('onFinalImageLoaded', () => {
+ it('ignores a late load from a session that finished before the current preview owner', () => {
+ const { eventB } = startTwoSessions();
+ // A completes first and hands the shared preview to the still-running B...
+ lifecycle.onTerminal(buildTerminalEvent({ item_id: 1, status: 'completed' }), true);
+ // ...then B completes and starts its own resolve illusion, retaining B's last frame.
+ lifecycle.onTerminal(buildTerminalEvent({ item_id: 2, status: 'completed' }), true);
+ expect(stores.$isProgressImageResolving.get()).toBe(true);
+ // A's final image only now finishes loading. Clearing here would cut B's illusion short.
+ lifecycle.onFinalImageLoaded('session-1');
+ expect(stores.$progressEvent.get()).toBe(eventB);
+ expect(stores.$progressImage.get()).toBe(eventB.image);
+ expect(stores.$isProgressImageResolving.get()).toBe(true);
+ // B's own final image ends the illusion.
+ lifecycle.onFinalImageLoaded('session-2');
+ expect(stores.$progressEvent.get()).toBeNull();
+ expect(stores.$progressImage.get()).toBeNull();
+ expect(stores.$isProgressImageResolving.get()).toBe(false);
+ });
+
+ it('clears the retained preview on a timeout when its final image never loads', () => {
+ // The retained session's image may never load: the load can fail (the viewer's
+ // reports errors through onError, not onLoad), or auto-switch may end up selecting a
+ // concurrently-completed session's image, so the retained session's image is never
+ // rendered. An ignored load may have been the last one coming, so the preview — an opaque
+ // overlay — must not be left covering the viewer until the next generation.
+ const { eventB } = startTwoSessions();
+ lifecycle.onTerminal(buildTerminalEvent({ item_id: 1, status: 'completed' }), true);
+ lifecycle.onTerminal(buildTerminalEvent({ item_id: 2, status: 'completed' }), true);
+ lifecycle.onFinalImageLoaded('session-1');
+ expect(stores.$progressEvent.get()).toBe(eventB);
+ vi.advanceTimersByTime(RESOLVE_TIMEOUT_MS);
+ expect(stores.$progressEvent.get()).toBeNull();
+ expect(stores.$progressImage.get()).toBeNull();
+ expect(stores.$isProgressImageResolving.get()).toBe(false);
+ });
+
+ it.each([
+ [
+ 'a still-running session takes the preview over',
+ () => {
+ lifecycle.recordProgress(
+ buildProgressEvent({ item_id: 2, session_id: 'session-2', image: buildProgressImage(2) })
+ );
+ },
+ ],
+ [
+ 'the final image loads first',
+ () => {
+ lifecycle.onFinalImageLoaded('session-1');
+ },
+ ],
+ [
+ 'the preview state is reset',
+ () => {
+ lifecycle.reset();
+ },
+ ],
+ ])('disarms the resolve timeout when %s', (_desc, takeOver) => {
+ const eventA = buildProgressEvent({ item_id: 1, image: buildProgressImage(1) });
+ lifecycle.recordProgress(eventA);
+ lifecycle.onTerminal(buildTerminalEvent({ item_id: 1, status: 'completed' }), true);
+ expect(vi.getTimerCount()).toBe(1);
+ takeOver();
+ // The timeout must never fire against a preview that has since been replaced or dropped.
+ // Asserting the timer is gone, not just that the state survives it, is what makes this bite
+ // for the takeovers that leave the stores null — there, a leaked timer would clear state
+ // that is already clear.
+ expect(vi.getTimerCount()).toBe(0);
+ const eventAfterTakeOver = stores.$progressEvent.get();
+ vi.advanceTimersByTime(RESOLVE_TIMEOUT_MS * 2);
+ expect(stores.$progressEvent.get()).toBe(eventAfterTakeOver);
+ });
+
+ it.each([
+ ['an image with no session (e.g. an upload)', null],
+ ['an image from a session this viewer never tracked', 'session-from-a-previous-visit'],
+ ])('still clears the retained preview on %s', (_desc, sessionId) => {
+ const eventA = buildProgressEvent({ item_id: 1, image: buildProgressImage(1) });
+ lifecycle.recordProgress(eventA);
+ lifecycle.onTerminal(buildTerminalEvent({ item_id: 1, status: 'completed' }), true);
+ expect(stores.$isProgressImageResolving.get()).toBe(true);
+ // Unattributable loads keep the safety net: a retained preview must not cover the viewer
+ // indefinitely just because the completed item's own image never loads.
+ lifecycle.onFinalImageLoaded(sessionId);
+ expect(stores.$progressEvent.get()).toBeNull();
+ expect(stores.$progressImage.get()).toBeNull();
+ expect(stores.$isProgressImageResolving.get()).toBe(false);
+ });
+
+ it('clears on a load from a session tracked before a reset', () => {
+ // Attributions are dropped along with the state they describe, so images generated before a
+ // disconnect or socket swap can still end a later session's illusion.
+ startTwoSessions();
+ lifecycle.reset();
+ const eventC = buildProgressEvent({ item_id: 3, session_id: 'session-3', image: buildProgressImage(3) });
+ lifecycle.recordProgress(eventC);
+ lifecycle.onTerminal(buildTerminalEvent({ item_id: 3, status: 'completed' }), true);
+ expect(stores.$isProgressImageResolving.get()).toBe(true);
+ lifecycle.onFinalImageLoaded('session-1');
+ expect(stores.$progressEvent.get()).toBeNull();
+ expect(stores.$isProgressImageResolving.get()).toBe(false);
+ });
+ });
+
+ describe('onQueueCleared', () => {
+ it.each([
+ ['an unscoped (admin or single-user) clear', null, 'user-1'],
+ ["the current user's scoped clear", 'user-1', 'user-1'],
+ ])('drops all previews and blocks trailing progress on %s', (_desc, clearedUserId, currentUserId) => {
+ startTwoSessions();
+ const applied = lifecycle.onQueueCleared(buildQueueClearedEvent(clearedUserId), currentUserId);
+ expect(applied).toBe(true);
+ expect(stores.$progressEvent.get()).toBeNull();
+ expect(stores.$progressImage.get()).toBeNull();
+ expect(stores.$progressData.get()).toEqual({});
+ // A worker claimed between the clear's cancellation pass and deletion is stopped only by
+ // this event — its trailing progress must not repopulate the preview.
+ expect(lifecycle.recordProgress(buildProgressEvent({ item_id: 1, image: buildProgressImage(1) }))).toBe(false);
+ expect(lifecycle.recordProgress(buildProgressEvent({ item_id: 2, image: buildProgressImage(2) }))).toBe(false);
+ });
+
+ it('blocks trailing progress from an item claimed after the clear cancelled the running ones', () => {
+ // The clear cancels the items already running before deleting the rows, so a worker that
+ // claims an item in between never gets a terminal event: only its in_progress claim, which
+ // precedes the deletion, tells the viewer the item exists. Its first progress event arrives
+ // seconds later — a preview for a deleted item that nothing would ever take down.
+ expect(lifecycle.onItemStarted(7)).toBe(true);
+ expect(lifecycle.onQueueCleared(buildQueueClearedEvent(null), 'user-1')).toBe(true);
+ expect(
+ lifecycle.recordProgress(
+ buildProgressEvent({ item_id: 7, session_id: 'session-7', image: buildProgressImage(7) })
+ )
+ ).toBe(false);
+ expect(stores.$progressEvent.get()).toBeNull();
+ expect(stores.$progressImage.get()).toBeNull();
+ });
+
+ it('blocks trailing progress from a session that had not produced a preview image yet', () => {
+ // Item 1 has only reported progress without an image, so it is absent from $progressData and
+ // does not own the shared globals once item 2 reports an image.
+ lifecycle.recordProgress(buildProgressEvent({ item_id: 1, session_id: 'session-1', image: null }));
+ lifecycle.recordProgress(
+ buildProgressEvent({ item_id: 2, session_id: 'session-2', image: buildProgressImage(2) })
+ );
+ expect(lifecycle.onQueueCleared(buildQueueClearedEvent(null), 'user-1')).toBe(true);
+ // Its first image event must not resurrect the preview the clear just dropped.
+ expect(lifecycle.recordProgress(buildProgressEvent({ item_id: 1, image: buildProgressImage(1) }))).toBe(false);
+ expect(stores.$progressEvent.get()).toBeNull();
+ expect(stores.$progressImage.get()).toBeNull();
+ expect(stores.$progressData.get()).toEqual({});
+ });
+
+ it.each([
+ ["another user's scoped clear", 'user-2'],
+ ['the sanitized broadcast of a foreign scoped clear', 'redacted'],
+ ])('leaves previews alone on %s', (_desc, clearedUserId) => {
+ const { eventA } = startTwoSessions();
+ const applied = lifecycle.onQueueCleared(buildQueueClearedEvent(clearedUserId), 'user-1');
+ expect(applied).toBe(false);
+ expect(stores.$progressEvent.get()).toBe(eventA);
+ expect(stores.$progressData.get()[1]?.itemId).toBe(1);
+ expect(stores.$progressData.get()[2]?.itemId).toBe(2);
+ });
+ });
+
+ describe('reset', () => {
+ it('drops all preview state without marking items finished', () => {
+ startTwoSessions();
+ stores.$isProgressImageResolving.set(true);
+ lifecycle.reset();
+ expect(stores.$progressEvent.get()).toBeNull();
+ expect(stores.$progressImage.get()).toBeNull();
+ expect(stores.$progressData.get()).toEqual({});
+ expect(stores.$isProgressImageResolving.get()).toBe(false);
+ // A new connection's events for a re-used id are not blocked — reset is not a cancel.
+ expect(lifecycle.recordProgress(buildProgressEvent({ item_id: 1, image: buildProgressImage(1) }))).toBe(true);
+ });
+ });
+});
diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.ts b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.ts
new file mode 100644
index 00000000000..e6c78cc7499
--- /dev/null
+++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.ts
@@ -0,0 +1,319 @@
+import type { ProgressImage as ProgressImageType } from 'features/nodes/types/common';
+import type { MapStore, WritableAtom } from 'nanostores';
+import type { S } from 'services/api/types';
+
+/** Live progress for a single in-flight session (queue item). Used to tile the viewer when several
+ * sessions run concurrently (multi-GPU). Only items that have produced a preview image are tracked.
+ * `seq` orders data by most recent update, so the shared single-image preview can be handed to the
+ * freshest remaining session when its current owner terminates. */
+export type ViewerProgressDatum = {
+ itemId: number;
+ seq: number;
+ progressEvent: S['InvocationProgressEvent'];
+ progressImage: ProgressImageType;
+};
+
+export type ViewerProgressDataMap = Record;
+
+/** The subset of LRUCache the lifecycle needs — kept minimal so tests can pass a plain Map. */
+type FinishedQueueItemIds = {
+ has: (itemId: number) => boolean;
+ set: (itemId: number, value: boolean) => unknown;
+};
+
+/** The subset of LRUCache the lifecycle needs — kept minimal so tests can pass a plain Map. */
+type ItemIdBySessionId = {
+ get: (sessionId: string) => number | undefined;
+ set: (sessionId: string, itemId: number) => unknown;
+ clear: () => unknown;
+};
+
+/**
+ * How long a completed session's retained preview may wait for its final image to load before it
+ * is cleared anyway. The "resolve" illusion normally ends on that image's load event, but nothing
+ * guarantees the event arrives: the load can fail (the viewer's reports errors through
+ * onError, not onLoad), or auto-switch can end up selecting a concurrently-completed session's
+ * image instead, in which case the retained session's image is never rendered at all. Without this
+ * bound the preview — an opaque overlay — could cover the viewer until the next generation.
+ */
+export const RESOLVE_TIMEOUT_MS = 3000;
+
+export type ViewerProgressStores = {
+ $progressEvent: WritableAtom;
+ $progressImage: WritableAtom;
+ /** Per-session progress, keyed by queue item id. Drives the tiled multi-session preview. */
+ $progressData: MapStore;
+ $isProgressImageResolving: WritableAtom;
+ /** Finished queue items, tracked so trailing progress events cannot repopulate the preview. */
+ finishedQueueItemIds: FinishedQueueItemIds;
+ /** Queue item id of each session we have seen progress for, keyed by session id. Outlives the
+ * item's terminal event so a late final-image load can be attributed to the session that
+ * produced it (see onFinalImageLoaded). */
+ itemIdBySessionId: ItemIdBySessionId;
+};
+
+const pickLatestDatum = (data: ViewerProgressDataMap): ViewerProgressDatum | null => {
+ let latest: ViewerProgressDatum | null = null;
+ for (const datum of Object.values(data)) {
+ if (datum !== undefined && (latest === null || datum.seq > latest.seq)) {
+ latest = datum;
+ }
+ }
+ return latest;
+};
+
+/**
+ * The store-side lifecycle of the image viewer's live-preview state, factored out of the React
+ * provider so it can be unit tested. The provider owns the socket subscriptions and the
+ * ownership/scope checks on incoming events; every store mutation happens here.
+ *
+ * The state it manages:
+ * - `$progressData`: one entry per session with a preview image (the tiled multi-session view).
+ * - `$progressEvent` / `$progressImage`: the shared single-image preview, owned by the session
+ * that most recently reported progress.
+ */
+export const createViewerProgressLifecycle = (stores: ViewerProgressStores) => {
+ const {
+ $progressEvent,
+ $progressImage,
+ $progressData,
+ $isProgressImageResolving,
+ finishedQueueItemIds,
+ itemIdBySessionId,
+ } = stores;
+ let seq = 0;
+ // The queue item whose retained preview the final gallery image's onLoad should clear — the tail
+ // end of the "resolve" illusion for a completed session (see onTerminal / onFinalImageLoaded).
+ // Null when no illusion is pending. Always written through setPendingResolve, which keeps the
+ // safety timeout in sync.
+ let pendingResolveItemId: number | null = null;
+ // Every item we have seen progress for and not yet seen terminate, including items that have not
+ // produced a preview image (those are absent from $progressData). A queue clear deletes items
+ // without emitting a per-item terminal event, so this is the set that must be marked finished
+ // there — otherwise an image-less session could later emit an image and resurrect the preview.
+ const unfinishedItemIds = new Set();
+ let resolveTimeoutId: ReturnType | null = null;
+
+ const clearRetainedPreview = (): void => {
+ $isProgressImageResolving.set(false);
+ $progressEvent.set(null);
+ $progressImage.set(null);
+ };
+
+ /**
+ * Arm (or, with null, disarm) the pending "resolve" illusion. Every write to
+ * `pendingResolveItemId` goes through here, so the safety timeout exists exactly while the
+ * illusion is pending: anything that ends the illusion — a load, a takeover by another session,
+ * a reset — also cancels the timeout, and it can never clear a preview that some other session
+ * has since taken over.
+ *
+ * The timeout bounds the illusion only. A preview left standing by a path that does not arm one
+ * (a progress event that carries no image replaces $progressEvent but not $progressImage, so the
+ * previous session's frame stays up while the next queue item spins up) is unbounded here, as it
+ * is today — that frame is taken down by the next preview image rather than by this timeout.
+ */
+ const setPendingResolve = (itemId: number | null): void => {
+ if (resolveTimeoutId !== null) {
+ clearTimeout(resolveTimeoutId);
+ resolveTimeoutId = null;
+ }
+ pendingResolveItemId = itemId;
+ if (itemId === null) {
+ return;
+ }
+ resolveTimeoutId = setTimeout(() => {
+ resolveTimeoutId = null;
+ pendingResolveItemId = null;
+ clearRetainedPreview();
+ }, RESOLVE_TIMEOUT_MS);
+ };
+
+ const clearAll = (): void => {
+ setPendingResolve(null);
+ unfinishedItemIds.clear();
+ // Session attributions describe state this reset just dropped. Keeping them would make later
+ // loads of those images look like another session's, suppressing clears they should perform.
+ itemIdBySessionId.clear();
+ clearRetainedPreview();
+ $progressData.set({});
+ };
+
+ /**
+ * A worker claimed this queue item (`in_progress`). Tracked so a queue clear can mark it
+ * finished: the clear cancels the items that were already running before it deletes the rows,
+ * but a worker that claims an item in between gets no terminal event at all — its row is gone —
+ * and its first progress event would otherwise appear seconds after the clear and put a preview
+ * for a deleted item on screen, with nothing left to ever take it down.
+ *
+ * Returns false if the item already finished (event ignored).
+ */
+ const onItemStarted = (itemId: number): boolean => {
+ if (finishedQueueItemIds.has(itemId)) {
+ return false;
+ }
+ unfinishedItemIds.add(itemId);
+ return true;
+ };
+
+ /** Record a progress event. Returns false if the item already finished (event ignored). */
+ const recordProgress = (data: S['InvocationProgressEvent']): boolean => {
+ if (finishedQueueItemIds.has(data.item_id)) {
+ return false;
+ }
+ unfinishedItemIds.add(data.item_id);
+ itemIdBySessionId.set(data.session_id, data.item_id);
+ setPendingResolve(null);
+ $isProgressImageResolving.set(false);
+ $progressEvent.set(data);
+ if (data.image) {
+ $progressImage.set(data.image);
+ // Track per-session so the viewer can tile concurrent sessions (multi-GPU).
+ $progressData.setKey(data.item_id, {
+ itemId: data.item_id,
+ seq: ++seq,
+ progressEvent: data,
+ progressImage: data.image,
+ });
+ }
+ return true;
+ };
+
+ /** Handle a terminal status for a queue item. Returns false if it already finished (ignored). */
+ const onTerminal = (data: S['QueueItemStatusChangedEvent'], autoSwitch: boolean): boolean => {
+ if (finishedQueueItemIds.has(data.item_id)) {
+ return false;
+ }
+ finishedQueueItemIds.set(data.item_id, true);
+ unfinishedItemIds.delete(data.item_id);
+ // Remove this session's tile from the multi-session preview as soon as it reaches a terminal
+ // state. The single-image "resolve" illusion below is handled separately via onLoadImage.
+ $progressData.setKey(data.item_id, undefined);
+ // The shared $progressEvent/$progressImage globals may currently hold a DIFFERENT session's
+ // latest preview (multi-GPU). Only the item that owns them may replace or clear them —
+ // otherwise canceling item A would blank item B's still-running preview until B's next image
+ // event.
+ const globalProgressEvent = $progressEvent.get();
+ if (globalProgressEvent !== null && globalProgressEvent.item_id !== data.item_id) {
+ return true;
+ }
+ const successor = pickLatestDatum($progressData.get());
+ if (successor !== null) {
+ // The terminated item owned the shared preview, but other sessions are still generating:
+ // hand the preview to the most recently updated one immediately. The tiled view only renders
+ // with more than one active session, so once a single session remains it is displayed
+ // through these globals — leaving them cleared (or parked on the finished session's stale
+ // frame via the resolve illusion) would hide a still-running preview. This applies to every
+ // terminal status, including successful completion with auto-switch.
+ setPendingResolve(null);
+ $isProgressImageResolving.set(false);
+ $progressEvent.set(successor.progressEvent);
+ $progressImage.set(successor.progressImage);
+ return true;
+ }
+ if (globalProgressEvent === null) {
+ // Nothing is retained (this item never reported progress), so there is no illusion to run —
+ // arming one would leave $isProgressImageResolving stuck on until the next generation.
+ setPendingResolve(null);
+ $isProgressImageResolving.set(false);
+ return true;
+ }
+ // Completed queue items have the progress event cleared by the onLoadImage callback. This allows the viewer to
+ // create the illusion of the progress image "resolving" into the final image. If we cleared the progress image
+ // now, there would be a flicker where the progress image disappears before the final image appears, and the
+ // last-selected gallery image should be shown for a brief moment.
+ //
+ // When gallery auto-switch is disabled, we do not need to create this illusion, because we are not going to
+ // switch to the final image automatically. In this case, we clear the progress image immediately.
+ //
+ // We also clear the progress image if the queue item is canceled or failed, as there is no final image to show.
+ if (
+ data.status === 'canceled' ||
+ data.status === 'failed' ||
+ !autoSwitch ||
+ // When the origin is 'canvas' and destination is 'canvas' (without a ':' suffix), that means the
+ // image is going to be added to the staging area. In this case, we need to clear the progress image else it
+ // will be stuck on the viewer.
+ (data.origin === 'canvas' && data.destination !== 'canvas')
+ ) {
+ setPendingResolve(null);
+ clearRetainedPreview();
+ } else {
+ setPendingResolve(data.item_id);
+ $isProgressImageResolving.set(true);
+ }
+ return true;
+ };
+
+ /**
+ * The final gallery image (or video) finished loading. If a completed session's "resolve"
+ * illusion is pending, this is its tail end: the retained preview is cleared so the final image
+ * shows. A no-op otherwise (e.g. when the preview was handed to a still-running session).
+ *
+ * `sessionId` identifies the item that was loaded (ImageDTO/VideoDTO `session_id`). Several
+ * sessions run concurrently under multi-GPU and auto-switch, so a load can arrive late, after a
+ * *different* session took over the retained preview: session A completes and hands the preview
+ * to B, B then completes and starts its own resolve illusion, and only then does A's final image
+ * finish loading. Clearing on A's load would cut B's illusion short — exactly the flicker the
+ * illusion exists to hide — so a load is ignored when it can be positively attributed to another
+ * session we tracked. Loads we cannot attribute (uploads, images from before this viewer
+ * mounted) still clear.
+ *
+ * Ignoring a load must never be the difference between the preview clearing and not clearing:
+ * the retained session's own image may never load at all (see RESOLVE_TIMEOUT_MS), and the
+ * ignored load may have been the last one coming. The timeout armed alongside the illusion is
+ * what bounds it — this check only decides whether the illusion ends early.
+ */
+ const onFinalImageLoaded = (sessionId: string | null): void => {
+ if (pendingResolveItemId === null) {
+ return;
+ }
+ if (sessionId !== null) {
+ const loadedItemId = itemIdBySessionId.get(sessionId);
+ if (loadedItemId !== undefined && loadedItemId !== pendingResolveItemId) {
+ return;
+ }
+ }
+ setPendingResolve(null);
+ clearRetainedPreview();
+ };
+
+ /**
+ * Handle a queue-cleared event. A clear deletes queue items without emitting a per-item terminal
+ * status event for every one of them (a worker claimed mid-clear is stopped only by this event),
+ * so the tracked previews must be dropped here. Which items were deleted depends on the event's
+ * scope (mirroring workflowExecutionCoordinator.onQueueCleared): an unscoped clear (user_id=null
+ * — an admin or single-user clear) deleted every item; a clear scoped to the current user
+ * deleted all of this client's items; another user's scoped clear — received in full by admins
+ * or as the sanitized user_id="redacted" broadcast by everyone else — deleted none of this
+ * client's items, and this store only ever tracks the client's own items.
+ *
+ * Returns whether the clear applied to this client's previews.
+ */
+ const onQueueCleared = (data: S['QueueClearedEvent'], currentUserId: string | null): boolean => {
+ const clearedUserId = data.user_id ?? null;
+ if (clearedUserId !== null && clearedUserId !== currentUserId) {
+ return false;
+ }
+ // Mark every session we have seen progress for and not yet seen terminate as finished, so a
+ // trailing invocation_progress event from a worker that the clear is still stopping cannot
+ // repopulate the preview. This must cover sessions that have not produced a preview image yet
+ // — they are absent from $progressData and may not own the shared globals, but their first
+ // image event would otherwise resurrect the preview after the clear.
+ for (const itemId of unfinishedItemIds) {
+ finishedQueueItemIds.set(itemId, true);
+ }
+ clearAll();
+ return true;
+ };
+
+ /**
+ * Drop all preview state without marking items finished. For socket disconnection and socket
+ * replacement (auth-token/user change): the tracked sessions belong to the old connection and
+ * will never emit another terminal event on this one.
+ */
+ const reset = (): void => {
+ clearAll();
+ };
+
+ return { onFinalImageLoaded, onItemStarted, onQueueCleared, onTerminal, recordProgress, reset };
+};