From 94ee82b58df03f098febb2fe040a435f6c7ecc8a Mon Sep 17 00:00:00 2001
From: Wojtek Figat <wojtek@figat.pl>
Date: Fri, 28 May 2021 12:09:26 +0200
Subject: [PATCH] Fix Editor UI tree control click to handle selection update
 before

#533
---
 Source/Editor/GUI/Tree/TreeNode.cs             | 18 ++++++++++++------
 .../Windows/Assets/PrefabWindow.Actions.cs     |  2 ++
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/Source/Editor/GUI/Tree/TreeNode.cs b/Source/Editor/GUI/Tree/TreeNode.cs
index f7e5f115a..7fdc1c746 100644
--- a/Source/Editor/GUI/Tree/TreeNode.cs
+++ b/Source/Editor/GUI/Tree/TreeNode.cs
@@ -723,12 +723,6 @@ namespace FlaxEditor.GUI.Tree
         /// <inheritdoc />
         public override bool OnMouseUp(Vector2 location, MouseButton button)
         {
-            // Check if mouse hits bar
-            if (button == MouseButton.Right && TestHeaderHit(ref location))
-            {
-                ParentTree.OnRightClickInternal(this, ref location);
-            }
-
             // Clear flag for left button
             if (button == MouseButton.Left)
             {
@@ -773,11 +767,23 @@ namespace FlaxEditor.GUI.Tree
                         Expand();
                 }
 
+                // Check if mouse hits bar
+                if (button == MouseButton.Right && TestHeaderHit(ref location))
+                {
+                    ParentTree.OnRightClickInternal(this, ref location);
+                }
+
                 // Handled
                 Focus();
                 return true;
             }
 
+            // Check if mouse hits bar
+            if (button == MouseButton.Right && TestHeaderHit(ref location))
+            {
+                ParentTree.OnRightClickInternal(this, ref location);
+            }
+
             // Base
             return base.OnMouseUp(location, button);
         }
diff --git a/Source/Editor/Windows/Assets/PrefabWindow.Actions.cs b/Source/Editor/Windows/Assets/PrefabWindow.Actions.cs
index af568cd3f..6687bc548 100644
--- a/Source/Editor/Windows/Assets/PrefabWindow.Actions.cs
+++ b/Source/Editor/Windows/Assets/PrefabWindow.Actions.cs
@@ -41,6 +41,8 @@ namespace FlaxEditor.Windows.Assets
                 var newRoot = Object.Find<Actor>(ref newRootId);
 
                 _window.Graph.MainActor = null;
+                _window.Viewport.Instance = null;
+
                 if (SceneGraphFactory.Nodes.TryGetValue(oldRootId, out var oldRootNode))
                     oldRootNode.Dispose();
                 if (SceneGraphFactory.Nodes.TryGetValue(newRootId, out var newRootNode))