Tweaks to existing "copilot completions" processing code to get it ready for auto-fixing copilot suggestions.

src/EditorFeatures/Core/Copilot/CopilotEditorUtilities.cs

@@ -0,0 +1,62 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System.Collections.Generic;
+using System.Collections.Immutable;
+using Microsoft.CodeAnalysis.Editor.Shared.Extensions;
+using Microsoft.CodeAnalysis.PooledObjects;
+using Microsoft.CodeAnalysis.Text;
+using Microsoft.VisualStudio.Language.Proposals;
+using Microsoft.VisualStudio.Text;
+
+namespace Microsoft.CodeAnalysis.Copilot;
+
+internal static class CopilotEditorUtilities
+{
+    /// <summary>
+    /// Returns the single roslyn solution snapshot that is affected by the edits in the proposal.
+    /// </summary>
+    /// <remarks>
+    /// Will fail in the event that the edits touch multiple workspaces, or if the edits touch
+    /// multiple versions of a solution in a workspace (for example, edits to different versions of
+    /// the same <see cref="ITextSnapshot"/>. Will also fail if this edits non-roslyn files, or files
+    /// we cannot process semantics for (like non C#/VB files).
+    /// </remarks>
+    public static Solution? TryGetAffectedSolution(ProposalBase proposal)
+    {
+        Solution? solution = null;
+        foreach (var edit in proposal.Edits)
+        {
+            var document = edit.Span.Snapshot.GetOpenDocumentInCurrentContextWithChanges();
+
+            // Edit touches a file roslyn doesn't know about.  Don't touch this.
+            if (document is null)
+                return null;
+
+            // Only bother for languages we can actually process semantics for.
+            if (document.SupportsSemanticModel)
+                return null;
+
+            var currentSolution = document.Project.Solution;
+
+            // Edit touches multiple solutions.  Don't bother with this for now for simplicity's sake.
+            if (solution != null && solution != currentSolution)
+                return null;
+
+            solution = currentSolution;
+        }
+
+        return solution;
+    }
+
+    /// <inheritdoc cref="CopilotUtilities.TryNormalizeCopilotTextChanges"/>
+    public static ImmutableArray<TextChange> TryGetNormalizedTextChanges(IEnumerable<ProposedEdit> edits)
+    {
+        using var _ = ArrayBuilder<TextChange>.GetInstance(out var textChanges);
+        foreach (var edit in edits)
+            textChanges.Add(new TextChange(edit.Span.Span.ToTextSpan(), edit.ReplacementText));
+
+        return CopilotUtilities.TryNormalizeCopilotTextChanges(textChanges);
+    }
+}

src/EditorFeatures/Core/Copilot/CopilotWpfTextCreationListener.cs

@@ -15,6 +15,7 @@
 using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
 using Microsoft.CodeAnalysis.Host.Mef;
 using Microsoft.CodeAnalysis.Options;
+using Microsoft.CodeAnalysis.Remote;
 using Microsoft.CodeAnalysis.Shared.TestHooks;
 using Microsoft.CodeAnalysis.Text;
 using Microsoft.CodeAnalysis.Threading;
@@ -106,6 +107,14 @@ private static async ValueTask ProcessCompletionEventAsync(
         const string featureId = "Completion";
         var proposalId = proposal.ProposalId;
 
+        var solution = CopilotEditorUtilities.TryGetAffectedSolution(proposal);
+        if (solution is null)
+            return;
+
+        // We're about to potentially make multiple calls to oop here.  So keep a session alive to avoid
+        // resyncing any data unnecessary.
+        using var _1 = await RemoteKeepAliveSession.CreateAsync(solution, cancellationToken).ConfigureAwait(false);
+
         foreach (var editGroup in proposal.Edits.GroupBy(e => e.Span.Snapshot))
         {
             cancellationToken.ThrowIfCancellationRequested();
@@ -116,12 +125,10 @@ private static async ValueTask ProcessCompletionEventAsync(
             if (document is null)
                 continue;
 
-            using var _ = PooledObjects.ArrayBuilder<TextChange>.GetInstance(out var textChanges);
-            foreach (var edit in editGroup)
-                textChanges.Add(new TextChange(edit.Span.Span.ToTextSpan(), edit.ReplacementText));
+            var normalizedEdits = CopilotEditorUtilities.TryGetNormalizedTextChanges(editGroup);
 
             await CopilotChangeAnalysisUtilities.AnalyzeCopilotChangeAsync(
-                document, accepted, featureId, proposalId, textChanges, cancellationToken).ConfigureAwait(false);
+                document, accepted, featureId, proposalId, normalizedEdits, cancellationToken).ConfigureAwait(false);
         }
     }
 }

