diff --git a/Source/Editor/GUI/ContextMenu/ContextMenuButton.cs b/Source/Editor/GUI/ContextMenu/ContextMenuButton.cs
index fee2332c4..680895d18 100644
--- a/Source/Editor/GUI/ContextMenu/ContextMenuButton.cs
+++ b/Source/Editor/GUI/ContextMenu/ContextMenuButton.cs
@@ -50,6 +50,11 @@ namespace FlaxEditor.GUI.ContextMenu
         /// </summary>
         public bool AutoCheck;
 
+        /// <summary>
+        /// Closes the context menu after clicking the button, otherwise menu will stay open.
+        /// </summary>
+        public bool CloseMenuOnClick = true;
+
         /// <summary>
         /// Initializes a new instance of the <see cref="ContextMenuButton"/> class.
         /// </summary>
@@ -90,14 +95,15 @@ namespace FlaxEditor.GUI.ContextMenu
         /// </summary>
         public void Click()
         {
-            // Close topmost context menu
-            ParentContextMenu?.TopmostCM.Hide();
+            if (CloseMenuOnClick)
+            {
+                // Close topmost context menu
+                ParentContextMenu?.TopmostCM.Hide();
+            }
 
             // Auto check logic
             if (AutoCheck)
-            {
                 Checked = !Checked;
-            }
 
             // Fire event
             Clicked?.Invoke();
diff --git a/Source/Editor/Viewport/EditorViewport.cs b/Source/Editor/Viewport/EditorViewport.cs
index 97a58e3de..108f7903f 100644
--- a/Source/Editor/Viewport/EditorViewport.cs
+++ b/Source/Editor/Viewport/EditorViewport.cs
@@ -491,12 +491,17 @@ namespace FlaxEditor.Viewport
                     {
                         var v = EditorViewportViewFlagsValues[i];
                         var button = viewFlags.AddButton(v.Name);
+                        button.CloseMenuOnClick = false;
                         button.Tag = v.Mode;
                     }
                     viewFlags.ButtonClicked += button =>
                     {
                         if (button.Tag != null)
-                            Task.ViewFlags ^= (ViewFlags)button.Tag;
+                        {
+                            var v = (ViewFlags)button.Tag;
+                            Task.ViewFlags ^= v;
+                            button.Icon = (Task.View.Flags & v) != 0 ? Style.Current.CheckBoxTick : SpriteHandle.Invalid;
+                        }
                     };
                     viewFlags.VisibleChanged += WidgetViewFlagsShowHide;
                 }