refactor(web): encapsulate agent version restore interaction (#42597)

web/features/agent-v2/agent-detail/configure/components/orchestrate/__tests__/publish-bar.spec.tsx

@@ -307,7 +307,10 @@ function renderPublishBar({
     queryClient.setQueryData(composerQueryKey, composerState)
   }
 
-  const renderPublishBarTree = (nextProps?: { isPublishing?: boolean }) => (
+  const renderPublishBarTree = (nextProps?: {
+    isPublishing?: boolean
+    selectedVersionSnapshot?: AgentConfigSnapshotSummaryResponse | null
+  }) => (
     <NuqsWrapper>
       <QueryClientProvider client={queryClient}>
         <JotaiProvider store={store}>
@@ -315,7 +318,11 @@ function renderPublishBar({
             agentId="agent-1"
             agentName="Iris"
             isPublishing={nextProps?.isPublishing ?? isPublishing}
-            selectedVersionSnapshot={selectedVersionSnapshot}
+            selectedVersionSnapshot={
+              nextProps?.selectedVersionSnapshot === undefined
+                ? selectedVersionSnapshot
+                : nextProps.selectedVersionSnapshot
+            }
             onPublish={onPublish}
             onExitVersions={onExitVersions}
             onOpenVersions={vi.fn()}
@@ -460,6 +467,43 @@ describe('AgentConfigurePublishBar', () => {
     expect(toastMock.success).toHaveBeenCalledWith('common.api.actionSuccess')
   })
 
+  it('starts a fresh restore confirmation after leaving and re-entering version view', async () => {
+    const user = userEvent.setup()
+    const { rerender, rerenderPublishBar } = renderPublishBar({
+      selectedVersionSnapshot: activeConfigSnapshot,
+    })
+    await user.click(
+      screen.getByRole('button', { name: /agentV2.agentDetail.versionHistory.restore/ }),
+    )
+    expect(await screen.findByRole('alertdialog')).toBeInTheDocument()
+
+    rerender(rerenderPublishBar({ selectedVersionSnapshot: null }))
+    await waitFor(() => expect(screen.queryByRole('alertdialog')).not.toBeInTheDocument())
+    rerender(
+      rerenderPublishBar({
+        selectedVersionSnapshot: {
+          ...activeConfigSnapshot,
+          id: 'snapshot-2',
+          version_note: 'Version two',
+        },
+      }),
+    )
+    expect(screen.queryByRole('alertdialog')).not.toBeInTheDocument()
+    expect(restoreVersionMutation).not.toHaveBeenCalled()
+
+    await user.click(
+      screen.getByRole('button', { name: /agentV2.agentDetail.versionHistory.restore/ }),
+    )
+    const dialog = await screen.findByRole('alertdialog', { name: /Version two/ })
+    await user.click(within(dialog).getByRole('button', { name: /restore/i }))
+    await waitFor(() =>
+      expect(restoreVersionMutation).toHaveBeenCalledWith(
+        { params: { agent_id: 'agent-1', version_id: 'snapshot-2' } },
+        expect.anything(),
+      ),
+    )
+  })
+
   it('shows the agent upgrade dialog before opening pricing on the sandbox plan', async () => {
     restoreAccess.plan = 'sandbox'
     const user = userEvent.setup()

web/features/agent-v2/agent-detail/configure/components/orchestrate/publish-bar/index.tsx

@@ -21,8 +21,7 @@ import { isAgentComposerDirtyAtom } from '@/features/agent-v2/agent-composer/sto
 import { useFormatTimeFromNow } from '@/hooks/use-format-time-from-now'
 import useTimestamp from '@/hooks/use-timestamp'
 import { consoleQuery } from '@/service/console'
-import { AgentVersionRestoreDialogs } from '../../preview/versions-panel/restore-dialogs'
-import { useAgentVersionRestore } from '../../preview/versions-panel/use-agent-version-restore'
+import { AgentVersionRestore } from '../../version-restore'
 import { AgentPublishImpactDetails } from './publish-impact-details'
 
 const PUBLISH_AGENT_HOTKEY = 'Mod+Shift+P' satisfies Hotkey
@@ -135,13 +134,6 @@ export function AgentConfigurePublishBar({
       enabled: publishIsAvailable && !isPublishing && !selectedVersionSnapshot,
     })
   useQuery(workflowReferencesQueryOptions)
-  const restore = useAgentVersionRestore({
-    agentId,
-    onRestored: async () => {
-      await onVersionRestored?.()
-      onExitVersions?.()
-    },
-  })
   const canPublish = publishIsAvailable && !isPublishing
 
   const handlePublish = async () => {
@@ -197,27 +189,25 @@ export function AgentConfigurePublishBar({
 
   if (selectedVersionSnapshot) {
     return (
-      <>
-        <AgentVersionRestoreBar
-          version={selectedVersionSnapshot}
-          isRestoring={restore.isPending}
-          restoreDisabled={restore.disabled}
-          onExitVersions={onExitVersions}
-          onRestoreVersion={
-            restore.canRestore ? () => restore.requestRestore(selectedVersionSnapshot) : undefined
-          }
-        />
-        <AgentVersionRestoreDialogs
-          version={restore.version}
-          isUpgradeOpen={restore.isUpgradeOpen}
-          onUpgradeClose={restore.closeUpgrade}
-          isConfirmOpen={restore.isConfirmOpen}
-          onConfirmOpenChange={restore.onConfirmOpenChange}
-          isPending={restore.isPending}
-          disabled={restore.disabled}
-          onConfirm={restore.confirmRestore}
-        />
-      </>
+      <AgentVersionRestore
+        agentId={agentId}
+        onRestored={async () => {
+          await onVersionRestored?.()
+          onExitVersions?.()
+        }}
+      >
+        {(restore) => (
+          <AgentVersionRestoreBar
+            version={selectedVersionSnapshot}
+            isRestoring={restore.isPending}
+            restoreDisabled={restore.disabled}
+            onExitVersions={onExitVersions}
+            onRestoreVersion={
+              restore.canRestore ? () => restore.requestRestore(selectedVersionSnapshot) : undefined
+            }
+          />
+        )}
+      </AgentVersionRestore>
     )
   }
 

web/features/agent-v2/agent-detail/configure/components/preview/versions-panel/index.tsx

@@ -15,10 +15,9 @@ import { userProfileQueryOptions } from '@/features/account-profile/client'
 import { useAgentPermissions } from '@/features/agent-v2/permissions'
 import { deploymentEditionAtom } from '@/features/system-features/state'
 import { consoleQuery } from '@/service/console'
+import { AgentVersionRestore } from '../../version-restore'
 import { CurrentDraftItem } from './current-draft-item'
 import { VersionFilter } from './filter'
-import { AgentVersionRestoreDialogs } from './restore-dialogs'
-import { useAgentVersionRestore } from './use-agent-version-restore'
 import { VersionFilterEmpty } from './version-filter-empty'
 import { VersionItem } from './version-item'
 
@@ -41,15 +40,6 @@ export function AgentPreviewVersionsPanel({
   onBeforeRestore,
   restoreDisabled,
 }: AgentPreviewVersionsPanelProps) {
-  const restore = useAgentVersionRestore({
-    agentId,
-    onBeforeRestore,
-    disabled: restoreDisabled,
-    onRestored: async () => {
-      await onVersionRestored?.()
-      onSelectVersion(null)
-    },
-  })
   const { t } = useTranslation('agentV2')
   const { t: tCommon } = useTranslation('common')
   const { t: tWorkflow } = useTranslation('workflow')
@@ -124,58 +114,60 @@ export function AgentPreviewVersionsPanel({
         </button>
       </div>
 
-      <div className="min-h-0 flex-1 overflow-y-auto px-3 py-2">
-        {versionsQuery.isPending && (
-          <div className="space-y-1">
-            <div className="h-10 animate-pulse rounded-lg bg-state-base-hover" />
-            <div className="h-18 animate-pulse rounded-lg bg-state-base-hover" />
-            <div className="h-10 animate-pulse rounded-lg bg-state-base-hover" />
-          </div>
-        )}
-        {!versionsQuery.isPending && versions.length === 0 && (
-          <div className="rounded-lg border border-components-panel-border bg-components-panel-on-panel-item-bg px-3 py-6 text-center system-sm-regular text-text-tertiary">
-            {t(($) => $['agentDetail.versionHistory.empty'])}
-          </div>
-        )}
-        {!versionsQuery.isPending && versions.length > 0 && (
-          <div className="flex flex-col gap-px">
-            <CurrentDraftItem
-              isActive={!activeVersionId}
-              isLast={filteredVersions.length === 0}
-              onSelect={() => onSelectVersion(null)}
-            />
-            {filteredVersions.length === 0 && isFiltering && (
-              <VersionFilterEmpty onReset={handleResetFilter} />
+      <AgentVersionRestore
+        agentId={agentId}
+        onBeforeRestore={onBeforeRestore}
+        disabled={restoreDisabled}
+        onRestored={async () => {
+          await onVersionRestored?.()
+          onSelectVersion(null)
+        }}
+      >
+        {(restore) => (
+          <div className="min-h-0 flex-1 overflow-y-auto px-3 py-2">
+            {versionsQuery.isPending && (
+              <div className="space-y-1">
+                <div className="h-10 animate-pulse rounded-lg bg-state-base-hover" />
+                <div className="h-18 animate-pulse rounded-lg bg-state-base-hover" />
+                <div className="h-10 animate-pulse rounded-lg bg-state-base-hover" />
+              </div>
+            )}
+            {!versionsQuery.isPending && versions.length === 0 && (
+              <div className="rounded-lg border border-components-panel-border bg-components-panel-on-panel-item-bg px-3 py-6 text-center system-sm-regular text-text-tertiary">
+                {t(($) => $['agentDetail.versionHistory.empty'])}
+              </div>
+            )}
+            {!versionsQuery.isPending && versions.length > 0 && (
+              <div className="flex flex-col gap-px">
+                <CurrentDraftItem
+                  isActive={!activeVersionId}
+                  isLast={filteredVersions.length === 0}
+                  onSelect={() => onSelectVersion(null)}
+                />
+                {filteredVersions.length === 0 && isFiltering && (
+                  <VersionFilterEmpty onReset={handleResetFilter} />
+                )}
+                {filteredVersions.map((version, index) => (
+                  <VersionItem
+                    key={version.id}
+                    version={version}
+                    activeVersionId={activeVersionId}
+                    isLatest={version.id === latestVersionId}
+                    isFirst={false}
+                    isLast={index === filteredVersions.length - 1}
+                    onSelect={onSelectVersion}
+                    onExport={canImportExportDSL ? handleExport : undefined}
+                    onRestore={restore.canRestore ? restore.requestRestore : undefined}
+                    restoreDisabled={restore.disabled}
+                    exportDisabled={exportDisabled}
+                    showUpgrade={shouldShowUpgrade}
+                  />
+                ))}
+              </div>
             )}
-            {filteredVersions.map((version, index) => (
-              <VersionItem
-                key={version.id}
-                version={version}
-                activeVersionId={activeVersionId}
-                isLatest={version.id === latestVersionId}
-                isFirst={false}
-                isLast={index === filteredVersions.length - 1}
-                onSelect={onSelectVersion}
-                onExport={canImportExportDSL ? handleExport : undefined}
-                onRestore={restore.canRestore ? restore.requestRestore : undefined}
-                restoreDisabled={restore.disabled}
-                exportDisabled={exportDisabled}
-                showUpgrade={shouldShowUpgrade}
-              />
-            ))}
           </div>
         )}
-      </div>
-      <AgentVersionRestoreDialogs
-        version={restore.version}
-        isUpgradeOpen={restore.isUpgradeOpen}
-        onUpgradeClose={restore.closeUpgrade}
-        isConfirmOpen={restore.isConfirmOpen}
-        onConfirmOpenChange={restore.onConfirmOpenChange}
-        isPending={restore.isPending}
-        disabled={restore.disabled}
-        onConfirm={restore.confirmRestore}
-      />
+      </AgentVersionRestore>
     </aside>
   )
 }

web/features/agent-v2/agent-detail/configure/components/preview/versions-panel/restore-dialogs.tsx

@@ -1,76 +0,0 @@
-import type { AgentConfigSnapshotSummaryResponse } from '@dify/contracts/api/console/agent/types.gen'
-import {
-  AlertDialog,
-  AlertDialogActions,
-  AlertDialogCancelButton,
-  AlertDialogConfirmButton,
-  AlertDialogContent,
-  AlertDialogDescription,
-  AlertDialogTitle,
-} from '@langgenius/dify-ui/alert-dialog'
-import { useTranslation } from 'react-i18next'
-import { PlanUpgradeModal } from '@/app/components/billing/plan-upgrade-modal'
-
-type AgentVersionRestoreDialogsProps = {
-  version: AgentConfigSnapshotSummaryResponse | null
-  isUpgradeOpen: boolean
-  onUpgradeClose: () => void
-  isConfirmOpen: boolean
-  onConfirmOpenChange: (open: boolean) => void
-  isPending: boolean
-  disabled: boolean
-  onConfirm: () => void
-}
-
-export function AgentVersionRestoreDialogs({
-  version,
-  isUpgradeOpen,
-  onUpgradeClose,
-  isConfirmOpen,
-  onConfirmOpenChange,
-  isPending,
-  disabled,
-  onConfirm,
-}: AgentVersionRestoreDialogsProps) {
-  const { t } = useTranslation('agentV2')
-  const versionLabel = version
-    ? version.version_note ||
-      t(($) => $['agentDetail.versionHistory.versionName'], { version: version.version })
-    : ''
-
-  return (
-    <>
-      <PlanUpgradeModal
-        show={isUpgradeOpen}
-        onClose={onUpgradeClose}
-        title={t(($) => $['upgrade.agentRestore.title'], { ns: 'billing' })}
-        description={t(($) => $['upgrade.agentRestore.description'], { ns: 'billing' })}
-      />
-      <AlertDialog open={isConfirmOpen} onOpenChange={onConfirmOpenChange}>
-        <AlertDialogContent>
-          <div className="flex flex-col gap-2 p-6 pb-4">
-            <AlertDialogTitle className="title-2xl-semi-bold text-text-primary">
-              {`${t(($) => $['agentDetail.versionHistory.restore'])} ${versionLabel}`}
-            </AlertDialogTitle>
-            <AlertDialogDescription className="system-md-regular text-text-secondary">
-              {t(($) => $['versionHistory.restorationTip'], { ns: 'workflow' })}
-            </AlertDialogDescription>
-          </div>
-          <AlertDialogActions>
-            <AlertDialogCancelButton variant="secondary" disabled={isPending}>
-              {t(($) => $['operation.cancel'], { ns: 'common' })}
-            </AlertDialogCancelButton>
-            <AlertDialogConfirmButton
-              tone="default"
-              disabled={disabled}
-              loading={isPending}
-              onClick={onConfirm}
-            >
-              {t(($) => $['agentDetail.versionHistory.restore'])}
-            </AlertDialogConfirmButton>
-          </AlertDialogActions>
-        </AlertDialogContent>
-      </AlertDialog>
-    </>
-  )
-}

web/features/agent-v2/agent-detail/configure/components/version-restore/index.tsx

@@ -1,19 +1,41 @@
+'use client'
+
 import type { AgentConfigSnapshotSummaryResponse } from '@dify/contracts/api/console/agent/types.gen'
+import type { ReactNode } from 'react'
+import {
+  AlertDialog,
+  AlertDialogActions,
+  AlertDialogCancelButton,
+  AlertDialogConfirmButton,
+  AlertDialogContent,
+  AlertDialogDescription,
+  AlertDialogTitle,
+} from '@langgenius/dify-ui/alert-dialog'
 import { useMutation, useQuery, useQueryClient } from '@tanstack/react-query'
 import { useAtomValue } from 'jotai'
 import { useState } from 'react'
 import { useTranslation } from 'react-i18next'
+import { PlanUpgradeModal } from '@/app/components/billing/plan-upgrade-modal'
 import { toast } from '@/app/notifications'
 import { useAgentPermissions } from '@/features/agent-v2/permissions'
 import { deploymentEditionAtom } from '@/features/system-features/state'
 import { consoleQuery } from '@/service/console'
 
-export function useAgentVersionRestore({
+type RestoreActions = {
+  requestRestore: (version: AgentConfigSnapshotSummaryResponse) => void
+  canRestore: boolean
+  disabled: boolean
+  isPending: boolean
+}
+
+export function AgentVersionRestore({
   agentId,
   onBeforeRestore,
   disabled: externallyDisabled = false,
   onRestored,
+  children,
 }: {
+  children: (actions: RestoreActions) => ReactNode
   agentId: string
   disabled?: boolean
   onBeforeRestore?: () => void | Promise<void>
@@ -29,9 +51,8 @@ export function useAgentVersionRestore({
       select: (data) => data.billing.subscription.plan,
     }),
   )
-  const [isUpgradeOpen, setIsUpgradeOpen] = useState(false)
+  const [stage, setStage] = useState<'closed' | 'upgrade' | 'confirm'>('closed')
   const [version, setVersion] = useState<AgentConfigSnapshotSummaryResponse | null>(null)
-  const [isConfirmOpen, setIsConfirmOpen] = useState(false)
   const showUpgrade = deploymentEdition === 'CLOUD' && plan === 'sandbox'
   const disabled =
     externallyDisabled || !canRestore || (deploymentEdition === 'CLOUD' && plan === undefined)
@@ -57,7 +78,7 @@ export function useAgentVersionRestore({
           }),
         ])
         await onRestored?.()
-        setIsConfirmOpen(false)
+        setStage('closed')
         toast.success(t(($) => $['api.actionSuccess'], { ns: 'common' }))
       },
       onError: () => {
@@ -69,36 +90,71 @@ export function useAgentVersionRestore({
   const requestRestore = (selectedVersion: AgentConfigSnapshotSummaryResponse) => {
     if (disabled || mutation.isPending) return
     if (showUpgrade) {
-      setIsUpgradeOpen(true)
+      setStage('upgrade')
       return
     }
     setVersion(selectedVersion)
-    setIsConfirmOpen(true)
+    setStage('confirm')
   }
 
   const confirmRestore = () => {
     if (!version || disabled || mutation.isPending) return
     if (showUpgrade) {
-      setIsConfirmOpen(false)
-      setIsUpgradeOpen(true)
+      setStage('upgrade')
       return
     }
     mutation.mutate({ params: { agent_id: agentId, version_id: version.id } })
   }
 
-  return {
-    requestRestore,
-    isPending: mutation.isPending,
-    disabled: disabled || mutation.isPending,
-    canRestore,
-    showUpgrade,
-    version,
-    isUpgradeOpen,
-    closeUpgrade: () => setIsUpgradeOpen(false),
-    isConfirmOpen,
-    onConfirmOpenChange: (open: boolean) => {
-      if (!mutation.isPending) setIsConfirmOpen(open)
-    },
-    confirmRestore,
-  }
+  const versionLabel = version
+    ? version.version_note ||
+      t(($) => $['agentDetail.versionHistory.versionName'], { version: version.version })
+    : ''
+
+  return (
+    <>
+      {children({
+        requestRestore,
+        isPending: mutation.isPending,
+        disabled: disabled || mutation.isPending,
+        canRestore,
+      })}
+      <PlanUpgradeModal
+        show={stage === 'upgrade'}
+        onClose={() => setStage('closed')}
+        title={t(($) => $['upgrade.agentRestore.title'], { ns: 'billing' })}
+        description={t(($) => $['upgrade.agentRestore.description'], { ns: 'billing' })}
+      />
+      <AlertDialog
+        open={stage === 'confirm'}
+        onOpenChange={(open) => {
+          if (!mutation.isPending) setStage(open ? 'confirm' : 'closed')
+        }}
+      >
+        <AlertDialogContent>
+          <div className="flex flex-col gap-2 p-6 pb-4">
+            <AlertDialogTitle className="title-2xl-semi-bold text-text-primary">
+              {`${t(($) => $['agentDetail.versionHistory.restore'])} ${versionLabel}`}
+            </AlertDialogTitle>
+            <AlertDialogDescription className="system-md-regular text-text-secondary">
+              {t(($) => $['versionHistory.restorationTip'], { ns: 'workflow' })}
+            </AlertDialogDescription>
+          </div>
+          <AlertDialogActions>
+            <AlertDialogCancelButton variant="secondary" disabled={mutation.isPending}>
+              {t(($) => $['operation.cancel'], { ns: 'common' })}
+            </AlertDialogCancelButton>
+            <AlertDialogConfirmButton
+              tone="default"
+              disabled={disabled || mutation.isPending}
+              loading={mutation.isPending}
+              onClick={confirmRestore}
+            >
+              {t(($) => $['agentDetail.versionHistory.restore'])}
+            </AlertDialogConfirmButton>
+          </AlertDialogActions>
+        </AlertDialogContent>
+      </AlertDialog>
+    </>
+  )
 }