diff --git a/Scripts/Editor/NodeEditorGUI.cs b/Scripts/Editor/NodeEditorGUI.cs
index 18d46d6..c19d276 100755
--- a/Scripts/Editor/NodeEditorGUI.cs
+++ b/Scripts/Editor/NodeEditorGUI.cs
@@ -23,23 +23,44 @@ namespace XNodeEditor
         /// <summary> Executed after all other window GUI. Useful if Zoom is ruining your day. Automatically resets after being run.</summary>
         public event Action onLateGUI;
         private static readonly Vector3[] polyLineTempArray = new Vector3[2];
+        private bool graphFindAttempted;
+        private bool editModeEntered;
 
         protected virtual void OnGUI()
         {
             Event e = Event.current;
             Matrix4x4 m = GUI.matrix;
-            if (graph == null)
+            if (graph == null && !graphFindAttempted)
             {
-                return;
+                graphFindAttempted = true;
+                if (!OnOpen(graphInstanceID, 0))
+                {
+                    return;
+                }
+            }
+            else
+            {
+                graphFindAttempted = false;
             }
 
             ValidateGraphEditor();
             Controls();
 
             DrawGrid(position, zoom, panOffset);
-            DrawConnections();
-            DrawDraggedConnection();
-            DrawNodes();
+            // Hack used to prevent connections flickering when exiting play mode.
+            if (!editModeEntered)
+            {
+                DrawNodes();
+                DrawConnections();
+                DrawDraggedConnection();
+            }
+            else
+            {
+                DrawConnections();
+                DrawDraggedConnection();
+                DrawNodes();
+            }
+
             DrawSelectionBox();
             DrawTooltip();
             graphEditor.OnGUI();
@@ -99,7 +120,8 @@ namespace XNodeEditor
 
             // Draw tiled background
             GUI.DrawTextureWithTexCoords(rect, gridTex, new Rect(tileOffset, tileAmount));
-            GUI.DrawTextureWithTexCoords(rect, crossTex, new Rect(tileOffset + new Vector2(0.5f, 0.5f), tileAmount));
+            GUI.DrawTextureWithTexCoords(rect, crossTex,
+                new Rect(tileOffset + new Vector2(0.5f, 0.5f), tileAmount));
         }
 
         public void DrawSelectionBox()
diff --git a/Scripts/Editor/NodeEditorWindow.cs b/Scripts/Editor/NodeEditorWindow.cs
index 90928d5..2cec68d 100644
--- a/Scripts/Editor/NodeEditorWindow.cs
+++ b/Scripts/Editor/NodeEditorWindow.cs
@@ -57,6 +57,7 @@ namespace XNodeEditor
         private void OnEnable()
         {
             Undo.undoRedoPerformed += OnUndoRedoPerformed;
+            EditorApplication.playModeStateChanged += PlaymodeStateChanged;
 
             // Reload portConnectionPoints if there are any
             int length = _references.Length;
@@ -76,6 +77,7 @@ namespace XNodeEditor
         private void OnDisable()
         {
             Undo.undoRedoPerformed -= OnUndoRedoPerformed;
+            EditorApplication.playModeStateChanged -= PlaymodeStateChanged;
 
             // Cache portConnectionPoints before serialization starts
             int count = portConnectionPoints.Count;
@@ -90,8 +92,32 @@ namespace XNodeEditor
             }
         }
 
+        private void PlaymodeStateChanged(PlayModeStateChange playModeStateChange)
+        {
+            if (playModeStateChange == PlayModeStateChange.EnteredEditMode)
+            {
+                // graph = EditorUtility.InstanceIDToObject(graphInstanceID) as NodeGraph;
+                // OnOpen(graphInstanceID, 0);
+                // Repaint();
+                editModeEntered = true;
+            }
+            else if (playModeStateChange == PlayModeStateChange.ExitingPlayMode)
+            {
+                // OnOpen(graphInstanceID, 0);
+                // Repaint();
+                editModeEntered = false;
+            }
+
+            // if (!EditorApplication.isPlaying)
+            // {
+            //     OnOpen(graphInstanceID, 0);
+            //     Repaint();
+            // }
+        }
+
         public Dictionary<Node, Vector2> nodeSizes { get; } = new Dictionary<Node, Vector2>();
         public NodeGraph graph;
+        public int graphInstanceID;
         public Vector2 panOffset
         {
             get => _panOffset;
@@ -304,6 +330,8 @@ namespace XNodeEditor
             NodeEditorWindow w = GetWindow(typeof(NodeEditorWindow), false, "xNode", true) as NodeEditorWindow;
             w.wantsMouseMove = true;
             w.graph = graph;
+            w.graphInstanceID = graph.GetInstanceID();
+
             return w;
         }