diff --git a/invokeai/app/api/sockets.py b/invokeai/app/api/sockets.py index 70ca97740cf..bdaf641c1c1 100644 --- a/invokeai/app/api/sockets.py +++ b/invokeai/app/api/sockets.py @@ -306,7 +306,10 @@ async def _handle_queue_event(self, event: FastAPIEvent[QueueEventBase]): per-session side effects. InvocationEventBase events stay private (owner + admins only). RecallParametersUpdatedEvent - is also private. QueueClearedEvent has no user identity and is broadcast to the queue room. + is also private. QueueClearedEvent is broadcast to the queue room when unscoped (an admin or + single-user clear that deleted every user's items); a user-scoped clear goes full to + owner + admins with a sanitized companion to the rest of the queue room, so other users + refresh their queue lists without treating the clear as their own. IMPORTANT: Check InvocationEventBase BEFORE QueueItemEventBase since InvocationEventBase inherits from QueueItemEventBase. The order of isinstance checks matters! @@ -445,9 +448,37 @@ async def _handle_queue_event(self, event: FastAPIEvent[QueueEventBase]): f"Emitted private queue_items_retried event to user rooms {event_data.user_ids} and admin room" ) + # QueueClearedEvent: an unscoped clear (user_id=None — admin or single-user mode) + # deleted every user's items, so everyone gets the full event. A user-scoped clear + # only deleted that user's rows: full event to owner+admin (single emit to a room + # list so a socket in both rooms receives it once), sanitized companion to the rest + # of the queue room so their queue lists and badge counts refetch without treating + # the clear as their own. + elif isinstance(event_data, QueueClearedEvent): + if event_data.user_id is None: + await self._sio.emit( + event=event_name, data=event_data.model_dump(mode="json"), room=event_data.queue_id + ) + logger.debug(f"Emitted unscoped queue_cleared to all subscribers in queue {event_data.queue_id}") + else: + user_room = f"user:{event_data.user_id}" + await self._sio.emit( + event=event_name, data=event_data.model_dump(mode="json"), room=[user_room, "admin"] + ) + sanitized = event_data.model_copy(update={"user_id": "redacted"}) + await self._sio.emit( + event=event_name, + data=sanitized.model_dump(mode="json"), + room=event_data.queue_id, + skip_sid=self._owner_and_admin_sids(event_data.user_id), + ) + logger.debug( + f"Emitted queue_cleared: full to {user_room}+admin, sanitized to queue {event_data.queue_id}" + ) + else: - # For remaining queue events (e.g. QueueClearedEvent) that do not - # carry user identity, emit to all subscribers in the queue room. + # For remaining queue events that do not carry user identity, + # emit to all subscribers in the queue room. await self._sio.emit( event=event_name, data=event_data.model_dump(mode="json"), room=event_data.queue_id ) diff --git a/invokeai/app/services/events/events_base.py b/invokeai/app/services/events/events_base.py index 0e1f71c2bc7..65be86f8399 100644 --- a/invokeai/app/services/events/events_base.py +++ b/invokeai/app/services/events/events_base.py @@ -113,9 +113,9 @@ def emit_queue_items_retried( """Emitted when a list of queue items are retried""" self.dispatch(QueueItemsRetriedEvent.build(retry_result, user_ids, retried_item_ids_by_user)) - def emit_queue_cleared(self, queue_id: str) -> None: - """Emitted when a queue is cleared""" - self.dispatch(QueueClearedEvent.build(queue_id)) + def emit_queue_cleared(self, queue_id: str, user_id: str | None = None) -> None: + """Emitted when a queue is cleared. `user_id` scopes the clear to one user's items; None means all.""" + self.dispatch(QueueClearedEvent.build(queue_id, user_id)) def emit_recall_parameters_updated(self, queue_id: str, user_id: str, parameters: dict) -> None: """Emitted when recall parameters are updated""" diff --git a/invokeai/app/services/events/events_common.py b/invokeai/app/services/events/events_common.py index 005b0e5435f..15b48bb4665 100644 --- a/invokeai/app/services/events/events_common.py +++ b/invokeai/app/services/events/events_common.py @@ -331,9 +331,14 @@ class QueueClearedEvent(QueueEventBase): __event_name__ = "queue_cleared" + user_id: str | None = Field( + default=None, + description="The ID of the user whose queue items were cleared, or None if all users' items were cleared", + ) + @classmethod - def build(cls, queue_id: str) -> "QueueClearedEvent": - return cls(queue_id=queue_id) + def build(cls, queue_id: str, user_id: str | None = None) -> "QueueClearedEvent": + return cls(queue_id=queue_id, user_id=user_id) class WorkflowEventBase(EventBase): diff --git a/invokeai/app/services/session_processor/workflow_call_runtime.py b/invokeai/app/services/session_processor/workflow_call_runtime.py index dba16650487..47576475101 100644 --- a/invokeai/app/services/session_processor/workflow_call_runtime.py +++ b/invokeai/app/services/session_processor/workflow_call_runtime.py @@ -8,7 +8,11 @@ ) from invokeai.app.invocations.workflow_return import WorkflowReturnOutput from invokeai.app.services.session_processor.workflow_call_batch import build_child_workflow_session_results -from invokeai.app.services.session_queue.session_queue_common import NodeFieldValue, SessionQueueItem +from invokeai.app.services.session_queue.session_queue_common import ( + NodeFieldValue, + SessionQueueItem, + SessionQueueItemNotFoundError, +) from invokeai.app.services.shared.graph import GraphExecutionState if TYPE_CHECKING: @@ -175,15 +179,19 @@ def fail_waiting_workflow_call(self, queue_item: SessionQueueItem, error_message error_traceback=error_message, ) - def _get_parent_queue_item(self, child_queue_item: SessionQueueItem) -> SessionQueueItem: + def _get_parent_queue_item(self, child_queue_item: SessionQueueItem) -> SessionQueueItem | None: parent_item_id = child_queue_item.parent_item_id if parent_item_id is None: raise ValueError("Child workflow queue item is missing parent_item_id metadata.") - return self._session_runner._services.session_queue.get_queue_item(parent_item_id) + try: + return self._session_runner._services.session_queue.get_queue_item(parent_item_id) + except SessionQueueItemNotFoundError: + # The parent may have been deleted while the child was running (e.g. the queue was cleared). + return None def _resume_parent_from_completed_child(self, child_queue_item: SessionQueueItem) -> None: parent_queue_item = self._get_parent_queue_item(child_queue_item) - if parent_queue_item.status in ("completed", "failed", "canceled"): + if parent_queue_item is None or parent_queue_item.status in ("completed", "failed", "canceled"): return try: output = self.get_child_workflow_return_output(child_queue_item.session) @@ -200,7 +208,12 @@ def _resume_parent_from_completed_child(self, child_queue_item: SessionQueueItem exclude_item_ids={child_queue_item.item_id}, ) self.fail_waiting_workflow_call(parent_queue_item, str(e)) - parent_queue_item = self._session_runner._services.session_queue.get_queue_item(parent_queue_item.item_id) + try: + parent_queue_item = self._session_runner._services.session_queue.get_queue_item( + parent_queue_item.item_id + ) + except SessionQueueItemNotFoundError: + return if getattr(parent_queue_item, "parent_item_id", None) is not None: self._fail_parent_from_failed_child(parent_queue_item) return @@ -229,7 +242,7 @@ def _resume_parent_from_completed_child(self, child_queue_item: SessionQueueItem def _fail_parent_from_failed_child(self, child_queue_item: SessionQueueItem) -> None: parent_queue_item = self._get_parent_queue_item(child_queue_item) - if parent_queue_item.status in ("completed", "failed", "canceled"): + if parent_queue_item is None or parent_queue_item.status in ("completed", "failed", "canceled"): return waiting_frame = parent_queue_item.session.waiting_workflow_call if waiting_frame is None: @@ -244,24 +257,38 @@ def _fail_parent_from_failed_child(self, child_queue_item: SessionQueueItem) -> f"The selected saved workflow '{waiting_frame.workflow_id}' failed during child execution." ) self.fail_waiting_workflow_call(parent_queue_item, child_error_message) - parent_queue_item = self._session_runner._services.session_queue.get_queue_item(parent_queue_item.item_id) + try: + parent_queue_item = self._session_runner._services.session_queue.get_queue_item(parent_queue_item.item_id) + except SessionQueueItemNotFoundError: + return if getattr(parent_queue_item, "parent_item_id", None) is not None: self._fail_parent_from_failed_child(parent_queue_item) def _cancel_parent_from_canceled_child(self, child_queue_item: SessionQueueItem) -> None: parent_queue_item = self._get_parent_queue_item(child_queue_item) - if parent_queue_item.status == "canceled": + if parent_queue_item is None or parent_queue_item.status == "canceled": return self._session_runner._services.session_queue.cancel_queue_item(parent_queue_item.item_id) def run_queue_item(self, queue_item: SessionQueueItem) -> None: self._session_runner.run(queue_item) - updated_queue_item = self._session_runner._services.session_queue.get_queue_item(queue_item.item_id) + try: + updated_queue_item = self._session_runner._services.session_queue.get_queue_item(queue_item.item_id) + except SessionQueueItemNotFoundError: + # The queue item was deleted while it was running (e.g. the queue was cleared or the current item + # was deleted). There is no parent lifecycle work to do for a deleted item. + return if getattr(updated_queue_item, "parent_item_id", None) is None: return - if updated_queue_item.status == "completed": - self._resume_parent_from_completed_child(updated_queue_item) - elif updated_queue_item.status == "failed": - self._fail_parent_from_failed_child(updated_queue_item) - elif updated_queue_item.status == "canceled": - self._cancel_parent_from_canceled_child(updated_queue_item) + try: + if updated_queue_item.status == "completed": + self._resume_parent_from_completed_child(updated_queue_item) + elif updated_queue_item.status == "failed": + self._fail_parent_from_failed_child(updated_queue_item) + elif updated_queue_item.status == "canceled": + self._cancel_parent_from_canceled_child(updated_queue_item) + except SessionQueueItemNotFoundError: + # An item in the parent chain was deleted after we looked it up but before we mutated it + # (the queue mutations re-read the row and raise when it has disappeared). The chain is + # being torn down; there is nothing left to update. + return diff --git a/invokeai/app/services/session_queue/session_queue_sqlite.py b/invokeai/app/services/session_queue/session_queue_sqlite.py index 53079322899..51980d68f3b 100644 --- a/invokeai/app/services/session_queue/session_queue_sqlite.py +++ b/invokeai/app/services/session_queue/session_queue_sqlite.py @@ -524,7 +524,7 @@ def clear(self, queue_id: str, user_id: Optional[str] = None) -> ClearResult: """, tuple(params), ) - self.__invoker.services.events.emit_queue_cleared(queue_id) + self.__invoker.services.events.emit_queue_cleared(queue_id, user_id) return ClearResult(deleted=count) def delete_queue_items_by_id(self, item_ids: list[int]) -> None: diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index 52ca3ecfcd5..d66ec1f7a01 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -59604,9 +59604,22 @@ "description": "The ID of the queue", "title": "Queue Id", "type": "string" + }, + "user_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "The ID of the user whose queue items were cleared, or None if all users' items were cleared", + "title": "User Id" } }, - "required": ["timestamp", "queue_id"], + "required": ["timestamp", "queue_id", "user_id"], "title": "QueueClearedEvent", "type": "object" }, diff --git a/invokeai/frontend/web/src/common/hooks/useGlobalHotkeys.ts b/invokeai/frontend/web/src/common/hooks/useGlobalHotkeys.ts index dd43c0b0947..070a4866a50 100644 --- a/invokeai/frontend/web/src/common/hooks/useGlobalHotkeys.ts +++ b/invokeai/frontend/web/src/common/hooks/useGlobalHotkeys.ts @@ -1,8 +1,8 @@ import { useAppStore } from 'app/store/storeHooks'; import { useDeleteImageModalApi } from 'features/deleteImageModal/store/state'; import { selectSelection } from 'features/gallery/store/gallerySelectors'; +import { useCancelCurrentQueueItem } from 'features/queue/hooks/useCancelCurrentQueueItem'; import { useClearQueue } from 'features/queue/hooks/useClearQueue'; -import { useDeleteCurrentQueueItem } from 'features/queue/hooks/useDeleteCurrentQueueItem'; import { useInvoke } from 'features/queue/hooks/useInvoke'; import { useRegisteredHotkeys } from 'features/system/components/HotkeysModal/useHotkeyData'; import { navigationApi } from 'features/ui/layouts/navigation-api'; @@ -37,17 +37,19 @@ export const useGlobalHotkeys = () => { dependencies: [queue], }); - const deleteCurrentQueueItem = useDeleteCurrentQueueItem(); + const cancelCurrentQueueItem = useCancelCurrentQueueItem(); useRegisteredHotkeys({ id: 'cancelQueueItem', category: 'app', - callback: deleteCurrentQueueItem.trigger, + callback: () => { + cancelCurrentQueueItem.trigger(); + }, options: { - enabled: !deleteCurrentQueueItem.isDisabled && !deleteCurrentQueueItem.isLoading, + enabled: !cancelCurrentQueueItem.isDisabled && !cancelCurrentQueueItem.isLoading, preventDefault: true, }, - dependencies: [deleteCurrentQueueItem], + dependencies: [cancelCurrentQueueItem], }); const clearQueue = useClearQueue(); diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index 3f0f07c07c2..2cec6cbaa8b 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -25206,6 +25206,12 @@ export type components = { * @description The ID of the queue */ queue_id: string; + /** + * User Id + * @description The ID of the user whose queue items were cleared, or None if all users' items were cleared + * @default null + */ + user_id: string | null; }; /** * QueueItemStatusChangedEvent diff --git a/invokeai/frontend/web/src/services/events/setEventListeners.tsx b/invokeai/frontend/web/src/services/events/setEventListeners.tsx index 30f705d670e..21f475f46b8 100644 --- a/invokeai/frontend/web/src/services/events/setEventListeners.tsx +++ b/invokeai/frontend/web/src/services/events/setEventListeners.tsx @@ -4,6 +4,7 @@ import { socketConnected } from 'app/store/middleware/listenerMiddleware/listene import type { AppStore } from 'app/store/store'; import { parseify } from 'common/util/serialize'; import { isNil, round } from 'es-toolkit/compat'; +import { selectCurrentUser } from 'features/auth/store/authSlice'; import { getDefaultRefImageConfig } from 'features/controlLayers/hooks/addLayerHooks'; import { allEntitiesDeleted, controlLayerRecalled } from 'features/controlLayers/store/canvasSlice'; import { canvasWorkflowIntegrationProcessingCompleted } from 'features/controlLayers/store/canvasWorkflowIntegrationSlice'; @@ -71,6 +72,7 @@ export const setEventListeners = ({ socket, store, setIsConnected }: SetEventLis clearCanvasWorkflowIntegrationProcessing: () => dispatch(canvasWorkflowIntegrationProcessingCompleted()), completedInvocationKeysByItemId, getAllNodeExecutionStates: () => $nodeExecutionStates.get(), + getCurrentUserId: () => selectCurrentUser(getState())?.user_id ?? null, getNodeExecutionState: (nodeId) => $nodeExecutionStates.get()[nodeId], logReconciliationError: (error, itemId) => { log.debug({ error: parseify(error) }, `Unable to reconcile workflow queue item ${itemId}`); @@ -538,6 +540,17 @@ export const setEventListeners = ({ socket, store, setIsConnected }: SetEventLis socket.on('queue_cleared', (data) => { log.debug({ data }, 'Queue cleared'); + // Clearing the queue deletes the in-progress item without emitting a per-item terminal status + // event, so the progress bar must be reset here — and the coordinator must mark the deleted + // tracked items terminal so a trailing invocation_progress event cannot repopulate the bar. + // The coordinator scopes a user-scoped clear (multiuser mode) to that user's items — on an + // admin client that may be a subset of the tracked items; on another user's client it is + // none of them — and reports whether the clear applied to any tracked item, so the progress + // bar is only reset when the clear could have deleted the item behind it. The queue tags + // below always need refreshing. + if (workflowExecutionCoordinator.onQueueCleared(data)) { + $lastProgressEvent.set(null); + } dispatch( queueApi.util.invalidateTags([ 'SessionQueueStatus', diff --git a/invokeai/frontend/web/src/services/events/workflowExecutionCoordinator.test.ts b/invokeai/frontend/web/src/services/events/workflowExecutionCoordinator.test.ts index 00c48c8a373..42e3e5edd25 100644 --- a/invokeai/frontend/web/src/services/events/workflowExecutionCoordinator.test.ts +++ b/invokeai/frontend/web/src/services/events/workflowExecutionCoordinator.test.ts @@ -30,6 +30,7 @@ const buildQueueStatusEvent = ( batch_id: 'batch-1', origin: 'workflows', destination: 'gallery', + user_id: 'user-1', status: 'completed', status_sequence: 1, batch_status: { @@ -71,6 +72,28 @@ const buildInvocationStartedEvent = ( ...overrides, }) as S['InvocationStartedEvent']; +const buildInvocationProgressEvent = ( + 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: 'prepared-node-1', + type: 'test_node', + }, + message: 'denoising', + percentage: 0.5, + image: null, + ...overrides, + }) as S['InvocationProgressEvent']; + const buildInvocationCompleteEvent = ( overrides: Partial = {} ): S['InvocationCompleteEvent'] => @@ -116,6 +139,13 @@ const buildInvocationErrorEvent = (overrides: Partial ...overrides, }) as S['InvocationErrorEvent']; +const buildQueueClearedEvent = (overrides: Partial = {}): S['QueueClearedEvent'] => + ({ + queue_id: 'default', + user_id: null, + ...overrides, + }) as S['QueueClearedEvent']; + const buildQueueItem = (status: S['SessionQueueItem']['status']): S['SessionQueueItem'] => ({ item_id: 1, @@ -143,7 +173,7 @@ const buildQueueItem = (status: S['SessionQueueItem']['status']): S['SessionQueu }, }) as unknown as S['SessionQueueItem']; -const createCoordinatorHarness = () => { +const createCoordinatorHarness = (currentUserId: string | null = 'user-1') => { const completedInvocationKeysByItemId = new Map>(); const nodeExecutionStates: Record = {}; const onInvocationComplete = vi.fn(); @@ -155,6 +185,7 @@ const createCoordinatorHarness = () => { clearCanvasWorkflowIntegrationProcessing, completedInvocationKeysByItemId, getAllNodeExecutionStates: () => nodeExecutionStates, + getCurrentUserId: () => currentUserId, getNodeExecutionState: (nodeId) => nodeExecutionStates[nodeId], logReconciliationError, onInvocationComplete, @@ -286,6 +317,146 @@ describe(createWorkflowExecutionCoordinator.name, () => { expect(queueItemRequests.size).toBe(1); }); + it('rejects trailing invocation progress after a queue item is canceled', () => { + const { coordinator } = createCoordinatorHarness(); + + coordinator.onQueueItemStatusChanged(buildQueueStatusEvent({ item_id: 1, status: 'in_progress', origin: null })); + expect(coordinator.onInvocationProgress(buildInvocationProgressEvent({ item_id: 1 }))).toBe(true); + + coordinator.onQueueItemStatusChanged(buildQueueStatusEvent({ item_id: 1, status: 'canceled', origin: null })); + + // A denoise step callback may race the cancelation and emit progress after the terminal status + // event. It must not be applied - it would repopulate the progress bar after it was cleared. + expect(coordinator.onInvocationProgress(buildInvocationProgressEvent({ item_id: 1 }))).toBe(false); + }); + + it('rejects trailing invocation progress after the queue is cleared', () => { + const { coordinator, queueItemRequests } = createCoordinatorHarness(); + + coordinator.onQueueItemStatusChanged(buildQueueStatusEvent({ item_id: 1, status: 'in_progress', origin: null })); + expect(coordinator.onInvocationProgress(buildInvocationProgressEvent({ item_id: 1 }))).toBe(true); + + // Clearing the queue deletes items without emitting per-item terminal status events, so the + // coordinator is told directly. A workflow item with a pending reconciliation is also tracked + // to verify the reconciliation is aborted. + coordinator.onQueueItemStatusChanged( + buildQueueStatusEvent({ item_id: 2, status: 'completed', origin: 'workflows' }) + ); + expect(coordinator.onQueueCleared(buildQueueClearedEvent({ user_id: null }))).toBe(true); + + expect(coordinator.onInvocationProgress(buildInvocationProgressEvent({ item_id: 1 }))).toBe(false); + expect(queueItemRequests.get(2)?.abort).toHaveBeenCalled(); + }); + + it('applies a queue clear scoped to the current user', () => { + const { coordinator, queueItemRequests } = createCoordinatorHarness('user-1'); + + coordinator.onQueueItemStatusChanged(buildQueueStatusEvent({ item_id: 1, status: 'in_progress', origin: null })); + expect(coordinator.onInvocationProgress(buildInvocationProgressEvent({ item_id: 1 }))).toBe(true); + coordinator.onQueueItemStatusChanged( + buildQueueStatusEvent({ item_id: 2, status: 'completed', origin: 'workflows' }) + ); + + expect(coordinator.onQueueCleared(buildQueueClearedEvent({ user_id: 'user-1' }))).toBe(true); + + expect(coordinator.onInvocationProgress(buildInvocationProgressEvent({ item_id: 1 }))).toBe(false); + expect(queueItemRequests.get(2)?.abort).toHaveBeenCalled(); + }); + + it("does not disturb this user's items when another user's scoped clear is broadcast", async () => { + // In multiuser mode a user-scoped clear only deletes that user's rows, but the queue_cleared + // event is broadcast to every queue subscriber (sanitized to user_id="redacted" for + // non-owner, non-admin recipients). This client's items were not touched: its in-flight + // reconciliation must complete and its in-progress item must keep accepting events. + const { coordinator, nodeExecutionStates, queueItemRequests } = createCoordinatorHarness('user-b'); + + coordinator.onQueueItemStatusChanged( + buildQueueStatusEvent({ item_id: 1, status: 'completed', origin: 'workflows', user_id: 'user-b' }) + ); + coordinator.onQueueItemStatusChanged( + buildQueueStatusEvent({ item_id: 2, status: 'in_progress', origin: null, user_id: 'user-b' }) + ); + expect(coordinator.onInvocationProgress(buildInvocationProgressEvent({ item_id: 2, user_id: 'user-b' }))).toBe( + true + ); + + expect(coordinator.onQueueCleared(buildQueueClearedEvent({ user_id: 'redacted' }))).toBe(false); + + expect(queueItemRequests.get(1)?.abort).not.toHaveBeenCalled(); + expect(coordinator.onInvocationProgress(buildInvocationProgressEvent({ item_id: 2, user_id: 'user-b' }))).toBe( + true + ); + + queueItemRequests.get(1)?.resolve(buildQueueItem('completed')); + await Promise.resolve(); + await Promise.resolve(); + + expect(nodeExecutionStates['node-1']?.status).toBe(zNodeStatus.enum.COMPLETED); + expect(nodeExecutionStates['node-1']?.outputs).toHaveLength(1); + }); + + it("applies another user's scoped clear to just that user's items on an admin client", async () => { + // Admins receive every user's events (and the full, non-redacted scoped queue_cleared event), + // so they may be tracking the cleared user's active item. The clear deletes that user's rows + // without per-item terminal events: the cleared user's tracked items must be marked terminal + // so trailing events are rejected, while the admin's own reconciliation and other users' live + // items are untouched. + const { coordinator, nodeExecutionStates, queueItemRequests } = createCoordinatorHarness('admin-1'); + + // The admin's own completed workflow item, reconciliation in flight. + coordinator.onQueueItemStatusChanged( + buildQueueStatusEvent({ item_id: 1, status: 'completed', origin: 'workflows', user_id: 'admin-1' }) + ); + // User A's in-progress item, actively emitting progress on the admin's client. + coordinator.onQueueItemStatusChanged( + buildQueueStatusEvent({ item_id: 2, status: 'in_progress', origin: null, user_id: 'user-a' }) + ); + expect(coordinator.onInvocationProgress(buildInvocationProgressEvent({ item_id: 2, user_id: 'user-a' }))).toBe( + true + ); + // User C's in-progress item, unaffected by user A's clear. + coordinator.onQueueItemStatusChanged( + buildQueueStatusEvent({ item_id: 3, status: 'in_progress', origin: null, user_id: 'user-c' }) + ); + expect(coordinator.onInvocationProgress(buildInvocationProgressEvent({ item_id: 3, user_id: 'user-c' }))).toBe( + true + ); + + expect(coordinator.onQueueCleared(buildQueueClearedEvent({ user_id: 'user-a' }))).toBe(true); + + // User A's item is gone: a trailing progress event must be rejected, not repopulate the bar. + expect(coordinator.onInvocationProgress(buildInvocationProgressEvent({ item_id: 2, user_id: 'user-a' }))).toBe( + false + ); + // User C's item is still live. + expect(coordinator.onInvocationProgress(buildInvocationProgressEvent({ item_id: 3, user_id: 'user-c' }))).toBe( + true + ); + + // The admin's own reconciliation was not aborted and completes normally. + expect(queueItemRequests.get(1)?.abort).not.toHaveBeenCalled(); + queueItemRequests.get(1)?.resolve(buildQueueItem('completed')); + await Promise.resolve(); + await Promise.resolve(); + + expect(nodeExecutionStates['node-1']?.status).toBe(zNodeStatus.enum.COMPLETED); + expect(nodeExecutionStates['node-1']?.outputs).toHaveLength(1); + }); + + it("aborts the cleared user's in-flight reconciliation on an admin client", () => { + // The scoped clear deleted the user's rows, so a reconciliation fetch for their terminal item + // can never succeed and must be aborted. + const { coordinator, queueItemRequests } = createCoordinatorHarness('admin-1'); + + coordinator.onQueueItemStatusChanged( + buildQueueStatusEvent({ item_id: 1, status: 'completed', origin: 'workflows', user_id: 'user-a' }) + ); + + expect(coordinator.onQueueCleared(buildQueueClearedEvent({ user_id: 'user-a' }))).toBe(true); + + expect(queueItemRequests.get(1)?.abort).toHaveBeenCalled(); + }); + it('still clears canvas workflow integration processing on late invocation errors', () => { const { clearCanvasWorkflowIntegrationProcessing, coordinator } = createCoordinatorHarness(); diff --git a/invokeai/frontend/web/src/services/events/workflowExecutionCoordinator.ts b/invokeai/frontend/web/src/services/events/workflowExecutionCoordinator.ts index fceb117d87d..3c443383aa7 100644 --- a/invokeai/frontend/web/src/services/events/workflowExecutionCoordinator.ts +++ b/invokeai/frontend/web/src/services/events/workflowExecutionCoordinator.ts @@ -15,6 +15,7 @@ import { } from 'services/events/nodeExecutionState'; import { createWorkflowExecutionState, + isTerminalQueueStatus, transitionWorkflowExecutionState, type WorkflowExecutionState, } from 'services/events/workflowExecutionState'; @@ -31,6 +32,7 @@ type WorkflowExecutionCoordinatorDeps = { clearCanvasWorkflowIntegrationProcessing: () => void; completedInvocationKeysByItemId: Map>; getAllNodeExecutionStates: () => Record; + getCurrentUserId: () => string | null; getNodeExecutionState: (nodeId: string) => NodeExecutionState | undefined; logReconciliationError: (error: unknown, itemId: number) => void; onInvocationComplete: (data: S['InvocationCompleteEvent']) => void; @@ -41,6 +43,12 @@ type WorkflowExecutionCoordinatorDeps = { export const createWorkflowExecutionCoordinator = (deps: WorkflowExecutionCoordinatorDeps) => { const workflowExecutionStates = new LRUCache({ max: 100 }); + // Each tracked item's owner, recorded from incoming events, lets a user-scoped queue clear be + // applied to just that user's items. Every event reaching the coordinator carries a real owner: + // invocation and status events are private to the owner and admins, and the sanitized + // (user_id="redacted") companions are filtered out before the coordinator. Sized to match + // workflowExecutionStates - an item without a tracked state needs no owner. + const itemOwnerUserIds = new LRUCache({ max: 100 }); const pendingWorkflowReconciliationRequests = new Map(); let activeWorkflowQueueItemId: number | null = null; @@ -59,7 +67,10 @@ export const createWorkflowExecutionCoordinator = (deps: WorkflowExecutionCoordi req?.abort?.(); req?.unsubscribe?.(); pendingWorkflowReconciliationRequests.delete(itemId); - workflowExecutionStates.delete(itemId); + // The workflow execution state entry is intentionally kept. A canceled queue item can emit a + // few trailing invocation events (e.g. a denoise step callback racing the cancelation), and the + // retained terminal state is what rejects them. Item ids are never reused, and the LRU cache + // bounds memory. clearCompletedInvocationKeysForQueueItem(deps.completedInvocationKeysByItemId, itemId); }; @@ -71,6 +82,48 @@ export const createWorkflowExecutionCoordinator = (deps: WorkflowExecutionCoordi pendingWorkflowReconciliationRequests.clear(); }; + const onQueueCleared = (data: S['QueueClearedEvent']): boolean => { + // Clearing the queue deletes its items without emitting per-item terminal status events, so the + // deleted tracked items must be marked terminal here for trailing invocation events to be + // rejected, and their pending reconciliations aborted. Which items were deleted depends on the + // event's scope: + // - 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 reaches admins in full — they track every user's items, so the + // cleared user's items must be terminated without disturbing other items or this client's + // own reconciliations — and reaches everyone else as a sanitized user_id="redacted" + // broadcast, which matches no tracked owner and applies to nothing. + // Returns whether the clear applied to any item this client tracks, so the caller knows + // whether to reset progress UI. + const clearedUserId = data.user_id ?? null; + const clearAppliesToAllItems = clearedUserId === null || clearedUserId === deps.getCurrentUserId(); + if (clearAppliesToAllItems) { + cancelPendingWorkflowReconciliations(); + } + let applied = clearAppliesToAllItems; + for (const itemId of [...workflowExecutionStates.keys()]) { + if (!clearAppliesToAllItems) { + if (itemOwnerUserIds.get(itemId) !== clearedUserId) { + continue; + } + const req = pendingWorkflowReconciliationRequests.get(itemId); + if (req) { + req.abort?.(); + req.unsubscribe?.(); + pendingWorkflowReconciliationRequests.delete(itemId); + applied = true; + } + } + const state = workflowExecutionStates.get(itemId); + if (!state || isTerminalQueueStatus(state.queueStatus)) { + continue; + } + workflowExecutionStates.set(itemId, { ...state, queueStatus: 'canceled' }); + applied = true; + } + return applied; + }; + const reconcileWorkflowQueueItemResults = (itemId: number, status: TerminalQueueStatus) => { const req = deps.reconcileQueueItem(itemId); pendingWorkflowReconciliationRequests.set(itemId, req); @@ -105,6 +158,7 @@ export const createWorkflowExecutionCoordinator = (deps: WorkflowExecutionCoordi }; const onInvocationStarted = (data: S['InvocationStartedEvent']) => { + itemOwnerUserIds.set(data.item_id, data.user_id); if ( !transitionWorkflowEvent(data.item_id, { type: 'invocation_started', @@ -126,6 +180,7 @@ export const createWorkflowExecutionCoordinator = (deps: WorkflowExecutionCoordi }; const onInvocationProgress = (data: S['InvocationProgressEvent']) => { + itemOwnerUserIds.set(data.item_id, data.user_id); if ( !transitionWorkflowEvent(data.item_id, { type: 'invocation_progress', @@ -151,6 +206,7 @@ export const createWorkflowExecutionCoordinator = (deps: WorkflowExecutionCoordi }; const onInvocationError = (data: S['InvocationErrorEvent']) => { + itemOwnerUserIds.set(data.item_id, data.user_id); if ( !transitionWorkflowEvent(data.item_id, { type: 'invocation_error', @@ -177,6 +233,7 @@ export const createWorkflowExecutionCoordinator = (deps: WorkflowExecutionCoordi }; const onInvocationComplete = (data: S['InvocationCompleteEvent']) => { + itemOwnerUserIds.set(data.item_id, data.user_id); if ( data.origin === 'workflows' && activeWorkflowQueueItemId !== null && @@ -196,6 +253,7 @@ export const createWorkflowExecutionCoordinator = (deps: WorkflowExecutionCoordi }; const onQueueItemStatusChanged = (data: S['QueueItemStatusChangedEvent']) => { + itemOwnerUserIds.set(data.item_id, data.user_id); if ( !transitionWorkflowEvent(data.item_id, { type: 'queue_item_status_changed', @@ -245,6 +303,7 @@ export const createWorkflowExecutionCoordinator = (deps: WorkflowExecutionCoordi onInvocationError, onInvocationProgress, onInvocationStarted, + onQueueCleared, onQueueItemStatusChanged, }; }; diff --git a/invokeai/frontend/web/src/services/events/workflowExecutionState.ts b/invokeai/frontend/web/src/services/events/workflowExecutionState.ts index bb571447ef0..76e679315d8 100644 --- a/invokeai/frontend/web/src/services/events/workflowExecutionState.ts +++ b/invokeai/frontend/web/src/services/events/workflowExecutionState.ts @@ -35,7 +35,8 @@ type WorkflowExecutionTransition = { const TERMINAL_QUEUE_STATUSES = new Set(['completed', 'failed', 'canceled']); const TERMINAL_INVOCATION_STATUSES = new Set(['completed', 'failed']); -const isTerminalQueueStatus = (status: QueueStatus | null) => status !== null && TERMINAL_QUEUE_STATUSES.has(status); +export const isTerminalQueueStatus = (status: QueueStatus | null) => + status !== null && TERMINAL_QUEUE_STATUSES.has(status); export const createWorkflowExecutionState = (): WorkflowExecutionState => ({ itemId: null, diff --git a/tests/app/routers/test_multiuser_authorization.py b/tests/app/routers/test_multiuser_authorization.py index f1ae8defa1a..5a7e673b18f 100644 --- a/tests/app/routers/test_multiuser_authorization.py +++ b/tests/app/routers/test_multiuser_authorization.py @@ -2060,9 +2060,9 @@ def test_queue_items_retried_payload_is_filtered_per_owner_room(self, socketio: assert admin_calls[0].kwargs["data"]["user_ids"] == ["owner-123", "owner-456"] assert admin_calls[0].kwargs["data"]["retried_item_ids_by_user"] == {"owner-123": [10], "owner-456": [20]} - def test_queue_cleared_still_broadcast(self, socketio: Any) -> None: - """QueueClearedEvent does not carry user identity and should still be broadcast - to all queue subscribers — this is a sanity check that we haven't over-scoped.""" + def test_unscoped_queue_cleared_still_broadcast(self, socketio: Any) -> None: + """An unscoped QueueClearedEvent (user_id=None — an admin or single-user clear that + deleted every user's items) should still be broadcast to all queue subscribers.""" import asyncio from unittest.mock import AsyncMock @@ -2075,8 +2075,50 @@ def test_queue_cleared_still_broadcast(self, socketio: Any) -> None: asyncio.run(socketio._handle_queue_event(("queue_cleared", event))) - rooms_emitted_to = [call.kwargs.get("room") for call in mock_emit.call_args_list] - assert "default" in rooms_emitted_to + assert len(mock_emit.call_args_list) == 1 + assert mock_emit.call_args_list[0].kwargs.get("room") == "default" + assert mock_emit.call_args_list[0].kwargs.get("data")["user_id"] is None + + def test_user_scoped_queue_cleared_routed_privately(self, socketio: Any) -> None: + """A user-scoped QueueClearedEvent only deleted that user's rows. The full event must + go to the owner + admin rooms; the rest of the queue room gets a sanitized companion + (user_id="redacted") so their queue lists refresh without treating the clear as their + own — otherwise another user's clear would abort their in-flight reconciliations and + mark their tracked items canceled.""" + import asyncio + from unittest.mock import AsyncMock + + from invokeai.app.services.events.events_common import QueueClearedEvent + + event = QueueClearedEvent.build(queue_id="default", user_id="owner-xyz") + + socketio._socket_users["sid-owner"] = {"user_id": "owner-xyz", "is_admin": False} + socketio._socket_users["sid-admin"] = {"user_id": "admin-1", "is_admin": True} + socketio._socket_users["sid-other"] = {"user_id": "other-user", "is_admin": False} + + mock_emit = AsyncMock() + socketio._sio.emit = mock_emit + + asyncio.run(socketio._handle_queue_event(("queue_cleared", event))) + + emits = [ + (c.kwargs.get("room"), c.kwargs.get("data"), c.kwargs.get("skip_sid")) for c in mock_emit.call_args_list + ] + + # Full event goes to owner + admin rooms in a single emit (room list deduplicates a + # socket that is in both rooms) + private_emits = [(p, s) for r, p, s in emits if r == ["user:owner-xyz", "admin"]] + assert len(private_emits) == 1 + assert private_emits[0][0]["user_id"] == "owner-xyz" + + # Sanitized companion goes to the queue room, skipping the owner's and admins' sids + queue_emits = [(p, s) for r, p, s in emits if r == "default"] + assert len(queue_emits) == 1, "expected exactly one sanitized emit to queue room" + sanitized_payload, skip_sid = queue_emits[0] + assert sanitized_payload["user_id"] == "redacted" + assert "sid-owner" in skip_sid + assert "sid-admin" in skip_sid + assert "sid-other" not in skip_sid def test_recall_parameters_emitted_once_to_owner_and_admin_rooms(self, socketio: Any) -> None: """RecallParametersUpdatedEvent must be delivered to the owner + admin rooms diff --git a/tests/app/services/test_workflow_call_runtime.py b/tests/app/services/test_workflow_call_runtime.py index 1edaaaae007..6f8a469a896 100644 --- a/tests/app/services/test_workflow_call_runtime.py +++ b/tests/app/services/test_workflow_call_runtime.py @@ -39,6 +39,26 @@ def test_run_does_not_fail_canceled_parent_after_child_return_error(monkeypatch) workflow_call_tests.test_run_does_not_fail_canceled_parent_after_child_return_error(monkeypatch) +def test_run_queue_item_tolerates_queue_item_deleted_mid_run(monkeypatch) -> None: + workflow_call_tests.test_run_queue_item_tolerates_queue_item_deleted_mid_run(monkeypatch) + + +def test_run_queue_item_tolerates_parent_deleted_while_child_runs(monkeypatch) -> None: + workflow_call_tests.test_run_queue_item_tolerates_parent_deleted_while_child_runs(monkeypatch) + + +def test_run_queue_item_tolerates_parent_deleted_before_completed_parent_mutation(monkeypatch) -> None: + workflow_call_tests.test_run_queue_item_tolerates_parent_deleted_before_completed_parent_mutation(monkeypatch) + + +def test_run_queue_item_tolerates_parent_deleted_before_failed_parent_mutation(monkeypatch) -> None: + workflow_call_tests.test_run_queue_item_tolerates_parent_deleted_before_failed_parent_mutation(monkeypatch) + + +def test_run_queue_item_tolerates_parent_deleted_before_canceled_parent_mutation(monkeypatch) -> None: + workflow_call_tests.test_run_queue_item_tolerates_parent_deleted_before_canceled_parent_mutation(monkeypatch) + + def test_workflow_call_coordinator_builds_child_queue_item_with_relationship_metadata(monkeypatch) -> None: workflow_call_tests.test_workflow_call_coordinator_builds_child_queue_item_with_relationship_metadata(monkeypatch) diff --git a/tests/app/services/workflow_call_test_utils.py b/tests/app/services/workflow_call_test_utils.py index cbe94707e93..8923d050bd9 100644 --- a/tests/app/services/workflow_call_test_utils.py +++ b/tests/app/services/workflow_call_test_utils.py @@ -23,6 +23,7 @@ WorkflowCallCoordinator, WorkflowCallQueueLifecycle, ) +from invokeai.app.services.session_queue.session_queue_common import SessionQueueItemNotFoundError from invokeai.app.services.shared.graph import Graph, GraphExecutionState, WorkflowCallFrame from invokeai.app.services.workflow_records.workflow_records_common import WorkflowCategory from tests.dangerously_run_function_in_subprocess import dangerously_run_function_in_subprocess @@ -887,6 +888,14 @@ def __init__(self) -> None: self.deleted_item_ids: list[int] = [] self.pending_count = 0 self.fail_enqueue_after: int | None = None + # Item ids whose mutation methods raise SessionQueueItemNotFoundError, simulating the row + # being deleted after a successful get_queue_item lookup. Like the SQLite implementation's + # mutations, which re-read the row and raise when it has disappeared. + self.not_found_item_ids: set[int] = set() + + def _raise_if_not_found(self, item_id: int) -> None: + if item_id in self.not_found_item_ids: + raise SessionQueueItemNotFoundError(f"No queue item with id {item_id}") def add_queue_item(self, queue_item): self.items[queue_item.item_id] = queue_item @@ -913,33 +922,48 @@ def _ensure_queue_item(self, item_id: int, session): return queue_item def get_queue_item(self, item_id: int): - return self.items[item_id] + queue_item = self.items.get(item_id) + if queue_item is None: + raise SessionQueueItemNotFoundError(f"No queue item with id {item_id}") + return queue_item def set_queue_item_session(self, item_id: int, session): + self._raise_if_not_found(item_id) queue_item = self._ensure_queue_item(item_id, session) queue_item.session = session self.session_updates.append((item_id, session)) return queue_item def suspend_queue_item(self, item_id: int): + self._raise_if_not_found(item_id) queue_item = self._ensure_queue_item(item_id, None) queue_item.status = "waiting" self.waiting_item_ids.append(item_id) return queue_item def resume_queue_item(self, item_id: int): + self._raise_if_not_found(item_id) queue_item = self._ensure_queue_item(item_id, None) queue_item.status = "pending" self.resumed_item_ids.append(item_id) return queue_item def complete_queue_item(self, item_id: int): + self._raise_if_not_found(item_id) queue_item = self._ensure_queue_item(item_id, None) queue_item.status = "completed" self.completed_item_ids.append(item_id) return queue_item + def cancel_queue_item(self, item_id: int): + self._raise_if_not_found(item_id) + queue_item = self._ensure_queue_item(item_id, None) + queue_item.status = "canceled" + self.canceled_item_ids.append(item_id) + return queue_item + def fail_queue_item(self, item_id: int, error_type: str, error_message: str, error_traceback: str): + self._raise_if_not_found(item_id) queue_item = self._ensure_queue_item(item_id, None) queue_item.status = "failed" queue_item.error_type = error_type @@ -1510,6 +1534,149 @@ def test_workflow_call_queue_lifecycle_resumes_parent_from_completed_child( assert parent_outputs[0].values == {"result": [3]} +def test_run_queue_item_tolerates_queue_item_deleted_mid_run(monkeypatch: pytest.MonkeyPatch) -> None: + session_queue = _DummySessionQueue() + runner, events, _workflow_records = _build_workflow_runner(monkeypatch, session_queue=session_queue) + lifecycle = WorkflowCallQueueLifecycle(runner) + + queue_item = SimpleNamespace( + item_id=1, + status="in_progress", + session=None, + session_id="session-id", + parent_item_id=None, + ) + session_queue.add_queue_item(queue_item) + + # Simulate the queue item being deleted while it is running (e.g. the queue was cleared or the + # current item was deleted via the API). + def run_and_delete(item) -> None: + session_queue.delete_queue_items_by_id([item.item_id]) + + monkeypatch.setattr(runner, "run", run_and_delete) + + lifecycle.run_queue_item(queue_item) + + assert session_queue.deleted_item_ids == [1] + assert session_queue.completed_item_ids == [] + assert session_queue.failed_item_ids == [] + assert events.errors == [] + + +def test_run_queue_item_tolerates_parent_deleted_while_child_runs(monkeypatch: pytest.MonkeyPatch) -> None: + for child_terminal_status in ("completed", "failed", "canceled"): + session_queue = _DummySessionQueue() + runner, events, _workflow_records = _build_workflow_runner(monkeypatch, session_queue=session_queue) + lifecycle = WorkflowCallQueueLifecycle(runner) + + child_queue_item = SimpleNamespace( + item_id=100, + status="in_progress", + session=None, + session_id="child-session-id", + parent_item_id=1, + ) + session_queue.add_queue_item(child_queue_item) + + # The parent (item id 1) does not exist in the queue - it was deleted while the child was running. + def run_to_terminal_status(item, status=child_terminal_status) -> None: + item.status = status + + monkeypatch.setattr(runner, "run", run_to_terminal_status) + + lifecycle.run_queue_item(child_queue_item) + + assert session_queue.canceled_item_ids == [] + assert session_queue.failed_item_ids == [] + assert events.errors == [] + + +def _setup_suspended_workflow_call_parent(monkeypatch: pytest.MonkeyPatch): + """Runs a workflow-call parent to its suspension point: parent (item 1) is waiting, child (item 100) is pending.""" + session_queue = _DummySessionQueue() + runner, events, _workflow_records = _build_workflow_runner(monkeypatch, session_queue=session_queue) + lifecycle = WorkflowCallQueueLifecycle(runner) + + graph = Graph() + graph.add_node(CallSavedWorkflowInvocation(id="call-node", workflow_id="workflow-a")) + session = GraphExecutionState(graph=graph) + queue_item = type( + "QueueItem", + (), + { + "item_id": 1, + "status": "in_progress", + "session": session, + "session_id": "session-id", + "user_id": "user-1", + "queue_id": "default", + "batch_id": "batch-1", + }, + )() + session_queue.add_queue_item(queue_item) + lifecycle.run_queue_item(queue_item) + assert session_queue.waiting_item_ids == [1] + assert session_queue.enqueued_child_item_ids == [100] + return session_queue, runner, lifecycle, events + + +def test_run_queue_item_tolerates_parent_deleted_before_completed_parent_mutation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + session_queue, runner, lifecycle, _events = _setup_suspended_workflow_call_parent(monkeypatch) + + # The parent is deleted after the child's completion handler has looked it up: get_queue_item + # still succeeds, but the parent mutations (set_queue_item_session et al.) raise. + session_queue.not_found_item_ids.add(1) + + child_queue_item = session_queue.dequeue() + assert child_queue_item is not None + lifecycle.run_queue_item(child_queue_item) + + assert 100 in session_queue.completed_item_ids + assert 1 not in session_queue.completed_item_ids + assert session_queue.resumed_item_ids == [] + + +def test_run_queue_item_tolerates_parent_deleted_before_failed_parent_mutation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + session_queue, runner, lifecycle, _events = _setup_suspended_workflow_call_parent(monkeypatch) + + session_queue.not_found_item_ids.add(1) + + child_queue_item = session_queue.dequeue() + assert child_queue_item is not None + + def run_to_failed(item) -> None: + item.status = "failed" + item.error_message = "child failed" + + monkeypatch.setattr(runner, "run", run_to_failed) + lifecycle.run_queue_item(child_queue_item) + + assert 1 not in session_queue.failed_item_ids + + +def test_run_queue_item_tolerates_parent_deleted_before_canceled_parent_mutation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + session_queue, runner, lifecycle, _events = _setup_suspended_workflow_call_parent(monkeypatch) + + session_queue.not_found_item_ids.add(1) + + child_queue_item = session_queue.dequeue() + assert child_queue_item is not None + + def run_to_canceled(item) -> None: + item.status = "canceled" + + monkeypatch.setattr(runner, "run", run_to_canceled) + lifecycle.run_queue_item(child_queue_item) + + assert session_queue.canceled_item_ids == [] + + def test_workflow_call_coordinator_cleans_up_enqueued_children_when_boundary_setup_fails( monkeypatch: pytest.MonkeyPatch, ) -> None: