Fix cohosting completion and code actions

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Extensions/ProjectExtensions.cs

@@ -11,8 +11,9 @@
 using Microsoft.AspNetCore.Razor;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.AspNetCore.Razor.PooledObjects;
-using Microsoft.CodeAnalysis.Razor;
+using Microsoft.AspNetCore.Razor.Threading;
 using Microsoft.CodeAnalysis.ExternalAccess.Razor;
+using Microsoft.CodeAnalysis.Razor;
 using Microsoft.CodeAnalysis.Razor.Telemetry;
 
 namespace Microsoft.CodeAnalysis;
@@ -84,23 +85,14 @@ internal static Document GetRequiredDocument(this Project project, DocumentId do
             ?? ThrowHelper.ThrowInvalidOperationException<Document>($"The document {documentId} did not exist in {project.Name}");
     }
 
-    public static bool TryGetCSharpDocument(this Project project, Uri csharpDocumentUri, [NotNullWhen(true)] out Document? document)
+    public static Task<SourceGeneratedDocument?> TryGetCSharpDocumentFromGeneratedDocumentUriAsync(this Project project, Uri generatedDocumentUri, CancellationToken cancellationToken)
     {
-        document = null;
-
-        var generatedDocumentIds = project.Solution.GetDocumentIdsWithUri(csharpDocumentUri);
-        var generatedDocumentId = generatedDocumentIds.FirstOrDefault(d => d.ProjectId == project.Id);
-        if (generatedDocumentId is null)
-        {
-            return false;
-        }
-
-        if (project.GetDocument(generatedDocumentId) is { } generatedDocument)
+        if (!TryGetHintNameFromGeneratedDocumentUri(project, generatedDocumentUri, out var hintName))
         {
-            document = generatedDocument;
+            return SpecializedTasks.Null<SourceGeneratedDocument>();
         }
 
-        return document is not null;
+        return TryGetSourceGeneratedDocumentFromHintNameAsync(project, hintName, cancellationToken);
     }
 
     /// <summary>

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Extensions/SolutionExtensions.cs

@@ -3,9 +3,11 @@
 
 using System;
 using System.Collections.Immutable;
+using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using Microsoft.AspNetCore.Razor;
+using Microsoft.CodeAnalysis.ExternalAccess.Razor;
 using Microsoft.CodeAnalysis.Razor;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;
 
