Merge pull request #48680 from sharwell/codeaction-ui

Avoid using the main thread within CodeAction.GetOperationsAsync

src/EditorFeatures/TestUtilities/Workspaces/TestWorkspace.cs

@@ -706,12 +706,12 @@ public void ChangeSolution(Solution solution)
         public override bool CanApplyParseOptionChange(ParseOptions oldOptions, ParseOptions newOptions, Project project)
             => true;
 
-        internal override async Task<bool> CanAddProjectReferenceAsync(ProjectId referencingProject, ProjectId referencedProject, CancellationToken cancellationToken)
+        internal override bool CanAddProjectReference(ProjectId referencingProject, ProjectId referencedProject)
         {
-            // VisualStudioWorkspace switches to the main thread for this call, so do the same thing here to catch tests
+            // VisualStudioWorkspace asserts the main thread for this call, so do the same thing here to catch tests
             // that fail to account for this possibility.
             var threadingContext = ExportProvider.GetExportedValue<IThreadingContext>();
-            await threadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(alwaysYield: true, cancellationToken);
+            Contract.ThrowIfFalse(threadingContext.HasMainThread && threadingContext.JoinableTaskContext.IsOnMainThread);
             return true;
         }
 

src/Features/Core/Portable/AddImport/CodeActions/AssemblyReferenceCodeAction.cs

@@ -7,6 +7,7 @@
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.CodeActions;
 using Microsoft.CodeAnalysis.Host;
+using Microsoft.CodeAnalysis.Shared.Utilities;
 using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.AddImport
@@ -23,17 +24,6 @@ public AssemblyReferenceCodeAction(
                 Contract.ThrowIfFalse(fixData.Kind == AddImportFixKind.ReferenceAssemblySymbol);
             }
 
-            private Task<string?> ResolvePathAsync(CancellationToken cancellationToken)
-            {
-                var assemblyResolverService = OriginalDocument.Project.Solution.Workspace.Services.GetRequiredService<IFrameworkAssemblyPathResolver>();
-
-                return assemblyResolverService.ResolveAssemblyPathAsync(
-                    OriginalDocument.Project.Id,
-                    FixData.AssemblyReferenceAssemblyName,
-                    FixData.AssemblyReferenceFullyQualifiedTypeName,
-                    cancellationToken);
-            }
-
             protected override Task<IEnumerable<CodeActionOperation>> ComputePreviewOperationsAsync(CancellationToken cancellationToken)
                 => ComputeOperationsAsync(isPreview: true, cancellationToken);
 
