diff --git a/ui/src/app/workflow/[workflowId]/stores/workflowStore.ts b/ui/src/app/workflow/[workflowId]/stores/workflowStore.ts
index 85d0c06..caa57ff 100644
--- a/ui/src/app/workflow/[workflowId]/stores/workflowStore.ts
+++ b/ui/src/app/workflow/[workflowId]/stores/workflowStore.ts
@@ -1,5 +1,5 @@
 import { ReactFlowInstance } from '@xyflow/react';
-import { NodeChange, EdgeChange } from '@xyflow/system';
+import { EdgeChange,NodeChange } from '@xyflow/system';
 import { create } from 'zustand';
 
 import { WorkflowError } from '@/client/types.gen';
diff --git a/ui/src/app/workflow/[workflowId]/utils/layoutNodes.ts b/ui/src/app/workflow/[workflowId]/utils/layoutNodes.ts
index 78e526f..868e7a3 100644
--- a/ui/src/app/workflow/[workflowId]/utils/layoutNodes.ts
+++ b/ui/src/app/workflow/[workflowId]/utils/layoutNodes.ts
@@ -7,8 +7,7 @@ export const layoutNodes = (
     nodes: FlowNode[],
     edges: FlowEdge[],
     rankdir: 'TB' | 'LR',
-    rfInstance: React.RefObject<ReactFlowInstance<FlowNode, FlowEdge> | null>,
-    saveWorkflow: (updateWorkflowDefinition: boolean) => Promise<void>
+    rfInstance: React.RefObject<ReactFlowInstance<FlowNode, FlowEdge> | null>
 ) => {
     const g = new dagre.graphlib.Graph();
     g.setGraph({ rankdir, nodesep: 250, ranksep: 250 });
diff --git a/ui/src/components/flow/edges/CustomEdge.tsx b/ui/src/components/flow/edges/CustomEdge.tsx
index f3b8c90..fb24ab9 100644
--- a/ui/src/components/flow/edges/CustomEdge.tsx
+++ b/ui/src/components/flow/edges/CustomEdge.tsx
@@ -3,6 +3,7 @@ import { AlertCircle, Pencil } from 'lucide-react';
 import { useCallback, useEffect, useState } from 'react';
 
 import { useWorkflow } from "@/app/workflow/[workflowId]/contexts/WorkflowContext";
+import { useWorkflowStore } from "@/app/workflow/[workflowId]/stores/workflowStore";
 import { Button } from "@/components/ui/button";
 import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog";
 import { Input } from "@/components/ui/input";
@@ -25,6 +26,14 @@ const EdgeDetailsDialog = ({ open, onOpenChange, data, onSave }: EdgeDetailsDial
     const [condition, setCondition] = useState(data?.condition ?? '');
     const [label, setLabel] = useState(data?.label ?? '');
 
+    // Update form state when data changes (e.g., from undo/redo)
+    useEffect(() => {
+        if (open) {
+            setCondition(data?.condition ?? '');
+            setLabel(data?.label ?? '');
+        }
+    }, [data, open]);
+
     const handleSave = () => {
         onSave({ condition: condition, label: label });
         onOpenChange(false);
@@ -87,8 +96,9 @@ interface CustomEdgeProps extends EdgeProps {
 export default function CustomEdge(props: CustomEdgeProps) {
     const { id, source, target, sourceX, sourceY, targetX, targetY, sourcePosition, targetPosition, data, style, selected } = props;
 
-    const { getEdges, setEdges, setNodes } = useReactFlow<FlowNode, FlowEdge>();
+    const { getEdges, setNodes } = useReactFlow<FlowNode, FlowEdge>();
     const { saveWorkflow } = useWorkflow();
+    const updateEdge = useWorkflowStore((state) => state.updateEdge);
     const [open, setOpen] = useState(false);
     const [isHovered, setIsHovered] = useState(false);
 
@@ -156,21 +166,14 @@ export default function CustomEdge(props: CustomEdgeProps) {
     }, [selected, isHovered, source, target, setNodes]);
 
     const handleSaveEdgeData = useCallback(async (updatedData: FlowEdgeData) => {
-        // Update the node data in the ReactFlow nodes state
-        setEdges((edges) => {
-            const updatedEdges = edges.map((edge) =>
-                edge.id === id
-                    ? { ...edge, data: updatedData }
-                    : edge
-            )
-            return updatedEdges;
-        }
-        );
+        // Use the workflow store's updateEdge method to properly track history
+        updateEdge(id, { data: updatedData });
+
         // Save the workflow after updating edge data with a small delay to ensure state is updated
         setTimeout(async () => {
             await saveWorkflow();
         }, 100);
-    }, [id, setEdges, saveWorkflow]);
+    }, [id, updateEdge, saveWorkflow]);
 
     return (
         <>
@@ -200,22 +203,23 @@ export default function CustomEdge(props: CustomEdgeProps) {
                     interactionWidth={20}
                 />
             </g>
-            {/* Show label when selected or hovered, positioned at edge center */}
-            {(selected || isHovered) && (
-                <EdgeLabelRenderer>
-                    <div
-                        style={{
-                            position: 'absolute',
-                            pointerEvents: 'all',
-                            transformOrigin: 'center',
-                            transform: `translate(-50%, -50%) translate(${labelX + offsetX}px, ${labelY + offsetY}px)`,
-                            zIndex: 1000,
-                        }}
-                        className="nodrag nopan"
-                        onMouseEnter={() => setIsHovered(true)}
-                        onMouseLeave={() => setIsHovered(false)}
-                        onDoubleClick={() => setOpen(true)}
-                    >
+            {/* Always show label, expand on select/hover */}
+            <EdgeLabelRenderer>
+                <div
+                    style={{
+                        position: 'absolute',
+                        pointerEvents: 'all',
+                        transformOrigin: 'center',
+                        transform: `translate(-50%, -50%) translate(${labelX + offsetX}px, ${labelY + offsetY}px)`,
+                        zIndex: 1000,
+                    }}
+                    className="nodrag nopan"
+                    onMouseEnter={() => setIsHovered(true)}
+                    onMouseLeave={() => setIsHovered(false)}
+                    onDoubleClick={() => setOpen(true)}
+                >
+                    {/* Show full EdgeLabel when selected or hovered, otherwise show simple label */}
+                    {(selected || isHovered) ? (
                         <div className={cn(
                             "flex flex-col gap-2 bg-white rounded-lg border-2 shadow-xl min-w-[200px]",
                             "animate-in fade-in zoom-in duration-200",
@@ -245,9 +249,19 @@ export default function CustomEdge(props: CustomEdgeProps) {
                                 </div>
                             </div>
                         </div>
-                    </div>
-                </EdgeLabelRenderer>
-            )}
+                    ) : (
+                        /* Simple label shown by default */
+                        <div className={cn(
+                            "px-2 py-1 bg-white rounded border shadow-sm",
+                            data?.invalid ? "border-red-400 text-red-600" : "border-gray-300 text-gray-700"
+                        )}>
+                            <div className="text-xs font-medium">
+                                {data?.label || data?.condition || 'No condition'}
+                            </div>
+                        </div>
+                    )}
+                </div>
+            </EdgeLabelRenderer>
             <EdgeDetailsDialog
                 open={open}
                 onOpenChange={setOpen}
diff --git a/ui/src/components/flow/nodes/AgentNode.tsx b/ui/src/components/flow/nodes/AgentNode.tsx
index 41f6134..a07016d 100644
--- a/ui/src/components/flow/nodes/AgentNode.tsx
+++ b/ui/src/components/flow/nodes/AgentNode.tsx
@@ -1,6 +1,6 @@
 import { NodeProps, NodeToolbar, Position } from "@xyflow/react";
 import { Edit, Headset, PlusIcon,Trash2Icon } from "lucide-react";
-import { memo, useState } from "react";
+import { memo, useEffect, useState } from "react";
 
 import { useWorkflow } from "@/app/workflow/[workflowId]/contexts/WorkflowContext";
 import { ExtractionVariable,FlowNodeData } from "@/components/flow/types";
@@ -83,6 +83,19 @@ export const AgentNode = memo(({ data, selected, id }: AgentNodeProps) => {
         setOpen(newOpen);
     };
 
+    // Update form state when data changes (e.g., from undo/redo)
+    useEffect(() => {
+        if (open) {
+            setPrompt(data.prompt);
+            setName(data.name);
+            setAllowInterrupt(data.allow_interrupt ?? true);
+            setExtractionEnabled(data.extraction_enabled ?? false);
+            setExtractionPrompt(data.extraction_prompt ?? "");
+            setVariables(data.extraction_variables ?? []);
+            setAddGlobalPrompt(data.add_global_prompt ?? true);
+        }
+    }, [data, open]);
+
     return (
         <>
             <NodeContent
diff --git a/ui/src/components/flow/nodes/EndCall.tsx b/ui/src/components/flow/nodes/EndCall.tsx
index cfe5d28..7efa4d7 100644
--- a/ui/src/components/flow/nodes/EndCall.tsx
+++ b/ui/src/components/flow/nodes/EndCall.tsx
@@ -1,6 +1,6 @@
 import { NodeProps, NodeToolbar, Position } from "@xyflow/react";
 import { Edit, OctagonX, PlusIcon, Trash2Icon } from "lucide-react";
-import { memo, useState } from "react";
+import { memo, useEffect, useState } from "react";
 
 import { useWorkflow } from "@/app/workflow/[workflowId]/contexts/WorkflowContext";
 import { ExtractionVariable, FlowNodeData } from "@/components/flow/types";
@@ -87,6 +87,19 @@ export const EndCall = memo(({ data, selected, id }: EndCallNodeProps) => {
         setOpen(newOpen);
     };
 
+    // Update form state when data changes (e.g., from undo/redo)
+    useEffect(() => {
+        if (open) {
+            setPrompt(data.prompt);
+            setIsStatic(data.is_static ?? true);
+            setName(data.name);
+            setExtractionEnabled(data.extraction_enabled ?? false);
+            setExtractionPrompt(data.extraction_prompt ?? "");
+            setVariables(data.extraction_variables ?? []);
+            setAddGlobalPrompt(data.add_global_prompt ?? true);
+        }
+    }, [data, open]);
+
     return (
         <>
             <NodeContent
diff --git a/ui/src/components/flow/nodes/GlobalNode.tsx b/ui/src/components/flow/nodes/GlobalNode.tsx
index f2a27bb..9df37c9 100644
--- a/ui/src/components/flow/nodes/GlobalNode.tsx
+++ b/ui/src/components/flow/nodes/GlobalNode.tsx
@@ -1,6 +1,6 @@
 import { NodeProps, NodeToolbar, Position } from "@xyflow/react";
 import { Edit, Headset, Trash2Icon } from "lucide-react";
-import { memo, useState } from "react";
+import { memo, useEffect, useState } from "react";
 
 import { useWorkflow } from "@/app/workflow/[workflowId]/contexts/WorkflowContext";
 import { FlowNodeData } from "@/components/flow/types";
@@ -56,6 +56,14 @@ export const GlobalNode = memo(({ data, selected, id }: GlobalNodeProps) => {
         setOpen(newOpen);
     };
 
+    // Update form state when data changes (e.g., from undo/redo)
+    useEffect(() => {
+        if (open) {
+            setPrompt(data.prompt);
+            setName(data.name);
+        }
+    }, [data, open]);
+
     return (
         <>
             <NodeContent
diff --git a/ui/src/components/flow/nodes/StartCall.tsx b/ui/src/components/flow/nodes/StartCall.tsx
index b0c5efe..f9ca488 100644
--- a/ui/src/components/flow/nodes/StartCall.tsx
+++ b/ui/src/components/flow/nodes/StartCall.tsx
@@ -1,6 +1,6 @@
 import { NodeProps, NodeToolbar, Position } from "@xyflow/react";
 import { Edit, Play } from "lucide-react";
-import { memo, useState } from "react";
+import { memo, useEffect, useState } from "react";
 
 import { useWorkflow } from "@/app/workflow/[workflowId]/contexts/WorkflowContext";
 import { FlowNodeData } from "@/components/flow/types";
@@ -95,6 +95,21 @@ export const StartCall = memo(({ data, selected, id }: StartCallNodeProps) => {
         setOpen(newOpen);
     };
 
+    // Update form state when data changes (e.g., from undo/redo)
+    useEffect(() => {
+        if (open) {
+            setPrompt(data.prompt ?? "");
+            setIsStatic(data.is_static ?? true);
+            setName(data.name);
+            setAllowInterrupt(data.allow_interrupt ?? true);
+            setAddGlobalPrompt(data.add_global_prompt ?? true);
+            setWaitForUserResponse(data.wait_for_user_response ?? false);
+            setDetectVoicemail(data.detect_voicemail ?? true);
+            setDelayedStart(data.delayed_start ?? false);
+            setDelayedStartDuration(data.delayed_start_duration ?? 3);
+        }
+    }, [data, open]);
+
     return (
         <>
             <NodeContent
diff --git a/ui/src/components/flow/nodes/common/useNodeHandlers.ts b/ui/src/components/flow/nodes/common/useNodeHandlers.ts
index 29e8cc7..122bf3c 100644
--- a/ui/src/components/flow/nodes/common/useNodeHandlers.ts
+++ b/ui/src/components/flow/nodes/common/useNodeHandlers.ts
@@ -1,7 +1,7 @@
-import { useReactFlow } from "@xyflow/react";
 import { useCallback, useState } from "react";
 
-import { FlowEdge, FlowNode, FlowNodeData } from "@/components/flow/types";
+import { useWorkflowStore } from "@/app/workflow/[workflowId]/stores/workflowStore";
+import { FlowNodeData } from "@/components/flow/types";
 
 interface UseNodeHandlersProps {
     id: string;
@@ -10,25 +10,26 @@ interface UseNodeHandlersProps {
 
 export const useNodeHandlers = ({ id, additionalData = {} }: UseNodeHandlersProps) => {
     const [open, setOpen] = useState(false);
-    const { setNodes } = useReactFlow<FlowNode, FlowEdge>();
+    const updateNode = useWorkflowStore((state) => state.updateNode);
+    const deleteNode = useWorkflowStore((state) => state.deleteNode);
+    const nodes = useWorkflowStore((state) => state.nodes);
 
     const handleSaveNodeData = useCallback(
         (updatedData: FlowNodeData) => {
-            setNodes((nodes) => {
-                const updatedNodes = nodes.map((node) =>
-                    node.id === id
-                        ? { ...node, data: { ...node.data, ...updatedData, ...additionalData } }
-                        : node
-                );
-                return updatedNodes;
-            });
+            // Find the current node to merge data properly
+            const currentNode = nodes.find(node => node.id === id);
+            if (currentNode) {
+                updateNode(id, {
+                    data: { ...currentNode.data, ...updatedData, ...additionalData }
+                });
+            }
         },
-        [id, setNodes, additionalData]
+        [id, updateNode, additionalData, nodes]
     );
 
     const handleDeleteNode = useCallback(() => {
-        setNodes((nodes) => nodes.filter((node) => node.id !== id));
-    }, [id, setNodes]);
+        deleteNode(id);
+    }, [id, deleteNode]);
 
     return {
         open,