From a23b5c3408bb320135f7b1427f021b976211db67 Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Mon, 7 Apr 2025 14:57:31 +1000 Subject: [PATCH] refactor(ui): make workflow published status server-side state Whether a workflow is published or not shouldn't be something stored on the client. It's properly server-side state. This change removes the `is_published` flag from redux and updates all references to the flag to use the getWorkflow query. It also updates the socket event listener that handles session complete events. When a validation run completes, we invalidate the tags for the getWorkflow query. We need to do a bit of juggling to avoid a race condition (documented in the code). Works well though. --- .../flow/panels/TopPanel/TopLeftPanel.tsx | 5 +- .../ActiveWorkflowNameAndActions.tsx | 5 +- .../sidePanel/WorkflowsTabLeftPanel.tsx | 6 +-- .../IsolatedWorkflowBuilderWatcher.tsx | 3 +- .../workflow/PublishWorkflowPanelContent.tsx | 8 ++- .../components/sidePanel/workflow/publish.ts | 30 ++++++++++-- .../nodes/hooks/useIsWorkflowEditorLocked.ts | 10 ++-- .../web/src/features/nodes/store/types.ts | 2 +- .../src/features/nodes/store/workflowSlice.ts | 12 ++--- .../nodes/util/workflow/buildWorkflow.ts | 1 - .../InvokeButtonTooltip.tsx | 5 +- .../SaveWorkflowMenuItem.tsx | 5 +- .../hooks/useCreateNewWorkflow.ts | 2 - .../hooks/useSaveOrSaveAsWorkflow.ts | 5 +- .../src/services/events/setEventListeners.tsx | 49 +++++++++++++++---- 15 files changed, 98 insertions(+), 50 deletions(-) diff --git a/invokeai/frontend/web/src/features/nodes/components/flow/panels/TopPanel/TopLeftPanel.tsx b/invokeai/frontend/web/src/features/nodes/components/flow/panels/TopPanel/TopLeftPanel.tsx index 59c191ccc7..7320c1fce7 100644 --- a/invokeai/frontend/web/src/features/nodes/components/flow/panels/TopPanel/TopLeftPanel.tsx +++ b/invokeai/frontend/web/src/features/nodes/components/flow/panels/TopPanel/TopLeftPanel.tsx @@ -1,22 +1,21 @@ import { Alert, AlertDescription, AlertIcon, AlertTitle, Box, Flex } from '@invoke-ai/ui-library'; import { useStore } from '@nanostores/react'; -import { useAppSelector } from 'app/store/storeHooks'; import AddNodeButton from 'features/nodes/components/flow/panels/TopPanel/AddNodeButton'; import UpdateNodesButton from 'features/nodes/components/flow/panels/TopPanel/UpdateNodesButton'; import { $isInPublishFlow, $isSelectingOutputNode, useIsValidationRunInProgress, + useIsWorkflowPublished, } from 'features/nodes/components/sidePanel/workflow/publish'; import { useIsWorkflowEditorLocked } from 'features/nodes/hooks/useIsWorkflowEditorLocked'; -import { selectWorkflowIsPublished } from 'features/nodes/store/workflowSlice'; import { memo } from 'react'; import { useTranslation } from 'react-i18next'; export const TopLeftPanel = memo(() => { const isLocked = useIsWorkflowEditorLocked(); const isInPublishFlow = useStore($isInPublishFlow); - const isPublished = useAppSelector(selectWorkflowIsPublished); + const isPublished = useIsWorkflowPublished(); const isValidationRunInProgress = useIsValidationRunInProgress(); const isSelectingOutputNode = useStore($isSelectingOutputNode); diff --git a/invokeai/frontend/web/src/features/nodes/components/sidePanel/WorkflowListMenu/ActiveWorkflowNameAndActions.tsx b/invokeai/frontend/web/src/features/nodes/components/sidePanel/WorkflowListMenu/ActiveWorkflowNameAndActions.tsx index eb0d851f28..025a23ea80 100644 --- a/invokeai/frontend/web/src/features/nodes/components/sidePanel/WorkflowListMenu/ActiveWorkflowNameAndActions.tsx +++ b/invokeai/frontend/web/src/features/nodes/components/sidePanel/WorkflowListMenu/ActiveWorkflowNameAndActions.tsx @@ -1,8 +1,9 @@ import { Flex, Spacer } from '@invoke-ai/ui-library'; import { useAppSelector } from 'app/store/storeHooks'; +import { useIsWorkflowPublished } from 'features/nodes/components/sidePanel/workflow/publish'; import { WorkflowListMenuTrigger } from 'features/nodes/components/sidePanel/WorkflowListMenu/WorkflowListMenuTrigger'; import { WorkflowViewEditToggleButton } from 'features/nodes/components/sidePanel/WorkflowViewEditToggleButton'; -import { selectWorkflowIsPublished, selectWorkflowMode } from 'features/nodes/store/workflowSlice'; +import { selectWorkflowMode } from 'features/nodes/store/workflowSlice'; import { WorkflowLibraryMenu } from 'features/workflowLibrary/components/WorkflowLibraryMenu/WorkflowLibraryMenu'; import { memo } from 'react'; @@ -10,7 +11,7 @@ import SaveWorkflowButton from './SaveWorkflowButton'; export const ActiveWorkflowNameAndActions = memo(() => { const mode = useAppSelector(selectWorkflowMode); - const isPublished = useAppSelector(selectWorkflowIsPublished); + const isPublished = useIsWorkflowPublished(); return ( diff --git a/invokeai/frontend/web/src/features/nodes/components/sidePanel/WorkflowsTabLeftPanel.tsx b/invokeai/frontend/web/src/features/nodes/components/sidePanel/WorkflowsTabLeftPanel.tsx index 9789c179c3..38b9043e45 100644 --- a/invokeai/frontend/web/src/features/nodes/components/sidePanel/WorkflowsTabLeftPanel.tsx +++ b/invokeai/frontend/web/src/features/nodes/components/sidePanel/WorkflowsTabLeftPanel.tsx @@ -3,18 +3,18 @@ import { useStore } from '@nanostores/react'; import { useAppSelector } from 'app/store/storeHooks'; import { EditModeLeftPanelContent } from 'features/nodes/components/sidePanel/EditModeLeftPanelContent'; import { PublishedWorkflowPanelContent } from 'features/nodes/components/sidePanel/PublishedWorkflowPanelContent'; -import { $isInPublishFlow } from 'features/nodes/components/sidePanel/workflow/publish'; +import { $isInPublishFlow, useIsWorkflowPublished } from 'features/nodes/components/sidePanel/workflow/publish'; import { PublishWorkflowPanelContent } from 'features/nodes/components/sidePanel/workflow/PublishWorkflowPanelContent'; import { ActiveWorkflowDescription } from 'features/nodes/components/sidePanel/WorkflowListMenu/ActiveWorkflowDescription'; import { ActiveWorkflowNameAndActions } from 'features/nodes/components/sidePanel/WorkflowListMenu/ActiveWorkflowNameAndActions'; -import { selectWorkflowIsPublished, selectWorkflowMode } from 'features/nodes/store/workflowSlice'; +import { selectWorkflowMode } from 'features/nodes/store/workflowSlice'; import { memo } from 'react'; import { ViewModeLeftPanelContent } from './viewMode/ViewModeLeftPanelContent'; const WorkflowsTabLeftPanel = () => { const mode = useAppSelector(selectWorkflowMode); - const isPublished = useAppSelector(selectWorkflowIsPublished); + const isPublished = useIsWorkflowPublished(); const isInPublishFlow = useStore($isInPublishFlow); return ( diff --git a/invokeai/frontend/web/src/features/nodes/components/sidePanel/workflow/IsolatedWorkflowBuilderWatcher.tsx b/invokeai/frontend/web/src/features/nodes/components/sidePanel/workflow/IsolatedWorkflowBuilderWatcher.tsx index b598277015..ba156d0441 100644 --- a/invokeai/frontend/web/src/features/nodes/components/sidePanel/workflow/IsolatedWorkflowBuilderWatcher.tsx +++ b/invokeai/frontend/web/src/features/nodes/components/sidePanel/workflow/IsolatedWorkflowBuilderWatcher.tsx @@ -48,8 +48,9 @@ const queryOptions = { if (!currentData) { return { serverWorkflowHash: null }; } + const { is_published: _is_published, ...serverWorkflow } = currentData.workflow; return { - serverWorkflowHash: stableHash(currentData.workflow), + serverWorkflowHash: stableHash(serverWorkflow), }; }, } satisfies Parameters[1]; diff --git a/invokeai/frontend/web/src/features/nodes/components/sidePanel/workflow/PublishWorkflowPanelContent.tsx b/invokeai/frontend/web/src/features/nodes/components/sidePanel/workflow/PublishWorkflowPanelContent.tsx index af71c39147..0b9fdb57ff 100644 --- a/invokeai/frontend/web/src/features/nodes/components/sidePanel/workflow/PublishWorkflowPanelContent.tsx +++ b/invokeai/frontend/web/src/features/nodes/components/sidePanel/workflow/PublishWorkflowPanelContent.tsx @@ -25,7 +25,7 @@ import { $isReadyToDoValidationRun, $isSelectingOutputNode, $outputNodeId, - $validationRunBatchId, + $validationRunData, usePublishInputs, } from 'features/nodes/components/sidePanel/workflow/publish'; import { useInputFieldTemplateTitleOrThrow } from 'features/nodes/hooks/useInputFieldTemplateTitleOrThrow'; @@ -237,7 +237,11 @@ const PublishWorkflowButton = memo(() => { duration: null, }); assert(result.value.enqueueResult.batch.batch_id); - $validationRunBatchId.set(result.value.enqueueResult.batch.batch_id); + assert(result.value.batchConfig.validation_run_data); + $validationRunData.set({ + batchId: result.value.enqueueResult.batch.batch_id, + workflowId: result.value.batchConfig.validation_run_data.workflow_id, + }); log.debug(parseify(result.value), 'Enqueued batch'); } }, [enqueue, projectUrl, t]); diff --git a/invokeai/frontend/web/src/features/nodes/components/sidePanel/workflow/publish.ts b/invokeai/frontend/web/src/features/nodes/components/sidePanel/workflow/publish.ts index 0f5b53ba5b..45afd55a60 100644 --- a/invokeai/frontend/web/src/features/nodes/components/sidePanel/workflow/publish.ts +++ b/invokeai/frontend/web/src/features/nodes/components/sidePanel/workflow/publish.ts @@ -5,13 +5,17 @@ import { useAppSelector } from 'app/store/storeHooks'; import { $templates } from 'features/nodes/store/nodesSlice'; import { selectNodesSlice } from 'features/nodes/store/selectors'; import type { Templates } from 'features/nodes/store/types'; -import { selectWorkflowFormNodeFieldFieldIdentifiersDeduped } from 'features/nodes/store/workflowSlice'; +import { + selectWorkflowFormNodeFieldFieldIdentifiersDeduped, + selectWorkflowId, +} from 'features/nodes/store/workflowSlice'; import type { FieldIdentifier } from 'features/nodes/types/field'; import { isBoardFieldType } from 'features/nodes/types/field'; import { isInvocationNode } from 'features/nodes/types/invocation'; import { atom, computed } from 'nanostores'; import { useMemo } from 'react'; import { useGetBatchStatusQuery } from 'services/api/endpoints/queue'; +import { useGetWorkflowQuery } from 'services/api/endpoints/workflows'; import { assert } from 'tsafe'; export const $isInPublishFlow = atom(false); @@ -23,12 +27,12 @@ export const $isReadyToDoValidationRun = computed( return isInPublishFlow && outputNodeId !== null && !isSelectingOutputNode; } ); -export const $validationRunBatchId = atom(null); +export const $validationRunData = atom<{ batchId: string; workflowId: string } | null>(null); export const useIsValidationRunInProgress = () => { - const validationRunBatchId = useStore($validationRunBatchId); + const validationRunData = useStore($validationRunData); const { isValidationRunInProgress } = useGetBatchStatusQuery( - validationRunBatchId ? { batch_id: validationRunBatchId } : skipToken, + validationRunData?.batchId ? { batch_id: validationRunData.batchId } : skipToken, { selectFromResult: ({ currentData }) => { if (!currentData) { @@ -41,7 +45,7 @@ export const useIsValidationRunInProgress = () => { }, } ); - return validationRunBatchId !== null || isValidationRunInProgress; + return validationRunData !== null || isValidationRunInProgress; }; export const selectFieldIdentifiersWithInvocationTypes = createSelector( @@ -88,3 +92,19 @@ export const usePublishInputs = () => { return fieldIdentifiers; }; + +const queryOptions = { + selectFromResult: ({ currentData }) => { + if (!currentData) { + return { isPublished: false }; + } + return { isPublished: currentData.is_published }; + }, +} satisfies Parameters[1]; + +export const useIsWorkflowPublished = () => { + const workflowId = useAppSelector(selectWorkflowId); + const { isPublished } = useGetWorkflowQuery(workflowId ?? skipToken, queryOptions); + + return isPublished; +}; diff --git a/invokeai/frontend/web/src/features/nodes/hooks/useIsWorkflowEditorLocked.ts b/invokeai/frontend/web/src/features/nodes/hooks/useIsWorkflowEditorLocked.ts index fae669adb7..2738dad04b 100644 --- a/invokeai/frontend/web/src/features/nodes/hooks/useIsWorkflowEditorLocked.ts +++ b/invokeai/frontend/web/src/features/nodes/hooks/useIsWorkflowEditorLocked.ts @@ -1,11 +1,13 @@ import { useStore } from '@nanostores/react'; -import { useAppSelector } from 'app/store/storeHooks'; -import { $isInPublishFlow, useIsValidationRunInProgress } from 'features/nodes/components/sidePanel/workflow/publish'; -import { selectWorkflowIsPublished } from 'features/nodes/store/workflowSlice'; +import { + $isInPublishFlow, + useIsValidationRunInProgress, + useIsWorkflowPublished, +} from 'features/nodes/components/sidePanel/workflow/publish'; export const useIsWorkflowEditorLocked = () => { const isInPublishFlow = useStore($isInPublishFlow); - const isPublished = useAppSelector(selectWorkflowIsPublished); + const isPublished = useIsWorkflowPublished(); const isValidationRunInProgress = useIsValidationRunInProgress(); const isLocked = isInPublishFlow || isPublished || isValidationRunInProgress; diff --git a/invokeai/frontend/web/src/features/nodes/store/types.ts b/invokeai/frontend/web/src/features/nodes/store/types.ts index 2a504edc1f..3170c9722b 100644 --- a/invokeai/frontend/web/src/features/nodes/store/types.ts +++ b/invokeai/frontend/web/src/features/nodes/store/types.ts @@ -21,7 +21,7 @@ export type NodesState = { export type WorkflowMode = 'edit' | 'view'; -export type WorkflowsState = Omit & { +export type WorkflowsState = Omit & { _version: 1; mode: WorkflowMode; formFieldInitialValues: Record; diff --git a/invokeai/frontend/web/src/features/nodes/store/workflowSlice.ts b/invokeai/frontend/web/src/features/nodes/store/workflowSlice.ts index c9d3a322e7..a18f7d7628 100644 --- a/invokeai/frontend/web/src/features/nodes/store/workflowSlice.ts +++ b/invokeai/frontend/web/src/features/nodes/store/workflowSlice.ts @@ -54,7 +54,7 @@ const formElementDataChangedReducer = ( element.data = { ...element.data, ...changes } as T['data']; }; -export const getBlankWorkflow = (): Omit => { +export const getBlankWorkflow = (): Omit => { return { name: '', author: '', @@ -69,7 +69,6 @@ export const getBlankWorkflow = (): Omit => { // Even though these values are `undefined`, the keys _must_ be present for the presistence layer to rehydrate // them correctly. It uses a merge strategy that relies on the keys being present. id: undefined, - is_published: null, }; }; @@ -116,9 +115,6 @@ export const workflowSlice = createSlice({ workflowIDChanged: (state, action: PayloadAction) => { state.id = action.payload; }, - workflowIsPublishedChanged(state, action: PayloadAction) { - state.is_published = action.payload; - }, formReset: (state) => { const rootElement = buildContainer('column', []); state.form = { @@ -175,7 +171,9 @@ export const workflowSlice = createSlice({ }, extraReducers: (builder) => { builder.addCase(workflowLoaded, (state, action): WorkflowState => { - const { nodes, edges: _edges, ...workflowExtra } = action.payload; + // nodes and edges are handled in the nodes slice + // is_published is server state + const { nodes, edges: _edges, is_published: _is_published, ...workflowExtra } = action.payload; const formFieldInitialValues = getFormFieldInitialValues(workflowExtra.form, nodes); @@ -245,7 +243,6 @@ export const { workflowVersionChanged, workflowContactChanged, workflowIDChanged, - workflowIsPublishedChanged, formReset, formElementAdded, formElementRemoved, @@ -309,7 +306,6 @@ export const selectWorkflowId = createWorkflowSelector((workflow) => workflow.id export const selectWorkflowMode = createWorkflowSelector((workflow) => workflow.mode); export const selectWorkflowDescription = createWorkflowSelector((workflow) => workflow.description); export const selectWorkflowForm = createWorkflowSelector((workflow) => workflow.form); -export const selectWorkflowIsPublished = createWorkflowSelector((workflow) => workflow.is_published); export const selectFormRootElementId = createWorkflowSelector((workflow) => { return workflow.form.rootElementId; diff --git a/invokeai/frontend/web/src/features/nodes/util/workflow/buildWorkflow.ts b/invokeai/frontend/web/src/features/nodes/util/workflow/buildWorkflow.ts index e410b89bb4..e060c58507 100644 --- a/invokeai/frontend/web/src/features/nodes/util/workflow/buildWorkflow.ts +++ b/invokeai/frontend/web/src/features/nodes/util/workflow/buildWorkflow.ts @@ -33,7 +33,6 @@ const workflowKeys = [ 'meta', 'id', 'form', - 'is_published', ] satisfies (keyof WorkflowV3)[]; type BuildWorkflowFunction = (arg: BuildWorkflowArg) => WorkflowV3; diff --git a/invokeai/frontend/web/src/features/queue/components/InvokeButtonTooltip/InvokeButtonTooltip.tsx b/invokeai/frontend/web/src/features/queue/components/InvokeButtonTooltip/InvokeButtonTooltip.tsx index d5eb9e3032..bec8f2212a 100644 --- a/invokeai/frontend/web/src/features/queue/components/InvokeButtonTooltip/InvokeButtonTooltip.tsx +++ b/invokeai/frontend/web/src/features/queue/components/InvokeButtonTooltip/InvokeButtonTooltip.tsx @@ -6,10 +6,9 @@ import { selectSendToCanvas } from 'features/controlLayers/store/canvasSettingsS import { selectIterations } from 'features/controlLayers/store/paramsSlice'; import { selectDynamicPromptsIsLoading } from 'features/dynamicPrompts/store/dynamicPromptsSlice'; import { selectAutoAddBoardId } from 'features/gallery/store/gallerySelectors'; -import { $isInPublishFlow } from 'features/nodes/components/sidePanel/workflow/publish'; +import { $isInPublishFlow, useIsWorkflowPublished } from 'features/nodes/components/sidePanel/workflow/publish'; import { selectNodesSlice } from 'features/nodes/store/selectors'; import type { NodesState } from 'features/nodes/store/types'; -import { selectWorkflowIsPublished } from 'features/nodes/store/workflowSlice'; import type { BatchSizeResult } from 'features/nodes/util/node/resolveBatchValue'; import { getBatchSize } from 'features/nodes/util/node/resolveBatchValue'; import type { Reason } from 'features/queue/store/readiness'; @@ -178,7 +177,7 @@ const IsReadyText = memo(({ isReady, prepend }: { isReady: boolean; prepend: boo const isLoadingDynamicPrompts = useAppSelector(selectDynamicPromptsIsLoading); const [_, enqueueMutation] = useEnqueueBatchMutation(enqueueMutationFixedCacheKeyOptions); const isInPublishFlow = useStore($isInPublishFlow); - const isPublished = useAppSelector(selectWorkflowIsPublished); + const isPublished = useIsWorkflowPublished(); const text = useMemo(() => { if (enqueueMutation.isLoading) { diff --git a/invokeai/frontend/web/src/features/workflowLibrary/components/WorkflowLibraryMenu/SaveWorkflowMenuItem.tsx b/invokeai/frontend/web/src/features/workflowLibrary/components/WorkflowLibraryMenu/SaveWorkflowMenuItem.tsx index 0ecb749ec6..26ce030296 100644 --- a/invokeai/frontend/web/src/features/workflowLibrary/components/WorkflowLibraryMenu/SaveWorkflowMenuItem.tsx +++ b/invokeai/frontend/web/src/features/workflowLibrary/components/WorkflowLibraryMenu/SaveWorkflowMenuItem.tsx @@ -1,7 +1,6 @@ import { MenuItem } from '@invoke-ai/ui-library'; -import { useAppSelector } from 'app/store/storeHooks'; import { useDoesWorkflowHaveUnsavedChanges } from 'features/nodes/components/sidePanel/workflow/IsolatedWorkflowBuilderWatcher'; -import { selectWorkflowIsPublished } from 'features/nodes/store/workflowSlice'; +import { useIsWorkflowPublished } from 'features/nodes/components/sidePanel/workflow/publish'; import { useSaveOrSaveAsWorkflow } from 'features/workflowLibrary/hooks/useSaveOrSaveAsWorkflow'; import { memo } from 'react'; import { useTranslation } from 'react-i18next'; @@ -11,7 +10,7 @@ const SaveWorkflowMenuItem = () => { const { t } = useTranslation(); const saveOrSaveAsWorkflow = useSaveOrSaveAsWorkflow(); const doesWorkflowHaveUnsavedChanges = useDoesWorkflowHaveUnsavedChanges(); - const isPublished = useAppSelector(selectWorkflowIsPublished); + const isPublished = useIsWorkflowPublished(); return ( { meta: { category }, } = data.workflow; dispatch(workflowIDChanged(id)); - dispatch(workflowIsPublishedChanged(false)); dispatch(workflowNameChanged(name)); dispatch(workflowCategoryChanged(category)); dispatch(newWorkflowSaved({ category })); diff --git a/invokeai/frontend/web/src/features/workflowLibrary/hooks/useSaveOrSaveAsWorkflow.ts b/invokeai/frontend/web/src/features/workflowLibrary/hooks/useSaveOrSaveAsWorkflow.ts index 8c58d7ba6c..b9d65bca89 100644 --- a/invokeai/frontend/web/src/features/workflowLibrary/hooks/useSaveOrSaveAsWorkflow.ts +++ b/invokeai/frontend/web/src/features/workflowLibrary/hooks/useSaveOrSaveAsWorkflow.ts @@ -1,5 +1,4 @@ -import { useAppSelector } from 'app/store/storeHooks'; -import { selectWorkflowIsPublished } from 'features/nodes/store/workflowSlice'; +import { useIsWorkflowPublished } from 'features/nodes/components/sidePanel/workflow/publish'; import { useBuildWorkflowFast } from 'features/nodes/util/workflow/buildWorkflow'; import { saveWorkflowAs } from 'features/workflowLibrary/components/SaveWorkflowAsDialog'; import { isLibraryWorkflow, useSaveLibraryWorkflow } from 'features/workflowLibrary/hooks/useSaveLibraryWorkflow'; @@ -12,7 +11,7 @@ import { useCallback } from 'react'; */ export const useSaveOrSaveAsWorkflow = () => { const buildWorkflow = useBuildWorkflowFast(); - const isPublished = useAppSelector(selectWorkflowIsPublished); + const isPublished = useIsWorkflowPublished(); const { saveWorkflow } = useSaveLibraryWorkflow(); const saveOrSaveAsWorkflow = useCallback(() => { diff --git a/invokeai/frontend/web/src/services/events/setEventListeners.tsx b/invokeai/frontend/web/src/services/events/setEventListeners.tsx index 5e1e529a02..e7d7e33b32 100644 --- a/invokeai/frontend/web/src/services/events/setEventListeners.tsx +++ b/invokeai/frontend/web/src/services/events/setEventListeners.tsx @@ -1,5 +1,7 @@ import { ExternalLink } from '@invoke-ai/ui-library'; +import { isAnyOf } from '@reduxjs/toolkit'; import { logger } from 'app/logging/logger'; +import { listenerMiddleware } from 'app/store/middleware/listenerMiddleware'; import { socketConnected } from 'app/store/middleware/listenerMiddleware/listeners/socketConnected'; import { $baseUrl } from 'app/store/nanostores/baseUrl'; import { $bulkDownloadId } from 'app/store/nanostores/bulkDownloadId'; @@ -9,10 +11,9 @@ import { deepClone } from 'common/util/deepClone'; import { $isInPublishFlow, $outputNodeId, - $validationRunBatchId, + $validationRunData, } from 'features/nodes/components/sidePanel/workflow/publish'; import { $nodeExecutionStates, upsertExecutionState } from 'features/nodes/hooks/useNodeExecutionState'; -import { workflowIsPublishedChanged } from 'features/nodes/store/workflowSlice'; import { zNodeStatus } from 'features/nodes/types/invocation'; import ErrorToastDescription, { getTitle } from 'features/toast/ErrorToastDescription'; import { toast } from 'features/toast/toast'; @@ -22,6 +23,7 @@ import type { ApiTagDescription } from 'services/api'; import { api, LIST_TAG } from 'services/api'; import { modelsApi } from 'services/api/endpoints/models'; import { queueApi, queueItemsAdapter } from 'services/api/endpoints/queue'; +import { workflowsApi } from 'services/api/endpoints/workflows'; import { buildOnInvocationComplete } from 'services/events/onInvocationComplete'; import { buildOnModelInstallError } from 'services/events/onModelInstallError'; import type { ClientToServerEvents, ServerToClientEvents } from 'services/events/types'; @@ -425,14 +427,43 @@ export const setEventListeners = ({ socket, store, setIsConnected }: SetEventLis $lastProgressEvent.set(null); // When a validation run is completed, we want to clear the validation run batch ID & set the workflow as published - if (batch_status.batch_id === $validationRunBatchId.get()) { - $validationRunBatchId.set(null); - if (status === 'completed') { - dispatch(workflowIsPublishedChanged(true)); - $isInPublishFlow.set(false); - $outputNodeId.set(null); - } + const validationRunData = $validationRunData.get(); + if (!validationRunData || batch_status.batch_id !== validationRunData.batchId || status !== 'completed') { + return; } + + // The published status of a workflow is server state, provided to the client in by the getWorkflow query. + // After successfully publishing a workflow, we need to invalidate the query cache so that the published status is + // seen throughout the app. We also need to reset the publish flow state. + // + // But, there is a race condition! If we invalidate the query cache and then immediately clear the publish flow state, + // between the time when the publish state is cleared and the query is re-fetched, we will render the wrong UI. + // + // So, we really need to wait for the query re-fetch to complete before clearing the publish flow state. This isn't + // possible using the `invalidateTags()` API. But we can fudge it by adding a once-off listener for that query. + + listenerMiddleware.startListening({ + matcher: isAnyOf( + workflowsApi.endpoints.getWorkflow.matchFulfilled, + workflowsApi.endpoints.getWorkflow.matchRejected + ), + effect: (action, listenerApi) => { + if (workflowsApi.endpoints.getWorkflow.matchFulfilled(action)) { + // If this query was re-fetching the workflow that was just published, we can clear the publish flow state and + // unsubscribe from the listener + if (action.payload.workflow_id === validationRunData.workflowId) { + listenerApi.unsubscribe(); + $validationRunData.set(null); + $isInPublishFlow.set(false); + $outputNodeId.set(null); + } + } else if (workflowsApi.endpoints.getWorkflow.matchRejected(action)) { + // If the query failed, we can unsubscribe from the listener + listenerApi.unsubscribe(); + } + }, + }); + dispatch(workflowsApi.util.invalidateTags([{ type: 'Workflow', id: validationRunData.workflowId }])); } });