diff --git a/src/Avalonia.Controls.TreeDataGrid/Primitives/TreeDataGridColumnHeader.cs b/src/Avalonia.Controls.TreeDataGrid/Primitives/TreeDataGridColumnHeader.cs
index 148c8f5e..8098a82a 100644
--- a/src/Avalonia.Controls.TreeDataGrid/Primitives/TreeDataGridColumnHeader.cs
+++ b/src/Avalonia.Controls.TreeDataGrid/Primitives/TreeDataGridColumnHeader.cs
@@ -65,6 +65,11 @@ public void Realize(IColumns columns, int columnIndex)
                 newInpc.PropertyChanged += OnModelPropertyChanged;
         }
 
+        public void UpdateColumnIndex(int columnIndex)
+        {
+            ColumnIndex = columnIndex;
+        }
+
         public void Unrealize()
         {
             if (_model is INotifyPropertyChanged oldInpc)
diff --git a/src/Avalonia.Controls.TreeDataGrid/Primitives/TreeDataGridColumnHeadersPresenter.cs b/src/Avalonia.Controls.TreeDataGrid/Primitives/TreeDataGridColumnHeadersPresenter.cs
index facb2bc9..085dcaba 100644
--- a/src/Avalonia.Controls.TreeDataGrid/Primitives/TreeDataGridColumnHeadersPresenter.cs
+++ b/src/Avalonia.Controls.TreeDataGrid/Primitives/TreeDataGridColumnHeadersPresenter.cs
@@ -33,6 +33,7 @@ protected override void RealizeElement(Control element, IColumn column, int inde
 
         protected override void UpdateElementIndex(Control element, int oldIndex, int newIndex)
         {
+            ((TreeDataGridColumnHeader)element).UpdateColumnIndex(newIndex);
             ChildIndexChanged?.Invoke(this, new ChildIndexChangedEventArgs(element, newIndex));
         }
 
diff --git a/tests/Avalonia.Controls.TreeDataGrid.Tests/TreeDataGridTests_Flat.cs b/tests/Avalonia.Controls.TreeDataGrid.Tests/TreeDataGridTests_Flat.cs
index 43c19570..083396e7 100644
--- a/tests/Avalonia.Controls.TreeDataGrid.Tests/TreeDataGridTests_Flat.cs
+++ b/tests/Avalonia.Controls.TreeDataGrid.Tests/TreeDataGridTests_Flat.cs
@@ -312,6 +312,35 @@ public void Does_Not_Realize_Columns_Outside_Viewport()
             Assert.True(double.IsNaN(columns[3].ActualWidth));
         }
 
+        [AvaloniaFact(Timeout = 10000)]
+        public void Header_Column_Indexes_Are_Updated_When_Columns_Are_Updated()
+        {
+            var (target, items) = CreateTarget(columns: new IColumn<Model>[]
+            {
+                new TextColumn<Model, int>("ID", x => x.Id, width: new GridLength(1, GridUnitType.Star)),
+                new TextColumn<Model, string?>("Title1", x => x.Title,  width: new GridLength(1, GridUnitType.Star)),
+                new TextColumn<Model, string?>("Title2", x => x.Title,  width: new GridLength(1, GridUnitType.Star)),
+                new TextColumn<Model, string?>("Title3", x => x.Title,  width: new GridLength(1, GridUnitType.Star)),
+            });
+
+            AssertColumnIndexes(target, 0, 4);
+
+            var source =(FlatTreeDataGridSource<Model>)target.Source!;
+
+            var movedColumn = source.Columns[1];
+            source.Columns.Remove(movedColumn);
+
+            AssertColumnIndexes(target, 0, 3);
+
+            source.Columns.Add(movedColumn);
+
+            var root = (TestWindow)target.GetVisualRoot()!;
+            root.UpdateLayout();
+            Dispatcher.UIThread.RunJobs();
+
+            AssertColumnIndexes(target, 0, 4);
+        }
+
         [AvaloniaFact(Timeout = 10000)]
         public void Columns_Are_Correctly_Sized_After_Changing_Source()
         {