invoke-ai · lstein · Aug 18, 2026 · Jul 27, 2026 · Aug 9, 2026 · Aug 17, 2026
@@ -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 && (
         <Flex w="full" h="full" position="absolute" alignItems="center" justifyContent="center">
-          <DndImage imageDTO={imageToRender} onLoad={onLoadImage} borderRadius="base" />
+          <DndImage imageDTO={imageToRender} onLoad={onLoadRenderedImage} borderRadius="base" />
         </Flex>
       )}
       {!imageToRender && <NoContentForViewer />}

@@ -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 <NoContentForViewer />;

@@ -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<number, ViewerProgressDatum | undefined>;
+export type { ViewerProgressDatum } from 'features/gallery/components/ImageViewer/viewerProgressLifecycle';
 
 type ImageViewerContextValue = {
   $progressEvent: Atom<S['InvocationProgressEvent'] | null>;
@@ -33,7 +31,11 @@ type ImageViewerContextValue = {
   $activeProgressData: Atom<ViewerProgressDatum[]>;
   $isProgressImageResolving: Atom<boolean>;
   $isTemporarilyShowingSelectedImage: WritableAtom<boolean>;
-  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<ImageViewerContextValue | null>(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<number, boolean>({ 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<string, number>({ 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 ':<session id>' 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(
     () => ({