From 4cbbb9dd7cf4eae477d42e8d60146adf1aa86fee Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sun, 26 Jul 2026 21:44:45 -0400 Subject: [PATCH 1/4] fix(ui): make multi-GPU viewer previews survive owner termination and queue lifecycle events Two related lifecycle gaps in the image viewer's multi-session preview state (ImageViewer/context.tsx): 1. Terminal-owner fallback: when the session owning the shared $progressEvent/$progressImage globals reached a terminal state, its tile was removed but the globals were cleared (or parked on the finished session's stale frame via the resolve illusion). Since the tiled view only renders with >1 active session, the remaining session's still-running preview disappeared until its next image event. The globals are now handed to the most recently updated remaining session immediately, for every terminal status. 2. Stale lifecycle: $progressData was cleaned only by per-item terminal events. It is now also cleared on queue_cleared (scoped like workflowExecutionCoordinator.onQueueCleared, and marking the cleared items finished so a trailing progress event from a worker stopped only by the clear cannot repopulate the preview), on socket disconnect, and on $socket replacement (auth-token/user change). The store logic is extracted into viewerProgressLifecycle.ts so it can be unit tested without rendering; the provider keeps the socket subscriptions and ownership/scope checks. 16 new vitest cases cover promotion across terminal statuses and auto-switch modes, non-owner termination, clear scoping (own/unscoped/foreign/sanitized), and disconnect resets. Follow-up to PR #9263 (JPPhoto review, 2026-07-25). Co-Authored-By: Claude Fable 5 --- .../components/ImageViewer/context.tsx | 151 ++++++------- .../viewerProgressLifecycle.test.ts | 208 ++++++++++++++++++ .../ImageViewer/viewerProgressLifecycle.ts | 207 +++++++++++++++++ 3 files changed, 482 insertions(+), 84 deletions(-) create mode 100644 invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.test.ts create mode 100644 invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.ts 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..4bb39f7f8aa 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; @@ -58,10 +56,20 @@ 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 })); + // 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, + }) + )[0]; useEffect(() => { if (!socket) { @@ -74,24 +82,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 +95,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 +111,14 @@ export const ImageViewerContextProvider = memo((props: PropsWithChildren) => { if (getEventScope(store.getState, data) !== 'own') { return; } - if (finishedQueueItemIds.has(data.item_id)) { + 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 +127,52 @@ export const ImageViewerContextProvider = memo((props: PropsWithChildren) => { return () => { socket.off('queue_item_status_changed', onQueueItemStatusChanged); }; - }, [ - $isProgressImageResolving, - $progressData, - $progressEvent, - $progressImage, - autoSwitch, - finishedQueueItemIds, - socket, - store, - ]); + }, [autoSwitch, lifecycle, socket, store]); - const onLoadImage = useCallback(() => { - if (!shouldClearProgressImageOnLoadRef.current) { + 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(() => { + lifecycle.onFinalImageLoaded(); + }, [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..29d5dc36f20 --- /dev/null +++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.test.ts @@ -0,0 +1,208 @@ +import type { ProgressImage as ProgressImageType } from 'features/nodes/types/common'; +import { atom, map } from 'nanostores'; +import type { S } from 'services/api/types'; +import { beforeEach, describe, expect, it } from 'vitest'; + +import type { ViewerProgressDataMap, ViewerProgressStores } from './viewerProgressLifecycle'; +import { createViewerProgressLifecycle } 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(() => { + stores = { + $progressEvent: atom(null), + $progressImage: atom(null), + $progressData: map({}), + $isProgressImageResolving: atom(false), + finishedQueueItemIds: new Map(), + }; + lifecycle = createViewerProgressLifecycle(stores); + }); + + 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(); + 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(); + expect(stores.$progressEvent.get()).toBeNull(); + expect(stores.$progressImage.get()).toBeNull(); + expect(stores.$isProgressImageResolving.get()).toBe(false); + }); + + 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('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.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..c1773abebe8 --- /dev/null +++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.ts @@ -0,0 +1,207 @@ +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; +}; + +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; +}; + +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 } = stores; + let seq = 0; + // Whether the final gallery image's onLoad should clear the retained preview — the tail end of + // the "resolve" illusion for a completed session (see onTerminal / onFinalImageLoaded). + let clearProgressOnFinalImageLoad = false; + + const clearAll = (): void => { + clearProgressOnFinalImageLoad = false; + $isProgressImageResolving.set(false); + $progressEvent.set(null); + $progressImage.set(null); + $progressData.set({}); + }; + + /** 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; + } + clearProgressOnFinalImageLoad = 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, + 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); + // 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. + clearProgressOnFinalImageLoad = false; + $isProgressImageResolving.set(false); + $progressEvent.set(successor.progressEvent); + $progressImage.set(successor.progressImage); + 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') + ) { + clearProgressOnFinalImageLoad = false; + $isProgressImageResolving.set(false); + $progressEvent.set(null); + $progressImage.set(null); + } else { + clearProgressOnFinalImageLoad = true; + $isProgressImageResolving.set(true); + } + return true; + }; + + /** + * The final gallery image 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). + */ + const onFinalImageLoaded = (): void => { + if (!clearProgressOnFinalImageLoad) { + return; + } + clearProgressOnFinalImageLoad = false; + $isProgressImageResolving.set(false); + $progressEvent.set(null); + $progressImage.set(null); + }; + + /** + * 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 tracked session finished so a trailing invocation_progress event from a worker + // that the clear is still stopping cannot repopulate the preview. + for (const datum of Object.values($progressData.get())) { + if (datum !== undefined) { + finishedQueueItemIds.set(datum.itemId, true); + } + } + const globalProgressEvent = $progressEvent.get(); + if (globalProgressEvent !== null) { + finishedQueueItemIds.set(globalProgressEvent.item_id, 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, onQueueCleared, onTerminal, recordProgress, reset }; +}; From 4581a92908e06e78a934998657bcb7a786399627 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sun, 16 Aug 2026 22:01:43 -0400 Subject: [PATCH 2/4] fix(ui): scope viewer preview clears to the session that owns them MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses JPPhoto's review of #9389. A queue clear marked only sessions that had produced a preview image (plus the owner of the shared globals) as finished. A session that had reported progress without an image yet was left unmarked, so its first image event after the clear resurrected the preview the clear had just dropped. Every item seen on progress is now tracked, image or not, and marked finished by the clear. The final-image load callback carried no session identity, so a late load from an already-completed session cleared whatever preview was retained at the time — including a *different* session's pending resolve illusion (A completes and hands the preview to B, B completes and starts its own illusion, then A's image finally loads). The callback now takes the loaded DTO's session_id and ignores loads it can attribute to another tracked session; unattributable loads (uploads, pre-mount images) still clear, so a retained preview cannot cover the viewer indefinitely. Co-Authored-By: Claude Opus 5 (1M context) --- .../ImageViewer/CurrentImagePreview.tsx | 8 +- .../ImageViewer/CurrentVideoPreview.tsx | 4 +- .../components/ImageViewer/context.tsx | 19 +++- .../viewerProgressLifecycle.test.ts | 57 +++++++++++- .../ImageViewer/viewerProgressLifecycle.ts | 90 +++++++++++++------ 5 files changed, 144 insertions(+), 34 deletions(-) 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 4bb39f7f8aa..d683e95180e 100644 --- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/context.tsx +++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/context.tsx @@ -31,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); @@ -59,6 +63,9 @@ export const ImageViewerContextProvider = memo((props: PropsWithChildren) => { // 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(() => @@ -68,6 +75,7 @@ export const ImageViewerContextProvider = memo((props: PropsWithChildren) => { $progressData, $isProgressImageResolving, finishedQueueItemIds, + itemIdBySessionId, }) )[0]; @@ -170,9 +178,12 @@ export const ImageViewerContextProvider = memo((props: PropsWithChildren) => { }; }, [lifecycle, socket]); - const onLoadImage = useCallback(() => { - lifecycle.onFinalImageLoaded(); - }, [lifecycle]); + 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 index 29d5dc36f20..7e282b8a1eb 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 @@ -58,6 +58,7 @@ describe('viewerProgressLifecycle', () => { $progressData: map({}), $isProgressImageResolving: atom(false), finishedQueueItemIds: new Map(), + itemIdBySessionId: new Map(), }; lifecycle = createViewerProgressLifecycle(stores); }); @@ -102,7 +103,7 @@ describe('viewerProgressLifecycle', () => { 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(); + lifecycle.onFinalImageLoaded('session-1'); expect(stores.$progressEvent.get()).toBe(eventB); expect(stores.$progressImage.get()).toBe(eventB.image); } @@ -148,7 +149,7 @@ describe('viewerProgressLifecycle', () => { // 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(); + lifecycle.onFinalImageLoaded('session-1'); expect(stores.$progressEvent.get()).toBeNull(); expect(stores.$progressImage.get()).toBeNull(); expect(stores.$isProgressImageResolving.get()).toBe(false); @@ -162,6 +163,43 @@ describe('viewerProgressLifecycle', () => { }); }); + 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.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); + }); + }); + describe('onQueueCleared', () => { it.each([ ['an unscoped (admin or single-user) clear', null, 'user-1'], @@ -179,6 +217,21 @@ describe('viewerProgressLifecycle', () => { expect(lifecycle.recordProgress(buildProgressEvent({ item_id: 2, image: buildProgressImage(2) }))).toBe(false); }); + 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'], diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.ts b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.ts index c1773abebe8..806ff0c9cba 100644 --- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.ts +++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.ts @@ -21,6 +21,12 @@ type FinishedQueueItemIds = { 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; +}; + export type ViewerProgressStores = { $progressEvent: WritableAtom; $progressImage: WritableAtom; @@ -29,6 +35,10 @@ export type ViewerProgressStores = { $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 => { @@ -52,14 +62,28 @@ const pickLatestDatum = (data: ViewerProgressDataMap): ViewerProgressDatum | nul * that most recently reported progress. */ export const createViewerProgressLifecycle = (stores: ViewerProgressStores) => { - const { $progressEvent, $progressImage, $progressData, $isProgressImageResolving, finishedQueueItemIds } = stores; + const { + $progressEvent, + $progressImage, + $progressData, + $isProgressImageResolving, + finishedQueueItemIds, + itemIdBySessionId, + } = stores; let seq = 0; - // Whether the final gallery image's onLoad should clear the retained preview — the tail end of - // the "resolve" illusion for a completed session (see onTerminal / onFinalImageLoaded). - let clearProgressOnFinalImageLoad = false; + // 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. + 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(); const clearAll = (): void => { - clearProgressOnFinalImageLoad = false; + pendingResolveItemId = null; + unfinishedItemIds.clear(); $isProgressImageResolving.set(false); $progressEvent.set(null); $progressImage.set(null); @@ -71,7 +95,9 @@ export const createViewerProgressLifecycle = (stores: ViewerProgressStores) => { if (finishedQueueItemIds.has(data.item_id)) { return false; } - clearProgressOnFinalImageLoad = false; + unfinishedItemIds.add(data.item_id); + itemIdBySessionId.set(data.session_id, data.item_id); + pendingResolveItemId = null; $isProgressImageResolving.set(false); $progressEvent.set(data); if (data.image) { @@ -93,6 +119,7 @@ export const createViewerProgressLifecycle = (stores: ViewerProgressStores) => { 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); @@ -112,7 +139,7 @@ export const createViewerProgressLifecycle = (stores: ViewerProgressStores) => { // 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. - clearProgressOnFinalImageLoad = false; + pendingResolveItemId = null; $isProgressImageResolving.set(false); $progressEvent.set(successor.progressEvent); $progressImage.set(successor.progressImage); @@ -136,27 +163,43 @@ export const createViewerProgressLifecycle = (stores: ViewerProgressStores) => { // will be stuck on the viewer. (data.origin === 'canvas' && data.destination !== 'canvas') ) { - clearProgressOnFinalImageLoad = false; + pendingResolveItemId = null; $isProgressImageResolving.set(false); $progressEvent.set(null); $progressImage.set(null); } else { - clearProgressOnFinalImageLoad = true; + pendingResolveItemId = data.item_id; $isProgressImageResolving.set(true); } return true; }; /** - * The final gallery image 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). + * 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, an unrelated image the user selected) still clear, keeping the safety net that stops + * a retained preview from covering the viewer indefinitely. */ - const onFinalImageLoaded = (): void => { - if (!clearProgressOnFinalImageLoad) { + const onFinalImageLoaded = (sessionId: string | null): void => { + if (pendingResolveItemId === null) { return; } - clearProgressOnFinalImageLoad = false; + if (sessionId !== null) { + const loadedItemId = itemIdBySessionId.get(sessionId); + if (loadedItemId !== undefined && loadedItemId !== pendingResolveItemId) { + return; + } + } + pendingResolveItemId = null; $isProgressImageResolving.set(false); $progressEvent.set(null); $progressImage.set(null); @@ -179,16 +222,13 @@ export const createViewerProgressLifecycle = (stores: ViewerProgressStores) => { if (clearedUserId !== null && clearedUserId !== currentUserId) { return false; } - // Mark every tracked session finished so a trailing invocation_progress event from a worker - // that the clear is still stopping cannot repopulate the preview. - for (const datum of Object.values($progressData.get())) { - if (datum !== undefined) { - finishedQueueItemIds.set(datum.itemId, true); - } - } - const globalProgressEvent = $progressEvent.get(); - if (globalProgressEvent !== null) { - finishedQueueItemIds.set(globalProgressEvent.item_id, true); + // 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; From 41a781f3bb689f2891fed085f91539c1d8992ed8 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sun, 16 Aug 2026 22:19:41 -0400 Subject: [PATCH 3/4] fix(ui): bound the viewer's retained preview and track claimed items MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to the previous commit, from an adversarial review of it. Ignoring a mismatched final-image load removed the only thing that ever took a retained preview down. The retained session's own image may never load: the viewer's reports load failures through onError, not onLoad, and with concurrent completions auto-switch can settle on another session's image, so the retained one is never rendered at all. In those cases the preview — an opaque overlay — covered the viewer until the next generation. The illusion now carries a timeout armed alongside it, so an ignored load can only end it early, never keep it up. Also from the same review: - A queue clear cancels the items already running before deleting the rows, so a worker that claims an item in between gets no terminal event and its first progress event lands after the clear. Only the in_progress claim event names that item, so the viewer now tracks it. - Session attributions outlived the state they described: a disconnect or socket swap dropped the previews but kept the session ids, so images from before the reset were still read as another session's and refused to end a later illusion. They are dropped with everything else. - A completed item that never reported progress armed an illusion with nothing retained, leaving the resolving flag on. Co-Authored-By: Claude Opus 5 (1M context) --- .../components/ImageViewer/context.tsx | 8 ++ .../viewerProgressLifecycle.test.ts | 95 +++++++++++++++- .../ImageViewer/viewerProgressLifecycle.ts | 102 ++++++++++++++---- 3 files changed, 185 insertions(+), 20 deletions(-) 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 d683e95180e..19530ffebb9 100644 --- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/context.tsx +++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/context.tsx @@ -119,6 +119,14 @@ export const ImageViewerContextProvider = memo((props: PropsWithChildren) => { if (getEventScope(store.getState, data) !== 'own') { return; } + 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; } 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 7e282b8a1eb..3e49f9d81e1 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 @@ -1,10 +1,10 @@ import type { ProgressImage as ProgressImageType } from 'features/nodes/types/common'; import { atom, map } from 'nanostores'; import type { S } from 'services/api/types'; -import { beforeEach, describe, expect, it } from 'vitest'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import type { ViewerProgressDataMap, ViewerProgressStores } from './viewerProgressLifecycle'; -import { createViewerProgressLifecycle } from './viewerProgressLifecycle'; +import { createViewerProgressLifecycle, RESOLVE_TIMEOUT_MS } from './viewerProgressLifecycle'; const buildProgressImage = (itemId: number): ProgressImageType => ({ @@ -52,6 +52,7 @@ describe('viewerProgressLifecycle', () => { let lifecycle: ReturnType; beforeEach(() => { + vi.useFakeTimers(); stores = { $progressEvent: atom(null), $progressImage: atom(null), @@ -63,6 +64,10 @@ describe('viewerProgressLifecycle', () => { 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) }); @@ -155,6 +160,13 @@ describe('viewerProgressLifecycle', () => { 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); @@ -183,6 +195,55 @@ describe('viewerProgressLifecycle', () => { 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); + takeOver(); + const eventAfterTakeOver = stores.$progressEvent.get(); + // The timeout must never fire against a preview that has since been replaced or dropped. + 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'], @@ -198,6 +259,20 @@ describe('viewerProgressLifecycle', () => { 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', () => { @@ -217,6 +292,22 @@ describe('viewerProgressLifecycle', () => { 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. diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.ts b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.ts index 806ff0c9cba..f34fd714baf 100644 --- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.ts +++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.ts @@ -25,8 +25,19 @@ type FinishedQueueItemIds = { 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; @@ -73,23 +84,71 @@ export const createViewerProgressLifecycle = (stores: ViewerProgressStores) => { 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. + // 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 clearAll = (): void => { - pendingResolveItemId = null; - unfinishedItemIds.clear(); + 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 is always in sync with it — + * whoever takes over or clears the shared preview also cancels the timeout, so it can never + * clear a preview that has since been handed to another session. + */ + 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)) { @@ -97,7 +156,7 @@ export const createViewerProgressLifecycle = (stores: ViewerProgressStores) => { } unfinishedItemIds.add(data.item_id); itemIdBySessionId.set(data.session_id, data.item_id); - pendingResolveItemId = null; + setPendingResolve(null); $isProgressImageResolving.set(false); $progressEvent.set(data); if (data.image) { @@ -139,12 +198,19 @@ export const createViewerProgressLifecycle = (stores: ViewerProgressStores) => { // 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. - pendingResolveItemId = null; + 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 @@ -163,12 +229,10 @@ export const createViewerProgressLifecycle = (stores: ViewerProgressStores) => { // will be stuck on the viewer. (data.origin === 'canvas' && data.destination !== 'canvas') ) { - pendingResolveItemId = null; - $isProgressImageResolving.set(false); - $progressEvent.set(null); - $progressImage.set(null); + setPendingResolve(null); + clearRetainedPreview(); } else { - pendingResolveItemId = data.item_id; + setPendingResolve(data.item_id); $isProgressImageResolving.set(true); } return true; @@ -186,8 +250,12 @@ export const createViewerProgressLifecycle = (stores: ViewerProgressStores) => { * 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, an unrelated image the user selected) still clear, keeping the safety net that stops - * a retained preview from covering the viewer indefinitely. + * 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) { @@ -199,10 +267,8 @@ export const createViewerProgressLifecycle = (stores: ViewerProgressStores) => { return; } } - pendingResolveItemId = null; - $isProgressImageResolving.set(false); - $progressEvent.set(null); - $progressImage.set(null); + setPendingResolve(null); + clearRetainedPreview(); }; /** @@ -243,5 +309,5 @@ export const createViewerProgressLifecycle = (stores: ViewerProgressStores) => { clearAll(); }; - return { onFinalImageLoaded, onQueueCleared, onTerminal, recordProgress, reset }; + return { onFinalImageLoaded, onItemStarted, onQueueCleared, onTerminal, recordProgress, reset }; }; From 80bb03edc0110f36315a3aef910a6e90c9e5273c Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sun, 16 Aug 2026 22:35:04 -0400 Subject: [PATCH 4/4] test(ui): assert the resolve timeout is disarmed, not just harmless MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two of the three disarm cases could not fail: after a load or a reset the stores are already null, so a leaked timer's clear was indistinguishable from no timer at all. Assert the timer count instead. Also correct the setPendingResolve docstring, which claimed more than the code does: a progress event carrying no image cancels the illusion without replacing the retained frame, so that frame is bounded by the next preview image rather than by this timeout — as it is today. Co-Authored-By: Claude Opus 5 (1M context) --- .../ImageViewer/viewerProgressLifecycle.test.ts | 7 ++++++- .../ImageViewer/viewerProgressLifecycle.ts | 12 +++++++++--- 2 files changed, 15 insertions(+), 4 deletions(-) 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 3e49f9d81e1..a3181382a7f 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 @@ -237,9 +237,14 @@ describe('viewerProgressLifecycle', () => { 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(); - const eventAfterTakeOver = stores.$progressEvent.get(); // 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); }); diff --git a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.ts b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.ts index f34fd714baf..e6c78cc7499 100644 --- a/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.ts +++ b/invokeai/frontend/web/src/features/gallery/components/ImageViewer/viewerProgressLifecycle.ts @@ -102,9 +102,15 @@ export const createViewerProgressLifecycle = (stores: ViewerProgressStores) => { /** * Arm (or, with null, disarm) the pending "resolve" illusion. Every write to - * `pendingResolveItemId` goes through here so the safety timeout is always in sync with it — - * whoever takes over or clears the shared preview also cancels the timeout, so it can never - * clear a preview that has since been handed to another session. + * `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) {