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

invokeai/frontend/web/src/features/gallery/components/ImageViewer/CurrentImagePreview.tsx

@@ -1,11 +1,13 @@
 import { Box, Flex } from '@invoke-ai/ui-library';
 import { useStore } from '@nanostores/react';
 import { useAppSelector } from 'app/store/storeHooks';
+import { useMediaUrl } from 'features/auth/store/mediaCookieRefresh';
 import { CanvasAlertsInvocationProgress } from 'features/controlLayers/components/CanvasAlerts/CanvasAlertsInvocationProgress';
 import { DndImage } from 'features/dnd/DndImage';
 import ImageMetadataViewer from 'features/gallery/components/ImageMetadataViewer/ImageMetadataViewer';
 import NextPrevItemButtons from 'features/gallery/components/NextPrevItemButtons';
 import { useNextPrevItemNavigation } from 'features/gallery/components/useNextPrevItemNavigation';
+import { autoSwitchedImages } from 'features/gallery/store/autoSwitchedImages';
 import { selectLastSelectedItem } from 'features/gallery/store/gallerySelectors';
 import { useRegisteredHotkeys } from 'features/system/components/HotkeysModal/useHotkeyData';
 import { navigationApi } from 'features/ui/layouts/navigation-api';
@@ -48,6 +50,20 @@ export const CurrentImagePreview = memo(({ imageDTO }: { imageDTO: ImageDTO | nu
   const previousRenderedImageNameRef = useRef<string | null>(null);
   const selectedImageRevealTimeoutId = useRef(0);
 
+  // The reveal gate below deliberately preloads the *thumbnail*, not the full-resolution image. The
+  // progress overlay covers this element until onLoadImage fires, so gating on the multi-megabyte
+  // `/full` response would hold a stale latent preview on screen for that entire download on a slow
+  // connection. The 256px thumbnail is roughly 100x smaller and is typically higher resolution than
+  // the preview it replaces; DndImage renders it via Chakra's `fallbackSrc` and swaps the full image
+  // in, in place, once that finishes loading.
+  //
+  // The URL must go through useMediaUrl so it is byte-identical to the one DndImage requests. The
+  // media cookie version is a query parameter, so a mismatch is a different key and the bytes are
+  // fetched twice (measured: 2 requests mismatched vs 1 matched). Note the reuse here is the
+  // document's list of available images, which is keyed by URL and is not the HTTP cache — it still
+  // holds in multiuser mode, where images are served `Cache-Control: private, no-store`.
+  const previewSrc = useMediaUrl(imageDTO?.thumbnail_url);
+
   useEffect(() => {
     if (!selectedImageName) {
       setImageToRender(null);
@@ -65,9 +81,13 @@ export const CurrentImagePreview = memo(({ imageDTO }: { imageDTO: ImageDTO | nu
         return;
       }
       setImageToRender(imageDTO);
+      // Resolve the progress overlay as soon as the thumbnail settles — on success *or* error.
+      // Relying on DndImage's onLoad alone leaves the overlay stuck whenever the image fails to
+      // load, because Chakra reports that as onError instead.
+      onLoadImage();
     };
 
-    if (typeof window === 'undefined') {
+    if (typeof window === 'undefined' || !previewSrc) {
       onReady();
       return;
     }
@@ -76,7 +96,7 @@ export const CurrentImagePreview = memo(({ imageDTO }: { imageDTO: ImageDTO | nu
 
     preloader.onload = onReady;
     preloader.onerror = onReady;
-    preloader.src = imageDTO.image_url;
+    preloader.src = previewSrc;
 
     if (preloader.complete) {
       onReady();
@@ -87,7 +107,7 @@ export const CurrentImagePreview = memo(({ imageDTO }: { imageDTO: ImageDTO | nu
       preloader.onload = null;
       preloader.onerror = null;
     };
-  }, [imageDTO, imageToRender?.image_name, selectedImageName]);
+  }, [imageDTO, imageToRender?.image_name, onLoadImage, previewSrc, selectedImageName]);
 
   const hasProgressImage = progressImage !== null;
 
@@ -96,6 +116,14 @@ export const CurrentImagePreview = memo(({ imageDTO }: { imageDTO: ImageDTO | nu
     const previousRenderedImageName = previousRenderedImageNameRef.current;
     previousRenderedImageNameRef.current = renderedImageName;
 
+    // Consume on every change of the rendered image, not only when the reveal conditions below
+    // hold — in the common case the auto-switched image renders with no progress showing, and an
+    // entry left behind would suppress a genuine user selection of the same image later.
+    const wasAutoSwitchedTo =
+      renderedImageName !== null &&
+      renderedImageName !== previousRenderedImageName &&
+      autoSwitchedImages.consume(renderedImageName);
+
     window.clearTimeout(selectedImageRevealTimeoutId.current);
 
     if (
@@ -113,6 +141,16 @@ export const CurrentImagePreview = memo(({ imageDTO }: { imageDTO: ImageDTO | nu
       return;
     }
 
+    // The reveal exists to make a mid-generation *user* selection visible. An auto-switch to a
+    // just-finished image can land here late — after the next generation's first progress event
+    // has already reset $isProgressImageResolving — and must not flash the previous result over
+    // the live preview. The set(false) is required: the clearTimeout above already cancelled any
+    // running reveal's timer, so returning with the atom still true would wedge the reveal on.
+    if (wasAutoSwitchedTo) {
+      $isTemporarilyShowingSelectedImage.set(false);
+      return;
+    }
+
     $isTemporarilyShowingSelectedImage.set(true);
     selectedImageRevealTimeoutId.current = window.setTimeout(() => {
       $isTemporarilyShowingSelectedImage.set(false);

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

@@ -0,0 +1,65 @@
+import { readFileSync } from 'node:fs';
+import { fileURLToPath } from 'node:url';
+
+import { describe, expect, it } from 'vitest';
+
+const read = (file: string) => readFileSync(fileURLToPath(new URL(file, import.meta.url)), 'utf8');
+
+// The behaviour of the deferred clear itself is covered by real tests in
+// progressImageResolution.test.ts. These are wiring checks only — this directory has no DOM test
+// environment, so the provider cannot be mounted. They assert that context.tsx routes through the
+// tested unit rather than reimplementing the state inline, which is what previously allowed the
+// armed flag and its timer to drift apart.
+describe('ImageViewer progress image wiring', () => {
+  const context = read('./context.tsx');
+  const currentImagePreview = read('./CurrentImagePreview.tsx');
+
+  it('resets the viewer progress atoms on every socket lifecycle transition', () => {
+    // socket.io has no event replay, so a drop spanning the terminal queue_item_status_changed
+    // loses that event permanently. Without these the overlay covers the finished image until the
+    // page is reloaded. setEventListeners already does the same for the global progress stores.
+    for (const event of ['connect', 'connect_error', 'disconnect']) {
+      expect(context).toContain(`socket.on('${event}', onSocketLifecycleChange)`);
+      expect(context).toContain(`socket.off('${event}', onSocketLifecycleChange)`);
+    }
+  });
+
+  it('keeps the armed flag and its backstop timer in one owned unit', () => {
+    // Both must come from createDeferredClear. A local boolean ref plus a separate timeout id is
+    // exactly the shape that let a stale timer outlive the generation that armed it.
+    expect(context).toContain('createDeferredClear()');
+    expect(context).toContain('deferredClear.arm(onResolveDeadline)');
+    expect(context).toContain('deferredClear.isArmed()');
+    expect(context).not.toContain('shouldClearProgressImageOnLoadRef');
+    expect(context).not.toContain('setTimeout');
+  });
+
+  it('supersedes a pending backstop when a new progress event arrives', () => {
+    // Otherwise: item N completes and arms the backstop, its final image never loads, the user
+    // starts item N+1, and N's deadline fires mid-generation and blanks N+1's live preview.
+    const progressHandler = context.slice(
+      context.indexOf('const onInvocationProgress ='),
+      context.indexOf("socket.on('invocation_progress'")
+    );
+    expect(progressHandler).toContain('disarmDeferredClear()');
+  });
+
+  it('gates the viewer reveal on the thumbnail rather than the full-resolution image', () => {
+    // Gating on `/full` holds a stale latent preview on screen for the whole multi-megabyte
+    // download on a slow connection.
+    expect(currentImagePreview).toContain('useMediaUrl(imageDTO?.thumbnail_url)');
+    expect(currentImagePreview).toContain('preloader.src = previewSrc');
+    expect(currentImagePreview).not.toMatch(/preloader\.src\s*=\s*imageDTO\.image_url/);
+  });
+
+  it('clears the progress overlay when the preload settles, including on error', () => {
+    // Chakra reports a failed load as onError, not onLoad, so DndImage's onLoad alone is not
+    // enough to guarantee the overlay is ever cleared.
+    expect(currentImagePreview).toContain('preloader.onerror = onReady');
+    const onReady = currentImagePreview.slice(
+      currentImagePreview.indexOf('const onReady ='),
+      currentImagePreview.indexOf('if (typeof window ===')
+    );
+    expect(onReady).toContain('onLoadImage()');
+  });
+});

invokeai/frontend/web/src/features/gallery/components/ImageViewer/context.tsx

@@ -6,13 +6,15 @@ import type { ProgressImage as ProgressImageType } from 'features/nodes/types/co
 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';
 
+import { createDeferredClear, getTerminalProgressAction } from './progressImageResolution';
+
 /** 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 = {
@@ -58,11 +60,45 @@ export const ImageViewerContextProvider = memo((props: PropsWithChildren) => {
   )[0];
   const $isProgressImageResolving = useState(() => atom(false))[0];
   const $isTemporarilyShowingSelectedImage = useState(() => atom(false))[0];
-  const shouldClearProgressImageOnLoadRef = useRef(false);
+  // Owns both the "clear on load" flag and its backstop timer, so no path can reset one and leak
+  // the other. See createDeferredClear.
+  const [deferredClear] = useState(() => createDeferredClear());
   // 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 }));
 
+  // Cancels a pending deferred clear without touching the preview itself. Every path that takes
+  // responsibility for the preview away from the armed onLoadImage must call this, or the backstop
+  // outlives the generation that armed it and blanks a later one's live preview.
+  const disarmDeferredClear = useCallback(() => {
+    deferredClear.disarm();
+    $isProgressImageResolving.set(false);
+  }, [$isProgressImageResolving, deferredClear]);
+
+  const clearProgressImage = useCallback(() => {
+    disarmDeferredClear();
+    $progressEvent.set(null);
+    $progressImage.set(null);
+  }, [disarmDeferredClear, $progressEvent, $progressImage]);
+
+  // Nulling $progressImage tears down the whole overlay, tiles included — $activeProgressData only
+  // renders while it is set. So when other sessions are still producing previews (multi-GPU), the
+  // backstop must not clear: the overlay has already stopped being this item's to own. Disarming is
+  // enough; those sessions clear it via their own terminal events.
+  const onResolveDeadline = useCallback(() => {
+    if ($activeProgressData.get().length > 0) {
+      disarmDeferredClear();
+      return;
+    }
+    clearProgressImage();
+  }, [$activeProgressData, clearProgressImage, disarmDeferredClear]);
+
+  useEffect(() => {
+    return () => {
+      deferredClear.disarm();
+    };
+  }, [deferredClear]);
+
   useEffect(() => {
     if (!socket) {
       return;
@@ -81,8 +117,10 @@ export const ImageViewerContextProvider = memo((props: PropsWithChildren) => {
         );
         return;
       }
-      shouldClearProgressImageOnLoadRef.current = false;
-      $isProgressImageResolving.set(false);
+      // A new preview supersedes any deferred clear still armed by the previous queue item, whose
+      // final image may never have loaded. Leaving its backstop running would blank this preview
+      // mid-generation.
+      disarmDeferredClear();
       $progressEvent.set(data);
       if (data.image) {
         $progressImage.set(data.image);
@@ -100,7 +138,7 @@ export const ImageViewerContextProvider = memo((props: PropsWithChildren) => {
     return () => {
       socket.off('invocation_progress', onInvocationProgress);
     };
-  }, [$isProgressImageResolving, $progressData, $progressEvent, $progressImage, finishedQueueItemIds, socket, store]);
+  }, [$progressData, $progressEvent, $progressImage, disarmDeferredClear, finishedQueueItemIds, socket, store]);
 
   useEffect(() => {
     if (!socket) {
@@ -128,39 +166,29 @@ export const ImageViewerContextProvider = memo((props: PropsWithChildren) => {
         // 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) {
+
+        // See getTerminalProgressAction for why each outcome is chosen. 'arm' defers the clear to
+        // onLoadImage so the viewer can create the illusion of the progress image "resolving" into
+        // the final image — clearing it here instead would flicker through the previously-selected
+        // gallery image before the final one appears.
+        const action = getTerminalProgressAction(data, {
+          autoSwitch,
+          globalProgressItemId: $progressEvent.get()?.item_id ?? null,
+        });
+
+        if (action === 'ignore') {
           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);
+
+        if (action === 'clear') {
+          clearProgressImage();
+          return;
         }
+
+        $isProgressImageResolving.set(true);
+        // onLoadImage is not guaranteed to fire — see PROGRESS_IMAGE_RESOLVE_TIMEOUT_MS. Without
+        // this deadline the overlay can cover the finished image until the page is reloaded.
+        deferredClear.arm(onResolveDeadline);
       }
     };
 
@@ -173,23 +201,58 @@ export const ImageViewerContextProvider = memo((props: PropsWithChildren) => {
     $isProgressImageResolving,
     $progressData,
     $progressEvent,
-    $progressImage,
     autoSwitch,
+    clearProgressImage,
+    deferredClear,
     finishedQueueItemIds,
+    onResolveDeadline,
     socket,
     store,
   ]);
 
+  // The viewer's progress atoms are separate stores from the global ones in services/events/stores,
+  // which setEventListeners already resets on every socket lifecycle transition. Without the same
+  // reset here the two diverge: socket.io has no event replay, so a drop spanning the terminal
+  // queue_item_status_changed loses that event permanently and nothing is left to clear the opaque
+  // overlay covering the finished image. Backgrounding a tab long enough for the connection to be
+  // torn down is the common way to hit this.
+  //
+  // Clearing on disconnect — not just on reconnect — matches the progress *bars*, which already
+  // vanish then. If the generation is in fact still running, the next invocation_progress event
+  // repopulates the preview within a step.
+  useEffect(() => {
+    if (!socket) {
+      return;
+    }
+
+    const onSocketLifecycleChange = () => {
+      clearProgressImage();
+      // connect_error fires once per reconnection attempt, i.e. roughly once a second while the
+      // server is down. `set` compares by reference, so an unconditional `set({})` would notify
+      // every subscriber on every attempt; only replace the map when it actually holds something.
+      if (Object.keys($progressData.get()).length > 0) {
+        $progressData.set({});
+      }
+    };
+
+    socket.on('connect', onSocketLifecycleChange);
+    socket.on('connect_error', onSocketLifecycleChange);
+    socket.on('disconnect', onSocketLifecycleChange);
+
+    return () => {
+      socket.off('connect', onSocketLifecycleChange);
+      socket.off('connect_error', onSocketLifecycleChange);
+      socket.off('disconnect', onSocketLifecycleChange);
+    };
+  }, [$progressData, clearProgressImage, socket]);
+
   const onLoadImage = useCallback(() => {
-    if (!shouldClearProgressImageOnLoadRef.current) {
+    if (!deferredClear.isArmed()) {
       return;
     }
 
-    shouldClearProgressImageOnLoadRef.current = false;
-    $isProgressImageResolving.set(false);
-    $progressEvent.set(null);
-    $progressImage.set(null);
-  }, [$isProgressImageResolving, $progressEvent, $progressImage]);
+    clearProgressImage();
+  }, [clearProgressImage, deferredClear]);
 
   const value = useMemo(
     () => ({