From f256b0670fd91bd519665343084d1a32e7a86acc Mon Sep 17 00:00:00 2001
From: Wojtek Figat <wojtek@figat.pl>
Date: Tue, 11 Mar 2025 13:07:53 +0100
Subject: [PATCH] Add mouse capture usage for new rubber band selection

#3151
---
 .../Editor/Gizmo/ViewportRubberBandSelector.cs  | 10 ++++++++--
 .../Editor/Viewport/MainEditorGizmoViewport.cs  | 17 +++++++++++++----
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/Source/Editor/Gizmo/ViewportRubberBandSelector.cs b/Source/Editor/Gizmo/ViewportRubberBandSelector.cs
index 5ad861820..8b0fb1e97 100644
--- a/Source/Editor/Gizmo/ViewportRubberBandSelector.cs
+++ b/Source/Editor/Gizmo/ViewportRubberBandSelector.cs
@@ -34,12 +34,15 @@ public class ViewportRubberBandSelector
     /// <summary>
     /// Triggers the start of a rubber band selection.
     /// </summary>
-    public void TryStartingRubberBandSelection()
+    /// <returns>True if selection started, otherwise false.</returns>
+    public bool TryStartingRubberBandSelection()
     {
         if (!_isRubberBandSpanning && !_owner.Gizmos.Active.IsControllingMouse && !_owner.IsRightMouseButtonDown)
         {
             _tryStartRubberBand = true;
+            return true;
         }
+        return false;
     }
 
     /// <summary>
@@ -248,9 +251,12 @@ public class ViewportRubberBandSelector
     /// <summary>
     /// Immediately stops the rubber band.
     /// </summary>
-    public void StopRubberBand()
+    /// <returns>True if rubber band was active before stopping.</returns>
+    public bool StopRubberBand()
     {
+        var result = _tryStartRubberBand;
         _isRubberBandSpanning = false;
         _tryStartRubberBand = false;
+        return result;
     }
 }
diff --git a/Source/Editor/Viewport/MainEditorGizmoViewport.cs b/Source/Editor/Viewport/MainEditorGizmoViewport.cs
index da3cf5bf6..a1dc61b4a 100644
--- a/Source/Editor/Viewport/MainEditorGizmoViewport.cs
+++ b/Source/Editor/Viewport/MainEditorGizmoViewport.cs
@@ -493,14 +493,18 @@ namespace FlaxEditor.Viewport
         public override void OnLostFocus()
         {
             base.OnLostFocus();
-            _rubberBandSelector.StopRubberBand();
+
+            if (_rubberBandSelector.StopRubberBand())
+                EndMouseCapture();
         }
 
         /// <inheritdoc />
         public override void OnMouseLeave()
         {
             base.OnMouseLeave();
-            _rubberBandSelector.StopRubberBand();
+
+            if (_rubberBandSelector.StopRubberBand())
+                EndMouseCapture();
         }
 
         /// <summary>
@@ -615,8 +619,11 @@ namespace FlaxEditor.Viewport
         protected override void OnLeftMouseButtonDown()
         {
             base.OnLeftMouseButtonDown();
-            
-            _rubberBandSelector.TryStartingRubberBandSelection();
+
+            if (_rubberBandSelector.TryStartingRubberBandSelection())
+            {
+                StartMouseCapture();
+            }
         }
 
         /// <inheritdoc />
@@ -629,6 +636,8 @@ namespace FlaxEditor.Viewport
             // Select rubberbanded rect actor nodes or pick with gizmo
             if (!_rubberBandSelector.ReleaseRubberBandSelection())
             {
+                EndMouseCapture();
+
                 // Try to pick something with the current gizmo
                 Gizmos.Active?.Pick();
             }