feat: refactor workflow definition handling and integrate undo/redo functionality in WorkflowEditor

web/app/dashboard/ai-agents/_components/config-workbench.tsx

@@ -70,7 +70,6 @@ import {
   IMConversationServiceMode,
   Status,
 } from "@/lib/generated/enums"
-import { useWorkflowDefinitionHistory } from "../../ai-workflows/_components/use-workflow-definition-history"
 import { WorkflowEditor } from "../../ai-workflows/_components/workflow-editor"
 
 type DirectToolItem = CreateAIAgentPayload["directTools"][number]
@@ -154,16 +153,8 @@ export function AIAgentConfigWorkbench({
   const [selectedSkillIds, setSelectedSkillIds] = useState<number[]>([])
   const [directTools, setDirectTools] = useState<DirectToolItem[]>([])
 
-  const {
-    definition,
-    revision: workflowRevision,
-    canUndo: canUndoWorkflow,
-    canRedo: canRedoWorkflow,
-    replace: replaceWorkflowHistory,
-    update: updateWorkflowDefinition,
-    undo: undoWorkflowDefinition,
-    redo: redoWorkflowDefinition,
-  } = useWorkflowDefinitionHistory(fallbackDefinition)
+  const [definition, setDefinition] = useState<AIWorkflowDefinition>(fallbackDefinition)
+  const [workflowRevision, setWorkflowRevision] = useState(0)
 
   const [aiConfigs, setAIConfigs] = useState<AIConfig[]>([])
   const [knowledgeBases, setKnowledgeBases] = useState<KnowledgeBase[]>([])
@@ -181,8 +172,9 @@ export function AIAgentConfigWorkbench({
   }, [agentId])
 
   const replaceWorkflowDefinition = useCallback((nextDefinition: AIWorkflowDefinition) => {
-    replaceWorkflowHistory(nextDefinition)
-  }, [replaceWorkflowHistory])
+    setDefinition(nextDefinition)
+    setWorkflowRevision((current) => current + 1)
+  }, [])
 
   const loadData = useCallback(async () => {
     setLoading(true)
@@ -823,11 +815,8 @@ export function AIAgentConfigWorkbench({
                   key={workflowRevision}
                   definition={definition}
                   nodeSpecs={nodeSpecs}
-                  onDefinitionChange={updateWorkflowDefinition}
-                  onUndo={undoWorkflowDefinition}
-                  undoDisabled={!canUndoWorkflow || savingWorkflow || loading}
-                  onRedo={redoWorkflowDefinition}
-                  redoDisabled={!canRedoWorkflow || savingWorkflow || loading}
+                  onDefinitionChange={setDefinition}
+                  historyDisabled={savingWorkflow || loading}
                   onRestoreDefault={restoreDefaultWorkflow}
                   restoreDefaultDisabled={savingWorkflow || loading}
                   onValidate={validateWorkflowDraft}

web/app/dashboard/ai-workflows/_components/use-workflow-definition-history.ts

@@ -1,89 +0,0 @@
-"use client"
-
-import { useCallback, useState } from "react"
-
-import type { AIWorkflowDefinition } from "@/lib/api/admin"
-
-type WorkflowDefinitionHistoryState = {
-  present: AIWorkflowDefinition
-  past: AIWorkflowDefinition[]
-  future: AIWorkflowDefinition[]
-  revision: number
-}
-
-export function useWorkflowDefinitionHistory(initialDefinition: AIWorkflowDefinition) {
-  const [state, setState] = useState<WorkflowDefinitionHistoryState>({
-    present: initialDefinition,
-    past: [],
-    future: [],
-    revision: 0,
-  })
-
-  const replace = useCallback((definition: AIWorkflowDefinition) => {
-    setState((current) => ({
-      present: definition,
-      past: [],
-      future: [],
-      revision: current.revision + 1,
-    }))
-  }, [])
-
-  const update = useCallback((definition: AIWorkflowDefinition) => {
-    setState((current) => {
-      if (sameWorkflowDefinition(current.present, definition)) {
-        return current
-      }
-      return {
-        present: definition,
-        past: [...current.past.slice(-49), current.present],
-        future: [],
-        revision: current.revision,
-      }
-    })
-  }, [])
-
-  const undo = useCallback(() => {
-    setState((current) => {
-      const previous = current.past[current.past.length - 1]
-      if (!previous) {
-        return current
-      }
-      return {
-        present: previous,
-        past: current.past.slice(0, -1),
-        future: [current.present, ...current.future.slice(0, 49)],
-        revision: current.revision + 1,
-      }
-    })
-  }, [])
-
-  const redo = useCallback(() => {
-    setState((current) => {
-      const next = current.future[0]
-      if (!next) {
-        return current
-      }
-      return {
-        present: next,
-        past: [...current.past.slice(-49), current.present],
-        future: current.future.slice(1),
-        revision: current.revision + 1,
-      }
-    })
-  }, [])
-
-  return {
-    definition: state.present,
-    revision: state.revision,
-    canUndo: state.past.length > 0,
-    canRedo: state.future.length > 0,
-    replace,
-    update,
-    undo,
-    redo,
-  }
-}
-
-function sameWorkflowDefinition(left: AIWorkflowDefinition, right: AIWorkflowDefinition) {
-  return JSON.stringify(left) === JSON.stringify(right)
-}

web/app/dashboard/ai-workflows/_components/workflow-editor.tsx

@@ -13,6 +13,7 @@ import {
   type WorkflowNodeJSON,
   type WorkflowPortEntity,
   useClientContext,
+  useUndoRedo,
   usePlaygroundTools,
   useService,
 } from "@flowgram.ai/free-layout-editor"
@@ -51,10 +52,7 @@ export function WorkflowEditor({
   onDefinitionChange,
   onRestoreDefault,
   restoreDefaultDisabled = false,
-  onUndo,
-  undoDisabled = false,
-  onRedo,
-  redoDisabled = false,
+  historyDisabled = false,
   onValidate,
   validateDisabled = false,
   onSaveDraft,
@@ -68,10 +66,7 @@ export function WorkflowEditor({
   onDefinitionChange: (definition: AIWorkflowDefinition) => void
   onRestoreDefault?: () => void
   restoreDefaultDisabled?: boolean
-  onUndo?: () => void
-  undoDisabled?: boolean
-  onRedo?: () => void
-  redoDisabled?: boolean
+  historyDisabled?: boolean
   onValidate?: () => void
   validateDisabled?: boolean
   onSaveDraft?: () => void
@@ -136,10 +131,7 @@ export function WorkflowEditor({
             }
           }}
           onSelectBranch={handleSelectBranch}
-          onUndo={onUndo}
-          undoDisabled={undoDisabled}
-          onRedo={onRedo}
-          redoDisabled={redoDisabled}
+          historyDisabled={historyDisabled}
           onRestoreDefault={onRestoreDefault}
           restoreDefaultDisabled={restoreDefaultDisabled}
           onValidate={onValidate}
@@ -164,10 +156,7 @@ function WorkflowEditorInner({
   onDefinitionChange,
   onSelectNode,
   onSelectBranch,
-  onUndo,
-  undoDisabled,
-  onRedo,
-  redoDisabled,
+  historyDisabled,
   onRestoreDefault,
   restoreDefaultDisabled,
   onValidate,
@@ -186,10 +175,7 @@ function WorkflowEditorInner({
   onDefinitionChange: (definition: AIWorkflowDefinition) => void
   onSelectNode: (nodeId: string) => void
   onSelectBranch: (branch: SelectedWorkflowBranch | null) => void
-  onUndo?: () => void
-  undoDisabled?: boolean
-  onRedo?: () => void
-  redoDisabled?: boolean
+  historyDisabled?: boolean
   onRestoreDefault?: () => void
   restoreDefaultDisabled?: boolean
   onValidate?: () => void
@@ -201,6 +187,7 @@ function WorkflowEditorInner({
 }) {
   const context = useClientContext()
   const playgroundTools = usePlaygroundTools()
+  const undoRedo = useUndoRedo()
   const workflowDocument = useService(WorkflowDocument)
   const linesManager = useService(WorkflowLinesManager)
   const selectService = useService(WorkflowSelectService)
@@ -229,6 +216,16 @@ function WorkflowEditorInner({
     onDefinitionChange(context.document.toJSON() as AIWorkflowDefinition)
   }
 
+  const undo = async () => {
+    await undoRedo.undo()
+    emitCurrentDefinition()
+  }
+
+  const redo = async () => {
+    await undoRedo.redo()
+    emitCurrentDefinition()
+  }
+
   const openNodeMenuFromPort = useCallback((request: WorkflowPortAddRequest) => {
     const rootRect = editorRootRef.current?.getBoundingClientRect()
     setNodeMenu({
@@ -336,10 +333,10 @@ function WorkflowEditorInner({
       >
         <WorkflowEditorToolbar
           toolbarExtra={toolbarExtra}
-          onUndo={onUndo}
-          undoDisabled={undoDisabled}
-          onRedo={onRedo}
-          redoDisabled={redoDisabled}
+          onUndo={() => void undo()}
+          undoDisabled={historyDisabled || !undoRedo.canUndo}
+          onRedo={() => void redo()}
+          redoDisabled={historyDisabled || !undoRedo.canRedo}
           onRestoreDefault={onRestoreDefault}
           restoreDefaultDisabled={restoreDefaultDisabled}
           onValidate={onValidate}