From 8de3b4d03348c45225aa0c5228740f60b55a42e9 Mon Sep 17 00:00:00 2001
From: Joel <iamjoel007@gmail.com>
Date: Mon, 27 Jul 2026 17:32:52 +0800
Subject: [PATCH] fix: prevent the agent build draft UI from flashing during
 apply (#39642)

---
 .../configure/__tests__/page.spec.tsx         | 91 ++++++++++++++++++-
 .../use-agent-configure-build-draft.ts        | 25 +++--
 2 files changed, 105 insertions(+), 11 deletions(-)

diff --git a/web/features/agent-v2/agent-detail/configure/__tests__/page.spec.tsx b/web/features/agent-v2/agent-detail/configure/__tests__/page.spec.tsx
index 4c971ea3cdb..08152cbd89e 100644
--- a/web/features/agent-v2/agent-detail/configure/__tests__/page.spec.tsx
+++ b/web/features/agent-v2/agent-detail/configure/__tests__/page.spec.tsx
@@ -356,12 +356,14 @@ vi.mock('../components/orchestrate/build-draft-bar', () => ({
     changeSummary?: unknown
     changesCount: number
     disabled?: boolean
+    isApplying?: boolean
     onApply: () => void
     onDiscard: () => void
   }) => (
     <div role="region" aria-label="build-draft-bar">
       <span>{`changes:${props.changesCount}`}</span>
-      <button type="button" disabled={props.disabled} onClick={props.onApply}>
+      <span>{`applying:${props.isApplying ? 'yes' : 'no'}`}</span>
+      <button type="button" disabled={props.disabled || props.isApplying} onClick={props.onApply}>
         apply build draft
       </button>
       <button type="button" disabled={props.disabled} onClick={props.onDiscard}>
@@ -2959,6 +2961,93 @@ describe('AgentConfigurePage', () => {
       })
     })
 
+    it('should keep the build draft UI while the applied normal draft is still refreshing', async () => {
+      const user = userEvent.setup()
+      const queryClient = new QueryClient()
+      const refetchComposerDeferred = createDeferredPromise<unknown>()
+      const refetchComposer = vi.fn(async () => {
+        const result = await refetchComposerDeferred.promise
+        mocks.queryState.composer = {
+          ...mocks.queryState.composer,
+          data: {
+            agent_soul: {
+              prompt: {
+                system_prompt: 'applied prompt',
+              },
+            },
+          },
+        }
+
+        return result
+      })
+      mocks.queryState.composer = {
+        data: {
+          agent_soul: {
+            prompt: {
+              system_prompt: 'old draft prompt',
+            },
+          },
+        },
+        isFetching: false,
+        isError: false,
+        isPending: false,
+        isSuccess: true,
+        refetch: refetchComposer,
+      }
+      mocks.queryState.buildDraft = {
+        data: {
+          agent_soul: {
+            prompt: {
+              system_prompt: 'build prompt',
+            },
+          },
+          draft: {},
+          variant: 'agent_app',
+        },
+        dataUpdatedAt: 1,
+        error: null,
+        isFetching: false,
+        isError: false,
+        isPending: false,
+        isSuccess: true,
+        refetch: vi.fn(),
+      }
+
+      render(
+        <QueryClientProvider client={queryClient}>
+          <AgentConfigurePage agentId="agent-1" />
+        </QueryClientProvider>,
+      )
+
+      await user.click(screen.getByRole('button', { name: 'apply build draft' }))
+      await waitFor(() => expect(refetchComposer).toHaveBeenCalled())
+
+      expect(screen.getByRole('region', { name: 'build-draft-bar' })).toHaveTextContent(
+        'applying:yes',
+      )
+      expect(screen.getByRole('button', { name: 'apply build draft' })).toBeDisabled()
+      expect(screen.getByRole('region', { name: 'orchestrate-panel' })).toHaveTextContent(
+        'buildDraft:yes',
+      )
+      expect(screen.getByRole('region', { name: 'orchestrate-panel' })).toHaveTextContent(
+        'prompt:build prompt',
+      )
+      expect(toastMock.success).not.toHaveBeenCalled()
+
+      await act(async () => {
+        refetchComposerDeferred.resolve({})
+        await refetchComposerDeferred.promise
+      })
+
+      await waitFor(() => {
+        expect(screen.queryByRole('region', { name: 'build-draft-bar' })).not.toBeInTheDocument()
+        expect(screen.getByRole('region', { name: 'orchestrate-panel' })).toHaveTextContent(
+          'prompt:applied prompt',
+        )
+      })
+      expect(toastMock.success).toHaveBeenCalled()
+    })
+
     it('should keep exiting build draft when debug conversation refresh fails after applying build draft', async () => {
       const user = userEvent.setup()
       const queryClient = new QueryClient()
diff --git a/web/features/agent-v2/agent-detail/configure/use-agent-configure-build-draft.ts b/web/features/agent-v2/agent-detail/configure/use-agent-configure-build-draft.ts
index ab83bf70f2a..65c974c7ff3 100644
--- a/web/features/agent-v2/agent-detail/configure/use-agent-configure-build-draft.ts
+++ b/web/features/agent-v2/agent-detail/configure/use-agent-configure-build-draft.ts
@@ -14,7 +14,7 @@ import type {
 import { toast } from '@langgenius/dify-ui/toast'
 import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'
 import isEqual from 'fast-deep-equal'
-import { useCallback, useEffect, useMemo, useRef } from 'react'
+import { useCallback, useEffect, useMemo, useRef, useState } from 'react'
 import { useTranslation } from 'react-i18next'
 import { agentSoulConfigToFormState } from '@/features/agent-v2/agent-composer/conversions'
 import { consoleQuery } from '@/service/client'
@@ -299,6 +299,7 @@ export function useAgentConfigureBuildDraftActions({
   const buildDraftRefreshTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null)
   const buildDraftRefreshGenerationRef = useRef(0)
   const forceCheckoutBeforeNextBuildRunRef = useRef(false)
+  const [isApplyingBuildDraftWorkflow, setIsApplyingBuildDraftWorkflow] = useState(false)
   const buildDraftQueryOptions = consoleQuery.agent.byAgentId.buildDraft.get.queryOptions({
     input: {
       params: {
@@ -320,7 +321,7 @@ export function useAgentConfigureBuildDraftActions({
   )
   const { mutateAsync: finalizeBuildChatRequest, isPending: isFinalizingBuildChat } =
     finalizeBuildChatMutation
-  const { mutateAsync: applyBuildDraftRequest, isPending: isApplyingBuildDraft } =
+  const { mutateAsync: applyBuildDraftRequest, isPending: isApplyingBuildDraftRequest } =
     applyBuildDraftMutation
   const { mutateAsync: discardBuildDraftRequest, isPending: isDiscardingBuildDraft } =
     discardBuildDraftMutation
@@ -402,17 +403,17 @@ export function useAgentConfigureBuildDraftActions({
     async (shouldRefetchComposer: boolean) => {
       cancelBuildDraftRefresh()
       await resetBuildChatSession().catch(() => undefined)
+      let nextAgentSoulConfig = normalAgentSoulConfig
+      if (shouldRefetchComposer) {
+        const result = await refetchComposer()
+        nextAgentSoulConfig = getAgentSoulConfigFromRefetchResult(result) ?? normalAgentSoulConfig
+      }
       setSoulSourceOverride('draft')
       queryClient.removeQueries({
         queryKey: buildDraftQueryOptions.queryKey,
       })
-      if (shouldRefetchComposer) {
-        const result = await refetchComposer()
-        rebaseComposerDraft(getAgentSoulConfigFromRefetchResult(result) ?? normalAgentSoulConfig)
-        onComposerRebased?.()
-      } else {
-        rebaseComposerDraft(normalAgentSoulConfig)
-      }
+      rebaseComposerDraft(nextAgentSoulConfig)
+      if (shouldRefetchComposer) onComposerRebased?.()
     },
     [
       buildDraftQueryOptions.queryKey,
@@ -428,6 +429,7 @@ export function useAgentConfigureBuildDraftActions({
   )
 
   const applyBuildDraft = async () => {
+    setIsApplyingBuildDraftWorkflow(true)
     try {
       await finalizeBuildChatRequest({
         params: {
@@ -449,6 +451,8 @@ export function useAgentConfigureBuildDraftActions({
       toast.success(tCommon(($) => $['api.actionSuccess']))
     } catch {
       toast.error(tCommon(($) => $['api.actionFailed']))
+    } finally {
+      setIsApplyingBuildDraftWorkflow(false)
     }
   }
 
@@ -478,7 +482,8 @@ export function useAgentConfigureBuildDraftActions({
     applyBuildDraft,
     cancelBuildDraftRefresh,
     discardBuildDraft,
-    isApplyingBuildDraft: isFinalizingBuildChat || isApplyingBuildDraft,
+    isApplyingBuildDraft:
+      isApplyingBuildDraftWorkflow || isFinalizingBuildChat || isApplyingBuildDraftRequest,
     isDiscardingBuildDraft,
     prepareBuildDraftBeforeRun: prepareBuildDraftRun,
     refreshBuildDraftAfterBuildChat,