Fix problematic things on IProjectSnapshot (#9871)

Part of #9519

As per our meeting the other day, there are a few things on
`IProjectSnapshot` that are problematic. This PR resolves all of the
issues in one of two ways:

* For `Configuration`, `GetProjectEngine()` and `ProjectWorkspaceState`
the cohost implementations now simply throw
* There was some usage of these inside cohosting, because it re-used
methods on `DocumentState` to do jobs that the source gen will
eventually take over. To make this clearer, I modified the helper
methods so they took the underlying data, rather than pulling it from
`IProjectSnapshot` directly
* Changes the `TagHelpers` property into a
`GetTagHelpersAsync(CancellationToken)` method
* Because of the virality of async, I recommend reviewing this PR
commit-at-a-time, just so this last commit is seen separately, for your
sanity.
* There were only two places where moving to async would have been very
very annoying, so I added a `GetTagHelpersSynchronously()` extension
method that validates that its not called from cohosting. It's only
called from legacy editor code.

Note that this isn't the proposed solution in the aforementioned
meeting, that change can still come later. I was hoping to take that
idea and use DI to enforce something better than just a runtime check,
but it turned out that wasn't possible because some MEF services need
these properties in non-cohosting. This PR at least introduces the
runtime check, so we can continue moving forward.

Integration tests are running and should show up in in the checks list.

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/CodeActions/Razor/ComponentAccessibilityCodeActionProvider.cs

@@ -25,24 +25,22 @@ namespace Microsoft.AspNetCore.Razor.LanguageServer.CodeActions;
 
 internal sealed class ComponentAccessibilityCodeActionProvider : IRazorCodeActionProvider
 {
-    private static readonly Task<IReadOnlyList<RazorVSInternalCodeAction>?> s_emptyResult = Task.FromResult<IReadOnlyList<RazorVSInternalCodeAction>?>(null);
-
     private readonly ITagHelperFactsService _tagHelperFactsService;
 
     public ComponentAccessibilityCodeActionProvider(ITagHelperFactsService tagHelperFactsService)
     {
         _tagHelperFactsService = tagHelperFactsService ?? throw new ArgumentNullException(nameof(tagHelperFactsService));
     }
 
-    public Task<IReadOnlyList<RazorVSInternalCodeAction>?> ProvideAsync(RazorCodeActionContext context, CancellationToken cancellationToken)
+    public async Task<IReadOnlyList<RazorVSInternalCodeAction>?> ProvideAsync(RazorCodeActionContext context, CancellationToken cancellationToken)
     {
         using var _ = ListPool<RazorVSInternalCodeAction>.GetPooledObject(out var codeActions);
 
         // Locate cursor
         var node = context.CodeDocument.GetSyntaxTree().Root.FindInnermostNode(context.Location.AbsoluteIndex);
         if (node is null)
         {
-            return s_emptyResult;
+            return null;
         }
 
         // Find start tag. We allow this code action to work from anywhere in the start tag, which includes
@@ -53,35 +51,35 @@ public ComponentAccessibilityCodeActionProvider(ITagHelperFactsService tagHelper
         var startTag = (IStartTagSyntaxNode?)node.FirstAncestorOrSelf<SyntaxNode>(n => n is IStartTagSyntaxNode);
         if (startTag is null)
         {
-            return s_emptyResult;
+            return null;
         }
 
         if (context.Location.AbsoluteIndex < startTag.SpanStart)
         {
             // Cursor is before the start tag, so we shouldn't show a light bulb. This can happen
             // in cases where the cursor is in whitespace at the beginning of the document
             // eg: $$ <Component></Component>
-            return s_emptyResult;
+            return null;
         }
 
         // Ignore if start tag has dots, as we only handle short tags
         if (startTag.Name.Content.Contains("."))
         {
-            return s_emptyResult;
+            return null;
         }
 
         if (!IsApplicableTag(startTag))
         {
-            return s_emptyResult;
+            return null;
         }
 
         if (IsTagUnknown(startTag, context))
         {
-            AddComponentAccessFromTag(context, startTag, codeActions);
+            await AddComponentAccessFromTagAsync(context, startTag, codeActions, cancellationToken).ConfigureAwait(false);
             AddCreateComponentFromTag(context, startTag, codeActions);
         }
 
-        return Task.FromResult<IReadOnlyList<RazorVSInternalCodeAction>?>(codeActions.ToArray());
+        return codeActions.ToArray();
     }
 
     private static bool IsApplicableTag(IStartTagSyntaxNode startTag)
@@ -136,7 +134,7 @@ private void AddCreateComponentFromTag(RazorCodeActionContext context, IStartTag
         container.Add(codeAction);
     }
 
-    private void AddComponentAccessFromTag(RazorCodeActionContext context, IStartTagSyntaxNode startTag, List<RazorVSInternalCodeAction> container)
+    private async Task AddComponentAccessFromTagAsync(RazorCodeActionContext context, IStartTagSyntaxNode startTag, List<RazorVSInternalCodeAction> container, CancellationToken cancellationToken)
     {
         var haveAddedNonQualifiedFix = false;
 
@@ -155,7 +153,7 @@ private void AddComponentAccessFromTag(RazorCodeActionContext context, IStartTag
             }
         }
 
-        var matching = FindMatchingTagHelpers(context, startTag);
+        var matching = await FindMatchingTagHelpersAsync(context, startTag, cancellationToken).ConfigureAwait(false);
 
         // For all the matches, add options for add @using and fully qualify
         foreach (var tagHelperPair in matching)
@@ -205,7 +203,7 @@ private void AddComponentAccessFromTag(RazorCodeActionContext context, IStartTag
         }
     }
 
-    private List<TagHelperPair> FindMatchingTagHelpers(RazorCodeActionContext context, IStartTagSyntaxNode startTag)
+    private async Task<List<TagHelperPair>> FindMatchingTagHelpersAsync(RazorCodeActionContext context, IStartTagSyntaxNode startTag, CancellationToken cancellationToken)
     {
         // Get all data necessary for matching
         var tagName = startTag.Name.Content;
@@ -224,7 +222,9 @@ private List<TagHelperPair> FindMatchingTagHelpers(RazorCodeActionContext contex
         // Find all matching tag helpers
         using var _ = DictionaryPool<string, TagHelperPair>.GetPooledObject(out var matching);
 
-        foreach (var tagHelper in context.DocumentSnapshot.Project.TagHelpers)
+        var tagHelpers = await context.DocumentSnapshot.Project.GetTagHelpersAsync(cancellationToken).ConfigureAwait(false);
+
+        foreach (var tagHelper in tagHelpers)
         {
             if (SatisfiesRules(tagHelper.TagMatchingRules, tagName.AsSpan(), parentTagName.AsSpan(), attributes, out var caseInsensitiveMatch))
             {
@@ -233,7 +233,7 @@ private List<TagHelperPair> FindMatchingTagHelpers(RazorCodeActionContext contex
         }
 
         // Iterate and find the fully qualified version
-        foreach (var tagHelper in context.DocumentSnapshot.Project.TagHelpers)
+        foreach (var tagHelper in tagHelpers)
         {
             if (matching.TryGetValue(tagHelper.Name, out var tagHelperPair))
             {

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/CodeActions/Razor/GenerateMethodCodeActionResolver.cs

@@ -116,7 +116,7 @@ public GenerateMethodCodeActionResolver(
 
         var codeBehindTextDocumentIdentifier = new OptionalVersionedTextDocumentIdentifier() { Uri = codeBehindUri };
 
-        var templateWithMethodSignature = PopulateMethodSignature(documentContext, actionParams);
+        var templateWithMethodSignature = await PopulateMethodSignatureAsync(documentContext, actionParams, cancellationToken).ConfigureAwait(false);
         var classLocationLineSpan = @class.GetLocation().GetLineSpan();
         var formattedMethod = FormattingUtilities.AddIndentationToMethod(
             templateWithMethodSignature,
@@ -160,7 +160,7 @@ private async Task<WorkspaceEdit> GenerateMethodInCodeBlockAsync(
         string? razorClassName,
         CancellationToken cancellationToken)
     {
-        var templateWithMethodSignature = PopulateMethodSignature(documentContext, actionParams);
+        var templateWithMethodSignature = await PopulateMethodSignatureAsync(documentContext, actionParams, cancellationToken).ConfigureAwait(false);
         var edits = CodeBlockService.CreateFormattedTextEdit(code, templateWithMethodSignature, _razorLSPOptionsMonitor.CurrentValue);
 
         // If there are 3 edits, this means that there is no existing @code block, so we have an edit for '@code {', the method stub, and '}'.
@@ -250,14 +250,15 @@ private async Task<WorkspaceEdit> GenerateMethodInCodeBlockAsync(
         return new WorkspaceEdit() { DocumentChanges = new[] { razorTextDocEdit } };
     }
 
-    private static string PopulateMethodSignature(VersionedDocumentContext documentContext, GenerateMethodCodeActionParams actionParams)
+    private static async Task<string> PopulateMethodSignatureAsync(VersionedDocumentContext documentContext, GenerateMethodCodeActionParams actionParams, CancellationToken cancellationToken)
     {
         var templateWithMethodSignature = s_generateMethodTemplate.Replace(s_methodName, actionParams.MethodName);
 
         var returnType = actionParams.IsAsync ? "global::System.Threading.Tasks.Task" : "void";
         templateWithMethodSignature = templateWithMethodSignature.Replace(s_returnType, returnType);
 
-        var eventTagHelper = documentContext.Project.TagHelpers
+        var tagHelpers = await documentContext.Project.GetTagHelpersAsync(cancellationToken).ConfigureAwait(false);
+        var eventTagHelper = tagHelpers
             .FirstOrDefault(th => th.Name == actionParams.EventName && th.IsEventHandlerTagHelper() && th.GetEventArgsType() is not null);
         var eventArgsType = eventTagHelper is null
             ? string.Empty // Couldn't find the params, generate no params instead.

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Cohost/CohostDocumentSnapshot.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
 using System.Diagnostics.CodeAnalysis;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.CodeAnalysis;
@@ -10,10 +11,10 @@
 
 namespace Microsoft.AspNetCore.Razor.LanguageServer.Cohost;
 
-internal class CohostDocumentSnapshot(TextDocument textDocument, IProjectSnapshot projectSnapshot) : IDocumentSnapshot
+internal class CohostDocumentSnapshot(TextDocument textDocument, CohostProjectSnapshot projectSnapshot) : IDocumentSnapshot
 {
     private readonly TextDocument _textDocument = textDocument;
-    private readonly IProjectSnapshot _projectSnapshot = projectSnapshot;
+    private readonly CohostProjectSnapshot _projectSnapshot = projectSnapshot;
 
     private RazorCodeDocument? _codeDocument;
 
@@ -53,8 +54,10 @@ public async Task<RazorCodeDocument> GetGeneratedOutputAsync()
         // and simply compiles when asked, and if a new document snapshot comes in, we compile again. This is presumably worse for perf
         // but since we don't expect users to ever use cohosting without source generators, it's fine for now.
 
-        var imports = await DocumentState.ComputedStateTracker.GetImportsAsync(this).ConfigureAwait(false);
-        _codeDocument = await DocumentState.ComputedStateTracker.GenerateCodeDocumentAsync(Project, this, imports).ConfigureAwait(false);
+        var projectEngine = _projectSnapshot.GetProjectEngine_CohostOnly();
+        var tagHelpers = await _projectSnapshot.GetTagHelpersAsync(CancellationToken.None).ConfigureAwait(false);
+        var imports = await DocumentState.ComputedStateTracker.GetImportsAsync(this, projectEngine).ConfigureAwait(false);
+        _codeDocument = await DocumentState.ComputedStateTracker.GenerateCodeDocumentAsync(tagHelpers, projectEngine, this, imports).ConfigureAwait(false);
 
         return _codeDocument;
     }

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Cohost/CohostProjectSnapshot.cs

@@ -7,6 +7,7 @@
 using System.IO;
 using System.Linq;
 using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.AspNetCore.Razor.PooledObjects;
 using Microsoft.AspNetCore.Razor.ProjectEngineHost;
@@ -17,7 +18,6 @@
 using Microsoft.CodeAnalysis.CSharp;
 using Microsoft.CodeAnalysis.Razor;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;
-using Microsoft.VisualStudio.Threading;
 
 namespace Microsoft.AspNetCore.Razor.LanguageServer.Cohost;
 
@@ -29,11 +29,11 @@ internal class CohostProjectSnapshot : IProjectSnapshot
     private readonly ProjectKey _projectKey;
     private readonly Lazy<RazorConfiguration> _lazyConfiguration;
     private readonly Lazy<RazorProjectEngine> _lazyProjectEngine;
-    private readonly AsyncLazy<ImmutableArray<TagHelperDescriptor>> _tagHelpersLazy;
-    private readonly Lazy<ProjectWorkspaceState> _projectWorkspaceStateLazy;
     private readonly Lazy<ImmutableDictionary<string, ImmutableArray<string>>> _importsToRelatedDocumentsLazy;
 
-    public CohostProjectSnapshot(Project project, DocumentSnapshotFactory documentSnapshotFactory, ITelemetryReporter telemetryReporter, JoinableTaskContext joinableTaskContext)
+    private ImmutableArray<TagHelperDescriptor>? _tagHelpers;
+
+    public CohostProjectSnapshot(Project project, DocumentSnapshotFactory documentSnapshotFactory, ITelemetryReporter telemetryReporter)
     {
         _project = project;
         _documentSnapshotFactory = documentSnapshotFactory;
@@ -44,7 +44,7 @@ public CohostProjectSnapshot(Project project, DocumentSnapshotFactory documentSn
         _lazyProjectEngine = new Lazy<RazorProjectEngine>(() =>
         {
             return ProjectEngineFactories.DefaultProvider.Create(
-                Configuration,
+                _lazyConfiguration.Value,
                 rootDirectoryPath: Path.GetDirectoryName(FilePath).AssumeNotNull(),
                 configure: builder =>
                 {
@@ -54,20 +54,12 @@ public CohostProjectSnapshot(Project project, DocumentSnapshotFactory documentSn
                 });
         });
 
-        _tagHelpersLazy = new AsyncLazy<ImmutableArray<TagHelperDescriptor>>(() =>
-        {
-            var resolver = new CompilationTagHelperResolver(_telemetryReporter);
-            return resolver.GetTagHelpersAsync(_project, GetProjectEngine(), CancellationToken.None).AsTask();
-        }, joinableTaskContext.Factory);
-
-        _projectWorkspaceStateLazy = new Lazy<ProjectWorkspaceState>(() => ProjectWorkspaceState.Create(TagHelpers, CSharpLanguageVersion));
-
         _importsToRelatedDocumentsLazy = new Lazy<ImmutableDictionary<string, ImmutableArray<string>>>(() =>
         {
             var importsToRelatedDocuments = ImmutableDictionary.Create<string, ImmutableArray<string>>(FilePathNormalizer.Comparer);
             foreach (var document in DocumentFilePaths)
             {
-                var importTargetPaths = ProjectState.GetImportDocumentTargetPaths(document, FileKinds.GetFileKindFromFilePath(document), GetProjectEngine());
+                var importTargetPaths = ProjectState.GetImportDocumentTargetPaths(document, FileKinds.GetFileKindFromFilePath(document), _lazyProjectEngine.Value);
                 importsToRelatedDocuments = ProjectState.AddToImportsToRelatedDocuments(importsToRelatedDocuments, document, importTargetPaths);
             }
 
@@ -77,7 +69,7 @@ public CohostProjectSnapshot(Project project, DocumentSnapshotFactory documentSn
 
     public ProjectKey Key => _projectKey;
 
-    public RazorConfiguration Configuration => _lazyConfiguration.Value;
+    public RazorConfiguration Configuration => throw new InvalidOperationException("Should not be called for cohosted projects.");
 
     public IEnumerable<string> DocumentFilePaths
         => _project.AdditionalDocuments
@@ -96,9 +88,18 @@ public IEnumerable<string> DocumentFilePaths
 
     public LanguageVersion CSharpLanguageVersion => ((CSharpParseOptions)_project.ParseOptions!).LanguageVersion;
 
-    public ImmutableArray<TagHelperDescriptor> TagHelpers => _tagHelpersLazy.GetValue();
+    public async ValueTask<ImmutableArray<TagHelperDescriptor>> GetTagHelpersAsync(CancellationToken cancellationToken)
+    {
+        if (_tagHelpers is null)
+        {
+            var resolver = new CompilationTagHelperResolver(_telemetryReporter);
+            _tagHelpers = await resolver.GetTagHelpersAsync(_project, _lazyProjectEngine.Value, cancellationToken).ConfigureAwait(false);
+        }
+
+        return _tagHelpers.GetValueOrDefault();
+    }
 
-    public ProjectWorkspaceState ProjectWorkspaceState => _projectWorkspaceStateLazy.Value;
+    public ProjectWorkspaceState ProjectWorkspaceState => throw new InvalidOperationException("Should not be called for cohosted projects.");
 
     public IDocumentSnapshot? GetDocument(string filePath)
     {
@@ -111,7 +112,13 @@ public IEnumerable<string> DocumentFilePaths
         return _documentSnapshotFactory.GetOrCreate(textDocument);
     }
 
-    public RazorProjectEngine GetProjectEngine() => _lazyProjectEngine.Value;
+    public RazorProjectEngine GetProjectEngine() => throw new InvalidOperationException("Should not be called for cohosted projects.");
+
+    /// <summary>
+    /// NOTE: To be called only from CohostDocumentSnapshot.GetGeneratedOutputAsync(). Will be removed when that method uses the source generator directly.
+    /// </summary>
+    /// <returns></returns>
+    internal RazorProjectEngine GetProjectEngine_CohostOnly() => _lazyProjectEngine.Value;
 
     public ImmutableArray<IDocumentSnapshot> GetRelatedDocuments(IDocumentSnapshot document)
     {

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Cohost/ProjectSnapshotFactory.cs

@@ -5,26 +5,23 @@
 using System.Runtime.CompilerServices;
 using Microsoft.AspNetCore.Razor.Telemetry;
 using Microsoft.CodeAnalysis;
-using Microsoft.CodeAnalysis.Razor.ProjectSystem;
-using Microsoft.VisualStudio.Threading;
 
 namespace Microsoft.AspNetCore.Razor.LanguageServer.Cohost;
 
 [Export(typeof(ProjectSnapshotFactory)), Shared]
 [method: ImportingConstructor]
-internal class ProjectSnapshotFactory(DocumentSnapshotFactory documentSnapshotFactory, ITelemetryReporter telemetryReporter, JoinableTaskContext joinableTaskContext)
+internal class ProjectSnapshotFactory(DocumentSnapshotFactory documentSnapshotFactory, ITelemetryReporter telemetryReporter)
 {
-    private static readonly ConditionalWeakTable<Project, IProjectSnapshot> _projectSnapshots = new();
+    private static readonly ConditionalWeakTable<Project, CohostProjectSnapshot> _projectSnapshots = new();
 
     private readonly DocumentSnapshotFactory _documentSnapshotFactory = documentSnapshotFactory;
     private readonly ITelemetryReporter _telemetryReporter = telemetryReporter;
-    private readonly JoinableTaskContext _joinableTaskContext = joinableTaskContext;
 
-    public IProjectSnapshot GetOrCreate(Project project)
+    public CohostProjectSnapshot GetOrCreate(Project project)
     {
         if (!_projectSnapshots.TryGetValue(project, out var projectSnapshot))
         {
-            projectSnapshot = new CohostProjectSnapshot(project, _documentSnapshotFactory, _telemetryReporter, _joinableTaskContext);
+            projectSnapshot = new CohostProjectSnapshot(project, _documentSnapshotFactory, _telemetryReporter);
             _projectSnapshots.Add(project, projectSnapshot);
         }