diff --git a/Scripts/NodeGraph.cs b/Scripts/NodeGraph.cs
index f71642b..1f4b419 100644
--- a/Scripts/NodeGraph.cs
+++ b/Scripts/NodeGraph.cs
@@ -18,46 +18,55 @@ namespace XNode {
 
         /// <summary> Add a node to the graph by type </summary>
         public virtual Node AddNode(Type type) {
+#if UNITY_EDITOR
+            UnityEditor.Undo.RecordObject(this, "Add Node " + type.ToString());
+#endif
             Node node = ScriptableObject.CreateInstance(type) as Node;
+            nodes.Add(node);
+            node.graph = this;
 #if UNITY_EDITOR
             if (!Application.isPlaying) {
                 UnityEditor.AssetDatabase.AddObjectToAsset(node, this);
-                UnityEditor.AssetDatabase.SaveAssets();
                 node.name = UnityEditor.ObjectNames.NicifyVariableName(node.name);
             }
+            UnityEditor.EditorUtility.SetDirty(this);
 #endif
-            nodes.Add(node);
-            node.graph = this;
             return node;
         }
 
         /// <summary> Creates a copy of the original node in the graph </summary>
         public virtual Node CopyNode(Node original) {
+#if UNITY_EDITOR
+            UnityEditor.Undo.RecordObject(this, "Copy Node");
+#endif
             Node node = ScriptableObject.Instantiate(original);
             node.ClearConnections();
+            nodes.Add(node);
+            node.graph = this;
 #if UNITY_EDITOR
             if (!Application.isPlaying) {
                 UnityEditor.AssetDatabase.AddObjectToAsset(node, this);
-                UnityEditor.AssetDatabase.SaveAssets();
                 node.name = UnityEditor.ObjectNames.NicifyVariableName(node.name);
             }
+            UnityEditor.EditorUtility.SetDirty(this);
 #endif
-            nodes.Add(node);
-            node.graph = this;
             return node;
         }
 
         /// <summary> Safely remove a node and all its connections </summary>
         /// <param name="node"></param>
         public void RemoveNode(Node node) {
+#if UNITY_EDITOR
+            UnityEditor.Undo.RecordObject(this, "Remove Node");
+#endif
             node.ClearConnections();
+            nodes.Remove(node);
 #if UNITY_EDITOR
             if (!Application.isPlaying) {
                 DestroyImmediate(node, true);
-                UnityEditor.AssetDatabase.SaveAssets();
             }
+            UnityEditor.EditorUtility.SetDirty(this);
 #endif
-            nodes.Remove(node);
         }
 
         /// <summary> Remove all nodes and connections from the graph </summary>