From 628367b97b300e7da63040886ab13ef0c94f3010 Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Wed, 4 Jun 2025 23:47:29 +1000 Subject: [PATCH] feat(ui): move socket events handling into ctx component --- .../components/SimpleSession/StagingArea.tsx | 2 - .../components/SimpleSession/context.tsx | 128 ++++++++++++++++-- .../SimpleSession/use-progress-events.ts | 48 ------- 3 files changed, 118 insertions(+), 60 deletions(-) delete mode 100644 invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/use-progress-events.ts diff --git a/invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/StagingArea.tsx b/invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/StagingArea.tsx index c4dacd07d1..d00f6a4ec8 100644 --- a/invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/StagingArea.tsx +++ b/invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/StagingArea.tsx @@ -7,14 +7,12 @@ import { getQueueItemElementId } from 'features/controlLayers/components/SimpleS import { StagingAreaContent } from 'features/controlLayers/components/SimpleSession/StagingAreaContent'; import { StagingAreaHeader } from 'features/controlLayers/components/SimpleSession/StagingAreaHeader'; import { StagingAreaNoItems } from 'features/controlLayers/components/SimpleSession/StagingAreaNoItems'; -import { useProgressEvents } from 'features/controlLayers/components/SimpleSession/use-progress-events'; import { useStagingAreaKeyboardNav } from 'features/controlLayers/components/SimpleSession/use-staging-keyboard-nav'; import { memo, useEffect } from 'react'; export const StagingArea = memo(() => { const ctx = useCanvasSessionContext(); const hasItems = useStore(ctx.$hasItems); - useProgressEvents(); useStagingAreaKeyboardNav(); useEffect(() => { diff --git a/invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/context.tsx b/invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/context.tsx index b2b85d21a0..64eae0b8f4 100644 --- a/invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/context.tsx +++ b/invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/context.tsx @@ -1,3 +1,4 @@ +import { useStore } from '@nanostores/react'; import { createSelector } from '@reduxjs/toolkit'; import { EMPTY_ARRAY } from 'app/store/constants'; import { useAppStore } from 'app/store/nanostores/store'; @@ -12,6 +13,7 @@ import type { PropsWithChildren } from 'react'; import { createContext, memo, useContext, useEffect, useMemo, useState } from 'react'; import { queueApi } from 'services/api/endpoints/queue'; import type { S } from 'services/api/types'; +import { $socket } from 'services/events/stores'; import { assert } from 'tsafe'; export type ProgressData = { @@ -134,13 +136,48 @@ const CanvasSessionContext = createContext(nul export const CanvasSessionContextProvider = memo( ({ session, children }: PropsWithChildren<{ session: SimpleSessionIdentifier | AdvancedSessionIdentifier }>) => { + /** + * For best performance and interop with the Canvas, which is outside react but needs to interact with the react + * app, all canvas session state is packaged as nanostores atoms. The trickiest part is syncing the queue items + * with a nanostores atom. + */ + + /** + * App store + */ const store = useAppStore(); - const [$items] = useState(() => atom([])); - const [$hasItems] = useState(() => computed([$items], (items) => items.length > 0)); - const [$autoSwitch] = useState(() => atom(true)); - const [$selectedItemId] = useState(() => atom(null)); - const [$progressData] = useState(() => atom>({})); - const [$selectedItem] = useState(() => + + const socket = useStore($socket); + + /** + * Manually-synced atom containing the queue items for the current session. + */ + const $items = useState(() => atom([]))[0]; + + /** + * Whether auto-switch is enabled. + */ + const $autoSwitch = useState(() => atom(true))[0]; + + /** + * An ephemeral store of progress events and images for all items in the current session. + */ + const $progressData = useState(() => atom>({}))[0]; + + /** + * The currently selected queue item's ID, or null if one is not selected. + */ + const $selectedItemId = useState(() => atom(null))[0]; + + /** + * Whether there are any items. Computed from the queue items array. + */ + const $hasItems = useState(() => computed([$items], (items) => items.length > 0))[0]; + + /** + * The currently selected queue item, or null if one is not selected. + */ + const $selectedItem = useState(() => computed([$items, $selectedItemId], (items, selectedItemId) => { if (items.length === 0) { return null; @@ -150,8 +187,12 @@ export const CanvasSessionContextProvider = memo( } return items.find(({ item_id }) => item_id === selectedItemId) ?? null; }) - ); - const [$selectedItemIndex] = useState(() => + )[0]; + + /** + * The currently selected queue item's index in the list of items, or null if one is not selected. + */ + const $selectedItemIndex = useState(() => computed([$items, $selectedItemId], (items, selectedItemId) => { if (items.length === 0) { return null; @@ -161,20 +202,59 @@ export const CanvasSessionContextProvider = memo( } return items.findIndex(({ item_id }) => item_id === selectedItemId) ?? null; }) - ); + )[0]; + /** + * A redux selector to select all queue items from the RTK Query cache. It's important that this returns stable + * references if possible to reduce re-renders. All derivations of the queue items (e.g. filtering out canceled + * items) should be done in a nanostores computed. + */ const selectQueueItems = useMemo( () => createSelector( queueApi.endpoints.listAllQueueItems.select({ destination: session.id }), - ({ data }) => data?.filter((item) => item.status !== 'canceled') ?? EMPTY_ARRAY + ({ data }) => data ?? EMPTY_ARRAY ), [session.id] ); + // Set up socket listeners useEffect(() => { + if (!socket) { + return; + } + + const onProgress = (data: S['InvocationProgressEvent']) => { + if (data.destination !== session.id) { + return; + } + setProgress($progressData, data); + }; + + const onQueueItemStatusChanged = (data: S['QueueItemStatusChangedEvent']) => { + if (data.destination !== session.id) { + return; + } + if (data.status === 'completed' && $autoSwitch.get()) { + $selectedItemId.set(data.item_id); + } + }; + + socket.on('invocation_progress', onProgress); + socket.on('queue_item_status_changed', onQueueItemStatusChanged); + + return () => { + socket.off('invocation_progress', onProgress); + socket.off('queue_item_status_changed', onQueueItemStatusChanged); + }; + }, [$autoSwitch, $progressData, $selectedItemId, session.id, socket]); + + // Set up state subscriptions and effects + useEffect(() => { + // Seed the $items atom with the initial query cache state $items.set(selectQueueItems(store.getState())); + // Manually keep the $items atom in sync as the query cache is updated const unsubReduxSyncToItemsAtom = store.subscribe(() => { const prevItems = $items.get(); const items = selectQueueItems(store.getState()); @@ -183,29 +263,57 @@ export const CanvasSessionContextProvider = memo( } }); + // Handle cases that could result in a nonexistent queue item being selected. const unsubEnsureSelectedItemIdExists = effect([$items, $selectedItemId], (items, selectedItemId) => { + // If there are no items, cannot have a selected item. if (items.length === 0) { $selectedItemId.set(null); return; } + // If there is no selected item but there are items, select the first one. if (selectedItemId === null && items.length > 0) { $selectedItemId.set(items[0]?.item_id ?? null); return; } + // If an item is selected and it is not in the list of items, un-set it. This effect will run again and we'll + // the above case, selecting the first item if there are any. if (selectedItemId !== null && items.findIndex(({ item_id }) => item_id === selectedItemId) === -1) { $selectedItemId.set(null); return; } }); + // Clean up the progress data when a queue item is discarded. + const unsubCleanUpProgressData = effect([$items, $progressData], (items, progressData) => { + const toDelete: string[] = []; + for (const datum of Object.values(progressData)) { + if (items.findIndex(({ session_id }) => session_id === datum.sessionId) === -1) { + toDelete.push(datum.sessionId); + } + } + if (toDelete.length === 0) { + return; + } + const newProgressData = { ...progressData }; + for (const sessionId of toDelete) { + delete newProgressData[sessionId]; + } + // This will re-trigger the effect - maybe this could just be a listener on $items? Brain hurt + $progressData.set(newProgressData); + }); + + // Create an RTK Query subscription. Without this, the query cache selector will never return anything bc RTK + // doesn't know we care about it. const { unsubscribe: unsubQueueItemsQuery } = store.dispatch( queueApi.endpoints.listAllQueueItems.initiate({ destination: session.id }) ); + // Clean up all subscriptions and top-level (i.e. non-computed/derived state) return () => { unsubQueueItemsQuery(); unsubReduxSyncToItemsAtom(); unsubEnsureSelectedItemIdExists(); + unsubCleanUpProgressData(); $items.set([]); $progressData.set({}); $selectedItemId.set(null); diff --git a/invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/use-progress-events.ts b/invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/use-progress-events.ts deleted file mode 100644 index 737ca3fdd0..0000000000 --- a/invokeai/frontend/web/src/features/controlLayers/components/SimpleSession/use-progress-events.ts +++ /dev/null @@ -1,48 +0,0 @@ -import { useStore } from '@nanostores/react'; -import { useCanvasSessionContext } from 'features/controlLayers/components/SimpleSession/context'; -import { useEffect } from 'react'; -import type { S } from 'services/api/types'; -import { $socket, setProgress } from 'services/events/stores'; - -export const useProgressEvents = () => { - const ctx = useCanvasSessionContext(); - const socket = useStore($socket); - useEffect(() => { - if (!socket) { - return; - } - - const onQueueItemStatusChanged = (data: S['QueueItemStatusChangedEvent']) => { - if (data.destination !== ctx.session.id) { - return; - } - if (data.status === 'completed' && ctx.$autoSwitch.get()) { - ctx.$selectedItemId.set(data.item_id); - } - }; - - socket.on('queue_item_status_changed', onQueueItemStatusChanged); - - return () => { - socket.off('queue_item_status_changed', onQueueItemStatusChanged); - }; - }, [ctx.$autoSwitch, ctx.$progressData, ctx.$selectedItemId, ctx.session.id, socket]); - - useEffect(() => { - if (!socket) { - return; - } - const onProgress = (data: S['InvocationProgressEvent']) => { - if (data.destination !== ctx.session.id) { - return; - } - // TODO: clear progress when done w/ it memory leak - setProgress(ctx.$progressData, data); - }; - socket.on('invocation_progress', onProgress); - - return () => { - socket.off('invocation_progress', onProgress); - }; - }, [ctx.$progressData, ctx.session.id, socket]); -};