Fix await completion in an async iterator

src/EditorFeatures/Test2/IntelliSense/CSharpCompletionCommandHandlerTests_AwaitCompletion.vb

@@ -1508,6 +1508,76 @@ public class C
         }
     }
 }
+", state.GetDocumentText())
+            End Using
+        End Function
+
+        <WpfFact, WorkItem("https://github.com/dotnet/roslyn/issues/78822")>
+        Public Async Function AwaitCompletionAddsAsync_AsyncEnumerableMethodDeclaration1() As Task
+            Using state = TestStateFactory.CreateCSharpTestState(
+                <Document><![CDATA[
+using System.Collections.Generic;
+using System.Threading.Tasks;
+
+public class C
+{
+    public static IAsyncEnumerable<string> Main()
+    {
+        $$
+    }
+}
+]]>
+                </Document>)
+                state.SendTypeChars("aw")
+                Await state.AssertSelectedCompletionItem(displayText:="await", isHardSelected:=True)
+
+                state.SendTab()
+                Assert.Equal("
+using System.Collections.Generic;
+using System.Threading.Tasks;
+
+public class C
+{
+    public static async IAsyncEnumerable<string> Main()
+    {
+        await
+    }
+}
+", state.GetDocumentText())
+            End Using
+        End Function
+
+        <WpfFact, WorkItem("https://github.com/dotnet/roslyn/issues/78822")>
+        Public Async Function AwaitCompletionAddsAsync_AsyncEnumerableMethodDeclaration2() As Task
+            Using state = TestStateFactory.CreateCSharpTestState(
+                <Document><![CDATA[
+using System.Collections.Generic;
+using System.Threading.Tasks;
+
+public class C
+{
+    public static async IAsyncEnumerable<string> Main()
+    {
+        $$
+    }
+}
+]]>
+                </Document>)
+                state.SendTypeChars("aw")
+                Await state.AssertSelectedCompletionItem(displayText:="await", isHardSelected:=True)
+
+                state.SendTab()
+                Assert.Equal("
+using System.Collections.Generic;
+using System.Threading.Tasks;
+
+public class C
+{
+    public static async IAsyncEnumerable<string> Main()
+    {
+        await
+    }
+}
 ", state.GetDocumentText())
             End Using
         End Function

src/Features/CSharp/Portable/Completion/CompletionProviders/AwaitCompletionProvider.cs

@@ -15,7 +15,6 @@
 using Microsoft.CodeAnalysis.CSharp.Syntax;
 using Microsoft.CodeAnalysis.Host.Mef;
 using Microsoft.CodeAnalysis.Shared.Extensions;
-using Microsoft.CodeAnalysis.Shared.Extensions.ContextQuery;
 using Microsoft.CodeAnalysis.Text;
 using Roslyn.Utilities;
 
@@ -31,9 +30,6 @@ internal sealed class AwaitCompletionProvider() : AbstractAwaitCompletionProvide
 
     public override ImmutableHashSet<char> TriggerCharacters => CompletionUtilities.CommonTriggerCharactersWithArgumentList;
 
-    protected override bool IsAwaitKeywordContext(SyntaxContext syntaxContext)
-        => base.IsAwaitKeywordContext(syntaxContext);
-
     /// <summary>
     /// Gets the span start where async keyword should go.
     /// </summary>
@@ -86,8 +82,13 @@ protected override int GetAsyncKeywordInsertionPosition(SyntaxNode declaration)
             // Don't change the return type if we don't understand it, or it already seems task-like.
             var taskLikeTypes = new KnownTaskTypes(semanticModel.Compilation);
             var returnType = semanticModel.GetTypeInfo(existingReturnType, cancellationToken).Type;
-            if (returnType is null or IErrorTypeSymbol || taskLikeTypes.IsTaskLike(returnType))
+            if (returnType is null or IErrorTypeSymbol ||
+                taskLikeTypes.IsTaskLike(returnType) ||
+                returnType.OriginalDefinition.Equals(taskLikeTypes.IAsyncEnumerableOfTType) ||
+                returnType.OriginalDefinition.Equals(taskLikeTypes.IAsyncEnumeratorOfTType))
+            {
                 return null;
+            }
 
             return $"{nameof(Task)}<{existingReturnType}>";
         }

src/Features/Core/Portable/Completion/Providers/AbstractAwaitCompletionProvider.cs

@@ -11,7 +11,6 @@
 using Microsoft.CodeAnalysis.Editing;
 using Microsoft.CodeAnalysis.LanguageService;
 using Microsoft.CodeAnalysis.PooledObjects;
-using Microsoft.CodeAnalysis.Shared.Collections;
 using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.Shared.Extensions.ContextQuery;
 using Microsoft.CodeAnalysis.Text;
@@ -65,9 +64,6 @@ protected AbstractAwaitCompletionProvider(ISyntaxFacts syntaxFacts)
     protected abstract SyntaxNode? GetExpressionToPlaceAwaitInFrontOf(SyntaxTree syntaxTree, int position, CancellationToken cancellationToken);
     protected abstract SyntaxToken? GetDotTokenLeftOfPosition(SyntaxTree syntaxTree, int position, CancellationToken cancellationToken);
 
-    protected virtual bool IsAwaitKeywordContext(SyntaxContext syntaxContext)
-        => syntaxContext.IsAwaitKeywordContext;
-
     private static bool IsConfigureAwaitable(Compilation compilation, ITypeSymbol symbol)
     {
         var originalDefinition = symbol.OriginalDefinition;
@@ -90,7 +86,7 @@ public sealed override async Task ProvideCompletionsAsync(CompletionContext cont
 
         var syntaxContext = await context.GetSyntaxContextWithExistingSpeculativeModelAsync(document, cancellationToken).ConfigureAwait(false);
 
-        var isAwaitKeywordContext = IsAwaitKeywordContext(syntaxContext);
+        var isAwaitKeywordContext = syntaxContext.IsAwaitKeywordContext;
         var dotAwaitContext = GetDotAwaitKeywordContext(syntaxContext, cancellationToken);
         if (!isAwaitKeywordContext && dotAwaitContext == DotAwaitContext.None)
             return;