diff --git a/Scripts/Editor/NodeEditorWindow.cs b/Scripts/Editor/NodeEditorWindow.cs
index f434e73..90928d5 100644
--- a/Scripts/Editor/NodeEditorWindow.cs
+++ b/Scripts/Editor/NodeEditorWindow.cs
@@ -54,23 +54,10 @@ namespace XNodeEditor
             }
         }
 
-        private void OnDisable()
-        {
-            // Cache portConnectionPoints before serialization starts
-            int count = portConnectionPoints.Count;
-            _references = new NodePortReference[count];
-            _rects = new Rect[count];
-            int index = 0;
-            foreach (var portConnectionPoint in portConnectionPoints)
-            {
-                _references[index] = new NodePortReference(portConnectionPoint.Key);
-                _rects[index] = portConnectionPoint.Value;
-                index++;
-            }
-        }
-
         private void OnEnable()
         {
+            Undo.undoRedoPerformed += OnUndoRedoPerformed;
+
             // Reload portConnectionPoints if there are any
             int length = _references.Length;
             if (length == _rects.Length)
@@ -86,6 +73,23 @@ namespace XNodeEditor
             }
         }
 
+        private void OnDisable()
+        {
+            Undo.undoRedoPerformed -= OnUndoRedoPerformed;
+
+            // Cache portConnectionPoints before serialization starts
+            int count = portConnectionPoints.Count;
+            _references = new NodePortReference[count];
+            _rects = new Rect[count];
+            int index = 0;
+            foreach (var portConnectionPoint in portConnectionPoints)
+            {
+                _references[index] = new NodePortReference(portConnectionPoint.Key);
+                _rects[index] = portConnectionPoint.Value;
+                index++;
+            }
+        }
+
         public Dictionary<Node, Vector2> nodeSizes { get; } = new Dictionary<Node, Vector2>();
         public NodeGraph graph;
         public Vector2 panOffset
@@ -154,6 +158,11 @@ namespace XNodeEditor
             }
         }
 
+        private void OnUndoRedoPerformed()
+        {
+            Repaint();
+        }
+
         /// <summary> Make sure the graph editor is assigned and to the right object </summary>
         private void ValidateGraphEditor()
         {