diff --git a/src/Avalonia.Controls.DataGrid/DataGrid.cs b/src/Avalonia.Controls.DataGrid/DataGrid.cs
index 021618c643..e185c8e201 100644
--- a/src/Avalonia.Controls.DataGrid/DataGrid.cs
+++ b/src/Avalonia.Controls.DataGrid/DataGrid.cs
@@ -3734,7 +3734,7 @@ namespace Avalonia.Controls
             if (sender is Control editingElement)
             {
                 editingElement.LostFocus -= EditingElement_LostFocus;
-                if (EditingRow != null && EditingColumnIndex != -1)
+                if (EditingRow != null && _editingColumnIndex != -1)
                 {
                     FocusEditingCell(true);
                 }
@@ -4042,7 +4042,7 @@ namespace Avalonia.Controls
             var editingRow = EditingRow;
             if (editingRow is null)
             {
-                return false;
+                return true;
             }
 
             Debug.Assert(_editingColumnIndex >= 0);
@@ -4431,8 +4431,7 @@ namespace Avalonia.Controls
                     dataGridCell.Focus();
                     success = dataGridCell.ContainsFocusedElement();
                 }
-                //TODO Check
-                //success = dataGridCell.ContainsFocusedElement() ? true : dataGridCell.Focus();
+
                 _focusEditingControl = !success;
             }
             return success;
diff --git a/src/Avalonia.Controls.DataGrid/Utils/TreeHelper.cs b/src/Avalonia.Controls.DataGrid/Utils/TreeHelper.cs
index 6aebf05d6b..d2c4a58a70 100644
--- a/src/Avalonia.Controls.DataGrid/Utils/TreeHelper.cs
+++ b/src/Avalonia.Controls.DataGrid/Utils/TreeHelper.cs
@@ -54,7 +54,7 @@ namespace Avalonia.Controls.Utils
         /// <returns>True if the currently focused element is within the visual tree of the parent</returns>
         internal static bool ContainsFocusedElement(this Visual element)
         {
-            return (element == null) ? false : element.ContainsChild(FocusManager.Instance.Current as Visual);
+            return element is InputElement { IsKeyboardFocusWithin: true };
         }
     }
 }