@@ -45,21 +35,82 @@ private async Task<IEnumerable<CodeActionOperation>> ComputeOperationsAsync(bool
                 var newDocument = await GetUpdatedDocumentAsync(cancellationToken).ConfigureAwait(false);
                 var newProject = newDocument.Project;
 
-                // Now add the actual assembly reference.
-                if (!isPreview)
+                if (isPreview)
                 {
-                    var resolvedPath = await ResolvePathAsync(cancellationToken).ConfigureAwait(false);
-                    if (!string.IsNullOrWhiteSpace(resolvedPath))
-                    {
-                        var service = OriginalDocument.Project.Solution.Workspace.Services.GetRequiredService<IMetadataService>();
-                        var reference = service.GetReference(resolvedPath, MetadataReferenceProperties.Assembly);
-                        newProject = newProject.WithMetadataReferences(
-                            newProject.MetadataReferences.Concat(reference));
-                    }
+                    // If this is a preview, just return an ApplyChangesOperation for the updated document
+                    var operation = new ApplyChangesOperation(newProject.Solution);
+                    return SpecializedCollections.SingletonEnumerable<CodeActionOperation>(operation);
                 }
+                else
+                {
+                    // Otherwise return an operation that can apply the text changes and add the reference
+                    var operation = new AddAssemblyReferenceCodeActionOperation(
+                        FixData.AssemblyReferenceAssemblyName,
+                        FixData.AssemblyReferenceFullyQualifiedTypeName,
+                        newProject);
+                    return SpecializedCollections.SingletonEnumerable<CodeActionOperation>(operation);
+                }
+            }
+
+            private sealed class AddAssemblyReferenceCodeActionOperation : CodeActionOperation
+            {
+                private readonly string _assemblyReferenceAssemblyName;
+                private readonly string _assemblyReferenceFullyQualifiedTypeName;
+                private readonly Project _newProject;
+
+                public AddAssemblyReferenceCodeActionOperation(
+                    string assemblyReferenceAssemblyName,
+                    string assemblyReferenceFullyQualifiedTypeName,
+                    Project newProject)
+                {
+                    _assemblyReferenceAssemblyName = assemblyReferenceAssemblyName;
+                    _assemblyReferenceFullyQualifiedTypeName = assemblyReferenceFullyQualifiedTypeName;
+                    _newProject = newProject;
+                }
+
+                internal override bool ApplyDuringTests => true;
+
+                public override void Apply(Workspace workspace, CancellationToken cancellationToken)
+                {
+                    var operation = GetApplyChangesOperation(workspace);
+                    if (operation is null)
+                        return;
 
-                var operation = new ApplyChangesOperation(newProject.Solution);
-                return SpecializedCollections.SingletonEnumerable<CodeActionOperation>(operation);
+                    operation.Apply(workspace, cancellationToken);
+                }
+
+                internal override bool TryApply(Workspace workspace, IProgressTracker progressTracker, CancellationToken cancellationToken)
+                {
+                    var operation = GetApplyChangesOperation(workspace);
+                    if (operation is null)
+                        return false;
+
+                    return operation.TryApply(workspace, progressTracker, cancellationToken);
+                }
+
+                private ApplyChangesOperation? GetApplyChangesOperation(Workspace workspace)
+                {
+                    var resolvedPath = ResolvePath(workspace);
+                    if (string.IsNullOrWhiteSpace(resolvedPath))
+                        return null;
+
+                    var service = workspace.Services.GetRequiredService<IMetadataService>();
+                    var reference = service.GetReference(resolvedPath, MetadataReferenceProperties.Assembly);
+                    var newProject = _newProject.WithMetadataReferences(
+                        _newProject.MetadataReferences.Concat(reference));
+
+                    return new ApplyChangesOperation(newProject.Solution);
+                }
+
+                private string? ResolvePath(Workspace workspace)
+                {
+                    var assemblyResolverService = workspace.Services.GetRequiredService<IFrameworkAssemblyPathResolver>();
+
+                    return assemblyResolverService.ResolveAssemblyPath(
+                        _newProject.Id,
+                        _assemblyReferenceAssemblyName,
+                        _assemblyReferenceFullyQualifiedTypeName);
+                }
             }
         }
     }

src/Features/Core/Portable/AddImport/CodeActions/MetadataSymbolReferenceCodeAction.cs

@@ -2,11 +2,11 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
-#nullable disable
-
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.CodeActions;
+using Microsoft.CodeAnalysis.Shared.Extensions;
 using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.AddImport
@@ -21,14 +21,14 @@ public MetadataSymbolReferenceCodeAction(Document originalDocument, AddImportFix
                 Contract.ThrowIfFalse(fixData.Kind == AddImportFixKind.MetadataSymbol);
             }
 
-            protected override Task<Project> UpdateProjectAsync(Project project, bool isPreview, CancellationToken cancellationToken)
+            protected override Task<CodeActionOperation?> UpdateProjectAsync(Project project, bool isPreview, CancellationToken cancellationToken)
             {
-                var projectWithReference = project.Solution.GetProject(FixData.PortableExecutableReferenceProjectId);
+                var projectWithReference = project.Solution.GetRequiredProject(FixData.PortableExecutableReferenceProjectId);
                 var reference = projectWithReference.MetadataReferences
                                                     .OfType<PortableExecutableReference>()
                                                     .First(pe => pe.FilePath == FixData.PortableExecutableReferenceFilePathToAdd);
 
-                return Task.FromResult(project.AddMetadataReference(reference));
+                return Task.FromResult<CodeActionOperation?>(new ApplyChangesOperation(project.AddMetadataReference(reference).Solution));
             }
         }
     }

src/Features/Core/Portable/AddImport/CodeActions/ProjectSymbolReferenceCodeAction.cs

@@ -2,10 +2,10 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
-#nullable disable
-
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.CodeActions;
+using Microsoft.CodeAnalysis.Shared.Utilities;
 using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.AddImport
