From 795527576b1ff810b17d00388aa2122b87c2e93d Mon Sep 17 00:00:00 2001
From: Wojtek Figat <wojtek@figat.pl>
Date: Mon, 18 Mar 2024 15:20:05 +0100
Subject: [PATCH] Add dynamic UI mode checking for Prefab Viewport

---
 Source/Editor/Viewport/PrefabWindowViewport.cs   | 7 ++++---
 Source/Editor/Viewport/Previews/PrefabPreview.cs | 3 +++
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/Source/Editor/Viewport/PrefabWindowViewport.cs b/Source/Editor/Viewport/PrefabWindowViewport.cs
index 31f27f708..1bf121ce2 100644
--- a/Source/Editor/Viewport/PrefabWindowViewport.cs
+++ b/Source/Editor/Viewport/PrefabWindowViewport.cs
@@ -44,6 +44,7 @@ namespace FlaxEditor.Viewport
         private sealed class PrefabUIEditorRoot : UIEditorRoot
         {
             private readonly PrefabWindowViewport _viewport;
+            private bool UI => _viewport._hasUILinkedCached;
 
             public PrefabUIEditorRoot(PrefabWindowViewport viewport)
             : base(true)
@@ -52,9 +53,9 @@ namespace FlaxEditor.Viewport
                 Parent = viewport;
             }
 
-            public override bool EnableInputs => false;
-            public override bool EnableSelecting => true;
-            public override bool EnableBackground => _viewport._hasUILinkedCached;
+            public override bool EnableInputs => !UI;
+            public override bool EnableSelecting => UI;
+            public override bool EnableBackground => UI;
             public override TransformGizmo TransformGizmo => _viewport.TransformGizmo;
         }
 
diff --git a/Source/Editor/Viewport/Previews/PrefabPreview.cs b/Source/Editor/Viewport/Previews/PrefabPreview.cs
index 2c7b60a1f..7df524a72 100644
--- a/Source/Editor/Viewport/Previews/PrefabPreview.cs
+++ b/Source/Editor/Viewport/Previews/PrefabPreview.cs
@@ -98,6 +98,9 @@ namespace FlaxEditor.Viewport.Previews
 
         private void UpdateLinkage()
         {
+            // Clear flag
+            _hasUILinked = false;
+
             // Link UI canvases to the preview (eg. after canvas added to the prefab)
             LinkCanvas(_instance);