Avalonia Focus Causing Focus Loss

I'm trying to setup a textBlock, where when someone clicks Edit in a context menu it swaps the textblock out for a textbox so the user can edit what's in that field.

My approach so far has been to have both controls in a grid and swap what is visible by binding them both to a bool in the ViewModel

<Grid RowDefinitions="Auto, Auto">
        <TextBox Text="{Binding Debug}"></TextBox>
        <TextBlock Grid.Row="1" Text="{Binding Greeting}" HorizontalAlignment="Center" VerticalAlignment="Center" IsVisible="{Binding !IsEditing}">
            <TextBlock.ContextMenu>
                <ContextMenu>
                    <MenuItem Header="Change" Click="MenuItem_OnClick"></MenuItem>
                </ContextMenu>
            </TextBlock.ContextMenu>
        </TextBlock>
        
        <TextBox Grid.Row="1" HorizontalAlignment="Center" VerticalAlignment="Center" PropertyChanged="AvaloniaObject_OnPropertyChanged" IsVisible="{Binding IsEditing}" LostFocus="InputElement_OnLostFocus">
            
        </TextBox>
    </Grid>

public class MainWindowViewModel : ViewModelBase
    {
        public string Greeting => "Welcome to Avalonia!";

        private bool _isEditing = false;
        public bool IsEditing
        {
            get => _isEditing;
            set => this.RaiseAndSetIfChanged(ref _isEditing, value);
        }

        private string _debug = "";

        public string Debug
        {
            get => _debug;
            set => this.RaiseAndSetIfChanged(ref _debug, value);
        }
    }

When someone clicks the menu item I set the IsEditing bool to True causing the TextBox to become visible

private void MenuItem_OnClick(object? sender, RoutedEventArgs e)
        {
            var vm = (MainWindowViewModel)this.DataContext;
            vm.IsEditing = true;
        }

I'm then watching for changes to the IsVisible property so I can force the focus to go to the textbox so the user can immediately start typing in it without having to click on it, that's where my trouble is.

private void AvaloniaObject_OnPropertyChanged(object? sender, AvaloniaPropertyChangedEventArgs e)
        {
            
            var vm = (MainWindowViewModel)this.DataContext;
            if (vm != null)
            {
                if (e.Property.Name.ToLower().Contains("focus"))
                {
                    vm.Debug += $"{e.Property.Name} changed to {e.NewValue} \n";
                }
            }
            
            if (e.Property.Name == "IsVisible")
            {
                if ((bool)e.NewValue)
                {
                    var tb = (TextBox)sender;
                    tb.Focus();
                }
            }
        }

When the focus on the textbox gets lost, I set the IsEditing bool back to false so the textblock shows back up.

private void InputElement_OnLostFocus(object? sender, RoutedEventArgs e)
        {
            var vm = (MainWindowViewModel)this.DataContext;
            if (vm != null)
            {
                vm.IsEditing = false;
                vm.Debug += "Focus was lost on the textbox \n";
            }
        }

I've got a little debug textbox to try to figure out what is happening...it looks like the focus gets set to the textbox and then immediately gets lost I'm not sure what could be causing the focus loss here.

This only happens when I force the focus to the textbox.

If I remove the line tb.Focus(); that immediate focus loss doesn't happen, but the focus isn't set to the textbox so a user could immediately start typing

I've got an example that reproduces the problem setup here