2 files changed, 48 insertions, 0 deletions

diff --git a/Scripts/Fold/Editor/FoldGraphView.cs b/Scripts/Fold/Editor/FoldGraphView.cs
index 3b8dd69..61eea86 100644
--- a/Scripts/Fold/Editor/FoldGraphView.cs
+++ b/Scripts/Fold/Editor/FoldGraphView.cs
@@ -34,6 +34,18 @@ public abstract class FoldNodeView : Node
         title = data.Title;
         viewDataKey = data.guid;
         SetPosition(new Rect(data.position, Vector2.zero));
+
+        // Allow shift+click to add to the current selection without clearing it.
+        RegisterCallback<MouseDownEvent>(evt =>
+        {
+            if (!evt.shiftKey)
+                return;
+
+            if (!graphView.selection.Contains(this))
+                graphView.AddToSelection(this);
+            evt.StopImmediatePropagation();
+            evt.PreventDefault();
+        }, TrickleDown.TrickleDown);
     }
 
     public override void SetPosition(Rect newPos)
@@ -41,6 +53,7 @@ public abstract class FoldNodeView : Node
         base.SetPosition(newPos);
         if (Data != null)
         {
+            graphView.RecordUndo("Move Node");
             Data.position = newPos.position;
             graphView.MarkDirty();
         }
@@ -95,6 +108,7 @@ public class FloatValueNodeView : FoldNodeView, IValueNode<float>
         var field = new FloatField("Value") { value = DataTyped.value };
         field.RegisterValueChangedCallback(evt =>
         {
+            graphView.RecordUndo("Edit Float");
             DataTyped.value = evt.newValue;
             graphView.MarkDirty();
         });
@@ -117,6 +131,7 @@ public class VectorValueNodeView : FoldNodeView, IValueNode<Vector4>
         var field = new Vector4Field("Value") { value = DataTyped.value };
         field.RegisterValueChangedCallback(evt =>
         {
+            graphView.RecordUndo("Edit Vector");
             DataTyped.value = evt.newValue;
             graphView.MarkDirty();
         });
@@ -144,6 +159,7 @@ public class GameObjectNodeView : FoldNodeView, IValueNode<GameObject>
         };
         field.RegisterValueChangedCallback(evt =>
         {
+            graphView.RecordUndo("Edit GameObject");
             DataTyped.output = evt.newValue as GameObject;
             graphView.MarkDirty();
         });
@@ -352,6 +368,7 @@ public class FoldGraphView : GraphView
 
         graphViewChanged = OnGraphViewChanged;
         RegisterCallback<KeyDownEvent>(OnKeyDown);
+        Undo.undoRedoPerformed += OnUndoRedo;
 
         Reload();
     }
@@ -420,6 +437,7 @@ public class FoldGraphView : GraphView
 
     void RemoveNode(FoldNodeView node)
     {
+        RecordUndo("Delete Node");
         graphAsset.nodes.Remove(node.Data);
         graphAsset.edges.RemoveAll(e => e.inputNodeGuid == node.Data.guid || e.outputNodeGuid == node.Data.guid);
         nodeLookup.Remove(node.Data.guid);
@@ -432,6 +450,7 @@ public class FoldGraphView : GraphView
             edge.output?.node is not FoldNodeView outputNode)
             return;
 
+        RecordUndo("Add Edge");
         graphAsset.edges.Add(new FoldEdge
         {
             inputNodeGuid = inputNode.Data.guid,
@@ -448,6 +467,7 @@ public class FoldGraphView : GraphView
             edge.output?.node is not FoldNodeView outputNode)
             return;
 
+        RecordUndo("Remove Edge");
         graphAsset.edges.RemoveAll(e =>
             e.inputNodeGuid == inputNode.Data.guid &&
             e.outputNodeGuid == outputNode.Data.guid &&
@@ -511,6 +531,7 @@ public class FoldGraphView : GraphView
             return;
 
         suppressGraphChanges = true;
+        RecordUndo("Duplicate Nodes");
 
         var cloneMap = new Dictionary<string, FoldNodeView>();
         foreach (var original in selectedNodes)
@@ -584,6 +605,7 @@ public class FoldGraphView : GraphView
 
     public T CreateNode<T>(Vector2 position) where T : FoldNodeData, new()
     {
+        RecordUndo("Create Node");
         var data = new T { position = position };
         graphAsset.nodes.Add(data);
         var view = CreateNode(data);
@@ -595,6 +617,7 @@ public class FoldGraphView : GraphView
 
     internal FoldNodeView CreateNode(Type dataType, Vector2 position)
     {
+        RecordUndo("Create Node");
         var data = (FoldNodeData)Activator.CreateInstance(dataType);
         data.position = position;
         graphAsset.nodes.Add(data);
@@ -739,6 +762,23 @@ public class FoldGraphView : GraphView
         window.ShowNotification(message);
     }
 
+    void OnUndoRedo()
+    {
+        Reload();
+    }
+
+    public void Dispose()
+    {
+        Undo.undoRedoPerformed -= OnUndoRedo;
+    }
+
+    internal void RecordUndo(string label)
+    {
+        if (suppressGraphChanges)
+            return;
+        Undo.RegisterCompleteObjectUndo(graphAsset, label);
+    }
+
     void ShowSearchWindow(Port startPort, Vector2 worldPosition)
     {
         var provider = ScriptableObject.CreateInstance<FoldSearchProvider>();
diff --git a/Scripts/Fold/Editor/FoldWindow.cs b/Scripts/Fold/Editor/FoldWindow.cs
index 44e373b..814225c 100644
--- a/Scripts/Fold/Editor/FoldWindow.cs
+++ b/Scripts/Fold/Editor/FoldWindow.cs
@@ -21,6 +21,7 @@ public class FoldWindow : EditorWindow
     {
         graphAsset = LoadOrCreateGraph();
         ConstructGraphView();
+        Undo.undoRedoPerformed += OnUndoRedo;
     }
 
     void OnDisable()
@@ -28,10 +29,12 @@ public class FoldWindow : EditorWindow
         if (graphView != null)
         {
             rootVisualElement.Remove(graphView);
+            graphView.Dispose();
             graphView = null;
         }
 
         AssetDatabase.SaveAssets();
+        Undo.undoRedoPerformed -= OnUndoRedo;
     }
 
     FoldGraph LoadOrCreateGraph()
@@ -60,6 +63,11 @@ public class FoldWindow : EditorWindow
         rootVisualElement.Add(graphView);
     }
 
+    void OnUndoRedo()
+    {
+        graphView?.Reload();
+    }
+
     public void ShowNotification(string message)
     {
         if (string.IsNullOrEmpty(message))