Always take ownership from UIThreadContext when rename commit

src/EditorFeatures/Core.Wpf/InlineRename/CommandHandlers/RenameCommandHandler.cs

@@ -94,11 +94,11 @@ protected override void SetAdornmentFocusToPreviousElement(ITextView textView)
             }
         }
 
-        protected override void Commit(InlineRenameSession activeSession, ITextView textView)
+        protected override void CommitAndSetFocus(InlineRenameSession activeSession, ITextView textView, IUIThreadOperationContext operationContext)
         {
             try
             {
-                base.Commit(activeSession, textView);
+                base.CommitAndSetFocus(activeSession, textView, operationContext);
             }
             catch (NotSupportedException ex)
             {

src/EditorFeatures/Core/InlineRename/CommandHandlers/AbstractRenameCommandHandler.cs

@@ -11,6 +11,8 @@
 using Microsoft.VisualStudio.Text;
 using Microsoft.VisualStudio.Text.Editor;
 using Microsoft.VisualStudio.Text.Editor.Commanding;
+using Microsoft.VisualStudio.Utilities;
+using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.Editor.Implementation.InlineRename;
 
@@ -55,7 +57,7 @@ private CommandState GetCommandState(Func<CommandState> nextHandler)
     private CommandState GetCommandState()
         => _renameService.ActiveSession != null ? CommandState.Available : CommandState.Unspecified;
 
-    private void HandlePossibleTypingCommand<TArgs>(TArgs args, Action nextHandler, Action<InlineRenameSession, SnapshotSpan> actionIfInsideActiveSpan)
+    private void HandlePossibleTypingCommand<TArgs>(TArgs args, Action nextHandler, IUIThreadOperationContext operationContext, Action<InlineRenameSession, IUIThreadOperationContext, SnapshotSpan> actionIfInsideActiveSpan)
         where TArgs : EditorCommandArgs
     {
         if (_renameService.ActiveSession == null)
@@ -78,14 +80,14 @@ private void HandlePossibleTypingCommand<TArgs>(TArgs args, Action nextHandler,
         if (_renameService.ActiveSession.TryGetContainingEditableSpan(singleSpan.Start, out var containingSpan) &&
             containingSpan.Contains(singleSpan))
         {
-            actionIfInsideActiveSpan(_renameService.ActiveSession, containingSpan);
+            actionIfInsideActiveSpan(_renameService.ActiveSession, operationContext, containingSpan);
         }
         else if (_renameService.ActiveSession.IsInOpenTextBuffer(singleSpan.Start))
         {
             // It's in a read-only area that is open, so let's commit the rename 
             // and then let the character go through
 
-            CommitIfActiveAndCallNextHandler(args, nextHandler);
+            CommitIfActiveAndCallNextHandler(args, nextHandler, operationContext);
         }
         else
         {
@@ -94,23 +96,32 @@ private void HandlePossibleTypingCommand<TArgs>(TArgs args, Action nextHandler,
         }
     }
 
-    private void CommitIfActive(EditorCommandArgs args)
+    private void CommitIfActive(EditorCommandArgs args, IUIThreadOperationContext operationContext)
     {
         if (_renameService.ActiveSession != null)
         {
             var selection = args.TextView.Selection.VirtualSelectedSpans.First();
 
-            _renameService.ActiveSession.Commit();
+            Commit(operationContext);
 
             var translatedSelection = selection.TranslateTo(args.TextView.TextBuffer.CurrentSnapshot);
             args.TextView.Selection.Select(translatedSelection.Start, translatedSelection.End);
             args.TextView.Caret.MoveTo(translatedSelection.End);
         }
     }
 
-    private void CommitIfActiveAndCallNextHandler(EditorCommandArgs args, Action nextHandler)
+    private void CommitIfActiveAndCallNextHandler(EditorCommandArgs args, Action nextHandler, IUIThreadOperationContext operationContext)
     {
-        CommitIfActive(args);
+        CommitIfActive(args, operationContext);
         nextHandler();
     }
+
+    private void Commit(IUIThreadOperationContext operationContext)
+    {
+        RoslynDebug.AssertNotNull(_renameService.ActiveSession);
+        // Prevent Editor's typing responsiveness auto canceling the rename operation.
+        // InlineRenameSession will call IUIThreadOperationExecutor to sets up our own IUIThreadOperationContext
+        operationContext.TakeOwnership();
+        _renameService.ActiveSession.Commit();
+    }
 }

src/EditorFeatures/Core/InlineRename/CommandHandlers/AbstractRenameCommandHandler_BackspaceDeleteHandler.cs

@@ -21,7 +21,7 @@ public CommandState GetCommandState(DeleteKeyCommandArgs args, Func<CommandState
 
     public void ExecuteCommand(BackspaceKeyCommandArgs args, Action nextHandler, CommandExecutionContext context)
     {
-        HandlePossibleTypingCommand(args, nextHandler, (activeSession, span) =>
+        HandlePossibleTypingCommand(args, nextHandler, context.OperationContext, (activeSession, _, span) =>
             {
                 var caretPoint = args.TextView.GetCaretPoint(args.SubjectBuffer);
                 if (!args.TextView.Selection.IsEmpty || caretPoint.Value != span.Start)
@@ -33,7 +33,7 @@ public void ExecuteCommand(BackspaceKeyCommandArgs args, Action nextHandler, Com
 
     public void ExecuteCommand(DeleteKeyCommandArgs args, Action nextHandler, CommandExecutionContext context)
     {
-        HandlePossibleTypingCommand(args, nextHandler, (activeSession, span) =>
+        HandlePossibleTypingCommand(args, nextHandler, context.OperationContext, (activeSession, _, span) =>
             {
                 var caretPoint = args.TextView.GetCaretPoint(args.SubjectBuffer);
                 if (!args.TextView.Selection.IsEmpty || caretPoint.Value != span.End)

src/EditorFeatures/Core/InlineRename/CommandHandlers/AbstractRenameCommandHandler_CutPasteHandler.cs

@@ -16,7 +16,7 @@ public CommandState GetCommandState(CutCommandArgs args, Func<CommandState> next
 
     public void ExecuteCommand(CutCommandArgs args, Action nextHandler, CommandExecutionContext context)
     {
-        HandlePossibleTypingCommand(args, nextHandler, (activeSession, span) =>
+        HandlePossibleTypingCommand(args, nextHandler, context.OperationContext, (activeSession, _, span) =>
         {
             nextHandler();
         });
@@ -27,7 +27,7 @@ public CommandState GetCommandState(PasteCommandArgs args, Func<CommandState> ne
 
     public void ExecuteCommand(PasteCommandArgs args, Action nextHandler, CommandExecutionContext context)
     {
-        HandlePossibleTypingCommand(args, nextHandler, (activeSession, span) =>
+        HandlePossibleTypingCommand(args, nextHandler, context.OperationContext, (activeSession, _, span) =>
         {
             nextHandler();
         });

src/EditorFeatures/Core/InlineRename/CommandHandlers/AbstractRenameCommandHandler_MoveSelectedLinesHandler.cs

@@ -15,7 +15,7 @@ public CommandState GetCommandState(MoveSelectedLinesUpCommandArgs args)
 
     public bool ExecuteCommand(MoveSelectedLinesUpCommandArgs args, CommandExecutionContext context)
     {
-        CommitIfActive(args);
+        CommitIfActive(args, context.OperationContext);
         return false;
     }
 
@@ -24,7 +24,7 @@ public CommandState GetCommandState(MoveSelectedLinesDownCommandArgs args)
 
     public bool ExecuteCommand(MoveSelectedLinesDownCommandArgs args, CommandExecutionContext context)
     {
-        CommitIfActive(args);
+        CommitIfActive(args, context.OperationContext);
         return false;
     }
 }

src/EditorFeatures/Core/InlineRename/CommandHandlers/AbstractRenameCommandHandler_OpenLineAboveHandler.cs

@@ -15,9 +15,9 @@ public CommandState GetCommandState(OpenLineAboveCommandArgs args, Func<CommandS
 
     public void ExecuteCommand(OpenLineAboveCommandArgs args, Action nextHandler, CommandExecutionContext context)
     {
-        HandlePossibleTypingCommand(args, nextHandler, (activeSession, span) =>
+        HandlePossibleTypingCommand(args, nextHandler, context.OperationContext, (activeSession, operationContext, span) =>
         {
-            activeSession.Commit();
+            Commit(operationContext);
             nextHandler();
         });
     }

src/EditorFeatures/Core/InlineRename/CommandHandlers/AbstractRenameCommandHandler_OpenLineBelowHandler.cs

@@ -15,9 +15,9 @@ public CommandState GetCommandState(OpenLineBelowCommandArgs args, Func<CommandS
 
     public void ExecuteCommand(OpenLineBelowCommandArgs args, Action nextHandler, CommandExecutionContext context)
     {
-        HandlePossibleTypingCommand(args, nextHandler, (activeSession, span) =>
+        HandlePossibleTypingCommand(args, nextHandler, context.OperationContext, (activeSession, operationContext, span) =>
         {
-            activeSession.Commit();
+            Commit(operationContext);
             nextHandler();
         });
     }

src/EditorFeatures/Core/InlineRename/CommandHandlers/AbstractRenameCommandHandler_RefactoringWithCommandHandler.cs

@@ -18,7 +18,7 @@ public CommandState GetCommandState(ReorderParametersCommandArgs args)
 
     public bool ExecuteCommand(ReorderParametersCommandArgs args, CommandExecutionContext context)
     {
-        CommitIfActive(args);
+        CommitIfActive(args, context.OperationContext);
         return false;
     }
 
@@ -27,7 +27,7 @@ public CommandState GetCommandState(RemoveParametersCommandArgs args)
 
     public bool ExecuteCommand(RemoveParametersCommandArgs args, CommandExecutionContext context)
     {
-        CommitIfActive(args);
+        CommitIfActive(args, context.OperationContext);
         return false;
     }
 
@@ -36,7 +36,7 @@ public CommandState GetCommandState(ExtractInterfaceCommandArgs args)
 
     public bool ExecuteCommand(ExtractInterfaceCommandArgs args, CommandExecutionContext context)
     {
-        CommitIfActive(args);
+        CommitIfActive(args, context.OperationContext);
         return false;
     }
 
@@ -45,7 +45,7 @@ public CommandState GetCommandState(EncapsulateFieldCommandArgs args)
 
     public bool ExecuteCommand(EncapsulateFieldCommandArgs args, CommandExecutionContext context)
     {
-        CommitIfActive(args);
+        CommitIfActive(args, context.OperationContext);
         return false;
     }
 }

src/EditorFeatures/Core/InlineRename/CommandHandlers/AbstractRenameCommandHandler_RenameHandler.cs

@@ -14,6 +14,7 @@
 using Microsoft.CodeAnalysis.Text;
 using Microsoft.VisualStudio.Commanding;
 using Microsoft.VisualStudio.Text.Editor.Commanding.Commands;
+using Microsoft.VisualStudio.Utilities;
 
 namespace Microsoft.CodeAnalysis.Editor.Implementation.InlineRename;
 
@@ -43,11 +44,11 @@ public bool ExecuteCommand(RenameCommandArgs args, CommandExecutionContext conte
         }
 
         var token = _listener.BeginAsyncOperation(nameof(ExecuteCommand));
-        _ = ExecuteCommandAsync(args).CompletesAsyncOperation(token);
+        _ = ExecuteCommandAsync(args, context.OperationContext).CompletesAsyncOperation(token);
         return true;
     }
 
-    private async Task ExecuteCommandAsync(RenameCommandArgs args)
+    private async Task ExecuteCommandAsync(RenameCommandArgs args, IUIThreadOperationContext editorOperationContext)
     {
         _threadingContext.ThrowIfNotOnUIThread();
 
@@ -63,12 +64,6 @@ private async Task ExecuteCommandAsync(RenameCommandArgs args)
             return;
         }
 
-        var backgroundWorkIndicatorFactory = workspace.Services.GetRequiredService<IBackgroundWorkIndicatorFactory>();
-        using var context = backgroundWorkIndicatorFactory.Create(
-            args.TextView,
-            args.TextView.GetTextElementSpan(caretPoint.Value),
-            EditorFeaturesResources.Finding_token_to_rename);
-
         // If there is already an active session, commit it first
         if (_renameService.ActiveSession != null)
         {
@@ -82,10 +77,16 @@ private async Task ExecuteCommandAsync(RenameCommandArgs args)
             else
             {
                 // Otherwise, commit the existing session and start a new one.
-                _renameService.ActiveSession.Commit();
+                Commit(editorOperationContext);
             }
         }
 
+        var backgroundWorkIndicatorFactory = workspace.Services.GetRequiredService<IBackgroundWorkIndicatorFactory>();
+        using var context = backgroundWorkIndicatorFactory.Create(
+            args.TextView,
+            args.TextView.GetTextElementSpan(caretPoint.Value),
+            EditorFeaturesResources.Finding_token_to_rename);
+
         var cancellationToken = context.UserCancellationToken;
 
         var document = await args

src/EditorFeatures/Core/InlineRename/CommandHandlers/AbstractRenameCommandHandler_ReturnHandler.cs

@@ -5,6 +5,7 @@
 using Microsoft.VisualStudio.Commanding;
 using Microsoft.VisualStudio.Text.Editor;
 using Microsoft.VisualStudio.Text.Editor.Commanding.Commands;
+using Microsoft.VisualStudio.Utilities;
 
 namespace Microsoft.CodeAnalysis.Editor.Implementation.InlineRename;
 
@@ -17,20 +18,16 @@ public bool ExecuteCommand(ReturnKeyCommandArgs args, CommandExecutionContext co
     {
         if (_renameService.ActiveSession != null)
         {
-            // Prevent Editor's typing responsiveness auto canceling the rename operation.
-            // InlineRenameSession will call IUIThreadOperationExecutor to sets up our own IUIThreadOperationContext
-            context.OperationContext.TakeOwnership();
-
-            Commit(_renameService.ActiveSession, args.TextView);
+            CommitAndSetFocus(_renameService.ActiveSession, args.TextView, context.OperationContext);
             return true;
         }
 
         return false;
     }
 
-    protected virtual void Commit(InlineRenameSession activeSession, ITextView textView)
+    protected virtual void CommitAndSetFocus(InlineRenameSession activeSession, ITextView textView, IUIThreadOperationContext operationContext)
     {
-        activeSession.Commit();
+        Commit(operationContext);
         SetFocusToTextView(textView);
     }
 }

src/EditorFeatures/Core/InlineRename/CommandHandlers/AbstractRenameCommandHandler_SaveHandler.cs

@@ -16,7 +16,7 @@ public bool ExecuteCommand(SaveCommandArgs args, CommandExecutionContext context
     {
         if (_renameService.ActiveSession != null)
         {
-            _renameService.ActiveSession.Commit();
+            Commit(context.OperationContext);
             SetFocusToTextView(args.TextView);
         }
 

src/EditorFeatures/Core/InlineRename/CommandHandlers/AbstractRenameCommandHandler_TabHandler.cs

@@ -26,7 +26,7 @@ public void ExecuteCommand(TabKeyCommandArgs args, Action nextHandler, CommandEx
             return;
         }
 
-        HandlePossibleTypingCommand(args, nextHandler, (activeSession, span) =>
+        HandlePossibleTypingCommand(args, nextHandler, context.OperationContext, (activeSession, _, span) =>
         {
             var spans = new NormalizedSnapshotSpanCollection(
                 activeSession.GetBufferManager(args.SubjectBuffer)

src/EditorFeatures/Core/InlineRename/CommandHandlers/AbstractRenameCommandHandler_TypeCharHandler.cs

@@ -19,7 +19,7 @@ public CommandState GetCommandState(TypeCharCommandArgs args, Func<CommandState>
 
     public void ExecuteCommand(TypeCharCommandArgs args, Action nextHandler, CommandExecutionContext context)
     {
-        HandlePossibleTypingCommand(args, nextHandler, (activeSession, span) =>
+        HandlePossibleTypingCommand(args, nextHandler, context.OperationContext, (activeSession, _, span) =>
         {
             var document = args.SubjectBuffer.CurrentSnapshot.GetOpenDocumentInCurrentContextWithChanges();
             if (document == null)

src/EditorFeatures/Core/InlineRename/InlineRenameSession.cs

@@ -723,6 +723,10 @@ void DismissUIAndRollbackEdits()
         }
     }
 
+    /// <remarks>
+    /// This method would create its own UIThreadOperationContext to handle cancellation.
+    /// Caller must make sure to call IUIThreadOperationContext.TakeOwnership() before calling this, otherwise it won't be able to cancel the rename operation.
+    /// </remarks>
     public void Commit(bool previewChanges = false)
         => CommitSynchronously(previewChanges);
 
@@ -738,6 +742,10 @@ private bool CommitSynchronously(bool previewChanges)
         return _threadingContext.JoinableTaskFactory.Run(() => CommitWorkerAsync(previewChanges, canUseBackgroundWorkIndicator: false));
     }
 
+    /// <remarks>
+    /// This method would create its own UIThreadOperationContext to handle cancellation.
+    /// Caller must make sure to call IUIThreadOperationContext.TakeOwnership() before calling this, otherwise it won't be able to cancel the rename operation.
+    /// </remarks>
     public async Task CommitAsync(bool previewChanges)
     {
         if (this.RenameService.GlobalOptions.GetOption(InlineRenameSessionOptionsStorage.RenameAsynchronously))