Merge pull request #14516 from CyrusNajmabadi/addPackageRollback

Rollback any document changes if adding a nuget package fails.

src/EditorFeatures/CSharpTest/Diagnostics/AddUsing/AddUsingTests_NuGet.cs

@@ -60,6 +60,8 @@ public async Task TestSearchPackageSingleName()
                 var installerServiceMock = new Mock<IPackageInstallerService>(MockBehavior.Loose);
                 installerServiceMock.SetupGet(i => i.IsEnabled).Returns(true);
                 installerServiceMock.SetupGet(i => i.PackageSources).Returns(NugetPackageSources);
+                installerServiceMock.Setup(s => s.TryInstallPackage(It.IsAny<Workspace>(), It.IsAny<DocumentId>(), It.IsAny<string>(), "NuGetPackage", It.IsAny<string>(), It.IsAny<CancellationToken>()))
+                                    .Returns(true);
 
                 var packageServiceMock = new Mock<ISymbolSearchService>();
                 packageServiceMock.Setup(s => s.FindPackagesWithTypeAsync(
@@ -89,6 +91,8 @@ public async Task TestSearchPackageMultipleNames()
                 var installerServiceMock = new Mock<IPackageInstallerService>(MockBehavior.Loose);
                 installerServiceMock.SetupGet(i => i.IsEnabled).Returns(true);
                 installerServiceMock.SetupGet(i => i.PackageSources).Returns(NugetPackageSources);
+                installerServiceMock.Setup(s => s.TryInstallPackage(It.IsAny<Workspace>(), It.IsAny<DocumentId>(), It.IsAny<string>(), "NuGetPackage", It.IsAny<string>(), It.IsAny<CancellationToken>()))
+                                    .Returns(true);
 
                 var packageServiceMock = new Mock<ISymbolSearchService>();
                 packageServiceMock.Setup(s => s.FindPackagesWithTypeAsync(
@@ -188,8 +192,8 @@ public async Task TestInstallGetsCalledNoVersion()
                 var installerServiceMock = new Mock<IPackageInstallerService>(MockBehavior.Loose);
                 installerServiceMock.SetupGet(i => i.IsEnabled).Returns(true);
                 installerServiceMock.SetupGet(i => i.PackageSources).Returns(NugetPackageSources);
-                installerServiceMock.Setup(s => s.TryInstallPackage(
-                    It.IsAny<Workspace>(), It.IsAny<DocumentId>(), It.IsAny<string>(), "NuGetPackage", /*versionOpt*/ null, It.IsAny<CancellationToken>()));
+                installerServiceMock.Setup(s => s.TryInstallPackage(It.IsAny<Workspace>(), It.IsAny<DocumentId>(), It.IsAny<string>(), "NuGetPackage", /*versionOpt*/ null, It.IsAny<CancellationToken>()))
+                                    .Returns(true);
 
                 var packageServiceMock = new Mock<ISymbolSearchService>();
                 packageServiceMock.Setup(s => s.FindPackagesWithTypeAsync(
@@ -220,8 +224,8 @@ public async Task TestInstallGetsCalledWithVersion()
                 installerServiceMock.SetupGet(i => i.PackageSources).Returns(NugetPackageSources);
                 installerServiceMock.Setup(s => s.GetInstalledVersions("NuGetPackage"))
                     .Returns(new[] { "1.0" });
-                installerServiceMock.Setup(s => s.TryInstallPackage(
-                    It.IsAny<Workspace>(), It.IsAny<DocumentId>(), It.IsAny<string>(), "NuGetPackage", "1.0", It.IsAny<CancellationToken>()));
+                installerServiceMock.Setup(s => s.TryInstallPackage(It.IsAny<Workspace>(), It.IsAny<DocumentId>(), It.IsAny<string>(), "NuGetPackage", "1.0", It.IsAny<CancellationToken>()))
+                                    .Returns(true);
 
                 var packageServiceMock = new Mock<ISymbolSearchService>();
                 packageServiceMock.Setup(s => s.FindPackagesWithTypeAsync(NugetOrgSource, "NuGetType", 0, It.IsAny<CancellationToken>()))
@@ -236,6 +240,36 @@ class C
 @"
 using NuGetNamespace;
 
+class C
+{
+    NuGetType n;
+}", systemSpecialCase: false, fixProviderData: new FixProviderData(installerServiceMock.Object, packageServiceMock.Object));
+                installerServiceMock.Verify();
+            }
+
+            [WorkItem(14516, "https://github.com/dotnet/roslyn/pull/14516")]
+            [Fact, Trait(Traits.Feature, Traits.Features.CodeActionsAddUsing)]
+            public async Task TestFailedInstallRollsBackFile()
+            {
+                var installerServiceMock = new Mock<IPackageInstallerService>(MockBehavior.Loose);
+                installerServiceMock.SetupGet(i => i.IsEnabled).Returns(true);
+                installerServiceMock.SetupGet(i => i.PackageSources).Returns(NugetPackageSources);
+                installerServiceMock.Setup(s => s.GetInstalledVersions("NuGetPackage"))
+                    .Returns(new[] { "1.0" });
+                installerServiceMock.Setup(s => s.TryInstallPackage(It.IsAny<Workspace>(), It.IsAny<DocumentId>(), It.IsAny<string>(), "NuGetPackage", "1.0", It.IsAny<CancellationToken>()))
+                                    .Returns(false);
+
+                var packageServiceMock = new Mock<ISymbolSearchService>();
+                packageServiceMock.Setup(s => s.FindPackagesWithTypeAsync(NugetOrgSource, "NuGetType", 0, It.IsAny<CancellationToken>()))
+                    .Returns(CreateSearchResult("NuGetPackage", "NuGetType", CreateNameParts("NuGetNamespace")));
+
+                await TestAsync(
+@"
+class C
+{
+    [|NuGetType|] n;
+}",
+@"
 class C
 {
     NuGetType n;

src/EditorFeatures/Core/Implementation/CodeActions/CodeActionEditHandlerService.cs

@@ -271,7 +271,7 @@ public ITextBufferAssociatedViewService AssociatedViewService
                     seenApplyChanges = true;
                 }
 
-                operation.Apply(workspace, cancellationToken);
+                operation.TryApply(workspace, progressTracker, cancellationToken);
             }
         }
 

src/EditorFeatures/Test/CodeActions/AbstractCodeActionOrUserDiagnosticTest.cs

@@ -505,7 +505,10 @@ private static async Task VerifyAgainstWorkspaceDefinitionAsync(string expectedT
                 }
                 else if (operation.ApplyDuringTests)
                 {
-                    operation.Apply(workspace, new ProgressTracker(), CancellationToken.None);
+                    var oldSolution = workspace.CurrentSolution;
+                    operation.TryApply(workspace, new ProgressTracker(), CancellationToken.None);
+                    var newSolution = workspace.CurrentSolution;
+                    result = Tuple.Create(oldSolution, newSolution);
                 }
             }
 

src/EditorFeatures/Test/CodeActions/AbstractCodeActionTest.cs

@@ -65,7 +65,7 @@ private async Task<IEnumerable<CodeRefactoring>> GetCodeRefactoringsAsync(TestWo
             await VerifyPreviewContents(workspace, expectedPreviewContents, operations);
 
             var applyChangesOperation = operations.OfType<ApplyChangesOperation>().First();
-            applyChangesOperation.Apply(workspace, new ProgressTracker(), CancellationToken.None);
+            applyChangesOperation.TryApply(workspace, new ProgressTracker(), CancellationToken.None);
 
             foreach (var document in workspace.Documents)
             {

src/EditorFeatures/Test/RenameTracking/RenameTrackingTestState.cs

@@ -198,7 +198,7 @@ public async Task AssertTag(string expectedFromName, string expectedToName, bool
                 var operations = (await actions[0].GetOperationsAsync(CancellationToken.None)).ToArray();
                 Assert.Equal(1, operations.Length);
 
-                operations[0].Apply(this.Workspace, new ProgressTracker(), CancellationToken.None);
+                operations[0].TryApply(this.Workspace, new ProgressTracker(), CancellationToken.None);
             }
         }
 

src/EditorFeatures/VisualBasicTest/Diagnostics/AddImport/AddImportTests_NuGet.vb

@@ -49,6 +49,8 @@ Namespace Microsoft.CodeAnalysis.Editor.VisualBasic.UnitTests.CodeActions.AddImp
                 Dim installerServiceMock = New Mock(Of IPackageInstallerService)(MockBehavior.Loose)
                 installerServiceMock.SetupGet(Function(i) i.IsEnabled).Returns(True)
                 installerServiceMock.SetupGet(Function(i) i.PackageSources).Returns(NugetPackageSources)
+                installerServiceMock.Setup(Function(s) s.TryInstallPackage(It.IsAny(Of Workspace), It.IsAny(Of DocumentId), It.IsAny(Of String), "NuGetPackage", It.IsAny(Of String), It.IsAny(Of CancellationToken))).
+                                     Returns(True)
 
                 Dim packageServiceMock = New Mock(Of ISymbolSearchService)()
                 packageServiceMock.Setup(Function(s) s.FindPackagesWithTypeAsync(NugetOrgSource, "NuGetType", 0, It.IsAny(Of CancellationToken)())).
@@ -74,6 +76,8 @@ End Class", fixProviderData:=New ProviderData(installerServiceMock.Object, packa
                 Dim installerServiceMock = New Mock(Of IPackageInstallerService)(MockBehavior.Loose)
                 installerServiceMock.SetupGet(Function(i) i.IsEnabled).Returns(True)
                 installerServiceMock.SetupGet(Function(i) i.PackageSources).Returns(NugetPackageSources)
+                installerServiceMock.Setup(Function(s) s.TryInstallPackage(It.IsAny(Of Workspace), It.IsAny(Of DocumentId), It.IsAny(Of String), "NuGetPackage", It.IsAny(Of String), It.IsAny(Of CancellationToken))).
+                                     Returns(True)
 
                 Dim packageServiceMock = New Mock(Of ISymbolSearchService)()
                 packageServiceMock.Setup(Function(s) s.FindPackagesWithTypeAsync(NugetOrgSource, "NuGetType", 0, It.IsAny(Of CancellationToken)())).
@@ -87,6 +91,31 @@ End Class",
 "
 Imports NS1.NS2
 
+Class C
+    Dim n As NuGetType
+End Class", fixProviderData:=New ProviderData(installerServiceMock.Object, packageServiceMock.Object))
+            End Function
+
+            <Fact, Trait(Traits.Feature, Traits.Features.CodeActionsAddImport)>
+            Public Async Function TestFailedInstallDoesNotChangeFile() As Task
+                ' Make a loose mock for the installer service.  We don't care what this test
+                ' calls on it.
+                Dim installerServiceMock = New Mock(Of IPackageInstallerService)(MockBehavior.Loose)
+                installerServiceMock.SetupGet(Function(i) i.IsEnabled).Returns(True)
+                installerServiceMock.SetupGet(Function(i) i.PackageSources).Returns(NugetPackageSources)
+                installerServiceMock.Setup(Function(s) s.TryInstallPackage(It.IsAny(Of Workspace), It.IsAny(Of DocumentId), It.IsAny(Of String), "NuGetPackage", It.IsAny(Of String), It.IsAny(Of CancellationToken))).
+                                     Returns(False)
+
+                Dim packageServiceMock = New Mock(Of ISymbolSearchService)()
+                packageServiceMock.Setup(Function(s) s.FindPackagesWithTypeAsync(NugetOrgSource, "NuGetType", 0, It.IsAny(Of CancellationToken)())).
+                    Returns(CreateSearchResult("NuGetPackage", "NuGetType", CreateNameParts("NS1", "NS2")))
+
+                Await TestAsync(
+"
+Class C
+    Dim n As [|NuGetType|]
+End Class",
+"
 Class C
     Dim n As NuGetType
 End Class", fixProviderData:=New ProviderData(installerServiceMock.Object, packageServiceMock.Object))
@@ -162,8 +191,8 @@ fixProviderData:=data)
                 Dim installerServiceMock = New Mock(Of IPackageInstallerService)(MockBehavior.Loose)
                 installerServiceMock.SetupGet(Function(i) i.IsEnabled).Returns(True)
                 installerServiceMock.SetupGet(Function(i) i.PackageSources).Returns(NugetPackageSources)
-                installerServiceMock.Setup(Function(s) s.TryInstallPackage(
-                    It.IsAny(Of Workspace), It.IsAny(Of DocumentId), It.IsAny(Of String), "NuGetPackage", Nothing, It.IsAny(Of CancellationToken)))
+                installerServiceMock.Setup(Function(s) s.TryInstallPackage(It.IsAny(Of Workspace), It.IsAny(Of DocumentId), It.IsAny(Of String), "NuGetPackage", Nothing, It.IsAny(Of CancellationToken))).
+                                     Returns(True)
 
                 Dim packageServiceMock = New Mock(Of ISymbolSearchService)()
                 packageServiceMock.Setup(Function(s) s.FindPackagesWithTypeAsync(NugetOrgSource, "NuGetType", 0, It.IsAny(Of CancellationToken)())).
@@ -190,8 +219,8 @@ End Class", fixProviderData:=New ProviderData(installerServiceMock.Object, packa
                 installerServiceMock.SetupGet(Function(i) i.PackageSources).Returns(NugetPackageSources)
                 installerServiceMock.Setup(Function(s) s.GetInstalledVersions("NuGetPackage")).
                     Returns({"1.0"})
-                installerServiceMock.Setup(Function(s) s.TryInstallPackage(
-                    It.IsAny(Of Workspace), It.IsAny(Of DocumentId), It.IsAny(Of String), "NuGetPackage", "1.0", It.IsAny(Of CancellationToken)))
+                installerServiceMock.Setup(Function(s) s.TryInstallPackage(It.IsAny(Of Workspace), It.IsAny(Of DocumentId), It.IsAny(Of String), "NuGetPackage", "1.0", It.IsAny(Of CancellationToken))).
+                                     Returns(True)
 
                 Dim packageServiceMock = New Mock(Of ISymbolSearchService)()
                 packageServiceMock.Setup(Function(s) s.FindPackagesWithTypeAsync(NugetOrgSource, "NuGetType", 0, It.IsAny(Of CancellationToken)())).

src/Features/Core/Portable/AddImport/AbstractAddImportCodeFixProvider.cs

@@ -3,11 +3,9 @@
 using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
-using System.Collections.Immutable;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
-using Microsoft.CodeAnalysis.CodeActions;
 using Microsoft.CodeAnalysis.Editing;
 using Microsoft.CodeAnalysis.Host;
 using Microsoft.CodeAnalysis.Internal.Log;
@@ -460,41 +458,5 @@ private static bool NotNull(SymbolReference reference)
         {
             return reference.SymbolResult.Symbol != null;
         }
-
-        private class OperationBasedCodeAction : CodeAction
-        {
-            private readonly string _title;
-            private readonly Glyph? _glyph;
-            private readonly CodeActionPriority _priority;
-            private readonly Func<CancellationToken, Task<ImmutableArray<CodeActionOperation>>> _getOperations;
-            private readonly Func<Workspace, bool> _isApplicable;
-
-            public override string Title => _title;
-            internal override int? Glyph => _glyph.HasValue ? (int)_glyph.Value : (int?)null;
-            public override string EquivalenceKey => _title;
-            internal override CodeActionPriority Priority => _priority;
-
-            public OperationBasedCodeAction(
-                string title, Glyph? glyph, CodeActionPriority priority,
-                Func<CancellationToken, Task<ImmutableArray<CodeActionOperation>>> getOperations,
-                Func<Workspace, bool> isApplicable)
-            {
-                _title = title;
-                _glyph = glyph;
-                _priority = priority;
-                _getOperations = getOperations;
-                _isApplicable = isApplicable;
-            }
-
-            protected override async Task<IEnumerable<CodeActionOperation>> ComputeOperationsAsync(CancellationToken cancellationToken)
-                => await _getOperations(cancellationToken).ConfigureAwait(false);
-
-            internal override bool PerformFinalApplicabilityCheck => _isApplicable != null;
-
-            internal override bool IsApplicable(Workspace workspace)
-            {
-                return _isApplicable == null ? true : _isApplicable(workspace);
-            }
-        }
     }
 }
\ No newline at end of file

src/Features/Core/Portable/AddImport/CodeActions/PackageReference.InstallNugetPackageOperation.cs

+// Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
+
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading;
+using Microsoft.CodeAnalysis.CodeActions;
+using Microsoft.CodeAnalysis.Packaging;
+using Microsoft.CodeAnalysis.Shared.Utilities;
+
+namespace Microsoft.CodeAnalysis.CodeFixes.AddImport
+{
+    internal abstract partial class AbstractAddImportCodeFixProvider<TSimpleNameSyntax>
+    {
+        private partial class PackageReference
+        {
+            /// <summary>
+            /// Operation responsible purely for installing a nuget package with a specific 
+            /// version, or a the latest version of a nuget package.  Is not responsible
+            /// for adding an import to user code.
+            /// </summary>
+            private class InstallNugetPackageOperation : CodeActionOperation
+            {
+                private readonly Document _document;
+                private readonly IPackageInstallerService _installerService;
+                private readonly string _source;
+                private readonly string _packageName;
+                private readonly string _versionOpt;
+                private readonly bool _isLocal;
+                private readonly List<string> _projectsWithMatchingVersion;
+
+                public InstallNugetPackageOperation(
+                    IPackageInstallerService installerService,
+                    Document document,
+                    string source,
+                    string packageName,
+                    string versionOpt,
+                    bool isLocal)
+                {
+                    _installerService = installerService;
+                    _document = document;
+                    _source = source;
+                    _packageName = packageName;
+                    _versionOpt = versionOpt;
+                    _isLocal = isLocal;
+                    if (versionOpt != null)
+                    {
+                        const int projectsToShow = 5;
+                        var otherProjects = installerService.GetProjectsWithInstalledPackage(
+                            _document.Project.Solution, packageName, versionOpt).ToList();
+                        _projectsWithMatchingVersion = otherProjects.Take(projectsToShow).Select(p => p.Name).ToList();
+                        if (otherProjects.Count > projectsToShow)
+                        {
+                            _projectsWithMatchingVersion.Add("...");
+                        }
+                    }
+                }
+
+                public override string Title => _versionOpt == null
+                    ? string.Format(FeaturesResources.Find_and_install_latest_version_of_0, _packageName)
+                    : _isLocal
+                        ? string.Format(FeaturesResources.Use_locally_installed_0_version_1_This_version_used_in_colon_2, _packageName, _versionOpt, string.Join(", ", _projectsWithMatchingVersion))
+                        : string.Format(FeaturesResources.Install_0_1, _packageName, _versionOpt);
+
+                internal override bool ApplyDuringTests => true;
+
+                internal override bool TryApply(Workspace workspace, IProgressTracker progressTracker, CancellationToken cancellationToken)
+                {
+                    return _installerService.TryInstallPackage(
+                        workspace, _document.Id, _source, _packageName, _versionOpt, cancellationToken);
+                }
+            }
+        }
+    }
+}
\ No newline at end of file

src/Features/Core/Portable/AddImport/CodeActions/PackageReference.InstallPackageAndAddImportCodeAction.cs

+// Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
+
+using System.Collections.Generic;
+using System.Collections.Immutable;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.CodeActions;
+using Microsoft.CodeAnalysis.Shared.Utilities;
+using Microsoft.CodeAnalysis.Text;
+using Roslyn.Utilities;
+
+namespace Microsoft.CodeAnalysis.CodeFixes.AddImport
+{
+    internal abstract partial class AbstractAddImportCodeFixProvider<TSimpleNameSyntax> : CodeFixProvider, IEqualityComparer<PortableExecutableReference>
+    {
+        private partial class PackageReference
+        {
+            private struct InstallPackageAndAddImportData
+            {
+                /// <summary>
+                /// The document before we added the import. Used so we can roll back if installing
+                /// the package failed.
+                /// </summary>
+                public readonly Document OldDocument;
+
+                /// <summary>
+                /// The document after the import has been added.
+                /// </summary>
+                public readonly Document NewDocument;
+
+                /// <summary>
+                /// The operation that will actually install the nuget package.
+                /// </summary>
+                public readonly InstallNugetPackageOperation InstallOperation;
+
+                public InstallPackageAndAddImportData(Document oldDocument, Document newDocument, InstallNugetPackageOperation installOperation)
+                {
+                    OldDocument = oldDocument;
+                    NewDocument = newDocument;
+                    InstallOperation = installOperation;
+                }
+            }
+
+            private class InstallPackageAndAddImportCodeAction : CodeAction
+            {
+                private readonly string _title;
+                private readonly CodeActionPriority _priority;
+                private readonly AsyncLazy<InstallPackageAndAddImportData> _installData;
+
+                public override string Title => _title;
+                public override string EquivalenceKey => _title;
+                internal override CodeActionPriority Priority => _priority;
+
+                public InstallPackageAndAddImportCodeAction(
+                    string title, CodeActionPriority priority,
+                    AsyncLazy<InstallPackageAndAddImportData> installData)
+                {
+                    _title = title;
+                    _priority = priority;
+                    _installData = installData;
+                }
+
+                /// <summary>
+                /// For preview purposes we return all the operations in a list.  This way the 
+                /// preview system stiches things together in the UI to make a suitable display.
+                /// i.e. if we have a SolutionChangedOperation and some other operation with a 
+                /// Title, then the UI will show that nicely to the user.
+                /// </summary>
+                protected override async Task<IEnumerable<CodeActionOperation>> ComputePreviewOperationsAsync(CancellationToken cancellationToken)
+                {
+                    var installData = await _installData.GetValueAsync(cancellationToken).ConfigureAwait(false);
+
+                    // Make a SolutionChangeAction.  This way we can let it generate the diff
+                    // preview appropriately.
+                    var solutionChangeAction = new SolutionChangeAction(
+                        "", c => Task.FromResult(installData.NewDocument.Project.Solution));
+
+                    var result = ArrayBuilder<CodeActionOperation>.GetInstance();
+                    result.AddRange(await solutionChangeAction.GetPreviewOperationsAsync(cancellationToken).ConfigureAwait(false));
+                    result.Add(installData.InstallOperation);
+                    return result.ToImmutableAndFree();
+                }
+
+                /// <summary>
+                /// However, for application purposes, we end up returning a single operation
+                /// that will then apply all our sub actions in order, stopping the moment
+                /// one of them fails.
+                /// </summary>
+                protected override async Task<IEnumerable<CodeActionOperation>> ComputeOperationsAsync(
+                    CancellationToken cancellationToken)
+                {
+                    var installData = await _installData.GetValueAsync(cancellationToken).ConfigureAwait(false);
+
+                    var oldText = await installData.OldDocument.GetTextAsync(cancellationToken).ConfigureAwait(false);
+                    var newText = await installData.NewDocument.GetTextAsync(cancellationToken).ConfigureAwait(false);
+
+                    return ImmutableArray.Create<CodeActionOperation>(
+                        new InstallPackageAndAddImportOperation(
+                            installData.OldDocument.Id, oldText, newText, installData.InstallOperation));
+                }
+            }
+
+            private class InstallPackageAndAddImportOperation : CodeActionOperation
+            {
+                private readonly DocumentId _changedDocumentId;
+                private readonly SourceText _oldText;
+                private readonly SourceText _newText;
+                private readonly InstallNugetPackageOperation _installNugetPackage;
+
+                public InstallPackageAndAddImportOperation(
+                    DocumentId changedDocumentId, 
+                    SourceText oldText,
+                    SourceText newText, 
+                    InstallNugetPackageOperation item2)
+                {
+                    _changedDocumentId = changedDocumentId;
+                    _oldText = oldText;
+                    _newText = newText;
+                    _installNugetPackage = item2;
+                }
+
+                internal override bool ApplyDuringTests => _installNugetPackage.ApplyDuringTests;
+                public override string Title => _installNugetPackage.Title;
+
+                internal override bool TryApply(Workspace workspace, IProgressTracker progressTracker, CancellationToken cancellationToken)
+                {
+                    var newSolution = workspace.CurrentSolution.WithDocumentText(
+                        _changedDocumentId, _newText);
+
+                    // First make the changes to add the import to the document.
+                    if (workspace.TryApplyChanges(newSolution, progressTracker))
+                    {
+                        if (_installNugetPackage.TryApply(workspace, progressTracker, cancellationToken))
+                        {
+                            return true;
+                        }
+
+                        // Installing the nuget package failed.  Roll back the workspace.
+                        var rolledBackSolution = workspace.CurrentSolution.WithDocumentText(
+                            _changedDocumentId, _oldText);
+                        workspace.TryApplyChanges(rolledBackSolution, progressTracker);
+                    }
+
+                    return false;
+                }
+            }
+        }
+    }
+}
\ No newline at end of file

src/Features/Core/Portable/AddImport/CodeActions/PackageReference.InstallWithPackageManagerCodeAction.cs

+// Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
+
+using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.CodeActions;
+using Roslyn.Utilities;
+
+namespace Microsoft.CodeAnalysis.CodeFixes.AddImport
+{
+    internal abstract partial class AbstractAddImportCodeFixProvider<TSimpleNameSyntax>
+    {
+        private partial class PackageReference : Reference
+        {
+            private class InstallWithPackageManagerCodeAction : CodeAction
+            {
+                private readonly PackageReference reference;
+
+                public InstallWithPackageManagerCodeAction(PackageReference reference)
+                {
+                    this.reference = reference;
+                }
+
+                public override string Title => FeaturesResources.Install_with_package_manager;
+
+                protected override Task<IEnumerable<CodeActionOperation>> ComputeOperationsAsync(CancellationToken cancellationToken)
+                {
+                    return Task.FromResult(SpecializedCollections.SingletonEnumerable<CodeActionOperation>(
+                        new InstallWithPackageManagerCodeActionOperation(reference)));
+                }
+
+                private class InstallWithPackageManagerCodeActionOperation : CodeActionOperation
+                {
+                    private readonly PackageReference reference;
+
+                    public InstallWithPackageManagerCodeActionOperation(PackageReference reference)
+                    {
+                        this.reference = reference;
+                    }
+
+                    public override string Title => FeaturesResources.Install_with_package_manager;
+
+                    public override void Apply(Workspace workspace, CancellationToken cancellationToken)
+                    {
+                        reference._installerService.ShowManagePackagesDialog(reference._packageName);
+                    }
+                }
+            }
+        }
+    }
+}

src/Features/Core/Portable/AddImport/PackageReference.cs → src/Features/Core/Portable/AddImport/CodeActions/PackageReference.ParentCodeAction.cs

-// Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
+// Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
 
 using System.Collections.Generic;
 using System.Collections.Immutable;
@@ -6,57 +6,20 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.CodeActions;
-using Microsoft.CodeAnalysis.Packaging;
 using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.CodeFixes.AddImport
 {
     internal abstract partial class AbstractAddImportCodeFixProvider<TSimpleNameSyntax>
     {
-        private class PackageReference : Reference
+        private partial class PackageReference : Reference
         {
-            private readonly IPackageInstallerService _installerService;
-            private readonly string _source;
-            private readonly string _packageName;
-            private readonly string _versionOpt;
-
-            public PackageReference(
-                AbstractAddImportCodeFixProvider<TSimpleNameSyntax> provider,
-                IPackageInstallerService installerService,
-                SearchResult searchResult,
-                string source,
-                string packageName,
-                string versionOpt)
-                : base(provider, searchResult)
-            {
-                _installerService = installerService;
-                _source = source;
-                _packageName = packageName;
-                _versionOpt = versionOpt;
-            }
-
-            public override Task<CodeAction> CreateCodeActionAsync(
-                Document document, SyntaxNode node, bool placeSystemNamespaceFirst, CancellationToken cancellationToken)
-            {
-                return Task.FromResult<CodeAction>(new PackageReferenceCodeAction(
-                    this, document, node, placeSystemNamespaceFirst));
-            }
-
-            public override bool Equals(object obj)
-            {
-                var reference = obj as PackageReference;
-                return base.Equals(obj) &&
-                    _packageName == reference._packageName &&
-                    _versionOpt == reference._versionOpt;
-            }
-
-            public override int GetHashCode()
-            {
-                return Hash.Combine(_versionOpt,
-                    Hash.Combine(_packageName, base.GetHashCode()));
-            }
-
-            private class PackageReferenceCodeAction : CodeAction
+            /// <summary>
+            /// This is the top level 'Install Nuget Package' code action we show in 
+            /// the lightbulb.  It will have children to 'Install Latest', 
+            /// 'Install Version 'X' ..., and 'Install with package manager'.
+            /// </summary>
+            private class ParentCodeAction : CodeAction
             {
                 private readonly PackageReference _reference;
                 private readonly string _title;
@@ -73,7 +36,7 @@ private class PackageReferenceCodeAction : CodeAction
                 // Adding a nuget reference is lower priority than other fixes..
                 internal override CodeActionPriority Priority => CodeActionPriority.Low;
 
-                public PackageReferenceCodeAction(
+                public ParentCodeAction(
                     PackageReference reference,
                     Document document,
                     SyntaxNode node,
@@ -120,14 +83,16 @@ private class PackageReferenceCodeAction : CodeAction
                             ? string.Format(FeaturesResources.Use_local_version_0, versionOpt)
                             : string.Format(FeaturesResources.Install_version_0, versionOpt);
 
+                    var installData = new AsyncLazy<InstallPackageAndAddImportData>(
+                        c => GetInstallDataAsync(versionOpt, isLocal, document, node, placeSystemNamespaceFirst, c),
+                        cacheResult: true);
+
                     // Nuget hits should always come after other results.
-                    return new OperationBasedCodeAction(
-                        title, glyph: null, priority: CodeActionPriority.Low,
-                        getOperations: c => GetOperationsAsync(versionOpt, isLocal, document, node, placeSystemNamespaceFirst, c),
-                        isApplicable: null);
+                    return new InstallPackageAndAddImportCodeAction(
+                        title, CodeActionPriority.Low, installData);
                 }
 
-                private async Task<ImmutableArray<CodeActionOperation>> GetOperationsAsync(
+                private async Task<InstallPackageAndAddImportData> GetInstallDataAsync(
                     string versionOpt, 
                     bool isLocal,
                     Document document, 
@@ -135,109 +100,26 @@ private class PackageReferenceCodeAction : CodeAction
                     bool placeSystemNamespaceFirst, 
                     CancellationToken cancellationToken)
                 {
+                    var oldDocument = document;
                     _reference.ReplaceNameNode(ref node, ref document, cancellationToken);
 
                     var newDocument = await _reference.provider.AddImportAsync(
                         node, _reference.SearchResult.NameParts, document, placeSystemNamespaceFirst, cancellationToken).ConfigureAwait(false);
-                    var newSolution = newDocument.Project.Solution;
 
-                    var operation1 = new ApplyChangesOperation(newSolution);
+                    // We're going to be manually applying this new document to the workspace
+                    // (so we can roll it back ourselves if installing the nuget package fails).
+                    // As such, we need to do the postprocessing ourselves of tihs document to 
+                    // ensure things like formatting/simplification happen to it.
+                    newDocument = await this.PostProcessChangesAsync(
+                        newDocument, cancellationToken).ConfigureAwait(false);
 
-                    var text = await document.GetTextAsync(cancellationToken).ConfigureAwait(false);
-                    var operation2 = new InstallNugetPackageOperation(
+                    var installOperation = new InstallNugetPackageOperation(
                         _reference._installerService, document, _reference._source, _reference._packageName, versionOpt, isLocal);
 
-                    var operations = ImmutableArray.Create<CodeActionOperation>(operation1, operation2);
-                    return operations;
-                }
-            }
-
-            private class InstallWithPackageManagerCodeAction : CodeAction
-            {
-                private readonly PackageReference reference;
-
-                public InstallWithPackageManagerCodeAction(PackageReference reference)
-                {
-                    this.reference = reference;
-                }
-
-                public override string Title => FeaturesResources.Install_with_package_manager;
-
-                protected override Task<IEnumerable<CodeActionOperation>> ComputeOperationsAsync(CancellationToken cancellationToken)
-                {
-                    return Task.FromResult(SpecializedCollections.SingletonEnumerable<CodeActionOperation>(
-                        new InstallWithPackageManagerCodeActionOperation(reference)));
-                }
-
-                private class InstallWithPackageManagerCodeActionOperation : CodeActionOperation
-                {
-                    private readonly PackageReference reference;
-
-                    public InstallWithPackageManagerCodeActionOperation(PackageReference reference)
-                    {
-                        this.reference = reference;
-                    }
-
-                    public override string Title => FeaturesResources.Install_with_package_manager;
-
-                    public override void Apply(Workspace workspace, CancellationToken cancellationToken)
-                    {
-                        reference._installerService.ShowManagePackagesDialog(reference._packageName);
-                    }
-                }
-            }
-
-            private class InstallNugetPackageOperation : CodeActionOperation
-            {
-                private readonly Document _document;
-                private readonly IPackageInstallerService _installerService;
-                private readonly string _source;
-                private readonly string _packageName;
-                private readonly string _versionOpt;
-                private readonly bool _isLocal;
-                private readonly List<string> _projectsWithMatchingVersion;
-
-                public InstallNugetPackageOperation(
-                    IPackageInstallerService installerService,
-                    Document document,
-                    string source,
-                    string packageName,
-                    string versionOpt,
-                    bool isLocal)
-                {
-                    _installerService = installerService;
-                    _document = document;
-                    _source = source;
-                    _packageName = packageName;
-                    _versionOpt = versionOpt;
-                    _isLocal = isLocal;
-                    if (versionOpt != null)
-                    {
-                        const int projectsToShow = 5;
-                        var otherProjects = installerService.GetProjectsWithInstalledPackage(
-                            _document.Project.Solution, packageName, versionOpt).ToList();
-                        _projectsWithMatchingVersion = otherProjects.Take(projectsToShow).Select(p => p.Name).ToList();
-                        if (otherProjects.Count > projectsToShow)
-                        {
-                            _projectsWithMatchingVersion.Add("...");
-                        }
-                    }
-                }
-
-                public override string Title => _versionOpt == null
-                    ? string.Format(FeaturesResources.Find_and_install_latest_version_of_0, _packageName)
-                    : _isLocal
-                        ? string.Format(FeaturesResources.Use_locally_installed_0_version_1_This_version_used_in_colon_2, _packageName, _versionOpt, string.Join(", ", _projectsWithMatchingVersion))
-                        : string.Format(FeaturesResources.Install_0_1, _packageName, _versionOpt);
-
-                internal override bool ApplyDuringTests => true;
-
-                public override void Apply(Workspace workspace, CancellationToken cancellationToken)
-                {
-                    _installerService.TryInstallPackage(
-                        workspace, _document.Id, _source, _packageName, _versionOpt, cancellationToken);
+                    return new InstallPackageAndAddImportData(
+                        oldDocument, newDocument, installOperation);
                 }
             }
         }
     }
-}
+}
\ No newline at end of file

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

+// Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
+
+using System;
+using System.Collections.Generic;
+using System.Collections.Immutable;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.CodeActions;
+using Microsoft.CodeAnalysis.Shared.Utilities;
+using Roslyn.Utilities;
+
+namespace Microsoft.CodeAnalysis.CodeFixes.AddImport
+{
+    internal abstract partial class AbstractAddImportCodeFixProvider<TSimpleNameSyntax>
+    {
+        private partial class SymbolReference
+        {
+            /// <summary>
+            /// Code action we use when just adding a using, possibly with a project or
+            /// metadata reference.  We don't use the standard code action types because
+            /// we want to do things like show a glyph if this will do more than just add
+            /// an import.
+            /// </summary>
+            private class SymbolReferenceCodeAction : CodeAction
+            {
+                private readonly string _title;
+                private readonly Glyph? _glyph;
+                private readonly CodeActionPriority _priority;
+                private readonly AsyncLazy<CodeActionOperation> _getOperation;
+                private readonly Func<Workspace, bool> _isApplicable;
+
+                public override string Title => _title;
+                internal override int? Glyph => _glyph.HasValue ? (int)_glyph.Value : (int?)null;
+                public override string EquivalenceKey => _title;
+                internal override CodeActionPriority Priority => _priority;
+
+                public SymbolReferenceCodeAction(
+                    string title, Glyph? glyph, CodeActionPriority priority,
+                    AsyncLazy<CodeActionOperation> getOperation,
+                    Func<Workspace, bool> isApplicable)
+                {
+                    _title = title;
+                    _glyph = glyph;
+                    _priority = priority;
+                    _getOperation = getOperation;
+                    _isApplicable = isApplicable;
+                }
+
+                protected override async Task<IEnumerable<CodeActionOperation>> ComputeOperationsAsync(CancellationToken cancellationToken)
+                    => ImmutableArray.Create(await _getOperation.GetValueAsync(cancellationToken).ConfigureAwait(false));
+
+                internal override bool PerformFinalApplicabilityCheck
+                    => _isApplicable != null;
+
+                internal override bool IsApplicable(Workspace workspace)
+                    => _isApplicable == null ? true : _isApplicable(workspace);
+            }
+        }
+    }
+}
\ No newline at end of file

src/Features/Core/Portable/AddImport/References/PackageReference.cs

+// Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
+
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.CodeActions;
+using Microsoft.CodeAnalysis.Packaging;
+using Roslyn.Utilities;
+
+namespace Microsoft.CodeAnalysis.CodeFixes.AddImport
+{
+    internal abstract partial class AbstractAddImportCodeFixProvider<TSimpleNameSyntax>
+    {
+        private partial class PackageReference : Reference
+        {
+            private readonly IPackageInstallerService _installerService;
+            private readonly string _source;
+            private readonly string _packageName;
+            private readonly string _versionOpt;
+
+            public PackageReference(
+                AbstractAddImportCodeFixProvider<TSimpleNameSyntax> provider,
+                IPackageInstallerService installerService,
+                SearchResult searchResult,
+                string source,
+                string packageName,
+                string versionOpt)
+                : base(provider, searchResult)
+            {
+                _installerService = installerService;
+                _source = source;
+                _packageName = packageName;
+                _versionOpt = versionOpt;
+            }
+
+            public override Task<CodeAction> CreateCodeActionAsync(
+                Document document, SyntaxNode node, bool placeSystemNamespaceFirst, CancellationToken cancellationToken)
+            {
+                return Task.FromResult<CodeAction>(new ParentCodeAction(
+                    this, document, node, placeSystemNamespaceFirst));
+            }
+
+            public override bool Equals(object obj)
+            {
+                var reference = obj as PackageReference;
+                return base.Equals(obj) &&
+                    _packageName == reference._packageName &&
+                    _versionOpt == reference._versionOpt;
+            }
+
+            public override int GetHashCode()
+            {
+                return Hash.Combine(_versionOpt,
+                    Hash.Combine(_packageName, base.GetHashCode()));
+            }
+        }
+    }
+}

src/Features/Core/Portable/AddImport/SymbolReference.cs → src/Features/Core/Portable/AddImport/References/SymbolReference.cs

 // Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
 
 using System;
-using System.Collections.Generic;
-using System.Collections.Immutable;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.CodeActions;
@@ -12,7 +10,7 @@ namespace Microsoft.CodeAnalysis.CodeFixes.AddImport
 {
     internal abstract partial class AbstractAddImportCodeFixProvider<TSimpleNameSyntax>
     {
-        private abstract class SymbolReference : Reference
+        private abstract partial class SymbolReference : Reference
         {
             public readonly SymbolResult<INamespaceOrTypeSymbol> SymbolResult;
 
@@ -43,7 +41,7 @@ public override int GetHashCode()
                 return Hash.Combine(this.SymbolResult.DesiredName, base.GetHashCode());
             }
 
-            private async Task<ImmutableArray<CodeActionOperation>> GetOperationsAsync(
+            private async Task<CodeActionOperation> GetOperationAsync(
                 Document document, SyntaxNode node, bool placeSystemNamespaceFirst,
                 CancellationToken cancellationToken)
             {
@@ -51,7 +49,7 @@ public override int GetHashCode()
                 var updatedSolution = GetUpdatedSolution(newDocument);
 
                 var operation = new ApplyChangesOperation(updatedSolution);
-                return ImmutableArray.Create<CodeActionOperation>(operation);
+                return operation;
             }
 
             protected virtual Solution GetUpdatedSolution(Document newDocument)
@@ -87,8 +85,13 @@ protected virtual Solution GetUpdatedSolution(Document newDocument)
                     description = $"{this.SearchResult.DesiredName} - {description}";
                 }
 
-                return new OperationBasedCodeAction(description, GetGlyph(document), GetPriority(document),
-                    c => this.GetOperationsAsync(document, node, placeSystemNamespaceFirst, c),
+                var getOperation = new AsyncLazy<CodeActionOperation>(
+                    c => this.GetOperationAsync(document, node, placeSystemNamespaceFirst, c),
+                    cacheResult: true);
+
+                return new SymbolReferenceCodeAction(
+                    description, GetGlyph(document), GetPriority(document),
+                    getOperation,
                     this.GetIsApplicableCheck(document.Project));
             }
 

src/Features/Core/Portable/Features.csproj

@@ -94,6 +94,9 @@
       <Link>Shared\Utilities\DesktopShim.cs</Link>
     </Compile>
     <Compile Include="CodeStyle\AbstractCodeStyleDiagnosticAnalyzer.cs" />
+    <Compile Include="AddImport\CodeActions\PackageReference.InstallPackageAndAddImportCodeAction.cs" />
+    <Compile Include="AddImport\CodeActions\PackageReference.InstallNugetPackageOperation.cs" />
+    <Compile Include="AddImport\CodeActions\PackageReference.InstallWithPackageManagerCodeAction.cs" />
     <Compile Include="NavigateTo\AbstractNavigateToSearchService.SearchResult.cs" />
     <Compile Include="AddMissingReference\AbstractAddMissingReferenceCodeFixProvider.cs" />
     <Compile Include="AddMissingReference\CodeAction.cs" />
@@ -110,8 +113,10 @@
     <Compile Include="Navigation\NavigableItemFactory.cs" />
     <Compile Include="Navigation\NavigableItemFactory.DeclaredSymbolNavigableItem.cs" />
     <Compile Include="Navigation\NavigableItemFactory.SymbolLocationNavigableItem.cs" />
+    <Compile Include="AddImport\CodeActions\PackageReference.ParentCodeAction.cs" />
     <Compile Include="Remote\RemoteArguments.cs" />
     <Compile Include="Structure\BlockStructureOptions.cs" />
+    <Compile Include="AddImport\CodeActions\SymbolReference.SymbolReferenceCodeAction.cs" />
     <Compile Include="UseThrowExpression\AbstractUseThrowExpressionDiagnosticAnalyzer.cs" />
     <Compile Include="UseThrowExpression\UseThrowExpressionCodeFixProvider.cs" />
     <Compile Include="UseThrowExpression\UseThrowExpressionCodeFixProvider.FixAllProvider.cs" />
@@ -137,13 +142,13 @@
     <Compile Include="ChangeSignature\IUnifiedArgumentSyntax.cs" />
     <Compile Include="ChangeSignature\ParameterConfiguration.cs" />
     <Compile Include="ChangeSignature\SignatureChange.cs" />
-    <Compile Include="AddImport\MetadataSymbolReference.cs" />
-    <Compile Include="AddImport\AssemblyReference.cs" />
-    <Compile Include="AddImport\PackageReference.cs" />
-    <Compile Include="AddImport\ProjectSymbolReference.cs" />
-    <Compile Include="AddImport\Reference.cs" />
+    <Compile Include="AddImport\References\MetadataSymbolReference.cs" />
+    <Compile Include="AddImport\References\AssemblyReference.cs" />
+    <Compile Include="AddImport\References\PackageReference.cs" />
+    <Compile Include="AddImport\References\ProjectSymbolReference.cs" />
+    <Compile Include="AddImport\References\Reference.cs" />
     <Compile Include="AddImport\SearchScope.cs" />
-    <Compile Include="AddImport\SymbolReference.cs" />
+    <Compile Include="AddImport\References\SymbolReference.cs" />
     <Compile Include="AddImport\SymbolReferenceFinder.cs" />
     <Compile Include="AddImport\SymbolResult.cs" />
     <Compile Include="CodeFixes\CodeFixContextExtensions.cs" />

src/Workspaces/Core/Portable/CodeActions/CodeAction.cs

@@ -259,7 +259,7 @@ protected async Task<Solution> PostProcessChangesAsync(Solution changedSolution,
         /// <param name="document">The document changed by the <see cref="CodeAction"/>.</param>
         /// <param name="cancellationToken">A cancellation token.</param>
         /// <returns>A document with the post processing changes applied.</returns>
-        protected virtual async Task<Document> PostProcessChangesAsync(Document document, CancellationToken cancellationToken)
+        protected async virtual Task<Document> PostProcessChangesAsync(Document document, CancellationToken cancellationToken)
         {
             if (document.SupportsSyntaxTree)
             {

src/Workspaces/Core/Portable/CodeActions/Operations/ApplyChangesOperation.cs

@@ -40,13 +40,14 @@ public ApplyChangesOperation(Solution changedSolution)
 
         public override void Apply(Workspace workspace, CancellationToken cancellationToken)
         {
-            this.Apply(workspace, new ProgressTracker(), cancellationToken);
+            this.TryApply(workspace, new ProgressTracker(), cancellationToken);
         }
 
-        internal override void Apply(
+        internal override bool TryApply(
             Workspace workspace, IProgressTracker progressTracker, CancellationToken cancellationToken)
         {
             workspace.TryApplyChanges(ChangedSolution, progressTracker);
+            return true;
         }
     }
 }
\ No newline at end of file

src/Workspaces/Core/Portable/CodeActions/Operations/CodeActionOperation.cs

@@ -13,10 +13,7 @@ public abstract class CodeActionOperation
         /// <summary>
         /// A short title describing of the effect of the operation.
         /// </summary>
-        public virtual string Title
-        {
-            get { return null; }
-        }
+        public virtual string Title => null;
 
         /// <summary>
         /// Called by the host environment to apply the effect of the operation.
@@ -26,9 +23,10 @@ public virtual void Apply(Workspace workspace, CancellationToken cancellationTok
         {
         }
 
-        internal virtual void Apply(Workspace workspace, IProgressTracker progressTracker, CancellationToken cancellationToken)
+        internal virtual bool TryApply(Workspace workspace, IProgressTracker progressTracker, CancellationToken cancellationToken)
         {
             this.Apply(workspace, cancellationToken);
+            return true;
         }
 
         /// <summary>