diff --git a/src/Avalonia.Controls.TreeDataGrid/Selection/TreeDataGridRowSelectionModel.cs b/src/Avalonia.Controls.TreeDataGrid/Selection/TreeDataGridRowSelectionModel.cs
index 5fb4f7bb..d84ea0aa 100644
--- a/src/Avalonia.Controls.TreeDataGrid/Selection/TreeDataGridRowSelectionModel.cs
+++ b/src/Avalonia.Controls.TreeDataGrid/Selection/TreeDataGridRowSelectionModel.cs
@@ -1,6 +1,7 @@
 using System;
 using System.Collections;
 using System.Collections.Generic;
+using System.Linq;
 using Avalonia.Controls.Models.TreeDataGrid;
 using Avalonia.Controls.Presenters;
 using Avalonia.Controls.Primitives;
@@ -366,6 +367,44 @@ void ITreeDataGridSelectionInteraction.OnTextInput(TreeDataGrid sender, TextInpu
             return null;
         }
 
+        protected internal override bool TryGetItemAt(IndexPath index, out TModel? result)
+        {
+            if (_source is FlatTreeDataGridSource<TModel> treeSource)
+            {
+                bool valid = index.Any(i => i >= treeSource.Items.Count()) == false;
+                result = valid ? treeSource.Items.ElementAt(index[0]) : null;
+                return valid;
+            }
+
+            var items = (IReadOnlyList<TModel>?)Root.ItemsView;
+            var count = index.Count;
+
+            for (var i = 0; i < count; ++i)
+            {
+                if (items is null)
+                {
+                    result = default;
+                    return false;
+                }
+
+                var j = index[i];
+
+                if (j < items.Count)
+                {
+                    if (i == count - 1)
+                    {
+                        result = items[j];
+                        return true;
+                    }
+                    else
+                        items = GetChildren(items[j]) as IReadOnlyList<TModel>;
+                }
+            }
+
+            result = default;
+            return false;
+        }
+
         protected override void OnSourceCollectionChangeFinished()
         {
             if (_raiseViewSelectionChanged)
diff --git a/src/Avalonia.Controls.TreeDataGrid/Selection/TreeSelectionModelBase.cs b/src/Avalonia.Controls.TreeDataGrid/Selection/TreeSelectionModelBase.cs
index b161d3eb..12b96c3a 100644
--- a/src/Avalonia.Controls.TreeDataGrid/Selection/TreeSelectionModelBase.cs
+++ b/src/Avalonia.Controls.TreeDataGrid/Selection/TreeSelectionModelBase.cs
@@ -202,36 +202,7 @@ public bool IsSelected(IndexPath index)
 
         protected internal abstract IEnumerable<T>? GetChildren(T node);
         
-        protected virtual bool TryGetItemAt(IndexPath index, out T? result)
-        {
-            var items = (IReadOnlyList<T>?)_root.ItemsView;
-            var count = index.Count;
-
-            for (var i = 0; i < count; ++i)
-            {
-                if (items is null)
-                {
-                    result = default;
-                    return false;
-                }
-
-                var j = index[i];
-
-                if (j < items.Count)
-                {
-                    if (i == count - 1)
-                    {
-                        result = items[j];
-                        return true;
-                    }
-                    else
-                        items = GetChildren(items[j]) as IReadOnlyList<T>;
-                }
-            }
-
-            result = default;
-            return false;
-        }
+        protected internal abstract bool TryGetItemAt(IndexPath index, out T? result);
 
         protected virtual void OnSourceCollectionChangeFinished()
         {
diff --git a/tests/Avalonia.Controls.TreeDataGrid.Tests/Selection/TreeSelectionModelBaseTests_Multiple.cs b/tests/Avalonia.Controls.TreeDataGrid.Tests/Selection/TreeSelectionModelBaseTests_Multiple.cs
index 382a6b0d..33546d37 100644
--- a/tests/Avalonia.Controls.TreeDataGrid.Tests/Selection/TreeSelectionModelBaseTests_Multiple.cs
+++ b/tests/Avalonia.Controls.TreeDataGrid.Tests/Selection/TreeSelectionModelBaseTests_Multiple.cs
@@ -1596,6 +1596,14 @@ public TestTreeSelectionModel(IList<Node> data)
                     node.Children = CreateNodes(node.Id, node.TargetDepth);
                 return node.Children;
             }
+
+            protected internal override bool TryGetItemAt(IndexPath index, out Node? result)
+            {
+                var nodes = (IList<Node>)Source!;
+                bool valid = index.Any(i => i >= nodes.Count) == false;
+                result = valid ? nodes[index[0]] : null;
+                return valid;
+            }
         }
     }
 }
diff --git a/tests/Avalonia.Controls.TreeDataGrid.Tests/Selection/TreeSelectionModelBaseTests_Single.cs b/tests/Avalonia.Controls.TreeDataGrid.Tests/Selection/TreeSelectionModelBaseTests_Single.cs
index 6279bc35..ed167108 100644
--- a/tests/Avalonia.Controls.TreeDataGrid.Tests/Selection/TreeSelectionModelBaseTests_Single.cs
+++ b/tests/Avalonia.Controls.TreeDataGrid.Tests/Selection/TreeSelectionModelBaseTests_Single.cs
@@ -78,7 +78,7 @@ public void Can_Set_SelectedIndex()
                 Assert.Equal("Node 0-2", target.SelectedItem!.Caption);
                 Assert.Equal("Node 0-2", target.SelectedItems.Single()!.Caption);
             }
-            
+
             [AvaloniaFact(Timeout = 10000)]
             public void Can_Set_Grandchild_SelectedIndex()
             {
@@ -329,7 +329,7 @@ public void PropertyChanged_Is_Raised_When_SelectedIndex_Changes()
                 Assert.Equal(1, raised);
             }
         }
-        
+
         public class Select
         {
             [AvaloniaFact(Timeout = 10000)]
@@ -1293,6 +1293,14 @@ public TestTreeSelectionModel(AvaloniaList<Node> data)
                     node.Children = CreateNodes(node.Id, node.TargetDepth);
                 return node.Children;
             }
+
+            protected internal override bool TryGetItemAt(IndexPath index, out Node? result)
+            {
+                var nodes = (IList<Node>)Source!;
+                bool valid = index.Any(i => i >= nodes.Count) == false;
+                result = valid ? nodes[index[0]] : null;
+                return valid;
+            }
         }
     }
 }