@@ -14,7 +16,10 @@ namespace Microsoft.CodeAnalysis;
 internal static class SolutionExtensions
 {
     public static ImmutableArray<DocumentId> GetDocumentIdsWithUri(this Solution solution, Uri uri)
-        => solution.GetDocumentIdsWithFilePath(uri.GetDocumentFilePath());
+    {
+        Debug.Assert(RazorUri.IsGeneratedDocumentUri(uri) == false, "This won't work with source generated Uris");
+        return solution.GetDocumentIdsWithFilePath(uri.GetDocumentFilePath());
+    }
 
     public static bool TryGetRazorDocument(this Solution solution, Uri razorDocumentUri, [NotNullWhen(true)] out TextDocument? razorDocument)
     {

src/Razor/src/Microsoft.VisualStudio.LanguageServices.Razor/LanguageClient/Cohost/CohostCodeActionsEndpoint.cs

@@ -98,7 +98,8 @@ public ImmutableArray<Registration> GetRegistrations(VSInternalClientCapabilitie
 
     private async Task<RazorVSInternalCodeAction[]> GetCSharpCodeActionsAsync(TextDocument razorDocument, VSCodeActionParams request, Guid correlationId, CancellationToken cancellationToken)
     {
-        if (!razorDocument.Project.TryGetCSharpDocument(request.TextDocument.Uri, out var generatedDocument))
+        var generatedDocument = await razorDocument.Project.TryGetCSharpDocumentFromGeneratedDocumentUriAsync(request.TextDocument.Uri, cancellationToken).ConfigureAwait(false);
+        if (generatedDocument is null)
         {
             return [];
         }

src/Razor/src/Microsoft.VisualStudio.LanguageServices.Razor/LanguageClient/Cohost/CohostCodeActionsResolveEndpoint.cs

@@ -102,7 +102,8 @@ private async Task<CodeAction> ResolveCSharpCodeActionAsync(TextDocument razorDo
 
             var uri = resolveParams.DelegatedDocumentUri.AssumeNotNull();
 
-            if (!razorDocument.Project.TryGetCSharpDocument(uri, out var generatedDocument))
+            var generatedDocument = await razorDocument.Project.TryGetCSharpDocumentFromGeneratedDocumentUriAsync(uri, cancellationToken).ConfigureAwait(false);
+            if (generatedDocument is null)
             {
                 return codeAction;
             }

src/Razor/src/Microsoft.VisualStudio.LanguageServices.Razor/LanguageClient/Cohost/CohostInlineCompletionEndpoint.cs

@@ -68,13 +68,13 @@ public ImmutableArray<Registration> GetRegistrations(VSInternalClientCapabilitie
             return null;
         }
 
-        if (!razorDocument.Project.TryGetCSharpDocument(generatedDocumentUri, out var generatedDocument))
+        var generatedDocument = await razorDocument.Project.TryGetCSharpDocumentFromGeneratedDocumentUriAsync(generatedDocumentUri, cancellationToken).ConfigureAwait(false);
+        if (generatedDocument is null)
         {
             return null;
         }
 
         var result = await Completion.GetInlineCompletionItemsAsync(context, generatedDocument, position, formattingOptions, cancellationToken).ConfigureAwait(false);
-
         if (result is null)
         {
             return null;

src/Razor/test/Microsoft.VisualStudio.LanguageServices.Razor.Test/Cohost/CodeActions/AddUsingTests.cs

@@ -3,14 +3,15 @@
 
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.ExternalAccess.Razor;
+using Microsoft.CodeAnalysis.Razor.Protocol;
 using Xunit;
 using Xunit.Abstractions;
 
 namespace Microsoft.VisualStudio.Razor.LanguageClient.Cohost.CodeActions;
 
 public class AddUsingTests(ITestOutputHelper testOutputHelper) : CohostCodeActionsEndpointTestBase(testOutputHelper)
 {
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact]
     public async Task FullyQualify()
     {
         var input = """
@@ -27,10 +28,10 @@ public async Task FullyQualify()
             }
             """;
 
-        await VerifyCodeActionAsync(input, expected, "System.Text.StringBuilder");
+        await VerifyCodeActionAsync(input, expected, LanguageServerConstants.CodeActions.FullyQualify);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact]
     public async Task FullyQualify_Multiple()
     {
         await VerifyCodeActionAsync(
@@ -54,11 +55,11 @@ public class StringBuilder
                     {
                     }
                     """)],
-            codeActionName: "Fully qualify 'StringBuilder'",
+            codeActionName: LanguageServerConstants.CodeActions.FullyQualify,
             childActionIndex: 0);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact]
     public async Task AddUsing()
     {
         var input = """
@@ -79,7 +80,7 @@ @using System.Text
         await VerifyCodeActionAsync(input, expected, RazorPredefinedCodeFixProviderNames.AddImport);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact]
     public async Task AddUsing_Typo()
     {
         var input = """
@@ -100,7 +101,7 @@ @using System.Text
         await VerifyCodeActionAsync(input, expected, RazorPredefinedCodeFixProviderNames.AddImport);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact]
     public async Task AddUsing_WithExisting()
     {
         var input = """

src/Razor/test/Microsoft.VisualStudio.LanguageServices.Razor.Test/Cohost/CodeActions/CSharpCodeActionTests.cs

@@ -10,7 +10,7 @@ namespace Microsoft.VisualStudio.Razor.LanguageClient.Cohost.CodeActions;
 
 public class CSharpCodeActionTests(ITestOutputHelper testOutputHelper) : CohostCodeActionsEndpointTestBase(testOutputHelper)
 {
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact]
     public async Task GenerateConstructor()
     {
         var input = """
@@ -45,7 +45,7 @@ public Goo()
         await VerifyCodeActionAsync(input, expected, RazorPredefinedCodeRefactoringProviderNames.GenerateConstructorFromMembers);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact]
     public async Task UseExpressionBodiedMember()
     {
         var input = """
@@ -78,7 +78,7 @@ @using System.Linq
         await VerifyCodeActionAsync(input, expected, RazorPredefinedCodeRefactoringProviderNames.UseExpressionBody);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact(Skip = "Roslyn code refactoring provider is not finding the expression")]
     public async Task IntroduceLocal()
     {
         var input = """
@@ -125,7 +125,7 @@ void M(string[] args)
         await VerifyCodeActionAsync(input, expected, RazorPredefinedCodeRefactoringProviderNames.IntroduceVariable);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact(Skip = "Roslyn code refactoring provider is not finding the expression")]
     public async Task IntroduceLocal_All()
     {
         var input = """
@@ -172,7 +172,7 @@ void M(string[] args)
         await VerifyCodeActionAsync(input, expected, RazorPredefinedCodeRefactoringProviderNames.IntroduceVariable, childActionIndex: 1);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact]
     public async Task ConvertConcatenationToInterpolatedString_CSharpStatement()
     {
         var input = """
@@ -190,7 +190,7 @@ public async Task ConvertConcatenationToInterpolatedString_CSharpStatement()
         await VerifyCodeActionAsync(input, expected, RazorPredefinedCodeRefactoringProviderNames.ConvertConcatenationToInterpolatedString);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact]
     public async Task ConvertConcatenationToInterpolatedString_ExplicitExpression()
     {
         var input = """
@@ -204,7 +204,7 @@ public async Task ConvertConcatenationToInterpolatedString_ExplicitExpression()
         await VerifyCodeActionAsync(input, expected, RazorPredefinedCodeRefactoringProviderNames.ConvertConcatenationToInterpolatedString);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact]
     public async Task ConvertConcatenationToInterpolatedString_CodeBlock()
     {
         var input = """
@@ -224,7 +224,7 @@ public async Task ConvertConcatenationToInterpolatedString_CodeBlock()
         await VerifyCodeActionAsync(input, expected, RazorPredefinedCodeRefactoringProviderNames.ConvertConcatenationToInterpolatedString);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact]
     public async Task ConvertBetweenRegularAndVerbatimInterpolatedString_CodeBlock()
     {
         var input = """
@@ -244,7 +244,7 @@ public async Task ConvertBetweenRegularAndVerbatimInterpolatedString_CodeBlock()
         await VerifyCodeActionAsync(input, expected, RazorPredefinedCodeRefactoringProviderNames.ConvertBetweenRegularAndVerbatimInterpolatedString);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact]
     public async Task ConvertBetweenRegularAndVerbatimInterpolatedString_CodeBlock2()
     {
         var input = """
@@ -264,7 +264,7 @@ public async Task ConvertBetweenRegularAndVerbatimInterpolatedString_CodeBlock2(
         await VerifyCodeActionAsync(input, expected, RazorPredefinedCodeRefactoringProviderNames.ConvertBetweenRegularAndVerbatimInterpolatedString);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact]
     public async Task ConvertBetweenRegularAndVerbatimString_CodeBlock()
     {
         var input = """
@@ -284,7 +284,7 @@ public async Task ConvertBetweenRegularAndVerbatimString_CodeBlock()
         await VerifyCodeActionAsync(input, expected, RazorPredefinedCodeRefactoringProviderNames.ConvertBetweenRegularAndVerbatimString);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact]
     public async Task ConvertBetweenRegularAndVerbatimString_CodeBlock2()
     {
         var input = """
@@ -304,7 +304,7 @@ public async Task ConvertBetweenRegularAndVerbatimString_CodeBlock2()
         await VerifyCodeActionAsync(input, expected, RazorPredefinedCodeRefactoringProviderNames.ConvertBetweenRegularAndVerbatimString);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact]
     public async Task ConvertPlaceholderToInterpolatedString_CodeBlock()
     {
         var input = """
@@ -324,7 +324,7 @@ public async Task ConvertPlaceholderToInterpolatedString_CodeBlock()
         await VerifyCodeActionAsync(input, expected, RazorPredefinedCodeRefactoringProviderNames.ConvertPlaceholderToInterpolatedString);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact]
     public async Task ConvertToInterpolatedString_CodeBlock()
     {
         var input = """
@@ -344,7 +344,7 @@ public async Task ConvertToInterpolatedString_CodeBlock()
         await VerifyCodeActionAsync(input, expected, RazorPredefinedCodeRefactoringProviderNames.ConvertToInterpolatedString);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact]
     public async Task AddDebuggerDisplay()
     {
         var input = """

src/Razor/test/Microsoft.VisualStudio.LanguageServices.Razor.Test/Cohost/CodeActions/GenerateEventHandlerTests.cs

@@ -10,7 +10,7 @@ namespace Microsoft.VisualStudio.Razor.LanguageClient.Cohost.CodeActions;
 
 public class GenerateEventHandlerTests(ITestOutputHelper testOutputHelper) : CohostCodeActionsEndpointTestBase(testOutputHelper)
 {
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact]
     public async Task NoCodeBlock()
     {
         var input = """
@@ -30,7 +30,7 @@ private void DoesNotExist(MouseEventArgs args)
         await VerifyCodeActionAsync(input, expected, LanguageServerConstants.CodeActions.GenerateEventHandler);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact]
     public async Task CodeBlock()
     {
         var input = """
@@ -56,7 +56,7 @@ private void DoesNotExist(MouseEventArgs args)
         await VerifyCodeActionAsync(input, expected, LanguageServerConstants.CodeActions.GenerateEventHandler);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact(Skip = "@bind- attribute tag helper is not being found")]
     public async Task BindSet()
     {
         var input = """
@@ -84,7 +84,7 @@ private void DoesNotExist(string args)
         await VerifyCodeActionAsync(input, expected, LanguageServerConstants.CodeActions.GenerateEventHandler);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact(Skip = "@bind- attribute tag helper is not being found")]
     public async Task BindAfter()
     {
         var input = """
@@ -112,7 +112,7 @@ private void DoesNotExist()
         await VerifyCodeActionAsync(input, expected, LanguageServerConstants.CodeActions.GenerateEventHandler);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact]
     public async Task Callback()
     {
         var input = """
@@ -138,7 +138,7 @@ private void DoesNotExist(InputFileChangeEventArgs args)
         await VerifyCodeActionAsync(input, expected, LanguageServerConstants.CodeActions.GenerateEventHandler);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact]
     public async Task AsyncCallback()
     {
         var input = """
@@ -164,7 +164,7 @@ private Task DoesNotExistAsync(string args)
         await VerifyCodeActionAsync(input, expected, LanguageServerConstants.CodeActions.GenerateAsyncEventHandler);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact]
     public async Task BadCodeBehind()
     {
         await VerifyCodeActionAsync(
@@ -192,7 +192,7 @@ public partial class NotAComponent
             codeActionName: LanguageServerConstants.CodeActions.GenerateEventHandler);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact]
     public async Task CodeBehind()
     {
         await VerifyCodeActionAsync(
@@ -231,7 +231,7 @@ private void DoesNotExist(Microsoft.AspNetCore.Components.Web.MouseEventArgs arg
             codeActionName: LanguageServerConstants.CodeActions.GenerateEventHandler);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact]
     public async Task EmptyCodeBehind()
     {
         await VerifyCodeActionAsync(
@@ -264,7 +264,7 @@ private void DoesNotExist(Microsoft.AspNetCore.Components.Web.MouseEventArgs arg
             codeActionName: LanguageServerConstants.CodeActions.GenerateEventHandler);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact]
     public async Task GenerateAsyncEventHandler_NoCodeBlock()
     {
         var input = """
@@ -284,7 +284,7 @@ private Task DoesNotExist(MouseEventArgs args)
         await VerifyCodeActionAsync(input, expected, LanguageServerConstants.CodeActions.GenerateAsyncEventHandler);
     }
 
-    [Fact(Skip = "Need to map uri back to source generated document")]
+    [Fact]
     public async Task GenerateAsyncEventHandler_CodeBlock()
     {
         var input = """