src/EditorFeatures/Core/ExternalAccess/UnitTestGenerator/Api/UnitTestGeneratorAddMissingImportsFeatureServiceAccessor.cs

@@ -13,34 +13,29 @@
 
 namespace Microsoft.CodeAnalysis.ExternalAccess.UnitTestGenerator.Api;
 
-[Export]
-[Shared]
+#pragma warning disable CA1822 // Mark members as static. Existing binary api with UnitTestGenerator.
+
+[Export, Shared]
 [method: ImportingConstructor]
 [method: Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
 internal class UnitTestGeneratorAddMissingImportsFeatureServiceAccessor()
 {
-#pragma warning disable CA1822 // Mark members as static
     internal async Task<Document> AddMissingImportsAsync(Document document, TextSpan textSpan, CancellationToken cancellationToken)
-#pragma warning restore CA1822
     {
         var service = document.Project.GetRequiredLanguageService<IAddMissingImportsFeatureService>();
 
         // Unfortunately, the unit testing system doesn't have a way to report progress.
         return await service.AddMissingImportsAsync(document, textSpan, CodeAnalysisProgress.None, cancellationToken).ConfigureAwait(false);
     }
 
-#pragma warning disable CA1822 // Mark members as static
     internal async Task<WrappedMissingImportsAnalysisResult> AnalyzeAsync(Document document, TextSpan textSpan, CancellationToken cancellationToken)
-#pragma warning restore CA1822
     {
         var service = document.Project.GetRequiredLanguageService<IAddMissingImportsFeatureService>();
-        var result = await service.AnalyzeAsync(document, textSpan, cancellationToken).ConfigureAwait(false);
+        var result = await service.AnalyzeAsync(document, textSpan, cleanupDocument: true, cancellationToken).ConfigureAwait(false);
         return new WrappedMissingImportsAnalysisResult(result.SelectAsArray(data => new WrappedAddImportFixData(data)));
     }
 
-#pragma warning disable CA1822 // Mark members as static
     internal async Task<Document> AddMissingImportsAsync(Document document, WrappedMissingImportsAnalysisResult analysisResult, CancellationToken cancellationToken)
-#pragma warning restore CA1822
     {
         var service = document.Project.GetRequiredLanguageService<IAddMissingImportsFeatureService>();
         var unwrappedResult = analysisResult.AddImportFixDatas.SelectAsArray(result => result.Underlying);

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

@@ -72,7 +72,8 @@ public sealed override async Task RegisterCodeFixesAsync(CodeFixContext context)
             searchOptions = searchOptions with { SearchNuGetPackages = false };
         }
 
-        var addImportOptions = await document.GetAddImportOptionsAsync(searchOptions, cancellationToken).ConfigureAwait(false);
+        var addImportOptions = await document.GetAddImportOptionsAsync(
+            searchOptions, cleanupDocument: true, cancellationToken).ConfigureAwait(false);
 
         var fixesForDiagnostic = await addImportService.GetFixesForDiagnosticsAsync(
             document, span, diagnostics, MaxResults, symbolSearchService, addImportOptions, packageSources, cancellationToken).ConfigureAwait(false);

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

@@ -112,7 +112,8 @@ private async Task<ImmutableArray<AddImportFixData>> GetFixesInCurrentProcessAsy
                         {
                             cancellationToken.ThrowIfCancellationRequested();
 
-                            var fixData = await reference.TryGetFixDataAsync(document, node, options.CleanupOptions, cancellationToken).ConfigureAwait(false);
+                            var fixData = await reference.TryGetFixDataAsync(
+                                document, node, options.CleanupDocument, options.CleanupOptions, cancellationToken).ConfigureAwait(false);
                             result.AddIfNotNull(fixData);
 
                             if (result.Count > maxResults)

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

@@ -16,21 +16,33 @@ namespace Microsoft.CodeAnalysis.AddImport;
 internal readonly record struct AddImportOptions(
     [property: DataMember(Order = 0)] SymbolSearchOptions SearchOptions,
     [property: DataMember(Order = 1)] CodeCleanupOptions CleanupOptions,
-    [property: DataMember(Order = 2)] MemberDisplayOptions MemberDisplayOptions);
+    [property: DataMember(Order = 2)] MemberDisplayOptions MemberDisplayOptions,
+    [property: DataMember(Order = 3)] bool CleanupDocument);
 
 internal static class AddImportOptionsProviders
 {
-    public static AddImportOptions GetAddImportOptions(this IOptionsReader options, LanguageServices languageServices, SymbolSearchOptions searchOptions, bool allowImportsInHiddenRegions)
+    public static AddImportOptions GetAddImportOptions(
+        this IOptionsReader options,
+        LanguageServices languageServices,
+        SymbolSearchOptions searchOptions,
+        bool allowImportsInHiddenRegions,
+        bool cleanupDocument)
         => new()
         {
             SearchOptions = searchOptions,
             CleanupOptions = options.GetCodeCleanupOptions(languageServices, allowImportsInHiddenRegions),
-            MemberDisplayOptions = options.GetMemberDisplayOptions(languageServices.Language)
+            MemberDisplayOptions = options.GetMemberDisplayOptions(languageServices.Language),
+            CleanupDocument = cleanupDocument,
         };
 
-    public static async ValueTask<AddImportOptions> GetAddImportOptionsAsync(this Document document, SymbolSearchOptions searchOptions, CancellationToken cancellationToken)
+    public static async ValueTask<AddImportOptions> GetAddImportOptionsAsync(
+        this Document document,
+        SymbolSearchOptions searchOptions,
+        bool cleanupDocument,
+        CancellationToken cancellationToken)
     {
         var configOptions = await document.GetHostAnalyzerConfigOptionsAsync(cancellationToken).ConfigureAwait(false);
-        return configOptions.GetAddImportOptions(document.Project.Services, searchOptions, document.AllowImportsInHiddenRegions());
+        return configOptions.GetAddImportOptions(
+            document.Project.Services, searchOptions, document.AllowImportsInHiddenRegions(), cleanupDocument);
     }
 }

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

@@ -23,9 +23,9 @@ private sealed partial class AssemblyReference(
         private readonly ReferenceAssemblyResult _referenceAssemblyWithType = referenceAssemblyWithType;
 
         public override async Task<AddImportFixData> TryGetFixDataAsync(
-            Document document, SyntaxNode node, CodeCleanupOptions options, CancellationToken cancellationToken)
+            Document document, SyntaxNode node, bool cleanDocument, CodeCleanupOptions options, CancellationToken cancellationToken)
         {
-            var textChanges = await GetTextChangesAsync(document, node, options, cancellationToken).ConfigureAwait(false);
+            var textChanges = await GetTextChangesAsync(document, node, cleanDocument, options, cancellationToken).ConfigureAwait(false);
 
             var title = $"{provider.GetDescription(SearchResult.NameParts)} ({string.Format(FeaturesResources.from_0, _referenceAssemblyWithType.AssemblyName)})";
             var fullyQualifiedTypeName = string.Join(

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

@@ -26,10 +26,10 @@ private sealed partial class PackageReference(
         private readonly string _versionOpt = versionOpt;
 
         public override async Task<AddImportFixData> TryGetFixDataAsync(
-            Document document, SyntaxNode node, CodeCleanupOptions options, CancellationToken cancellationToken)
+            Document document, SyntaxNode node, bool cleanupDocument, CodeCleanupOptions options, CancellationToken cancellationToken)
         {
             var textChanges = await GetTextChangesAsync(
-                document, node, options, cancellationToken).ConfigureAwait(false);
+                document, node, cleanupDocument, options, cancellationToken).ConfigureAwait(false);
 
             return AddImportFixData.CreateForPackageSymbol(
                 textChanges, _source, _packageName, _versionOpt);

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

@@ -98,10 +98,10 @@ public override int GetHashCode()
         }
 
         public abstract Task<AddImportFixData> TryGetFixDataAsync(
-            Document document, SyntaxNode node, CodeCleanupOptions options, CancellationToken cancellationToken);
+            Document document, SyntaxNode node, bool cleanupDocument, CodeCleanupOptions options, CancellationToken cancellationToken);
 
         protected async Task<ImmutableArray<TextChange>> GetTextChangesAsync(
-            Document document, SyntaxNode node, CodeCleanupOptions options, CancellationToken cancellationToken)
+            Document document, SyntaxNode node, bool cleanupDocument, CodeCleanupOptions options, CancellationToken cancellationToken)
         {
             // Within an import, we're only adding a package/nuget reference (and we're not going to add another
             // using/import as we're already in one).  So no need for text changes.
@@ -116,13 +116,21 @@ protected async Task<ImmutableArray<TextChange>> GetTextChangesAsync(
             var newDocument = await provider.AddImportAsync(
                 node, SearchResult.NameParts, document, options.AddImportOptions, cancellationToken).ConfigureAwait(false);
 
-            var cleanedDocument = await CodeAction.CleanupDocumentAsync(
-                newDocument, options, cancellationToken).ConfigureAwait(false);
+            var cleanedDocument = await CleanDocumentAsync(newDocument, cleanupDocument, options, cancellationToken).ConfigureAwait(false);
 
             var textChanges = await cleanedDocument.GetTextChangesAsync(
                 originalDocument, cancellationToken).ConfigureAwait(false);
 
             return [.. textChanges];
         }
+
+        protected static async Task<Document> CleanDocumentAsync(Document newDocument, bool cleanupDocument, CodeCleanupOptions options, CancellationToken cancellationToken)
+        {
+            // We always need to at least cleanup syntax (this ensures spacing of elastic trivia on the new import).
+            // Optionally cleanup the rest of the document if that is requested.
+            return cleanupDocument
+                ? await CodeAction.CleanupDocumentAsync(newDocument, options, cancellationToken).ConfigureAwait(false)
+                : await CodeAction.CleanupSyntaxAsync(newDocument, options, cancellationToken).ConfigureAwait(false);
+        }
     }
 }

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

