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 bdaf641c1c1..02c0c5dc7d6 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. @@ -490,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/onInvocationComplete.test.ts b/invokeai/frontend/web/src/services/events/onInvocationComplete.test.ts new file mode 100644 index 00000000000..a9594d84f14 --- /dev/null +++ b/invokeai/frontend/web/src/services/events/onInvocationComplete.test.ts @@ -0,0 +1,149 @@ +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 { $lastProgressEvent } from 'services/events/stores'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +import { buildOnInvocationComplete } from './onInvocationComplete'; + +vi.mock('app/logging/logger', () => ({ + logger: () => ({ + debug: vi.fn(), + trace: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }), +})); + +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.clearAllMocks(); + vi.mocked(getImageDTOSafe).mockResolvedValue(imageDTO); + }); + + 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()); + + 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 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 () => { + const { handler, dispatch } = buildHandler(null); + + 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 ea6a237d4b2..1bc607ab61c 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, @@ -44,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})`); @@ -166,6 +176,10 @@ export const buildOnInvocationComplete = ( return; } + if (!isOwnEvent(data)) { + 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); @@ -271,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/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 new file mode 100644 index 00000000000..33c71f0cdaa --- /dev/null +++ b/tests/app/test_invocation_event_socketio.py @@ -0,0 +1,169 @@ +"""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"], + ) + + +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})