@@ -31,22 +31,58 @@ public ProjectSymbolReferenceCodeAction(
             private bool ShouldAddProjectReference()
                 => FixData.ProjectReferenceToAdd != null && FixData.ProjectReferenceToAdd != OriginalDocument.Project.Id;
 
-            protected override async Task<Project> UpdateProjectAsync(Project project, bool isPreview, CancellationToken cancellationToken)
+            protected override Task<CodeActionOperation?> UpdateProjectAsync(Project project, bool isPreview, CancellationToken cancellationToken)
             {
                 if (!ShouldAddProjectReference())
                 {
-                    return project;
+                    return SpecializedTasks.Null<CodeActionOperation>();
                 }
 
-                if (!isPreview)
+                var projectWithAddedReference = project.AddProjectReference(new ProjectReference(FixData.ProjectReferenceToAdd));
+                var applyOperation = new ApplyChangesOperation(projectWithAddedReference.Solution);
+                if (isPreview)
+                {
+                    return Task.FromResult<CodeActionOperation?>(applyOperation);
+                }
+
+                return Task.FromResult<CodeActionOperation?>(new AddProjectReferenceCodeActionOperation(OriginalDocument.Project.Id, FixData.ProjectReferenceToAdd, applyOperation));
+            }
+
+            private sealed class AddProjectReferenceCodeActionOperation : CodeActionOperation
+            {
+                private readonly ProjectId _referencingProject;
+                private readonly ProjectId _referencedProject;
+                private readonly ApplyChangesOperation _applyOperation;
+
+                public AddProjectReferenceCodeActionOperation(ProjectId referencingProject, ProjectId referencedProject, ApplyChangesOperation applyOperation)
                 {
-                    if (!await project.Solution.Workspace.CanAddProjectReferenceAsync(OriginalDocument.Project.Id, FixData.ProjectReferenceToAdd, cancellationToken).ConfigureAwait(false))
-                    {
-                        return project;
-                    }
+                    _referencingProject = referencingProject;
+                    _referencedProject = referencedProject;
+                    _applyOperation = applyOperation;
                 }
 
-                return project.AddProjectReference(new ProjectReference(FixData.ProjectReferenceToAdd));
+                internal override bool ApplyDuringTests => true;
+
+                public override void Apply(Workspace workspace, CancellationToken cancellationToken)
+                {
+                    if (!CanApply(workspace))
+                        return;
+
+                    _applyOperation.Apply(workspace, cancellationToken);
+                }
+
+                internal override bool TryApply(Workspace workspace, IProgressTracker progressTracker, CancellationToken cancellationToken)
+                {
+                    if (!CanApply(workspace))
+                        return false;
+
+                    return _applyOperation.TryApply(workspace, progressTracker, cancellationToken);
+                }
+
+                private bool CanApply(Workspace workspace)
+                {
+                    return workspace.CanAddProjectReference(_referencingProject, _referencedProject);
+                }
             }
         }
     }

src/Features/Core/Portable/AddImport/CodeActions/SymbolReference.SymbolReferenceCodeAction.cs

@@ -2,13 +2,12 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
-#nullable disable
-
 using System;
 using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.CodeActions;
+using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.AddImport
 {
@@ -31,32 +30,38 @@ protected SymbolReferenceCodeAction(
 
             protected override async Task<IEnumerable<CodeActionOperation>> ComputePreviewOperationsAsync(CancellationToken cancellationToken)
             {
-                var changedSolution = await GetChangedSolutionAsync(isPreview: true, cancellationToken).ConfigureAwait(false);
-                if (changedSolution == null)
+                var operation = await GetChangeSolutionOperationAsync(isPreview: true, cancellationToken).ConfigureAwait(false);
+                if (operation is null)
                 {
                     return Array.Empty<CodeActionOperation>();
                 }
 
-                return new CodeActionOperation[] { new ApplyChangesOperation(changedSolution) };
+                return SpecializedCollections.SingletonEnumerable(operation);
             }
 
-            protected override Task<Solution> GetChangedSolutionAsync(CancellationToken cancellationToken)
+            protected override async Task<IEnumerable<CodeActionOperation>> ComputeOperationsAsync(CancellationToken cancellationToken)
             {
-                return GetChangedSolutionAsync(isPreview: false, cancellationToken);
+                var operation = await GetChangeSolutionOperationAsync(isPreview: false, cancellationToken).ConfigureAwait(false);
+                if (operation is null)
+                {
+                    return Array.Empty<CodeActionOperation>();
+                }
+
+                return SpecializedCollections.SingletonEnumerable(operation);
             }
 
-            private async Task<Solution> GetChangedSolutionAsync(bool isPreview, CancellationToken cancellationToken)
+            private async Task<CodeActionOperation?> GetChangeSolutionOperationAsync(bool isPreview, CancellationToken cancellationToken)
             {
                 var updatedDocument = await GetUpdatedDocumentAsync(cancellationToken).ConfigureAwait(false);
 
-                // Defer to subtype to add any p2p or metadata refs as appropriate.
+                // Defer to subtype to add any p2p or metadata refs as appropriate. If no changes to project references
+                // are necessary, the call to 'UpdateProjectAsync' will return null, in which case we fall back to just
+                // returning the updated document with its text changes.
                 var updatedProject = await UpdateProjectAsync(updatedDocument.Project, isPreview, cancellationToken).ConfigureAwait(false);
-
-                var updatedSolution = updatedProject.Solution;
-                return updatedSolution;
+                return updatedProject ?? new ApplyChangesOperation(updatedDocument.Project.Solution);
             }
 
-            protected abstract Task<Project> UpdateProjectAsync(Project project, bool isPreview, CancellationToken cancellationToken);
+            protected abstract Task<CodeActionOperation?> UpdateProjectAsync(Project project, bool isPreview, CancellationToken cancellationToken);
         }
     }
 }