@@ -47,7 +47,7 @@ public override int GetHashCode()
 
         private async Task<ImmutableArray<TextChange>> GetTextChangesAsync(
             Document document, SyntaxNode contextNode,
-            CodeCleanupOptions options, bool hasExistingImport,
+            bool cleanupDocument, CodeCleanupOptions options, bool hasExistingImport,
             CancellationToken cancellationToken)
         {
             // Defer to the language to add the actual import/using.
@@ -63,8 +63,8 @@ private async Task<ImmutableArray<TextChange>> GetTextChangesAsync(
                 newContextNode, SymbolResult.Symbol, newDocument,
                 options.AddImportOptions, cancellationToken).ConfigureAwait(false);
 
-            var cleanedDocument = await CodeAction.CleanupDocumentAsync(
-                updatedDocument, options, cancellationToken).ConfigureAwait(false);
+            var cleanedDocument = await CleanDocumentAsync(
+                updatedDocument, cleanupDocument, options, cancellationToken).ConfigureAwait(false);
 
             var textChanges = await cleanedDocument.GetTextChangesAsync(
                 document, cancellationToken).ConfigureAwait(false);
@@ -73,8 +73,7 @@ private async Task<ImmutableArray<TextChange>> GetTextChangesAsync(
         }
 
         public sealed override async Task<AddImportFixData> TryGetFixDataAsync(
-            Document document, SyntaxNode node,
-            CodeCleanupOptions options, CancellationToken cancellationToken)
+            Document document, SyntaxNode node, bool cleanupDocument, CodeCleanupOptions options, CancellationToken cancellationToken)
         {
             var semanticModel = await document.GetSemanticModelAsync(cancellationToken).ConfigureAwait(false);
             var (description, hasExistingImport) = GetDescription(document, options, node, semanticModel, cancellationToken);
@@ -108,7 +107,7 @@ public sealed override async Task<AddImportFixData> TryGetFixDataAsync(
             }
 
             var textChanges = await GetTextChangesAsync(
-                document, node, options, hasExistingImport, cancellationToken).ConfigureAwait(false);
+                document, node, cleanupDocument, options, hasExistingImport, cancellationToken).ConfigureAwait(false);
 
             return GetFixData(
                 document, textChanges, description,

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

@@ -32,7 +32,7 @@ internal abstract class AbstractAddMissingImportsFeatureService : IAddMissingImp
     protected abstract ImmutableArray<AbstractFormattingRule> GetFormatRules(SourceText text);
 
     public async Task<ImmutableArray<AddImportFixData>> AnalyzeAsync(
-        Document document, TextSpan textSpan, CancellationToken cancellationToken)
+        Document document, TextSpan textSpan, bool cleanDocument, CancellationToken cancellationToken)
     {
         // Get the diagnostics that indicate a missing import.
         var addImportFeatureService = document.GetRequiredLanguageService<IAddImportFeatureService>();
@@ -53,6 +53,7 @@ public async Task<ImmutableArray<AddImportFixData>> AnalyzeAsync(
                 SearchReferenceAssemblies = false,
                 SearchNuGetPackages = false,
             },
+            cleanDocument,
             cancellationToken).ConfigureAwait(false);
 
         var unambiguousFixes = await addImportFeatureService.GetUniqueFixesAsync(

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

@@ -31,14 +31,13 @@ public override async Task ComputeRefactoringsAsync(CodeRefactoringContext conte
         // Currently this refactoring requires the SourceTextContainer to have a pasted text span.
         var sourceText = await document.GetValueTextAsync(cancellationToken).ConfigureAwait(false);
         if (!pasteTrackingService.TryGetPastedTextSpan(sourceText.Container, out var textSpan))
-        {
             return;
-        }
 
         // Check pasted text span for missing imports
         var addMissingImportsService = document.GetRequiredLanguageService<IAddMissingImportsFeatureService>();
 
-        var fixData = await addMissingImportsService.AnalyzeAsync(document, textSpan, cancellationToken).ConfigureAwait(false);
+        var fixData = await addMissingImportsService.AnalyzeAsync(
+            document, textSpan, cleanupDocument: true, cancellationToken).ConfigureAwait(false);
         if (fixData.IsDefaultOrEmpty)
             return;