feat(gallery): refresh when other clients complete generations

invokeai/frontend/webv2/src/app/GalleryRealtimeRuntime.tsx

@@ -1,16 +1,23 @@
 import { createGalleryRealtimeRuntime, invalidateGallery } from '@features/gallery/queries';
+import { getSessionUserId } from '@features/identity';
+import { isWebv2QueueItemOrigin } from '@features/queue/contracts';
 import { useMountEffect } from '@platform/react/useMountEffect';
 import { socketHub } from '@platform/transport/socketHub';
 import { useQueryClient } from '@tanstack/react-query';
 
-/** App-owned composition: uploads the backend announces over the socket refresh the Gallery read model. */
+/**
+ * App-owned composition: uploads and other clients' completed generations the backend announces over the socket
+ * refresh the Gallery read model.
+ */
 export const GalleryRealtimeRuntime = () => {
   const queryClient = useQueryClient();
 
   useMountEffect(() => {
     const runtime = createGalleryRealtimeRuntime({
       backend: socketHub,
+      getSessionUserId,
       invalidate: () => invalidateGallery(queryClient),
+      isClientQueueItemOrigin: isWebv2QueueItemOrigin,
     });
 
     runtime.start();

invokeai/frontend/webv2/src/app/RecallParametersRuntime.tsx

@@ -1,4 +1,4 @@
-import { getAuthSession } from '@features/identity';
+import { getSessionUserId } from '@features/identity';
 import { useMountEffect } from '@platform/react/useMountEffect';
 import { socketHub } from '@platform/transport/socketHub';
 import { isEditableHotkeyTarget } from '@workbench/hotkeys/keys';
@@ -10,12 +10,6 @@ import { useWorkbenchInternalStore } from '@workbench/WorkbenchContext';
 import { useWorkbenchWidgetRegistry } from '@workbench/WorkbenchWidgetRegistryContext';
 import { useTranslation } from 'react-i18next';
 
-/** Single-user servers tag events with a system user id, so only multi-user sessions can be fenced. */
-const getSessionUserId = (): string | null => {
-  const session = getAuthSession();
-  return session.multiuserEnabled ? (session.user?.user_id ?? null) : null;
-};
-
 /**
  * App-owned composition: external `POST /api/v1/recall` updates land in the active project's Generate panel, and
  * `POST /api/v1/recall/video` updates in its Video panel; each brings its panel to the front.

invokeai/frontend/webv2/src/features/gallery/data/realtimeRuntime.test.ts

@@ -28,13 +28,31 @@ describe('gallery realtime runtime', () => {
     }),
   };
   const invalidate = vi.fn();
+  let sessionUserId: string | null = 'alice';
+  const getSessionUserId = () => sessionUserId;
+  const CLIENT_ORIGIN = 'this-client';
+  const createRuntime = () =>
+    createGalleryRealtimeRuntime({
+      backend,
+      getSessionUserId,
+      invalidate,
+      isClientQueueItemOrigin: (origin) => origin === CLIENT_ORIGIN,
+    });
+  const emitStatusChanged = (event: { origin?: string | null; status?: string; user_id?: string }): void =>
+    handlers.get('queue_item_status_changed')?.({
+      origin: 'chatbot',
+      status: 'completed',
+      user_id: 'alice',
+      ...event,
+    } as never);
 
   beforeEach(() => {
     vi.useFakeTimers();
     handlers.clear();
     detachers.length = 0;
     connectionHandler = null;
     initialStatus = 'connected';
+    sessionUserId = 'alice';
     backend.on.mockClear();
     backend.onConnectionChange.mockClear();
     invalidate.mockReset();
@@ -43,15 +61,15 @@ describe('gallery realtime runtime', () => {
   afterEach(() => vi.useRealTimers());
 
   it('coalesces a burst of upload events into one invalidation and detaches on dispose', () => {
-    const runtime = createGalleryRealtimeRuntime({ backend, invalidate });
+    const runtime = createRuntime();
 
     runtime.start();
     runtime.start();
     handlers.get('image_uploaded')?.({ image_name: 'a.png' } as never);
     handlers.get('image_uploaded')?.({ image_name: 'b.png' } as never);
     handlers.get('video_uploaded')?.({ video_name: 'c.mp4' } as never);
 
-    expect(backend.on).toHaveBeenCalledTimes(2);
+    expect(backend.on).toHaveBeenCalledTimes(3);
     expect(backend.onConnectionChange).toHaveBeenCalledTimes(1);
     expect(invalidate).not.toHaveBeenCalled();
 
@@ -68,7 +86,7 @@ describe('gallery realtime runtime', () => {
   });
 
   it('holds passes a minimum interval apart under a stream of uploads and still runs a trailing pass', () => {
-    const runtime = createGalleryRealtimeRuntime({ backend, invalidate });
+    const runtime = createRuntime();
 
     runtime.start();
 
@@ -93,7 +111,7 @@ describe('gallery realtime runtime', () => {
 
   it('refetches after an outage but not on the boot-time connect', () => {
     initialStatus = 'connecting';
-    const runtime = createGalleryRealtimeRuntime({ backend, invalidate });
+    const runtime = createRuntime();
 
     runtime.start();
     connectionHandler?.('connected');
@@ -111,7 +129,7 @@ describe('gallery realtime runtime', () => {
 
   it('treats a connect after mounting on a downed backend as a reconnect', () => {
     initialStatus = 'disconnected';
-    const runtime = createGalleryRealtimeRuntime({ backend, invalidate });
+    const runtime = createRuntime();
 
     runtime.start();
     connectionHandler?.('connected');
@@ -122,7 +140,7 @@ describe('gallery realtime runtime', () => {
   });
 
   it('drops a pending invalidation on dispose', () => {
-    const runtime = createGalleryRealtimeRuntime({ backend, invalidate });
+    const runtime = createRuntime();
 
     runtime.start();
     handlers.get('image_uploaded')?.({ image_name: 'a.png' } as never);
@@ -138,7 +156,7 @@ describe('gallery realtime runtime', () => {
   });
 
   it('ignores socket callbacks after its account scope expires', () => {
-    const runtime = createGalleryRealtimeRuntime({ backend, invalidate });
+    const runtime = createRuntime();
 
     runtime.start();
     accountLifecycle.invalidate();
@@ -148,4 +166,66 @@ describe('gallery realtime runtime', () => {
     expect(invalidate).not.toHaveBeenCalled();
     runtime.dispose();
   });
+
+  it("refreshes once for a burst of the session user's completions enqueued by other clients", () => {
+    const runtime = createRuntime();
+
+    runtime.start();
+    emitStatusChanged({});
+    emitStatusChanged({ origin: null });
+    vi.advanceTimersByTime(250);
+
+    expect(invalidate).toHaveBeenCalledTimes(1);
+    runtime.dispose();
+  });
+
+  it('ignores completions that cannot add to this gallery or that Queue already refreshes for', () => {
+    const runtime = createRuntime();
+
+    runtime.start();
+    // Another user's item, as an admin receives it, and as everyone else receives it.
+    emitStatusChanged({ user_id: 'bob' });
+    emitStatusChanged({ origin: null, user_id: 'redacted' });
+    // An item this client enqueued, which Queue refreshes for in the tab tracking it.
+    emitStatusChanged({ origin: CLIENT_ORIGIN });
+
+    for (const status of ['pending', 'in_progress', 'waiting', 'failed', 'canceled'] as const) {
+      emitStatusChanged({ status });
+    }
+
+    vi.advanceTimersByTime(1000);
+
+    expect(invalidate).not.toHaveBeenCalled();
+    runtime.dispose();
+  });
+
+  it("refreshes for the server's system user when a single-user session cannot fence by user", () => {
+    sessionUserId = null;
+    const runtime = createRuntime();
+
+    runtime.start();
+    emitStatusChanged({ user_id: 'system' });
+    vi.advanceTimersByTime(250);
+
+    expect(invalidate).toHaveBeenCalledTimes(1);
+    runtime.dispose();
+  });
+
+  it('attributes each completion to the session user at the time it arrives', () => {
+    sessionUserId = 'alice';
+    const runtime = createRuntime();
+
+    runtime.start();
+    sessionUserId = 'bob';
+    emitStatusChanged({ user_id: 'alice' });
+    vi.advanceTimersByTime(1000);
+
+    expect(invalidate).not.toHaveBeenCalled();
+
+    emitStatusChanged({ user_id: 'bob' });
+    vi.advanceTimersByTime(250);
+
+    expect(invalidate).toHaveBeenCalledTimes(1);
+    runtime.dispose();
+  });
 });

invokeai/frontend/webv2/src/features/gallery/data/realtimeRuntime.ts

@@ -7,24 +7,38 @@ export interface GalleryRealtimeBackend {
   onConnectionChange(handler: (status: BackendConnectionStatus) => void): () => void;
 }
 
+/** The fields of the backend's `queue_item_status_changed` payload that Gallery reads. */
+interface QueueItemStatusChangedPayload {
+  origin: string | null;
+  status: string;
+  user_id: string;
+}
+
 export interface GalleryRealtimeRuntime {
   dispose(): void;
   start(): void;
 }
 
 /**
- * Refresh external uploads and reconnects at least minIntervalMs apart. Coalesce pending events and guarantee a
- * trailing pass without repeatedly cancelling page fetches.
+ * Refresh external uploads, the session user's completions enqueued by other clients, and reconnects at least
+ * minIntervalMs apart. Coalesce pending events and guarantee a trailing pass without repeatedly cancelling page
+ * fetches.
  */
 export const createGalleryRealtimeRuntime = ({
   backend,
   coalesceMs = 250,
+  getSessionUserId,
   invalidate,
+  isClientQueueItemOrigin,
   minIntervalMs = 1000,
 }: {
   backend: GalleryRealtimeBackend;
   coalesceMs?: number;
+  /** Null when events cannot be attributed to a user (single-user servers tag every event `system`). */
+  getSessionUserId: () => string | null;
   invalidate: () => void | Promise<void>;
+  /** True for origins of items this client enqueues; Queue refreshes for those in the tab that tracks them. */
+  isClientQueueItemOrigin: (origin: string | null) => boolean;
   minIntervalMs?: number;
 }): GalleryRealtimeRuntime => {
   const owner = captureAccountScope();
@@ -53,6 +67,22 @@ export const createGalleryRealtimeRuntime = ({
     }, delay);
   };
 
+  const handleQueueItemStatusChanged = (event: QueueItemStatusChangedPayload): void => {
+    // Admins also receive other users' full events and everyone else's arrive redacted, so only the session user's
+    // completions can add to this gallery. Items this client enqueued are left to Queue, which refreshes only in the
+    // tab tracking the item: another tab or device of the same user is not refreshed for them. Known limitation,
+    // left intentionally: refreshing here too would double-refresh the enqueuing tab.
+    const sessionUserId = getSessionUserId();
+
+    if (
+      event.status === 'completed' &&
+      !isClientQueueItemOrigin(event.origin) &&
+      (sessionUserId === null || event.user_id === sessionUserId)
+    ) {
+      scheduleInvalidation();
+    }
+  };
+
   const start = (): void => {
     if (isStarted || !isAccountScopeCurrent(owner)) {
       return;
@@ -67,6 +97,7 @@ export const createGalleryRealtimeRuntime = ({
     detachers.push(
       backend.on('image_uploaded', scheduleInvalidation),
       backend.on('video_uploaded', scheduleInvalidation),
+      backend.on('queue_item_status_changed', handleQueueItemStatusChanged),
       backend.onConnectionChange((status) => {
         const isReconnect = previousStatus === 'disconnected' && status === 'connected';
 

invokeai/frontend/webv2/src/features/identity/index.ts

@@ -6,6 +6,7 @@ export {
   ensureAuthSession,
   ensureReadyAuthSession,
   getAuthSession,
+  getSessionUserId,
   getUserStorageScope,
   identityTransportAuthAdapter,
   isLoginAttemptSupersededError,

invokeai/frontend/webv2/src/features/identity/session.ts

@@ -80,6 +80,15 @@ export const useAuthSession = (): AuthSession => store.useSnapshot();
 /** Workbench imperative reads follow route-guard resolution; changing users remounts the route. */
 export const getAuthSession = (): AuthSession => store.getSnapshot();
 
+/**
+ * The user id that owns this session's socket events. Single-user servers tag events with a system user id, so
+ * only multi-user sessions can be fenced.
+ */
+export const getSessionUserId = (): string | null => {
+  const session = getAuthSession();
+  return session.multiuserEnabled ? (session.user?.user_id ?? null) : null;
+};
+
 /** Stable subscription for App-composed capability read ports. */
 export const subscribeAuthSession = (listener: () => void): (() => void) => store.subscribe(listener);
 

invokeai/frontend/webv2/src/features/queue/contracts.ts

@@ -63,6 +63,7 @@ export {
   buildUtilityQueueItemOrigin,
   isTerminalBackendStatus,
   isUtilityQueueItemOrigin,
+  isWebv2QueueItemOrigin,
   parseQueueItemOrigin,
   parseQueueItemOriginProjectId,
   type BackendSocketEvents,

invokeai/frontend/webv2/src/features/queue/data/events.test.ts

@@ -0,0 +1,16 @@
+import { describe, expect, it } from 'vitest';
+
+import { buildQueueItemOrigin, buildUtilityQueueItemOrigin, isWebv2QueueItemOrigin } from './events';
+
+describe('isWebv2QueueItemOrigin', () => {
+  it('recognizes every origin this client builds and no other client', () => {
+    // Gallery skips these completions; an unrecognized form would double-refresh, a false match would miss one.
+    expect(isWebv2QueueItemOrigin(buildQueueItemOrigin('local-1'))).toBe(true);
+    expect(isWebv2QueueItemOrigin(buildQueueItemOrigin('local-1', 'project-1'))).toBe(true);
+    expect(isWebv2QueueItemOrigin(buildUtilityQueueItemOrigin('filter-1'))).toBe(true);
+
+    expect(isWebv2QueueItemOrigin('canvas')).toBe(false);
+    expect(isWebv2QueueItemOrigin('chatbot:webv2:1')).toBe(false);
+    expect(isWebv2QueueItemOrigin(null)).toBe(false);
+  });
+});

invokeai/frontend/webv2/src/features/queue/data/events.ts

@@ -152,6 +152,10 @@ export const buildQueueItemOrigin = (localQueueItemId: string, projectId?: strin
 export const buildUtilityQueueItemOrigin = (utilityId: string): string =>
   `${UTILITY_QUEUE_ITEM_ORIGIN_PREFIX}${utilityId}`;
 
+/** True for every origin this client builds: local, project and utility queue items. */
+export const isWebv2QueueItemOrigin = (origin: string | null | undefined): boolean =>
+  origin?.startsWith(QUEUE_ITEM_ORIGIN_PREFIX) ?? false;
+
 /** True when `origin` belongs to the utility queue (never a project/local queue item). */
 export const isUtilityQueueItemOrigin = (origin: string | null | undefined): boolean =>
   origin?.startsWith(UTILITY_QUEUE_ITEM_ORIGIN_PREFIX) ?? false;