From f576d80799ef9987e79d95cb3cecc8225b68a1a0 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sun, 12 Jul 2026 21:27:08 -0400 Subject: [PATCH 1/3] fix(ui): don't auto-switch boards on other users' generations In multiuser mode, invocation-complete events are emitted to the admin room as well as the owner's room so that admins' gallery caches stay fresh. The gallery auto-switch handler treated every received event as the current user's own generation, so an admin's selected board would jump to another user's board whenever that user generated an image. Gate the auto-switch (board + image selection) on the event's user_id matching the logged-in user. Cache updates (board totals, image name lists) still run for all received events. Single-user mode, where there is no current user, is unaffected. Co-Authored-By: Claude Fable 5 --- .../events/onInvocationComplete.test.ts | 144 ++++++++++++++++++ .../services/events/onInvocationComplete.tsx | 9 ++ 2 files changed, 153 insertions(+) create mode 100644 invokeai/frontend/web/src/services/events/onInvocationComplete.test.ts 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..48256d66338 --- /dev/null +++ b/invokeai/frontend/web/src/services/events/onInvocationComplete.test.ts @@ -0,0 +1,144 @@ +import { boardIdSelected, imageSelected } from 'features/gallery/store/gallerySlice'; +import { boardsApi } from 'services/api/endpoints/boards'; +import { getImageDTOSafe } from 'services/api/endpoints/images'; +import type { ImageDTO, S } from 'services/api/types'; +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(), + }), +})); + +vi.mock('services/api/endpoints/images', () => ({ + getImageDTOSafe: vi.fn(), + imagesApi: { + util: { + updateQueryData: vi.fn(() => ({ type: 'imagesApi/updateQueryData' })), + invalidateTags: vi.fn(() => ({ type: 'imagesApi/invalidateTags' })), + }, + }, +})); + +vi.mock('services/api/endpoints/boards', () => ({ + boardsApi: { + endpoints: { + getBoardImagesTotal: { + select: vi.fn(() => () => ({ data: undefined })), + }, + }, + util: { + upsertQueryEntries: vi.fn(() => ({ type: 'boardsApi/upsertQueryEntries' })), + updateQueryData: vi.fn(() => ({ type: 'boardsApi/updateQueryData' })), + }, + }, +})); + +vi.mock('services/api/endpoints/queue', () => ({ + queueApi: { + util: { + invalidateTags: vi.fn(() => ({ type: 'queueApi/invalidateTags' })), + }, + }, +})); + +vi.mock('features/gallery/store/gallerySelectors', () => ({ + selectAutoSwitch: () => true, + selectGalleryView: () => 'images', + selectGetImageNamesQueryArgs: () => ({ search_term: '', starred_first: true, order_dir: 'DESC' }), + selectListBoardsQueryArgs: () => ({}), + selectSelectedBoardId: () => 'admin-board', +})); + +vi.mock('features/nodes/hooks/useNodeExecutionState', () => ({ + $nodeExecutionStates: { get: () => ({}) }, + upsertExecutionState: vi.fn(), +})); + +vi.mock('features/controlLayers/store/canvasWorkflowIntegrationSlice', () => ({ + canvasWorkflowIntegrationProcessingCompleted: vi.fn(() => ({ + type: 'canvasWorkflowIntegration/processingCompleted', + })), +})); + +vi.mock('services/events/nodeExecutionState', () => ({ + getUpdatedNodeExecutionStateOnInvocationComplete: vi.fn(() => undefined), +})); + +vi.mock('services/events/stores', () => ({ + $lastProgressEvent: { set: vi.fn() }, +})); + +const OWNER_USER_ID = 'user-1'; +const OTHER_USER_ID = 'admin-1'; + +const imageDTO = { + image_name: 'img-1.png', + board_id: 'user-board', + is_intermediate: false, + image_category: 'general', +} as ImageDTO; + +const buildEvent = (): S['InvocationCompleteEvent'] => + ({ + item_id: 1, + batch_id: 'batch-1', + session_id: 'session-1', + queue_id: 'default', + user_id: OWNER_USER_ID, + origin: null, + destination: null, + invocation: { id: 'node-1', type: 'l2i' }, + invocation_source_id: 'node-1', + result: { image: { image_name: imageDTO.image_name } }, + }) as unknown as S['InvocationCompleteEvent']; + +const buildHandler = (currentUser: { user_id: string; is_admin: boolean } | null) => { + const dispatch = vi.fn(); + const getState = vi.fn(() => ({ auth: { user: currentUser } })); + const handler = buildOnInvocationComplete(getState as never, dispatch as never, new Map()); + return { handler, dispatch }; +}; + +const dispatchedActionTypes = (dispatch: ReturnType) => + dispatch.mock.calls.map(([action]) => (action as { type?: string }).type); + +describe('buildOnInvocationComplete gallery auto-switch', () => { + beforeEach(() => { + vi.mocked(getImageDTOSafe).mockResolvedValue(imageDTO); + }); + + it("does not switch boards for another user's generation, but still updates gallery caches", async () => { + const { handler, dispatch } = buildHandler({ user_id: OTHER_USER_ID, is_admin: true }); + + await handler(buildEvent()); + + const types = dispatchedActionTypes(dispatch); + expect(types).not.toContain(boardIdSelected.type); + expect(types).not.toContain(imageSelected.type); + // Cache updates still happen so the admin's view of the other user's board stays fresh. + expect(boardsApi.util.upsertQueryEntries).toHaveBeenCalled(); + }); + + it("switches boards for the current user's own generation", async () => { + const { handler, dispatch } = buildHandler({ user_id: OWNER_USER_ID, is_admin: false }); + + await handler(buildEvent()); + + expect(dispatchedActionTypes(dispatch)).toContain(boardIdSelected.type); + }); + + it('switches boards in single-user mode (no current user)', async () => { + const { handler, dispatch } = buildHandler(null); + + await handler(buildEvent()); + + expect(dispatchedActionTypes(dispatch)).toContain(boardIdSelected.type); + }); +}); diff --git a/invokeai/frontend/web/src/services/events/onInvocationComplete.tsx b/invokeai/frontend/web/src/services/events/onInvocationComplete.tsx index ea6a237d4b2..0bd062208d7 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, @@ -166,6 +167,14 @@ export const buildOnInvocationComplete = ( return; } + // In multiuser mode, admins receive invocation events for all users' generations so the cache + // updates above keep their gallery fresh - but the selected board/image must only follow this + // user's own generations. In single-user mode there is no current user and no gating is needed. + const currentUser = selectCurrentUser(getState()); + if (currentUser && data.user_id !== currentUser.user_id) { + return; + } + // Finally, we may need to autoswitch to the new image. We'll only do it for the last image in the list. const lastImageDTO = imageDTOs.at(-1); From 22aed0ec0c3d9d1edcce289e61080c846c1c9ba2 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sun, 12 Jul 2026 21:38:46 -0400 Subject: [PATCH 2/3] fix(events): keep other users' progress out of admins' personal UI Invocation progress events drive personal UI only (the global progress bar, the workflow editor's current-image node, the image viewer's progress image). Emitting them to the admin room made an admin's progress display animate and sawtooth in response to other users' generations. No admin-facing feature consumes other users' progress, so emit progress events to the owner's room only. The remaining invocation events (started, complete, error) still go to admins since the gallery cache updates depend on them, but now via a single emit to [user room, admin room]. python-socketio dedupes recipients across a room list, so an admin who owns the queue item - including the "system" user in single-user mode, which is also an admin - receives exactly one copy instead of two. The old double delivery meant the second, image-stripped copy of a progress event clobbered the owner's own progress image, and completion events were processed twice. On the frontend, an invocation-complete event now only clears the progress indicator when the event belongs to the current user, sharing the ownership check introduced for gallery auto-switch. Co-Authored-By: Claude Fable 5 --- invokeai/app/api/sockets.py | 27 +++-- .../events/onInvocationComplete.test.ts | 9 +- .../services/events/onInvocationComplete.tsx | 20 +++- tests/app/test_invocation_event_socketio.py | 109 ++++++++++++++++++ 4 files changed, 147 insertions(+), 18 deletions(-) create mode 100644 tests/app/test_invocation_event_socketio.py diff --git a/invokeai/app/api/sockets.py b/invokeai/app/api/sockets.py index bdaf641c1c1..7cfaf69c400 100644 --- a/invokeai/app/api/sockets.py +++ b/invokeai/app/api/sockets.py @@ -325,18 +325,25 @@ async def _handle_queue_event(self, event: FastAPIEvent[QueueEventBase]): if isinstance(event_data, InvocationEventBase) and hasattr(event_data, "user_id"): user_room = f"user:{event_data.user_id}" - # Emit to the user's room - await self._sio.emit(event=event_name, data=event_data.model_dump(mode="json"), room=user_room) - - # Also emit to admin room so admins can see all events, but strip image preview data - # from InvocationProgressEvent to prevent admins from seeing other users' image content if isinstance(event_data, InvocationProgressEvent): - admin_event_data = event_data.model_copy(update={"image": None}) - await self._sio.emit(event=event_name, data=admin_event_data.model_dump(mode="json"), room="admin") + # Progress events only drive personal UI (the global progress bar and progress + # image previews) and are high-frequency. No admin UI consumes other users' + # progress, so emit to the owner only. This also keeps other users' progress + # from hijacking an admin's progress bar and image previews. + await self._sio.emit(event=event_name, data=event_data.model_dump(mode="json"), room=user_room) + logger.debug(f"Emitted invocation progress event to user room {user_room}") else: - await self._sio.emit(event=event_name, data=event_data.model_dump(mode="json"), room="admin") - - logger.debug(f"Emitted private invocation event {event_name} to user room {user_room} and admin room") + # started/complete/error also feed admins' gallery cache updates, so admins + # receive them for all users. Emit to the union of owner + admin rooms in a + # SINGLE call so an admin owner receives exactly one copy (see the + # RecallParametersUpdatedEvent note below on python-socketio's recipient + # dedup across a room list). + await self._sio.emit( + event=event_name, data=event_data.model_dump(mode="json"), room=[user_room, "admin"] + ) + logger.debug( + f"Emitted private invocation event {event_name} to user room {user_room} and admin room" + ) # QueueItemStatusChangedEvent: full to owner+admin, sanitized to everyone else in # the queue room so their queue list, badge, and item caches refresh. diff --git a/invokeai/frontend/web/src/services/events/onInvocationComplete.test.ts b/invokeai/frontend/web/src/services/events/onInvocationComplete.test.ts index 48256d66338..a9594d84f14 100644 --- a/invokeai/frontend/web/src/services/events/onInvocationComplete.test.ts +++ b/invokeai/frontend/web/src/services/events/onInvocationComplete.test.ts @@ -2,6 +2,7 @@ import { boardIdSelected, imageSelected } from 'features/gallery/store/gallerySl import { boardsApi } from 'services/api/endpoints/boards'; import { getImageDTOSafe } from 'services/api/endpoints/images'; import type { ImageDTO, S } from 'services/api/types'; +import { $lastProgressEvent } from 'services/events/stores'; import { beforeEach, describe, expect, it, vi } from 'vitest'; import { buildOnInvocationComplete } from './onInvocationComplete'; @@ -111,10 +112,11 @@ const dispatchedActionTypes = (dispatch: ReturnType) => describe('buildOnInvocationComplete gallery auto-switch', () => { beforeEach(() => { + vi.clearAllMocks(); vi.mocked(getImageDTOSafe).mockResolvedValue(imageDTO); }); - it("does not switch boards for another user's generation, but still updates gallery caches", async () => { + it("does not switch boards or clear the progress indicator for another user's generation, but still updates gallery caches", async () => { const { handler, dispatch } = buildHandler({ user_id: OTHER_USER_ID, is_admin: true }); await handler(buildEvent()); @@ -122,16 +124,18 @@ describe('buildOnInvocationComplete gallery auto-switch', () => { const types = dispatchedActionTypes(dispatch); expect(types).not.toContain(boardIdSelected.type); expect(types).not.toContain(imageSelected.type); + expect($lastProgressEvent.set).not.toHaveBeenCalled(); // Cache updates still happen so the admin's view of the other user's board stays fresh. expect(boardsApi.util.upsertQueryEntries).toHaveBeenCalled(); }); - it("switches boards for the current user's own generation", async () => { + it("switches boards and clears the progress indicator for the current user's own generation", async () => { const { handler, dispatch } = buildHandler({ user_id: OWNER_USER_ID, is_admin: false }); await handler(buildEvent()); expect(dispatchedActionTypes(dispatch)).toContain(boardIdSelected.type); + expect($lastProgressEvent.set).toHaveBeenCalledWith(null); }); it('switches boards in single-user mode (no current user)', async () => { @@ -140,5 +144,6 @@ describe('buildOnInvocationComplete gallery auto-switch', () => { await handler(buildEvent()); expect(dispatchedActionTypes(dispatch)).toContain(boardIdSelected.type); + expect($lastProgressEvent.set).toHaveBeenCalledWith(null); }); }); diff --git a/invokeai/frontend/web/src/services/events/onInvocationComplete.tsx b/invokeai/frontend/web/src/services/events/onInvocationComplete.tsx index 0bd062208d7..1bc607ab61c 100644 --- a/invokeai/frontend/web/src/services/events/onInvocationComplete.tsx +++ b/invokeai/frontend/web/src/services/events/onInvocationComplete.tsx @@ -45,6 +45,15 @@ export const buildOnInvocationComplete = ( dispatch: AppDispatch, completedInvocationKeysByItemId: Map> ) => { + // In multiuser mode, admins receive invocation events for all users' generations so the cache + // updates below keep their gallery fresh - but personal UI (board/image selection, the progress + // indicator) must only follow this user's own generations. In single-user mode there is no + // current user and every event is the user's own. + const isOwnEvent = (data: S['InvocationCompleteEvent']) => { + const currentUser = selectCurrentUser(getState()); + return !currentUser || data.user_id === currentUser.user_id; + }; + const addImagesToGallery = async (data: S['InvocationCompleteEvent']) => { if (nodeTypeDenylist.includes(data.invocation.type)) { log.trace(`Skipping denylisted node type (${data.invocation.type})`); @@ -167,11 +176,7 @@ export const buildOnInvocationComplete = ( return; } - // In multiuser mode, admins receive invocation events for all users' generations so the cache - // updates above keep their gallery fresh - but the selected board/image must only follow this - // user's own generations. In single-user mode there is no current user and no gating is needed. - const currentUser = selectCurrentUser(getState()); - if (currentUser && data.user_id !== currentUser.user_id) { + if (!isOwnEvent(data)) { return; } @@ -280,6 +285,9 @@ export const buildOnInvocationComplete = ( // Add images to gallery (canvas workflow integration results go to staging area automatically) await addImagesToGallery(data); - $lastProgressEvent.set(null); + // Another user's completion must not clear this user's progress indicator. + if (isOwnEvent(data)) { + $lastProgressEvent.set(null); + } }; }; diff --git a/tests/app/test_invocation_event_socketio.py b/tests/app/test_invocation_event_socketio.py new file mode 100644 index 00000000000..2066bbc056d --- /dev/null +++ b/tests/app/test_invocation_event_socketio.py @@ -0,0 +1,109 @@ +"""Tests for socket routing of invocation events in multiuser mode. + +Invocation progress events drive personal UI (the global progress bar and progress image +previews) and must be delivered only to the owner - admins receiving other users' progress +would see their own progress display hijacked. The other invocation events (started, +complete, error) also feed admins' gallery cache updates, so they go to the owner and the +admin room - but in a single emit so that an admin who owns the queue item (which includes +the "system" user in single-user mode) receives exactly one copy. +""" + +from unittest.mock import AsyncMock + +import pytest +from fastapi import FastAPI + +from invokeai.app.api.sockets import SocketIO +from invokeai.app.services.events.events_common import ( + InvocationCompleteEvent, + InvocationErrorEvent, + InvocationProgressEvent, + InvocationStartedEvent, +) + + +@pytest.fixture +def anyio_backend() -> str: + return "asyncio" + + +_COMMON_FIELDS = { + "queue_id": "default", + "item_id": 1, + "batch_id": "batch-1", + "user_id": "owner-1", + "session_id": "session-1", + "invocation": {"type": "add", "id": "node-1", "a": 1, "b": 2}, + "invocation_source_id": "node-1", +} + + +@pytest.mark.anyio +async def test_progress_event_is_emitted_only_to_owner() -> None: + socketio = SocketIO(FastAPI()) + socketio._sio.emit = AsyncMock() + + event = InvocationProgressEvent(**_COMMON_FIELDS, message="denoising", percentage=0.5) + + await socketio._handle_queue_event(("invocation_progress", event)) + + socketio._sio.emit.assert_awaited_once_with( + event="invocation_progress", + data=event.model_dump(mode="json"), + room="user:owner-1", + ) + + +@pytest.mark.anyio +async def test_complete_event_is_emitted_once_to_owner_and_admin_rooms() -> None: + socketio = SocketIO(FastAPI()) + socketio._sio.emit = AsyncMock() + + event = InvocationCompleteEvent(**_COMMON_FIELDS, result={"type": "integer_output", "value": 3}) + + await socketio._handle_queue_event(("invocation_complete", event)) + + # A single emit to the union of rooms - python-socketio dedupes recipients across a room + # list, so an admin owner (or the single-user "system" user) receives exactly one copy. + socketio._sio.emit.assert_awaited_once_with( + event="invocation_complete", + data=event.model_dump(mode="json"), + room=["user:owner-1", "admin"], + ) + + +@pytest.mark.anyio +async def test_started_event_is_emitted_once_to_owner_and_admin_rooms() -> None: + socketio = SocketIO(FastAPI()) + socketio._sio.emit = AsyncMock() + + event = InvocationStartedEvent(**_COMMON_FIELDS) + + await socketio._handle_queue_event(("invocation_started", event)) + + socketio._sio.emit.assert_awaited_once_with( + event="invocation_started", + data=event.model_dump(mode="json"), + room=["user:owner-1", "admin"], + ) + + +@pytest.mark.anyio +async def test_error_event_is_emitted_once_to_owner_and_admin_rooms() -> None: + socketio = SocketIO(FastAPI()) + socketio._sio.emit = AsyncMock() + + event = InvocationErrorEvent( + **_COMMON_FIELDS, + error_type="ValueError", + error_message="oops", + error_traceback="Traceback (most recent call last): ...", + ) + + await socketio._handle_queue_event(("invocation_error", event)) + + socketio._sio.emit.assert_awaited_once_with( + event="invocation_error", + data=event.model_dump(mode="json"), + room=["user:owner-1", "admin"], + ) From c13e1e4d20cf3608378a06fbf06aec9c66e0eacb Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sun, 12 Jul 2026 22:15:35 -0400 Subject: [PATCH 3/3] fix(events): scope activity indicators to the user's own work in multiuser mode Three related fixes so that one user's generation activity no longer animates other users' (including admins') personal UI: 1. The queue status endpoint now returns per-user counts for every caller, admins included. A new is_admin flag on get_queue_status() disables current-item redaction for admins so they keep visibility of other users' current item while gaining their own user_pending/user_in_progress counts. Redaction remains fail-closed for callers that pass user_id without the flag. The frontend progress bar, busy favicon/tab title, and floating invoke button spinner now prefer the per-user counts over the global counts, so they only react to the user's own queue items. Queue status embedded in socket events never carries user counts; when patching the cache from an event, the previously-fetched per-user counts are retained instead of being nulled out. 2. Model load events now carry the user whose action triggered the load (threaded from the invocation context's queue item and the utility/convert endpoints) and are routed to that user's socket room instead of being broadcast to every connected client. These events only feed the loading-models spinner, which is personal UI. Model install/download events remain broadcast. schema.ts and openapi.json regenerated for the new event field. 3. On a queue item reaching a terminal status, the failure toast and progress indicator reset now only fire for the current user's own items - admins receive these events for all users and previously got other users' error toasts and had their in-flight progress display cleared. Single-user mode is unaffected throughout: there is no current user in the frontend auth state, the sole "system" user is in its own user room, and its per-user counts equal the global counts. Co-Authored-By: Claude Fable 5 --- invokeai/app/api/routers/model_manager.py | 2 +- invokeai/app/api/routers/session_queue.py | 12 ++-- invokeai/app/api/routers/utilities.py | 10 ++-- invokeai/app/api/sockets.py | 17 +++++- invokeai/app/services/events/events_base.py | 10 ++-- invokeai/app/services/events/events_common.py | 14 +++-- .../services/model_load/model_load_base.py | 9 ++- .../services/model_load/model_load_default.py | 13 +++- .../session_queue/session_queue_base.py | 9 ++- .../session_queue/session_queue_sqlite.py | 9 ++- .../app/services/shared/invocation_context.py | 6 +- invokeai/frontend/web/openapi.json | 18 +++++- .../app/hooks/useSyncFaviconQueueStatus.ts | 7 ++- .../system/components/ProgressBar.tsx | 7 ++- .../components/FloatingLeftPanelButtons.tsx | 4 +- .../frontend/web/src/services/api/schema.ts | 19 +++++- .../services/events/queueStatusEvents.test.ts | 14 +++++ .../src/services/events/queueStatusEvents.ts | 20 ++++++- .../src/services/events/setEventListeners.tsx | 8 ++- .../routers/test_multiuser_authorization.py | 8 ++- .../test_session_queue_status_user_scoping.py | 26 ++++++++ tests/app/test_invocation_event_socketio.py | 60 +++++++++++++++++++ 22 files changed, 257 insertions(+), 45 deletions(-) diff --git a/invokeai/app/api/routers/model_manager.py b/invokeai/app/api/routers/model_manager.py index bdd2e406444..0321a254296 100644 --- a/invokeai/app/api/routers/model_manager.py +++ b/invokeai/app/api/routers/model_manager.py @@ -1155,7 +1155,7 @@ async def convert_model( with TemporaryDirectory(dir=ApiDependencies.invoker.services.configuration.models_path) as tmpdir: convert_path = pathlib.Path(tmpdir) / pathlib.Path(model_config.path).stem - converted_model = loader.load_model(model_config) + converted_model = loader.load_model(model_config, user_id=current_admin.user_id) # write the converted file to the convert path raw_model = converted_model.model assert hasattr(raw_model, "save_pretrained") diff --git a/invokeai/app/api/routers/session_queue.py b/invokeai/app/api/routers/session_queue.py index 89d2f9a7395..9850364af23 100644 --- a/invokeai/app/api/routers/session_queue.py +++ b/invokeai/app/api/routers/session_queue.py @@ -468,12 +468,14 @@ async def get_queue_status( current_user: CurrentUserOrDefault, queue_id: str = Path(description="The queue id to perform this operation on"), ) -> SessionQueueAndProcessorStatus: - """Gets the status of the session queue. Returns global counts; non-admin users additionally - get their own pending/in_progress counts (so the UI can show an X/Y badge) and cannot see the - current item's identifiers unless they own it.""" + """Gets the status of the session queue. Returns global counts; every user additionally gets + their own pending/in_progress counts (so the UI can show an X/Y badge and scope personal UI + like the progress bar to the user's own activity). Non-admin users cannot see the current + item's identifiers unless they own it.""" try: - user_id = None if current_user.is_admin else current_user.user_id - queue = ApiDependencies.invoker.services.session_queue.get_queue_status(queue_id, user_id=user_id) + queue = ApiDependencies.invoker.services.session_queue.get_queue_status( + queue_id, user_id=current_user.user_id, is_admin=current_user.is_admin + ) processor = ApiDependencies.invoker.services.session_processor.get_status() return SessionQueueAndProcessorStatus(queue=queue, processor=processor) except Exception as e: diff --git a/invokeai/app/api/routers/utilities.py b/invokeai/app/api/routers/utilities.py index 3ba8fee292c..4541aadb622 100644 --- a/invokeai/app/api/routers/utilities.py +++ b/invokeai/app/api/routers/utilities.py @@ -103,7 +103,7 @@ def _resolve_model_path(model_config_path: str) -> Path: return (base_models_path / model_path).resolve() -def _run_expand_prompt(prompt: str, model_key: str, max_tokens: int, system_prompt: str | None) -> str: +def _run_expand_prompt(prompt: str, model_key: str, max_tokens: int, system_prompt: str | None, user_id: str) -> str: """Run text LLM inference synchronously (called from thread).""" model_manager = ApiDependencies.invoker.services.model_manager model_config = model_manager.store.get_model(model_key) @@ -112,7 +112,7 @@ def _run_expand_prompt(prompt: str, model_key: str, max_tokens: int, system_prom raise ValueError(f"Model '{model_key}' is not a TextLLM model (got {model_config.type})") with _model_load_lock: - loaded_model = model_manager.load.load_model(model_config) + loaded_model = model_manager.load.load_model(model_config, user_id=user_id) with torch.no_grad(), loaded_model.model_on_device() as (_, model): model_abs_path = _resolve_model_path(model_config.path) @@ -147,6 +147,7 @@ async def expand_prompt(current_user: CurrentUserOrDefault, body: ExpandPromptRe body.model_key, body.max_tokens, body.system_prompt, + current_user.user_id, ) return ExpandPromptResponse(expanded_prompt=expanded) except UnknownModelException: @@ -172,7 +173,7 @@ class ImageToPromptResponse(BaseModel): error: str | None = None -def _run_image_to_prompt(image_name: str, model_key: str, instruction: str) -> str: +def _run_image_to_prompt(image_name: str, model_key: str, instruction: str, user_id: str) -> str: """Run LLaVA OneVision inference synchronously (called from thread).""" model_manager = ApiDependencies.invoker.services.model_manager model_config = model_manager.store.get_model(model_key) @@ -181,7 +182,7 @@ def _run_image_to_prompt(image_name: str, model_key: str, instruction: str) -> s raise ValueError(f"Model '{model_key}' is not a LLaVA OneVision model (got {model_config.type})") with _model_load_lock: - loaded_model = model_manager.load.load_model(model_config) + loaded_model = model_manager.load.load_model(model_config, user_id=user_id) # Load the image from InvokeAI's image store image = ApiDependencies.invoker.services.images.get_pil_image(image_name) @@ -229,6 +230,7 @@ async def image_to_prompt(current_user: CurrentUserOrDefault, body: ImageToPromp body.image_name, body.model_key, body.instruction, + current_user.user_id, ) return ImageToPromptResponse(prompt=prompt) except UnknownModelException: diff --git a/invokeai/app/api/sockets.py b/invokeai/app/api/sockets.py index 7cfaf69c400..02c0c5dc7d6 100644 --- a/invokeai/app/api/sockets.py +++ b/invokeai/app/api/sockets.py @@ -497,7 +497,22 @@ async def _handle_queue_event(self, event: FastAPIEvent[QueueEventBase]): logger.error(f"Error handling queue event {event[0]}: {e}", exc_info=True) async def _handle_model_event(self, event: FastAPIEvent[ModelEventBase | DownloadEventBase]) -> None: - await self._sio.emit(event=event[0], data=event[1].model_dump(mode="json")) + event_name, event_data = event + + # Model load events only drive personal UI (the loading-models spinner that puts the + # progress bar into indeterminate mode), so they are routed to the user whose action + # triggered the load. Broadcasting them made every user's progress bar animate whenever + # any user's generation loaded a model. In single-user mode the owner is "system" and + # every socket is in user:system, preserving the old behavior. + if isinstance(event_data, (ModelLoadStartedEvent, ModelLoadCompleteEvent)): + await self._sio.emit( + event=event_name, data=event_data.model_dump(mode="json"), room=f"user:{event_data.user_id}" + ) + return + + # Model install / download events remain broadcast to all connected sockets - they feed + # the model manager UI, which is not per-user. + await self._sio.emit(event=event_name, data=event_data.model_dump(mode="json")) async def _handle_bulk_image_download_event(self, event: FastAPIEvent[BulkDownloadEventBase]) -> None: event_name, event_data = event diff --git a/invokeai/app/services/events/events_base.py b/invokeai/app/services/events/events_base.py index 65be86f8399..0288e934716 100644 --- a/invokeai/app/services/events/events_base.py +++ b/invokeai/app/services/events/events_base.py @@ -176,15 +176,17 @@ def emit_download_error(self, job: "DownloadJob") -> None: # region Model loading - def emit_model_load_started(self, config: "AnyModelConfig", submodel_type: Optional["SubModelType"] = None) -> None: + def emit_model_load_started( + self, config: "AnyModelConfig", submodel_type: Optional["SubModelType"] = None, user_id: str = "system" + ) -> None: """Emitted when a model load is started.""" - self.dispatch(ModelLoadStartedEvent.build(config, submodel_type)) + self.dispatch(ModelLoadStartedEvent.build(config, submodel_type, user_id)) def emit_model_load_complete( - self, config: "AnyModelConfig", submodel_type: Optional["SubModelType"] = None + self, config: "AnyModelConfig", submodel_type: Optional["SubModelType"] = None, user_id: str = "system" ) -> None: """Emitted when a model load is complete.""" - self.dispatch(ModelLoadCompleteEvent.build(config, submodel_type)) + self.dispatch(ModelLoadCompleteEvent.build(config, submodel_type, user_id)) # endregion diff --git a/invokeai/app/services/events/events_common.py b/invokeai/app/services/events/events_common.py index 15b48bb4665..94c4fcfeb90 100644 --- a/invokeai/app/services/events/events_common.py +++ b/invokeai/app/services/events/events_common.py @@ -510,10 +510,13 @@ class ModelLoadStartedEvent(ModelEventBase): config: AnyModelConfig = Field(description="The model's config") submodel_type: Optional[SubModelType] = Field(default=None, description="The submodel type, if any") + user_id: str = Field(default="system", description="The ID of the user whose action triggered the load") @classmethod - def build(cls, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> "ModelLoadStartedEvent": - return cls(config=config, submodel_type=submodel_type) + def build( + cls, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None, user_id: str = "system" + ) -> "ModelLoadStartedEvent": + return cls(config=config, submodel_type=submodel_type, user_id=user_id) @payload_schema.register @@ -524,10 +527,13 @@ class ModelLoadCompleteEvent(ModelEventBase): config: AnyModelConfig = Field(description="The model's config") submodel_type: Optional[SubModelType] = Field(default=None, description="The submodel type, if any") + user_id: str = Field(default="system", description="The ID of the user whose action triggered the load") @classmethod - def build(cls, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> "ModelLoadCompleteEvent": - return cls(config=config, submodel_type=submodel_type) + def build( + cls, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None, user_id: str = "system" + ) -> "ModelLoadCompleteEvent": + return cls(config=config, submodel_type=submodel_type, user_id=user_id) @payload_schema.register diff --git a/invokeai/app/services/model_load/model_load_base.py b/invokeai/app/services/model_load/model_load_base.py index 87a405b4ea4..a06f869cbf2 100644 --- a/invokeai/app/services/model_load/model_load_base.py +++ b/invokeai/app/services/model_load/model_load_base.py @@ -15,12 +15,19 @@ class ModelLoadServiceBase(ABC): """Wrapper around AnyModelLoader.""" @abstractmethod - def load_model(self, model_config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> LoadedModel: + def load_model( + self, + model_config: AnyModelConfig, + submodel_type: Optional[SubModelType] = None, + user_id: Optional[str] = None, + ) -> LoadedModel: """ Given a model's configuration, load it and return the LoadedModel object. :param model_config: Model configuration record (as returned by ModelRecordBase.get_model()) :param submodel: For main (pipeline models), the submodel to fetch. + :param user_id: The user whose action triggered the load, threaded into the model load + events so they can be routed to that user's UI (defaults to the system user). """ @property diff --git a/invokeai/app/services/model_load/model_load_default.py b/invokeai/app/services/model_load/model_load_default.py index 2e2d2ae219d..480c88dd74e 100644 --- a/invokeai/app/services/model_load/model_load_default.py +++ b/invokeai/app/services/model_load/model_load_default.py @@ -50,18 +50,25 @@ def ram_cache(self) -> ModelCache: """Return the RAM cache used by this loader.""" return self._ram_cache - def load_model(self, model_config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> LoadedModel: + def load_model( + self, + model_config: AnyModelConfig, + submodel_type: Optional[SubModelType] = None, + user_id: Optional[str] = None, + ) -> LoadedModel: """ Given a model's configuration, load it and return the LoadedModel object. :param model_config: Model configuration record (as returned by ModelRecordBase.get_model()) :param submodel: For main (pipeline models), the submodel to fetch. + :param user_id: The user whose action triggered the load, threaded into the model load + events so they can be routed to that user's UI (defaults to the system user). """ # We don't have an invoker during testing # TODO(psyche): Mock this method on the invoker in the tests if hasattr(self, "_invoker"): - self._invoker.services.events.emit_model_load_started(model_config, submodel_type) + self._invoker.services.events.emit_model_load_started(model_config, submodel_type, user_id or "system") implementation, model_config, submodel_type = self._registry.get_implementation(model_config, submodel_type) # type: ignore loaded_model: LoadedModel = implementation( @@ -71,7 +78,7 @@ def load_model(self, model_config: AnyModelConfig, submodel_type: Optional[SubMo ).load_model(model_config, submodel_type) if hasattr(self, "_invoker"): - self._invoker.services.events.emit_model_load_complete(model_config, submodel_type) + self._invoker.services.events.emit_model_load_complete(model_config, submodel_type, user_id or "system") return loaded_model diff --git a/invokeai/app/services/session_queue/session_queue_base.py b/invokeai/app/services/session_queue/session_queue_base.py index a07abf10ee1..2737816d713 100644 --- a/invokeai/app/services/session_queue/session_queue_base.py +++ b/invokeai/app/services/session_queue/session_queue_base.py @@ -79,18 +79,25 @@ def get_queue_status( queue_id: str, user_id: Optional[str] = None, acting_user_id: Optional[str] = None, + is_admin: bool = False, ) -> SessionQueueStatus: """Gets the status of the queue. Aggregate counts (pending/in_progress/.../total) are always global across all users. If user_id is provided, the requesting user's own counts are additionally returned in - the user_pending/user_in_progress fields (left None otherwise). + the user_pending/user_in_progress fields (left None otherwise). Admin callers should + also pass their user_id so personal UI (e.g. the progress bar) can distinguish their + own activity from other users'. acting_user_id is independent of user_id and controls only current-item redaction: when set, the returned status omits item_id/session_id/batch_id unless the currently-running item belongs to acting_user_id. The redaction is decided from the same get_current() snapshot used to embed those identifiers, so it cannot race against a concurrent state change. + + is_admin disables current-item redaction entirely: admins may see the identifiers of + any user's current item. Redaction stays fail-closed - a caller that passes user_id + without is_admin gets the non-admin behavior. """ pass diff --git a/invokeai/app/services/session_queue/session_queue_sqlite.py b/invokeai/app/services/session_queue/session_queue_sqlite.py index 51980d68f3b..c36d9537287 100644 --- a/invokeai/app/services/session_queue/session_queue_sqlite.py +++ b/invokeai/app/services/session_queue/session_queue_sqlite.py @@ -1143,6 +1143,7 @@ def get_queue_status( queue_id: str, user_id: Optional[str] = None, acting_user_id: Optional[str] = None, + is_admin: bool = False, ) -> SessionQueueStatus: with self._db.transaction() as cursor: # Aggregate counts are always global (across all users). This lets a non-admin's @@ -1190,11 +1191,13 @@ def get_queue_status( # so a concurrent transition (e.g. B finishing while A's status changes) cannot leave # stale identifiers in the result. The aggregate counts stay global; only the current # item's identifiers are gated. acting_user_id (event path) takes precedence over - # user_id (API path) when deciding the redaction owner; either being None means an - # admin/global caller who may see the current item. + # user_id (API path) when deciding the redaction owner; either being None means a + # global caller who may see the current item. is_admin disables redaction outright so + # admin callers can pass their user_id (for the per-user counts) without losing + # visibility of other users' current item. owner_user_id = user_id if acting_user_id is None else acting_user_id show_current_item = current_item is not None and ( - owner_user_id is None or current_item.user_id == owner_user_id + is_admin or owner_user_id is None or current_item.user_id == owner_user_id ) return SessionQueueStatus( diff --git a/invokeai/app/services/shared/invocation_context.py b/invokeai/app/services/shared/invocation_context.py index e38766d5ba2..9c142359bc5 100644 --- a/invokeai/app/services/shared/invocation_context.py +++ b/invokeai/app/services/shared/invocation_context.py @@ -394,7 +394,7 @@ def load( if submodel_type: message += f" ({submodel_type.value})" self._util.signal_progress(message) - return self._services.model_manager.load.load_model(model, submodel_type) + return self._services.model_manager.load.load_model(model, submodel_type, user_id=self._data.queue_item.user_id) def load_by_attrs( self, name: str, base: BaseModelType, type: ModelType, submodel_type: Optional[SubModelType] = None @@ -424,7 +424,9 @@ def load_by_attrs( if submodel_type: message += f" ({submodel_type.value})" self._util.signal_progress(message) - return self._services.model_manager.load.load_model(configs[0], submodel_type) + return self._services.model_manager.load.load_model( + configs[0], submodel_type, user_id=self._data.queue_item.user_id + ) @staticmethod def _raise_if_external(model: AnyModelConfig) -> None: diff --git a/invokeai/frontend/web/openapi.json b/invokeai/frontend/web/openapi.json index e2801e9e39a..f258a06b1d0 100644 --- a/invokeai/frontend/web/openapi.json +++ b/invokeai/frontend/web/openapi.json @@ -7738,7 +7738,7 @@ "get": { "tags": ["queue"], "summary": "Get Queue Status", - "description": "Gets the status of the session queue. Returns global counts; non-admin users additionally\nget their own pending/in_progress counts (so the UI can show an X/Y badge) and cannot see the\ncurrent item's identifiers unless they own it.", + "description": "Gets the status of the session queue. Returns global counts; every user additionally gets\ntheir own pending/in_progress counts (so the UI can show an X/Y badge and scope personal UI\nlike the progress bar to the user's own activity). Non-admin users cannot see the current\nitem's identifiers unless they own it.", "operationId": "get_queue_status", "security": [ { @@ -57112,9 +57112,15 @@ ], "default": null, "description": "The submodel type, if any" + }, + "user_id": { + "default": "system", + "description": "The ID of the user whose action triggered the load", + "title": "User Id", + "type": "string" } }, - "required": ["timestamp", "config", "submodel_type"], + "required": ["timestamp", "config", "submodel_type", "user_id"], "title": "ModelLoadCompleteEvent", "type": "object" }, @@ -57422,9 +57428,15 @@ ], "default": null, "description": "The submodel type, if any" + }, + "user_id": { + "default": "system", + "description": "The ID of the user whose action triggered the load", + "title": "User Id", + "type": "string" } }, - "required": ["timestamp", "config", "submodel_type"], + "required": ["timestamp", "config", "submodel_type", "user_id"], "title": "ModelLoadStartedEvent", "type": "object" }, diff --git a/invokeai/frontend/web/src/app/hooks/useSyncFaviconQueueStatus.ts b/invokeai/frontend/web/src/app/hooks/useSyncFaviconQueueStatus.ts index 7bd55f25f4f..9fa58ebc949 100644 --- a/invokeai/frontend/web/src/app/hooks/useSyncFaviconQueueStatus.ts +++ b/invokeai/frontend/web/src/app/hooks/useSyncFaviconQueueStatus.ts @@ -7,8 +7,13 @@ const invokeLogoSVG = 'assets/images/invoke-favicon.svg'; const invokeAlertLogoSVG = 'assets/images/invoke-alert-favicon.svg'; const queryOptions = { + // The busy favicon/title reflect the user's own activity. In multiuser mode the global + // counts include other users' generations, so prefer the per-user counts. selectFromResult: (res) => ({ - queueSize: res.data ? res.data.queue.pending + res.data.queue.in_progress : 0, + queueSize: res.data + ? (res.data.queue.user_pending ?? res.data.queue.pending) + + (res.data.queue.user_in_progress ?? res.data.queue.in_progress) + : 0, }), } satisfies Parameters[1]; diff --git a/invokeai/frontend/web/src/features/system/components/ProgressBar.tsx b/invokeai/frontend/web/src/features/system/components/ProgressBar.tsx index 5a4abdd4d28..a65b010f1a8 100644 --- a/invokeai/frontend/web/src/features/system/components/ProgressBar.tsx +++ b/invokeai/frontend/web/src/features/system/components/ProgressBar.tsx @@ -12,6 +12,9 @@ const ProgressBar = (props: ProgressProps) => { const isConnected = useStore($isConnected); const lastProgressEvent = useStore($lastProgressEvent); const loadingModelsCount = useStore($loadingModelsCount); + // The progress bar reflects the user's own activity. In multiuser mode the global + // in_progress count includes other users' generations, so prefer the per-user count. + const inProgressCount = queueStatus?.queue.user_in_progress ?? queueStatus?.queue.in_progress; const value = useMemo(() => { if (!lastProgressEvent) { return 0; @@ -28,7 +31,7 @@ const ProgressBar = (props: ProgressProps) => { return true; } - if (!queueStatus?.queue.in_progress) { + if (!inProgressCount) { return false; } @@ -45,7 +48,7 @@ const ProgressBar = (props: ProgressProps) => { } return false; - }, [isConnected, lastProgressEvent, queueStatus?.queue.in_progress, loadingModelsCount]); + }, [isConnected, lastProgressEvent, inProgressCount, loadingModelsCount]); return ( { if (!data) { return { isProcessing: false }; } - return { isProcessing: data.queue.in_progress > 0 }; + // The spinner reflects the user's own activity. In multiuser mode the global + // in_progress count includes other users' generations, so prefer the per-user count. + return { isProcessing: (data.queue.user_in_progress ?? data.queue.in_progress) > 0 }; }, }); diff --git a/invokeai/frontend/web/src/services/api/schema.ts b/invokeai/frontend/web/src/services/api/schema.ts index f938b7f0f2c..5eaa7ee096c 100644 --- a/invokeai/frontend/web/src/services/api/schema.ts +++ b/invokeai/frontend/web/src/services/api/schema.ts @@ -2197,9 +2197,10 @@ export type paths = { }; /** * Get Queue Status - * @description Gets the status of the session queue. Returns global counts; non-admin users additionally - * get their own pending/in_progress counts (so the UI can show an X/Y badge) and cannot see the - * current item's identifiers unless they own it. + * @description Gets the status of the session queue. Returns global counts; every user additionally gets + * their own pending/in_progress counts (so the UI can show an X/Y badge and scope personal UI + * like the progress bar to the user's own activity). Non-admin users cannot see the current + * item's identifiers unless they own it. */ get: operations["get_queue_status"]; put?: never; @@ -24168,6 +24169,12 @@ export type components = { * @default null */ submodel_type: components["schemas"]["SubModelType"] | null; + /** + * User Id + * @description The ID of the user whose action triggered the load + * @default system + */ + user_id: string; }; /** * ModelLoadStartedEvent @@ -24189,6 +24196,12 @@ export type components = { * @default null */ submodel_type: components["schemas"]["SubModelType"] | null; + /** + * User Id + * @description The ID of the user whose action triggered the load + * @default system + */ + user_id: string; }; /** * ModelLoaderOutput diff --git a/invokeai/frontend/web/src/services/events/queueStatusEvents.test.ts b/invokeai/frontend/web/src/services/events/queueStatusEvents.test.ts index 9c7dab37c90..8167b85335a 100644 --- a/invokeai/frontend/web/src/services/events/queueStatusEvents.test.ts +++ b/invokeai/frontend/web/src/services/events/queueStatusEvents.test.ts @@ -65,6 +65,20 @@ describe(getUpdatedQueueStatusOnQueueItemStatusChanged.name, () => { processor: current.processor, }); }); + + it('retains the cached per-user counts, which event queue statuses never carry', () => { + // e.g. an admin with no generations of their own observing another user's event: the event's + // queue_status has null user counts, but the admin's own counts (0) must not be lost, else + // personal UI falls back to the global counts and reacts to the other user's activity. + const current = buildQueueStatus({ user_pending: 0, user_in_progress: 0 }); + const event = buildQueueStatusChangedEvent(); + + const updated = getUpdatedQueueStatusOnQueueItemStatusChanged(current, event); + + expect(updated.queue.user_pending).toBe(0); + expect(updated.queue.user_in_progress).toBe(0); + expect(updated.queue.completed).toBe(1); + }); }); describe(getQueueStatusWithObservedEvents.name, () => { diff --git a/invokeai/frontend/web/src/services/events/queueStatusEvents.ts b/invokeai/frontend/web/src/services/events/queueStatusEvents.ts index 01b635a3ad3..8ef57fe2513 100644 --- a/invokeai/frontend/web/src/services/events/queueStatusEvents.ts +++ b/invokeai/frontend/web/src/services/events/queueStatusEvents.ts @@ -41,12 +41,28 @@ const recordQueueItemStatusChangedEvent = (event: QueueItemStatusChangedEvent): latestObservedQueueStatus = event.queue_status; }; +/** + * Queue status embedded in socket events is built without a requesting user, so its + * user_pending/user_in_progress are always null. When replacing a cached/fetched status with an + * event's status, retain the previous per-user counts - for a non-owner recipient (e.g. an admin + * observing another user's event) they are still accurate, and for the owner the tag invalidation + * that accompanies every status change refetches fresh per-user counts immediately. + */ +const withRetainedUserCounts = ( + next: S['SessionQueueStatus'], + previous: S['SessionQueueStatus'] +): S['SessionQueueStatus'] => ({ + ...next, + user_pending: next.user_pending ?? previous.user_pending, + user_in_progress: next.user_in_progress ?? previous.user_in_progress, +}); + export const getQueueStatusWithObservedEvents = (queueStatus: QueueStatusResponse): QueueStatusResponse => { const currentItemId = queueStatus.queue.item_id; if (latestObservedQueueStatus && currentItemId !== null && observedTerminalItemIds.has(currentItemId)) { return { ...queueStatus, - queue: latestObservedQueueStatus, + queue: withRetainedUserCounts(latestObservedQueueStatus, queueStatus.queue), }; } @@ -60,7 +76,7 @@ export const getUpdatedQueueStatusOnQueueItemStatusChanged = ( recordQueueItemStatusChangedEvent(event); return { ...queueStatus, - queue: event.queue_status, + queue: withRetainedUserCounts(event.queue_status, queueStatus.queue), }; }; diff --git a/invokeai/frontend/web/src/services/events/setEventListeners.tsx b/invokeai/frontend/web/src/services/events/setEventListeners.tsx index 21f475f46b8..3a83eb56496 100644 --- a/invokeai/frontend/web/src/services/events/setEventListeners.tsx +++ b/invokeai/frontend/web/src/services/events/setEventListeners.tsx @@ -522,7 +522,13 @@ export const setEventListeners = ({ socket, store, setIsConnected }: SetEventLis } dispatch(queueApi.util.invalidateTags(tagsToInvalidate)); - if (status === 'completed' || status === 'failed' || status === 'canceled') { + // Admins receive full events for all users' queue items. Personal UI side effects (the + // failure toast and the progress indicator) must only follow the current user's own items. + // In single-user mode there is no current user and every event is the user's own. + const currentUser = getState().auth.user; + const isOwnItem = !currentUser || data.user_id === currentUser.user_id; + + if (isOwnItem && (status === 'completed' || status === 'failed' || status === 'canceled')) { if (status === 'failed' && error_type) { toast({ id: `INVOCATION_ERROR_${error_type}`, diff --git a/tests/app/routers/test_multiuser_authorization.py b/tests/app/routers/test_multiuser_authorization.py index 5a7e673b18f..81a1503c379 100644 --- a/tests/app/routers/test_multiuser_authorization.py +++ b/tests/app/routers/test_multiuser_authorization.py @@ -1440,14 +1440,16 @@ def test_get_queue_status_route_returns_global_and_user_counts( assert queue_status["user_pending"] == 2 assert queue_status["user_in_progress"] == 0 - # Admin caller sees the same global total. Admins query with user_id=None, so no - # per-user counts are computed (the badge falls back to the global total for admins). + # Admin caller sees the same global total plus their own per-user counts (zero here - + # the admin owns no items), so personal UI like the progress bar can distinguish the + # admin's own activity from other users'. r = client.get("/api/v1/queue/default/status", headers=_auth(admin_tok)) assert r.status_code == 200 queue_status = r.json()["queue"] assert queue_status["pending"] == 3 assert queue_status["total"] == 3 - assert queue_status["user_pending"] is None + assert queue_status["user_pending"] == 0 + assert queue_status["user_in_progress"] == 0 # =========================================================================== diff --git a/tests/app/services/session_queue/test_session_queue_status_user_scoping.py b/tests/app/services/session_queue/test_session_queue_status_user_scoping.py index d58933879f8..c9e2147b0e8 100644 --- a/tests/app/services/session_queue/test_session_queue_status_user_scoping.py +++ b/tests/app/services/session_queue/test_session_queue_status_user_scoping.py @@ -107,6 +107,32 @@ def test_status_current_item_redacted_for_non_owner_but_counts_global(session_qu assert status.user_pending == 1 # A's single pending item +def test_status_admin_caller_gets_user_subcounts_and_current_item_identifiers( + session_queue: SqliteSessionQueue, +) -> None: + """An admin caller (user_id set, is_admin=True) gets their own subcounts - so personal UI + like the progress bar can distinguish the admin's own activity from other users' - while + still seeing the identifiers of another user's current item.""" + admin_id = "admin-1" + user_b = "user-b" + b_item_id = _insert_queue_item(session_queue, user_id=user_b) + _insert_queue_item(session_queue, user_id=admin_id) + + in_progress = session_queue.dequeue() + assert in_progress is not None and in_progress.item_id == b_item_id + + status = session_queue.get_queue_status(queue_id="default", user_id=admin_id, is_admin=True) + + # Admins are not subject to current-item redaction. + assert status.item_id == b_item_id + assert status.session_id is not None + assert status.batch_id is not None + # Global counts plus the admin's own subcounts. + assert status.in_progress == 1 # B's item, counted globally + assert status.user_in_progress == 0 # the admin owns none in progress + assert status.user_pending == 1 # the admin's single pending item + + def test_get_queue_item_ids_returns_all_users_ids(session_queue: SqliteSessionQueue) -> None: """get_queue_item_ids returns ids for every user so the virtualized list can show the (redacted) entries belonging to other users. Redaction happens at hydration time.""" diff --git a/tests/app/test_invocation_event_socketio.py b/tests/app/test_invocation_event_socketio.py index 2066bbc056d..33c71f0cdaa 100644 --- a/tests/app/test_invocation_event_socketio.py +++ b/tests/app/test_invocation_event_socketio.py @@ -107,3 +107,63 @@ async def test_error_event_is_emitted_once_to_owner_and_admin_rooms() -> None: data=event.model_dump(mode="json"), room=["user:owner-1", "admin"], ) + + +def _make_model_config(): + from invokeai.backend.model_manager.configs.factory import AnyModelConfigValidator + + return AnyModelConfigValidator.validate_python( + { + "key": "model-key-1", + "hash": "hash-1", + "path": "/models/some-model", + "name": "some-model", + "base": "sd-1", + "type": "vae", + "format": "diffusers", + "source": "/models/some-model", + "source_type": "path", + "file_size": 1024, + } + ) + + +@pytest.mark.anyio +async def test_model_load_events_are_emitted_only_to_triggering_user() -> None: + from invokeai.app.services.events.events_common import ModelLoadCompleteEvent, ModelLoadStartedEvent + + socketio = SocketIO(FastAPI()) + socketio._sio.emit = AsyncMock() + + config = _make_model_config() + started = ModelLoadStartedEvent.build(config, user_id="owner-1") + complete = ModelLoadCompleteEvent.build(config, user_id="owner-1") + + await socketio._handle_model_event(("model_load_started", started)) + await socketio._handle_model_event(("model_load_complete", complete)) + + socketio._sio.emit.assert_any_await( + event="model_load_started", data=started.model_dump(mode="json"), room="user:owner-1" + ) + socketio._sio.emit.assert_any_await( + event="model_load_complete", data=complete.model_dump(mode="json"), room="user:owner-1" + ) + assert socketio._sio.emit.await_count == 2 + + +@pytest.mark.anyio +async def test_model_install_events_remain_broadcast() -> None: + from invokeai.app.services.events.events_common import ModelInstallStartedEvent + + socketio = SocketIO(FastAPI()) + socketio._sio.emit = AsyncMock() + + from types import SimpleNamespace + + event = SimpleNamespace(model_dump=lambda mode="json": {"id": 1}) + # Not a load event, so it takes the broadcast path (no room argument). + assert not isinstance(event, ModelInstallStartedEvent) + + await socketio._handle_model_event(("model_install_started", event)) + + socketio._sio.emit.assert_awaited_once_with(event="model_install_started", data={"id": 1})