diff --git a/invokeai/frontend/web/src/services/events/onInvocationComplete.test.ts b/invokeai/frontend/web/src/services/events/onInvocationComplete.test.ts new file mode 100644 index 00000000000..512c6bce203 --- /dev/null +++ b/invokeai/frontend/web/src/services/events/onInvocationComplete.test.ts @@ -0,0 +1,105 @@ +import type { AppDispatch, AppGetState } from 'app/store/store'; +import type { S } from 'services/api/types'; +import { $lastProgressEvent } from 'services/events/stores'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +import { buildOnInvocationComplete } from './onInvocationComplete'; + +vi.mock('app/logging/logger', () => ({ + logger: () => ({ + debug: vi.fn(), + trace: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }), +})); + +// Spies for the workflow-editor node execution side effects. Kept lazy so vi.mock hoisting is safe. +const mockUpsertExecutionState = vi.fn(); +let mockNodeExecutionStates: Record = {}; +vi.mock('features/nodes/hooks/useNodeExecutionState', () => ({ + $nodeExecutionStates: { get: () => mockNodeExecutionStates }, + upsertExecutionState: (...args: unknown[]) => mockUpsertExecutionState(...args), +})); + +const mockGetUpdatedNodeExecutionState = vi.fn(); +vi.mock('services/events/nodeExecutionState', () => ({ + getUpdatedNodeExecutionStateOnInvocationComplete: (...args: unknown[]) => mockGetUpdatedNodeExecutionState(...args), +})); + +const buildEvent = (overrides: Partial = {}): S['InvocationCompleteEvent'] => + ({ + invocation: { type: 'some_node', id: 'inv-1' }, + invocation_source_id: 'node-1', + result: {}, + origin: 'workflows', + destination: 'canvas', + user_id: 'owner-1', + item_id: 1, + ...overrides, + }) as unknown as S['InvocationCompleteEvent']; + +const makeGetState = (user: { user_id: string } | null): AppGetState => + (() => ({ auth: { user } })) as unknown as AppGetState; + +const seedProgressEvent = () => { + const progress = { queue_id: 'default', item_id: 1 } as unknown as S['InvocationProgressEvent']; + $lastProgressEvent.set(progress); + return progress; +}; + +beforeEach(() => { + mockUpsertExecutionState.mockReset(); + mockGetUpdatedNodeExecutionState.mockReset(); + mockNodeExecutionStates = {}; + $lastProgressEvent.set(null); +}); + +describe('buildOnInvocationComplete cross-user isolation', () => { + it("ignores another user's completion event entirely (admin receiving via admin room)", async () => { + // An admin has a local workflow node with the same source id in progress, and a live progress event. + mockNodeExecutionStates = { 'node-1': { nodeId: 'node-1', status: 'IN_PROGRESS' } }; + mockGetUpdatedNodeExecutionState.mockReturnValue({ nodeId: 'node-1', status: 'COMPLETED' }); + const seededProgress = seedProgressEvent(); + + const dispatch = vi.fn() as unknown as AppDispatch; + const handler = buildOnInvocationComplete(makeGetState({ user_id: 'admin-1' }), dispatch, new Map()); + + // The event belongs to a different user. + await handler(buildEvent({ user_id: 'other-user', origin: 'canvas_workflow_integration' })); + + // No node execution state read/upsert, no canvas/gallery dispatches, progress event untouched. + expect(mockGetUpdatedNodeExecutionState).not.toHaveBeenCalled(); + expect(mockUpsertExecutionState).not.toHaveBeenCalled(); + expect(dispatch).not.toHaveBeenCalled(); + expect($lastProgressEvent.get()).toBe(seededProgress); + }); + + it("processes the client's own completion event", async () => { + mockNodeExecutionStates = { 'node-1': { nodeId: 'node-1', status: 'IN_PROGRESS' } }; + mockGetUpdatedNodeExecutionState.mockReturnValue({ nodeId: 'node-1', status: 'COMPLETED' }); + seedProgressEvent(); + + const dispatch = vi.fn() as unknown as AppDispatch; + const handler = buildOnInvocationComplete(makeGetState({ user_id: 'owner-1' }), dispatch, new Map()); + + await handler(buildEvent({ user_id: 'owner-1' })); + + expect(mockUpsertExecutionState).toHaveBeenCalledWith('node-1', expect.objectContaining({ nodeId: 'node-1' })); + // The handler clears the progress event at the end when it processes an event it owns. + expect($lastProgressEvent.get()).toBeNull(); + }); + + it('processes events in single-user mode where there is no authenticated user', async () => { + mockNodeExecutionStates = { 'node-1': { nodeId: 'node-1', status: 'IN_PROGRESS' } }; + mockGetUpdatedNodeExecutionState.mockReturnValue({ nodeId: 'node-1', status: 'COMPLETED' }); + + const dispatch = vi.fn() as unknown as AppDispatch; + const handler = buildOnInvocationComplete(makeGetState(null), dispatch, new Map()); + + await handler(buildEvent({ user_id: 'system' })); + + expect(mockUpsertExecutionState).toHaveBeenCalledWith('node-1', expect.objectContaining({ nodeId: 'node-1' })); + }); +}); diff --git a/invokeai/frontend/web/src/services/events/onInvocationComplete.tsx b/invokeai/frontend/web/src/services/events/onInvocationComplete.tsx index ea6a237d4b2..33e97452d80 100644 --- a/invokeai/frontend/web/src/services/events/onInvocationComplete.tsx +++ b/invokeai/frontend/web/src/services/events/onInvocationComplete.tsx @@ -1,5 +1,6 @@ import { logger } from 'app/logging/logger'; import type { AppDispatch, AppGetState } from 'app/store/store'; +import { selectCurrentUser } from 'features/auth/store/authSlice'; import { canvasWorkflowIntegrationProcessingCompleted } from 'features/controlLayers/store/canvasWorkflowIntegrationSlice'; import { selectAutoSwitch, @@ -247,6 +248,22 @@ export const buildOnInvocationComplete = ( return async (data: S['InvocationCompleteEvent']) => { log.debug({ data } as JsonObject, `Invocation complete (${data.invocation.type}, ${data.invocation_source_id})`); + // In multiuser mode, admins are subscribed to the "admin" socket room and therefore receive + // invocation events for *every* user, not just their own. Another user's completion event must + // not touch this client's state: it could mark our workflow-editor nodes complete and append the + // other user's outputs to matching node IDs, stop our canvas workflow integration spinner, insert + // their images into our gallery, or clear our own in-progress progress event. Bail out entirely + // before any of those side effects run. + // + // Only gate this when we actually know who is logged in. In single-user mode there is no + // authenticated user (selectCurrentUser is null) and every event is the local user's own, so + // we fall through and preserve the original behavior. + const currentUser = selectCurrentUser(getState()); + if (currentUser && data.user_id !== currentUser.user_id) { + log.trace(`Skipping invocation complete for another user (${data.user_id})`); + return; + } + const nodeExecutionState = $nodeExecutionStates.get()[data.invocation_source_id]; const updatedNodeExecutionState = getUpdatedNodeExecutionStateOnInvocationComplete( nodeExecutionState, diff --git a/invokeai/frontend/web/src/services/events/setEventListeners.tsx b/invokeai/frontend/web/src/services/events/setEventListeners.tsx index 30f705d670e..dbaf71ffeca 100644 --- a/invokeai/frontend/web/src/services/events/setEventListeners.tsx +++ b/invokeai/frontend/web/src/services/events/setEventListeners.tsx @@ -536,6 +536,16 @@ export const setEventListeners = ({ socket, store, setIsConnected }: SetEventLis } }); + socket.on('queue_counts_changed', (data) => { + // A content-free broadcast that the queue's global counts may have changed — typically + // because *another* user enqueued or one of their jobs changed status. We never receive + // those users' private queue item events, so this is the only signal that lets a non-admin's + // badge keep its global total (the "/Y" in "X/Y") and the in_progress-driven progress bar + // current. Only the redacted SessionQueueStatus is refetched; no per-user/batch caches. + log.trace({ data }, 'Queue counts changed'); + dispatch(queueApi.util.invalidateTags(['SessionQueueStatus'])); + }); + socket.on('queue_cleared', (data) => { log.debug({ data }, 'Queue cleared'); dispatch( diff --git a/invokeai/frontend/web/src/services/events/types.ts b/invokeai/frontend/web/src/services/events/types.ts index 72239986fb3..89c2bcb5d3c 100644 --- a/invokeai/frontend/web/src/services/events/types.ts +++ b/invokeai/frontend/web/src/services/events/types.ts @@ -30,6 +30,11 @@ export type ServerToClientEvents = { model_install_cancelled: (payload: S['ModelInstallCancelledEvent']) => void; model_load_complete: (payload: S['ModelLoadCompleteEvent']) => void; queue_item_status_changed: (payload: S['QueueItemStatusChangedEvent']) => void; + // Content-free broadcast to the whole queue room: the global queue counts may have changed + // (some user enqueued or a job changed status). Carries only queue_id — no per-user data — + // so every subscriber can refetch the redacted queue status. Emitted by the backend socket + // layer, not the event bus, so it has no generated `S[...]` schema type. + queue_counts_changed: (payload: { queue_id: string }) => void; queue_cleared: (payload: S['QueueClearedEvent']) => void; batch_enqueued: (payload: S['BatchEnqueuedEvent']) => void; queue_items_retried: (payload: S['QueueItemsRetriedEvent']) => void;