src/Features/Core/Portable/CodeRefactorings/AddMissingImports/AbstractAddMissingImportsFeatureService.cs

@@ -237,8 +237,35 @@ private static async Task<Document> CleanUpNewLinesAsync(Document document, Text
             IDocumentTextDifferencingService textDiffingService,
             CancellationToken cancellationToken)
         {
-            var newSolution = await codeAction.GetChangedSolutionAsync(
-                progressTracker, cancellationToken: cancellationToken).ConfigureAwait(false);
+            // CodeAction.GetChangedSolutionAsync is only implemented for code actions that can fully compute the new
+            // solution without deferred computation or taking a dependency on the main thread. In other cases, the
+            // implementation of GetChangedSolutionAsync will throw an exception and the code action application is
+            // expected to apply the changes by executing the operations in GetOperationsAsync (which may have other
+            // side effects). This code cannot assume the input CodeAction supports GetChangedSolutionAsync, so it first
+            // attempts to apply text changes obtained from GetOperationsAsync. Two forms are supported:
+            //
+            // 1. GetOperationsAsync returns an empty list of operations (i.e. no changes are required)
+            // 2. GetOperationsAsync returns a list of operations, where the first change is an ApplyChangesOperation to
+            //    change the text in the solution, and any remaining changes are deferred computation changes.
+            //
+            // If GetOperationsAsync does not adhere to one of these patterns, the code falls back to calling
+            // GetChangedSolutionAsync since there is no clear way to apply the changes otherwise.
+            var operations = await codeAction.GetOperationsAsync(cancellationToken).ConfigureAwait(false);
+            Solution newSolution;
+            if (operations.Length == 0)
+            {
+                newSolution = document.Project.Solution;
+            }
+            else if (operations.Length == 1 && operations[0] is ApplyChangesOperation applyChangesOperation)
+            {
+                newSolution = applyChangesOperation.ChangedSolution;
+            }
+            else
+            {
+                newSolution = await codeAction.GetChangedSolutionAsync(
+                    progressTracker, cancellationToken: cancellationToken).ConfigureAwait(false);
+            }
+
             var newDocument = newSolution.GetDocument(document.Id);
 
             // Use Line differencing to reduce the possibility of changes that overwrite existing code.

src/VisualStudio/Core/Def/Implementation/ProjectSystem/MetadataReferences/VisualStudioFrameworkAssemblyPathResolverFactory.cs

@@ -7,8 +7,6 @@
 using System.IO;
 using System.Reflection;
 using System.Runtime.Versioning;
-using System.Threading;
-using System.Threading.Tasks;
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
 using Microsoft.CodeAnalysis.Host;
@@ -48,13 +46,12 @@ public Service(IThreadingContext threadingContext, VisualStudioWorkspace? worksp
                 _serviceProvider = serviceProvider;
             }
 
-            public async Task<string?> ResolveAssemblyPathAsync(
+            public string? ResolveAssemblyPath(
                 ProjectId projectId,
                 string assemblyName,
-                string? fullyQualifiedTypeName,
-                CancellationToken cancellationToken)
+                string? fullyQualifiedTypeName)
             {
-                await ThreadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
+                AssertIsForeground();
 
                 var assembly = ResolveAssembly(projectId, assemblyName);
                 if (assembly != null)