Generate Documentation - Add loading state

src/EditorFeatures/Core.Wpf/QuickInfo/OnTheFlyDocsView.xaml.cs

@@ -14,6 +14,7 @@
 using System.Windows.Controls;
 using Microsoft.CodeAnalysis.Classification;
 using Microsoft.CodeAnalysis.Copilot;
+using Microsoft.CodeAnalysis.Editor.Copilot;
 using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
 using Microsoft.CodeAnalysis.ErrorReporting;
 using Microsoft.CodeAnalysis.Internal.Log;
@@ -198,7 +199,7 @@ private async Task SetResultTextAsync(ICopilotCodeAnalysisService copilotService
                         // https://dev.azure.com/devdiv/DevDiv/_wiki/wikis/DevDiv.wiki/45121/Free-SKU-Handling-Guidance-and-Recommendations
                         var uiShell = _serviceProvider.GetServiceOnMainThread<SVsUIShell, IVsUIShell>();
                         uiShell.PostExecCommand(
-                            new Guid("39B0DEDE-D931-4A92-9AA2-3447BC4998DC"),
+                            CopilotConstants.CopilotQuotaExceededGuid,
                             0x3901,
                             nCmdexecopt: 0,
                             pvaIn: null);

src/EditorFeatures/Core/Copilot/CopilotConstants.cs

@@ -0,0 +1,25 @@
+// 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;
+using Microsoft.VisualStudio.Language.Suggestions;
+
+namespace Microsoft.CodeAnalysis.Editor.Copilot;
+
+internal static class CopilotConstants
+{
+    public const int CopilotIconLogoId = 1;
+    public const int CopilotIconSparkleId = 2;
+    public const int CopilotIconSparkleBlueId = 3;
+
+    /// <summary>
+    /// This flag is used to indicate that a thinking state tip should be shown.
+    /// This will eventually be defined in Microsoft.VisualStudio.Language.Suggestions.TipStyle once the behavior is finalized.
+    /// </summary>
+    // Copied from https://devdiv.visualstudio.com/DevDiv/_git/IntelliCode-VS?path=%2Fsrc%2FVSIX%2FIntelliCode.VSIX%2FSuggestionService%2FImplementation%2FSuggestionSession.cs&amp
+    public const TipStyle ShowThinkingStateTipStyle = (TipStyle)0x20;
+
+    public static readonly Guid CopilotIconMonikerGuid = new("{4515B9BD-70A1-45FA-9545-D4536417C596}");
+    public static readonly Guid CopilotQuotaExceededGuid = new("39B0DEDE-D931-4A92-9AA2-3447BC4998DC");
+}

src/EditorFeatures/Core/DocumentationComments/CopilotGenerateDocumentationCommentManager.cs

@@ -67,7 +67,7 @@ private async Task GenerateDocumentationCommentProposalsAsync(Document document,
             var provider = textView.Properties.GetOrCreateSingletonProperty(typeof(CopilotGenerateDocumentationCommentProvider),
                 () => new CopilotGenerateDocumentationCommentProvider(_threadingContext, copilotService));
 
-            await provider!.InitializeAsync(textView, _suggestionServiceBase!, cancellationToken).ConfigureAwait(false);
+            await provider.InitializeAsync(textView, _suggestionServiceBase!, cancellationToken).ConfigureAwait(false);
 
             return provider;
         }

src/EditorFeatures/Core/DocumentationComments/CopilotGenerateDocumentationCommentProvider.cs

@@ -5,6 +5,7 @@
 using System;
 using System.Collections.Generic;
 using System.Collections.Immutable;
+using System.Diagnostics.CodeAnalysis;
 using System.Text;
 using System.Text.RegularExpressions;
 using System.Threading;
@@ -13,22 +14,23 @@
 using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
 using Microsoft.CodeAnalysis.PooledObjects;
 using Microsoft.CodeAnalysis.Text;
-using Microsoft.CodeAnalysis.Threading;
 using Microsoft.VisualStudio.Language.Proposals;
 using Microsoft.VisualStudio.Language.Suggestions;
 using Microsoft.VisualStudio.Text;
 using Microsoft.VisualStudio.Text.Editor;
+using Microsoft.VisualStudio.Threading;
 using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.DocumentationComments
 {
-    internal class CopilotGenerateDocumentationCommentProvider : SuggestionProviderBase
+    internal sealed class CopilotGenerateDocumentationCommentProvider : SuggestionProviderBase
     {
         private SuggestionManagerBase? _suggestionManager;
         private readonly ICopilotCodeAnalysisService _copilotService;
 
         public readonly IThreadingContext ThreadingContext;
 
+        [MemberNotNullWhen(true, nameof(_suggestionManager))]
         public bool Enabled => _enabled && (_suggestionManager != null);
         private bool _enabled = true;
 
@@ -46,37 +48,37 @@ public async Task InitializeAsync(ITextView textView, SuggestionServiceBase sugg
         public async Task GenerateDocumentationProposalAsync(DocumentationCommentSnippet snippet,
             ITextSnapshot oldSnapshot, VirtualSnapshotPoint oldCaret, CancellationToken cancellationToken)
         {
-            await Task.Yield().ConfigureAwait(false);
-
+            // Checks to see if the feature is enabled and if the suggestionManager is available
             if (!Enabled)
             {
                 return;
             }
 
-            // MemberNode is not null at this point, checked when determining if the file is exluded.
+            await TaskScheduler.Default;
+
+            // MemberNode is not null at this point, checked when determining if the file is excluded.
             var snippetProposal = GetSnippetProposal(snippet.SnippetText, snippet.MemberNode!, snippet.Position, snippet.CaretOffset);
 
             if (snippetProposal is null)
             {
                 return;
             }
 
-            // Do not do IntelliCode line completions if we're about to generate a documentation comment
-            // so that won't have interfering grey text.
-            var intellicodeLineCompletionsDisposable = await _suggestionManager!.DisableProviderAsync(SuggestionServiceNames.IntelliCodeLineCompletions, cancellationToken).ConfigureAwait(false);
-
-            var proposalEdits = await GetProposedEditsAsync(snippetProposal, _copilotService, oldSnapshot, snippet.IndentText, cancellationToken).ConfigureAwait(false);
-
-            var proposal = Proposal.TryCreateProposal(null, proposalEdits, oldCaret, flags: ProposalFlags.SingleTabToAccept);
+            // We need to disable IntelliCode Line Completions when starting a documentation comment session. Our suggestions take precedence to line completions in the documentation comment case.
+            // It needs to be disposed of in any case we have left the session, either after a user has accepted grey text or if we needed to bail out earlier in the process.
+            // Disposing of the provider is necessary to reenable the provider.
+            var intelliCodeLineCompletionsDisposable = await _suggestionManager.DisableProviderAsync(SuggestionServiceNames.IntelliCodeLineCompletions, cancellationToken).ConfigureAwait(false);
 
-            if (proposal is null)
+            var suggestion = new DocumentationCommentSuggestion(this, _suggestionManager, intelliCodeLineCompletionsDisposable);
+            Func<CancellationToken, Task<ProposalBase?>> generateProposal = async (cancellationToken) =>
             {
-                await intellicodeLineCompletionsDisposable.DisposeAsync().ConfigureAwait(false);
-                return;
-            }
+                var proposalEdits = await GetProposedEditsAsync(
+                    snippetProposal, _copilotService, oldSnapshot, snippet.IndentText, cancellationToken).ConfigureAwait(false);
+
+                return Proposal.TryCreateProposal(description: null, proposalEdits, oldCaret, flags: ProposalFlags.ShowCommitHighlight);
+            };
 
-            var suggestion = new DocumentationCommentSuggestion(this, proposal, _suggestionManager, intellicodeLineCompletionsDisposable);
-            await suggestion.TryDisplaySuggestionAsync(cancellationToken).ConfigureAwait(false);
+            await suggestion.StartSuggestionSessionWithProposalAsync(generateProposal, cancellationToken).ConfigureAwait(false);
         }
 
         /// <summary>

src/EditorFeatures/Core/DocumentationComments/DocumentationCommentSuggestion.cs

@@ -4,25 +4,27 @@
 
 using System;
 using System.ComponentModel;
+using System.Drawing;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.Editor.Copilot;
 using Microsoft.CodeAnalysis.Internal.Log;
+using Microsoft.CodeAnalysis.Shared.TestHooks;
 using Microsoft.VisualStudio.Language.Proposals;
 using Microsoft.VisualStudio.Language.Suggestions;
 using Microsoft.VisualStudio.Text;
+using Microsoft.VisualStudio.Threading;
 
 namespace Microsoft.CodeAnalysis.DocumentationComments
 {
-    internal class DocumentationCommentSuggestion(CopilotGenerateDocumentationCommentProvider providerInstance, ProposalBase proposal,
-        SuggestionManagerBase suggestionManager, VisualStudio.Threading.IAsyncDisposable? intellicodeLineCompletionsDisposable) : SuggestionBase
+    internal sealed class DocumentationCommentSuggestion(CopilotGenerateDocumentationCommentProvider providerInstance,
+        SuggestionManagerBase suggestionManager, VisualStudio.Threading.IAsyncDisposable? intelliCodeLineCompletionsDisposable) : SuggestionBase
     {
-        public ProposalBase Proposal { get; } = proposal;
-
         public SuggestionManagerBase SuggestionManager { get; } = suggestionManager;
 
-        public VisualStudio.Threading.IAsyncDisposable? IntellicodeLineCompletionsDisposable { get; set; } = intellicodeLineCompletionsDisposable;
+        public VisualStudio.Threading.IAsyncDisposable? IntelliCodeLineCompletionsDisposable { get; set; } = intelliCodeLineCompletionsDisposable;
 
-        public override TipStyle TipStyle => TipStyle.AlwaysShowTip;
+        public override TipStyle TipStyle => TipStyle.AlwaysShowTip | CopilotConstants.ShowThinkingStateTipStyle;
 
         public override EditDisplayStyle EditStyle => EditDisplayStyle.GrayText;
 
@@ -37,7 +39,7 @@ public override async Task OnAcceptedAsync(SuggestionSessionBase session, Propos
             var threadingContext = providerInstance.ThreadingContext;
 
             await threadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(cancel);
-            await DisposeAsync().ConfigureAwait(false);
+            await DisposeIntelliCodeCompletionsDisposableAsync().ConfigureAwait(false);
             Logger.Log(FunctionId.Copilot_Generate_Documentation_Accepted, logLevel: LogLevel.Information);
         }
 
@@ -65,34 +67,121 @@ public override Task OnProposalUpdatedAsync(SuggestionSessionBase session, Propo
             return Task.CompletedTask;
         }
 
-        public async Task TryDisplaySuggestionAsync(CancellationToken cancellationToken)
+        public async Task StartSuggestionSessionWithProposalAsync(
+            Func<CancellationToken, Task<ProposalBase?>> generateProposal, CancellationToken cancellationToken)
         {
-            _suggestionSession = await SuggestionManager.TryDisplaySuggestionAsync(this, cancellationToken).ConfigureAwait(false);
+            var sessionStarted = await StartSuggestionSessionAsync(cancellationToken).ConfigureAwait(false);
+            if (!sessionStarted)
+            {
+                return;
+            }
 
-            if (_suggestionSession != null)
+            var proposal = await generateProposal(cancellationToken).ConfigureAwait(false);
+            if (proposal is null)
             {
-                var success = await TryDisplayProposalAsync(_suggestionSession, cancellationToken).ConfigureAwait(false);
-                if (success)
-                {
-                    Logger.Log(FunctionId.Copilot_Generate_Documentation_Displayed, logLevel: LogLevel.Information);
-                }
+                await DismissSuggestionSessionAsync(cancellationToken).ConfigureAwait(false);
+                return;
             }
+
+            await TryDisplayDocumentationSuggestionAsync(proposal, cancellationToken).ConfigureAwait(false);
         }
 
-        private async Task<bool> TryDisplayProposalAsync(SuggestionSessionBase session, CancellationToken cancellationToken)
+        /// <summary>
+        /// Starts the Suggestion Session. The TryDisplaySuggestion call doesn't display any grey text, but starts the session such that we have the
+        /// exclusive right to display grey text later.
+        /// </summary>
+        /// <returns>If true, user will see the thinking state as long as the Suggestion Session is active and replace with grey text if a call to DisplayProposal succeeds.
+        /// If unable to retrieve the session, the caller should bail out.
+        /// </returns>
+        private async Task<bool> StartSuggestionSessionAsync(CancellationToken cancellationToken)
+        {
+            _suggestionSession = await RunWithEnqueueActionAsync(
+                "StartWork",
+                async () => await SuggestionManager.TryDisplaySuggestionAsync(this, cancellationToken).ConfigureAwait(false),
+            cancellationToken).ConfigureAwait(false);
+
+            if (_suggestionSession is null)
+            {
+                await DisposeIntelliCodeCompletionsDisposableAsync().ConfigureAwait(false);
+                return false;
+            }
+
+            return true;
+        }
+
+        /// <summary>
+        /// This is where we actually try to display the grey-text from the proposal
+        /// we created.
+        /// </summary>
+        public async Task TryDisplayDocumentationSuggestionAsync(ProposalBase proposal, CancellationToken cancellationToken)
         {
             try
             {
-                await providerInstance.ThreadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
-                await session.DisplayProposalAsync(Proposal, cancellationToken).ConfigureAwait(false);
-                return true;
+                await RunWithEnqueueActionAsync<bool>(
+                    "DisplayProposal",
+                    async () =>
+                    {
+                        await _suggestionSession!.DisplayProposalAsync(proposal, cancellationToken).ConfigureAwait(false);
+                        return true;
+                    },
+                    cancellationToken).ConfigureAwait(false);
+
+                Logger.Log(FunctionId.Copilot_Generate_Documentation_Displayed, logLevel: LogLevel.Information);
             }
             catch (OperationCanceledException)
             {
                 Logger.Log(FunctionId.Copilot_Generate_Documentation_Canceled, logLevel: LogLevel.Information);
             }
+        }
+
+        /// <summary>
+        /// Dismisses the session if the proposal we generated was invalid.
+        /// Needs to dispose of the IntelliCodeCompletionsDisposable so we no longer have exclusive right to
+        /// display any grey text.
+        /// </summary>
+        private async Task DismissSuggestionSessionAsync(CancellationToken cancellationToken)
+        {
+            await RunWithEnqueueActionAsync<bool>(
+                "DismissSuggestionSession",
+                async () =>
+                {
+                    await ClearSuggestionAsync(ReasonForDismiss.DismissedDueToInvalidProposal, cancellationToken).ConfigureAwait(false);
+                    return true;
+                },
+                cancellationToken).ConfigureAwait(false);
+        }
+
+        /// <summary>
+        /// In general, calls to a SuggestionManager or SuggestionSession need to be wrapped in an EnqueueAction.
+        /// This is the pattern recommended by VS Platform to avoid races.
+        /// Pattern from platform shown here:
+        /// https://devdiv.visualstudio.com/DevDiv/_git/IntelliCode-VS?path=/src/VSIX/IntelliCode.VSIX/SuggestionService/AmbientAI/SuggestionProviderForAmbientAI.cs
+        /// </summary>
+        private async Task<T> RunWithEnqueueActionAsync<T>(string description, Func<Task<T>> action, CancellationToken cancellationToken)
+        {
+            Assumes.NotNull(SuggestionManager);
+
+            var taskCompletionSource = new TaskCompletionSource<T>(TaskCreationOptions.RunContinuationsAsynchronously);
+
+            await providerInstance.ThreadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
+            SuggestionManager.EnqueueAction(description, async () =>
+            {
+                try
+                {
+                    var result = await action().ConfigureAwaitRunInline();
+                    taskCompletionSource.TrySetResult(result);
+                }
+                catch (OperationCanceledException operationCanceledException)
+                {
+                    taskCompletionSource.TrySetCanceled(operationCanceledException.CancellationToken);
+                }
+                catch (Exception exception)
+                {
+                    taskCompletionSource.TrySetException(exception);
+                }
+            });
 
-            return false;
+            return await taskCompletionSource.Task.WithCancellation(cancellationToken).ConfigureAwait(false);
         }
 
         private async Task ClearSuggestionAsync(ReasonForDismiss reason, CancellationToken cancellationToken)
@@ -103,15 +192,19 @@ private async Task ClearSuggestionAsync(ReasonForDismiss reason, CancellationTok
             }
 
             _suggestionSession = null;
-            await DisposeAsync().ConfigureAwait(false);
+            await DisposeIntelliCodeCompletionsDisposableAsync().ConfigureAwait(false);
         }
 
-        private async Task DisposeAsync()
+        /// <summary>
+        /// The IntelliCodeLineCompletionDisposable needs to be disposed any time we exit the SuggestionSession so that
+        /// line completions can be shown again.
+        /// </summary>
+        private async Task DisposeIntelliCodeCompletionsDisposableAsync()
         {
-            if (IntellicodeLineCompletionsDisposable != null)
+            if (IntelliCodeLineCompletionsDisposable != null)
             {
-                await IntellicodeLineCompletionsDisposable.DisposeAsync().ConfigureAwait(false);
-                IntellicodeLineCompletionsDisposable = null;
+                await IntelliCodeLineCompletionsDisposable.DisposeAsync().ConfigureAwait(false);
+                IntelliCodeLineCompletionsDisposable = null;
             }
         }
     }