Always use designtime document for formatting with FUSE (#10640)

Separated commits to functional change and code moving. This moves the
feature flag check into tooling instead of relying on the compiler to
know and generate documents differently.

That also allows for correctly doing document generation of
designtime/runtime from the tooling side. If FUSE is enabled tooling
will use runtime in all places except formatting. This is to keep
formatting working as it does today without FUSE and start A/B testing
faster.

src/Compiler/Microsoft.CodeAnalysis.Razor.Compiler/src/Language/RazorProjectEngine.cs

@@ -52,7 +52,7 @@ public class RazorProjectEngine
 
     public RazorCodeDocument Process(
         RazorSourceDocument source,
-        string fileKind,
+        string? fileKind,
         ImmutableArray<RazorSourceDocument> importSources,
         IReadOnlyList<TagHelperDescriptor> tagHelpers)
     {
@@ -109,9 +109,7 @@ public class RazorProjectEngine
             throw new ArgumentNullException(nameof(projectItem));
         }
 
-        var codeDocument = Configuration.LanguageServerFlags?.ForceRuntimeCodeGeneration == true
-            ? CreateCodeDocumentCore(projectItem)
-            : CreateCodeDocumentDesignTimeCore(projectItem);
+        var codeDocument = CreateCodeDocumentDesignTimeCore(projectItem);
         ProcessCore(codeDocument);
         return codeDocument;
     }
@@ -127,9 +125,7 @@ public class RazorProjectEngine
             throw new ArgumentNullException(nameof(source));
         }
 
-        var codeDocument = Configuration.LanguageServerFlags?.ForceRuntimeCodeGeneration == true
-            ? CreateCodeDocumentCore(source, fileKind, importSources, tagHelpers, configureParser: null, configureCodeGeneration: null)
-            : CreateCodeDocumentDesignTimeCore(source, fileKind, importSources, tagHelpers, configureParser: null, configureCodeGeneration: null);
+        var codeDocument = CreateCodeDocumentDesignTimeCore(source, fileKind, importSources, tagHelpers, configureParser: null, configureCodeGeneration: null);
         ProcessCore(codeDocument);
         return codeDocument;
     }

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Formatting/FormattingContext.cs

@@ -266,7 +266,7 @@ internal class FormattingContext : IDisposable
 
         var changedSnapshot = OriginalSnapshot.WithText(changedText);
 
-        var codeDocument = await changedSnapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
+        var codeDocument = await changedSnapshot.GetFormatterCodeDocumentAsync().ConfigureAwait(false);
 
         DEBUG_ValidateComponents(CodeDocument, codeDocument);
 

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Formatting/RazorFormattingService.cs

@@ -8,6 +8,7 @@ using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.AspNetCore.Razor.TextDifferencing;
+using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;
 using Microsoft.CodeAnalysis.Razor.Protocol;
 using Microsoft.CodeAnalysis.Razor.Workspaces;
@@ -40,7 +41,7 @@ internal class RazorFormattingService : IRazorFormattingService
         FormattingOptions options,
         CancellationToken cancellationToken)
     {
-        var codeDocument = await documentContext.Snapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
+        var codeDocument = await documentContext.Snapshot.GetFormatterCodeDocumentAsync().ConfigureAwait(false);
 
         // Range formatting happens on every paste, and if there are Razor diagnostics in the file
         // that can make some very bad results. eg, given:

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

@@ -59,4 +59,30 @@ internal static class IDocumentSnapshotExtensions
         var fileName = Path.GetFileNameWithoutExtension(documentSnapshot.FilePath);
         return fileName.AsSpan().Equals(path.Span, FilePathComparison.Instance);
     }
+
+    public static Task<RazorCodeDocument> GetFormatterCodeDocumentAsync(this IDocumentSnapshot documentSnapshot)
+    {
+        var forceRuntimeCodeGeneration = documentSnapshot.Project.Configuration.LanguageServerFlags?.ForceRuntimeCodeGeneration ?? false;
+        if (!forceRuntimeCodeGeneration)
+        {
+            return documentSnapshot.GetGeneratedOutputAsync();
+        }
+
+        // if forceRuntimeCodeGeneration is on, GetGeneratedOutputAsync will get runtime code. As of now
+        // the formatting service doesn't expect the form of code generated to be what the compiler does with
+        // runtime. For now force usage of design time and avoid the cache. There may be a slight perf hit
+        // but either the user is typing (which will invalidate the cache) or the user is manually attempting to
+        // format. We expect formatting to invalidate the cache if it changes things and consider this an
+        // acceptable overhead for now.
+        return GetDesignTimeDocumentAsync(documentSnapshot);
+    }
+
+    private static async Task<RazorCodeDocument> GetDesignTimeDocumentAsync(IDocumentSnapshot documentSnapshot)
+    {
+        var project = documentSnapshot.Project;
+        var tagHelpers = await project.GetTagHelpersAsync(CancellationToken.None).ConfigureAwait(false);
+        var projectEngine = project.GetProjectEngine();
+        var imports = await DocumentState.GetImportsAsync(documentSnapshot, projectEngine).ConfigureAwait(false);
+        return await DocumentState.GenerateCodeDocumentAsync(tagHelpers, project.GetProjectEngine(), documentSnapshot, imports, false).ConfigureAwait(false);
+    }
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/ProjectSystem/DocumentState.ComputedStateTracker.cs

@@ -0,0 +1,252 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT license. See License.txt in the project root for license information.
+
+using System;
+using System.Diagnostics.CodeAnalysis;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Razor.Language;
+
+namespace Microsoft.CodeAnalysis.Razor.ProjectSystem;
+
+internal partial class DocumentState
+{
+    private class ComputedStateTracker
+    {
+        private readonly object _lock;
+
+        private ComputedStateTracker? _older;
+
+        // We utilize a WeakReference here to avoid bloating committed memory. If pieces request document output inbetween GC collections
+        // then we will provide the weak referenced task; otherwise we require any state requests to be re-computed.
+        private WeakReference<Task<(RazorCodeDocument, VersionStamp)>>? _taskUnsafeReference;
+
+        private ComputedOutput? _computedOutput;
+
+        public ComputedStateTracker(DocumentState state, ComputedStateTracker? older = null)
+        {
+            _lock = state._lock;
+            _older = older;
+        }
+
+        public bool IsResultAvailable
+        {
+            get
+            {
+                if (_computedOutput?.TryGetCachedOutput(out _, out _) == true)
+                {
+                    return true;
+                }
+
+                if (_taskUnsafeReference is null)
+                {
+                    return false;
+                }
+
+                if (_taskUnsafeReference.TryGetTarget(out var taskUnsafe))
+                {
+                    return taskUnsafe.IsCompleted;
+                }
+
+                return false;
+            }
+        }
+
+        public async Task<(RazorCodeDocument, VersionStamp)> GetGeneratedOutputAndVersionAsync(ProjectSnapshot project, IDocumentSnapshot document)
+        {
+            if (_computedOutput?.TryGetCachedOutput(out var cachedCodeDocument, out var cachedInputVersion) == true)
+            {
+                return (cachedCodeDocument, cachedInputVersion);
+            }
+
+            var (codeDocument, inputVersion) = await GetMemoizedGeneratedOutputAndVersionAsync(project, document).ConfigureAwait(false);
+
+            _computedOutput = new ComputedOutput(codeDocument, inputVersion);
+            return (codeDocument, inputVersion);
+        }
+
+        private Task<(RazorCodeDocument, VersionStamp)> GetMemoizedGeneratedOutputAndVersionAsync(ProjectSnapshot project, IDocumentSnapshot document)
+        {
+            if (project is null)
+            {
+                throw new ArgumentNullException(nameof(project));
+            }
+
+            if (document is null)
+            {
+                throw new ArgumentNullException(nameof(document));
+            }
+
+            if (_taskUnsafeReference is null ||
+                !_taskUnsafeReference.TryGetTarget(out var taskUnsafe))
+            {
+                TaskCompletionSource<(RazorCodeDocument, VersionStamp)>? tcs = null;
+
+                lock (_lock)
+                {
+                    if (_taskUnsafeReference is null ||
+                        !_taskUnsafeReference.TryGetTarget(out taskUnsafe))
+                    {
+                        // So this is a bit confusing. Instead of directly calling the Razor parser inside of this lock we create an indirect TaskCompletionSource
+                        // to represent when it completes. The reason behind this is that there are several scenarios in which the Razor parser will run synchronously
+                        // (mostly all in VS) resulting in this lock being held for significantly longer than expected. To avoid threads queuing up repeatedly on the
+                        // above lock and blocking we can allow those threads to await asynchronously for the completion of the original parse.
+
+                        tcs = new(TaskCreationOptions.RunContinuationsAsynchronously);
+                        taskUnsafe = tcs.Task;
+                        _taskUnsafeReference = new WeakReference<Task<(RazorCodeDocument, VersionStamp)>>(taskUnsafe);
+                    }
+                }
+
+                if (tcs is null)
+                {
+                    // There's no task completion source created meaning a value was retrieved from cache, just return it.
+                    return taskUnsafe;
+                }
+
+                // Typically in VS scenarios this will run synchronously because all resources are readily available.
+                var outputTask = ComputeGeneratedOutputAndVersionAsync(project, document);
+                if (outputTask.IsCompleted)
+                {
+                    // Compiling ran synchronously, lets just immediately propagate to the TCS
+                    PropagateToTaskCompletionSource(outputTask, tcs);
+                }
+                else
+                {
+                    // Task didn't run synchronously (most likely outside of VS), lets allocate a bit more but utilize ContinueWith
+                    // to properly connect the output task and TCS
+                    _ = outputTask.ContinueWith(
+                        static (task, state) =>
+                        {
+                            Assumes.NotNull(state);
+                            var tcs = (TaskCompletionSource<(RazorCodeDocument, VersionStamp)>)state;
+
+                            PropagateToTaskCompletionSource(task, tcs);
+                        },
+                        tcs,
+                        CancellationToken.None,
+                        TaskContinuationOptions.ExecuteSynchronously,
+                        TaskScheduler.Default);
+                }
+            }
+
+            return taskUnsafe;
+
+            static void PropagateToTaskCompletionSource(
+                Task<(RazorCodeDocument, VersionStamp)> targetTask,
+                TaskCompletionSource<(RazorCodeDocument, VersionStamp)> tcs)
+            {
+                if (targetTask.Status == TaskStatus.RanToCompletion)
+                {
+#pragma warning disable VSTHRD103 // Call async methods when in an async method
+                    tcs.SetResult(targetTask.Result);
+#pragma warning restore VSTHRD103 // Call async methods when in an async method
+                }
+                else if (targetTask.Status == TaskStatus.Canceled)
+                {
+                    tcs.SetCanceled();
+                }
+                else if (targetTask.Status == TaskStatus.Faulted)
+                {
+                    // Faulted tasks area always aggregate exceptions so we need to extract the "true" exception if it's available:
+                    Assumes.NotNull(targetTask.Exception);
+                    var exception = targetTask.Exception.InnerException ?? targetTask.Exception;
+                    tcs.SetException(exception);
+                }
+            }
+        }
+
+        private async Task<(RazorCodeDocument, VersionStamp)> ComputeGeneratedOutputAndVersionAsync(ProjectSnapshot project, IDocumentSnapshot document)
+        {
+            // We only need to produce the generated code if any of our inputs is newer than the
+            // previously cached output.
+            //
+            // First find the versions that are the inputs:
+            // - The project + computed state
+            // - The imports
+            // - This document
+            //
+            // All of these things are cached, so no work is wasted if we do need to generate the code.
+            var configurationVersion = project.State.ConfigurationVersion;
+            var projectWorkspaceStateVersion = project.State.ProjectWorkspaceStateVersion;
+            var documentCollectionVersion = project.State.DocumentCollectionVersion;
+            var imports = await GetImportsAsync(document, project.GetProjectEngine()).ConfigureAwait(false);
+            var documentVersion = await document.GetTextVersionAsync().ConfigureAwait(false);
+
+            // OK now that have the previous output and all of the versions, we can see if anything
+            // has changed that would require regenerating the code.
+            var inputVersion = documentVersion;
+            if (inputVersion.GetNewerVersion(configurationVersion) == configurationVersion)
+            {
+                inputVersion = configurationVersion;
+            }
+
+            if (inputVersion.GetNewerVersion(projectWorkspaceStateVersion) == projectWorkspaceStateVersion)
+            {
+                inputVersion = projectWorkspaceStateVersion;
+            }
+
+            if (inputVersion.GetNewerVersion(documentCollectionVersion) == documentCollectionVersion)
+            {
+                inputVersion = documentCollectionVersion;
+            }
+
+            foreach (var import in imports)
+            {
+                var importVersion = import.Version;
+                if (inputVersion.GetNewerVersion(importVersion) == importVersion)
+                {
+                    inputVersion = importVersion;
+                }
+            }
+
+            if (_older?._taskUnsafeReference != null &&
+                _older._taskUnsafeReference.TryGetTarget(out var taskUnsafe))
+            {
+                var (olderOutput, olderInputVersion) = await taskUnsafe.ConfigureAwait(false);
+                if (inputVersion.GetNewerVersion(olderInputVersion) == olderInputVersion)
+                {
+                    // Nothing has changed, we can use the cached result.
+                    lock (_lock)
+                    {
+                        _taskUnsafeReference = _older._taskUnsafeReference;
+                        _older = null;
+                        return (olderOutput, olderInputVersion);
+                    }
+                }
+            }
+
+            var tagHelpers = await project.GetTagHelpersAsync(CancellationToken.None).ConfigureAwait(false);
+            var forceRuntimeCodeGeneration = project.Configuration.LanguageServerFlags?.ForceRuntimeCodeGeneration ?? false;
+            var codeDocument = await GenerateCodeDocumentAsync(tagHelpers, project.GetProjectEngine(), document, imports, forceRuntimeCodeGeneration).ConfigureAwait(false);
+            return (codeDocument, inputVersion);
+        }
+
+        private class ComputedOutput
+        {
+            private readonly VersionStamp _inputVersion;
+            private readonly WeakReference<RazorCodeDocument> _codeDocumentReference;
+
+            public ComputedOutput(RazorCodeDocument codeDocument, VersionStamp inputVersion)
+            {
+                _codeDocumentReference = new WeakReference<RazorCodeDocument>(codeDocument);
+                _inputVersion = inputVersion;
+            }
+
+            public bool TryGetCachedOutput([NotNullWhen(true)] out RazorCodeDocument? codeDocument, out VersionStamp inputVersion)
+            {
+                // The goal here is to capture a weak reference to the code document so if there's ever a sub-system that's still utilizing it
+                // our computed output maintains its cache.
+
+                if (_codeDocumentReference.TryGetTarget(out codeDocument))
+                {
+                    inputVersion = _inputVersion;
+                    return true;
+                }
+
+                inputVersion = default;
+                return false;
+            }
+        }
+    }
+}

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/ProjectSystem/DocumentState.ImportItem.cs

@@ -0,0 +1,12 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT license. See License.txt in the project root for license information.
+
+namespace Microsoft.CodeAnalysis.Razor.ProjectSystem;
+
+internal partial class DocumentState
+{
+    internal record struct ImportItem(string? FilePath, VersionStamp Version, IDocumentSnapshot Document)
+    {
+        public readonly string? FileKind => Document.FileKind;
+    }
+}

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/ProjectSystem/DocumentState.cs

@@ -4,7 +4,6 @@
 using System;
 using System.Collections.Immutable;
 using System.Diagnostics.CodeAnalysis;
-using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor;
 using Microsoft.AspNetCore.Razor.Language;
@@ -13,7 +12,7 @@ using Microsoft.CodeAnalysis.Text;
 
 namespace Microsoft.CodeAnalysis.Razor.ProjectSystem;
 
-internal class DocumentState
+internal partial class DocumentState
 {
     private static readonly TextAndVersion s_emptyText = TextAndVersion.Create(
         SourceText.From(string.Empty),
@@ -246,6 +245,11 @@ internal class DocumentState
 
         foreach (var item in importItems)
         {
+            if (item is NotFoundProjectItem)
+            {
+                continue;
+            }
+
             if (item.PhysicalPath is null)
             {
                 // This is a default import.
@@ -261,284 +265,48 @@ internal class DocumentState
         return imports.ToImmutable();
     }
 
-    // Internal, because we are temporarily sharing code with CohostDocumentSnapshot
-    internal class ComputedStateTracker
+    internal static async Task<RazorCodeDocument> GenerateCodeDocumentAsync(ImmutableArray<TagHelperDescriptor> tagHelpers, RazorProjectEngine projectEngine, IDocumentSnapshot document, ImmutableArray<ImportItem> imports, bool forceRuntimeCodeGeneration)
     {
-        private readonly object _lock;
-
-        private ComputedStateTracker? _older;
-
-        // We utilize a WeakReference here to avoid bloating committed memory. If pieces request document output inbetween GC collections
-        // then we will provide the weak referenced task; otherwise we require any state requests to be re-computed.
-        private WeakReference<Task<(RazorCodeDocument, VersionStamp)>>? _taskUnsafeReference;
-
-        private ComputedOutput? _computedOutput;
-
-        public ComputedStateTracker(DocumentState state, ComputedStateTracker? older = null)
+        // OK we have to generate the code.
+        using var importSources = new PooledArrayBuilder<RazorSourceDocument>(imports.Length);
+        foreach (var item in imports)
         {
-            _lock = state._lock;
-            _older = older;
+            var importProjectItem = item.FilePath is null ? null : projectEngine.FileSystem.GetItem(item.FilePath, item.FileKind);
+            var sourceDocument = await GetRazorSourceDocumentAsync(item.Document, importProjectItem).ConfigureAwait(false);
+            importSources.Add(sourceDocument);
         }
 
-        public bool IsResultAvailable
+        var projectItem = document.FilePath is null ? null : projectEngine.FileSystem.GetItem(document.FilePath, document.FileKind);
+        var documentSource = await GetRazorSourceDocumentAsync(document, projectItem).ConfigureAwait(false);
+
+        if (forceRuntimeCodeGeneration)
         {
-            get
-            {
-                if (_computedOutput?.TryGetCachedOutput(out _, out _) == true)
-                {
-                    return true;
-                }
-
-                if (_taskUnsafeReference is null)
-                {
-                    return false;
-                }
-
-                if (_taskUnsafeReference.TryGetTarget(out var taskUnsafe))
-                {
-                    return taskUnsafe.IsCompleted;
-                }
-
-                return false;
-            }
+            return projectEngine.Process(documentSource, fileKind: document.FileKind, importSources.DrainToImmutable(), tagHelpers);
         }
 
-        public async Task<(RazorCodeDocument, VersionStamp)> GetGeneratedOutputAndVersionAsync(ProjectSnapshot project, IDocumentSnapshot document)
+        return projectEngine.ProcessDesignTime(documentSource, fileKind: document.FileKind, importSources.DrainToImmutable(), tagHelpers);
+    }
+
+    internal static Task<RazorCodeDocument> GenerateFormattingCodeDocumentAsync(ImmutableArray<TagHelperDescriptor> tagHelpers, RazorProjectEngine projectEngine, IDocumentSnapshot document, ImmutableArray<ImportItem> imports)
+        => GenerateCodeDocumentAsync(tagHelpers, projectEngine, document, imports, forceRuntimeCodeGeneration: false);
+
+    internal static async Task<ImmutableArray<ImportItem>> GetImportsAsync(IDocumentSnapshot document, RazorProjectEngine projectEngine)
+    {
+        var imports = GetImportsCore(document.Project, projectEngine, document.FilePath.AssumeNotNull(), document.FileKind.AssumeNotNull());
+        using var result = new PooledArrayBuilder<ImportItem>(imports.Length);
+
+        foreach (var snapshot in imports)
         {
-            if (_computedOutput?.TryGetCachedOutput(out var cachedCodeDocument, out var cachedInputVersion) == true)
-            {
-                return (cachedCodeDocument, cachedInputVersion);
-            }
-
-            var (codeDocument, inputVersion) = await GetMemoizedGeneratedOutputAndVersionAsync(project, document).ConfigureAwait(false);
-
-            _computedOutput = new ComputedOutput(codeDocument, inputVersion);
-            return (codeDocument, inputVersion);
+            var versionStamp = await snapshot.GetTextVersionAsync().ConfigureAwait(false);
+            result.Add(new ImportItem(snapshot.FilePath, versionStamp, snapshot));
         }
 
-        private Task<(RazorCodeDocument, VersionStamp)> GetMemoizedGeneratedOutputAndVersionAsync(ProjectSnapshot project, IDocumentSnapshot document)
-        {
-            if (project is null)
-            {
-                throw new ArgumentNullException(nameof(project));
-            }
+        return result.DrainToImmutable();
+    }
 
-            if (document is null)
-            {
-                throw new ArgumentNullException(nameof(document));
-            }
-
-            if (_taskUnsafeReference is null ||
-                !_taskUnsafeReference.TryGetTarget(out var taskUnsafe))
-            {
-                TaskCompletionSource<(RazorCodeDocument, VersionStamp)>? tcs = null;
-
-                lock (_lock)
-                {
-                    if (_taskUnsafeReference is null ||
-                        !_taskUnsafeReference.TryGetTarget(out taskUnsafe))
-                    {
-                        // So this is a bit confusing. Instead of directly calling the Razor parser inside of this lock we create an indirect TaskCompletionSource
-                        // to represent when it completes. The reason behind this is that there are several scenarios in which the Razor parser will run synchronously
-                        // (mostly all in VS) resulting in this lock being held for significantly longer than expected. To avoid threads queuing up repeatedly on the
-                        // above lock and blocking we can allow those threads to await asynchronously for the completion of the original parse.
-
-                        tcs = new(TaskCreationOptions.RunContinuationsAsynchronously);
-                        taskUnsafe = tcs.Task;
-                        _taskUnsafeReference = new WeakReference<Task<(RazorCodeDocument, VersionStamp)>>(taskUnsafe);
-                    }
-                }
-
-                if (tcs is null)
-                {
-                    // There's no task completion source created meaning a value was retrieved from cache, just return it.
-                    return taskUnsafe;
-                }
-
-                // Typically in VS scenarios this will run synchronously because all resources are readily available.
-                var outputTask = ComputeGeneratedOutputAndVersionAsync(project, document);
-                if (outputTask.IsCompleted)
-                {
-                    // Compiling ran synchronously, lets just immediately propagate to the TCS
-                    PropagateToTaskCompletionSource(outputTask, tcs);
-                }
-                else
-                {
-                    // Task didn't run synchronously (most likely outside of VS), lets allocate a bit more but utilize ContinueWith
-                    // to properly connect the output task and TCS
-                    _ = outputTask.ContinueWith(
-                        static (task, state) =>
-                        {
-                            Assumes.NotNull(state);
-                            var tcs = (TaskCompletionSource<(RazorCodeDocument, VersionStamp)>)state;
-
-                            PropagateToTaskCompletionSource(task, tcs);
-                        },
-                        tcs,
-                        CancellationToken.None,
-                        TaskContinuationOptions.ExecuteSynchronously,
-                        TaskScheduler.Default);
-                }
-            }
-
-            return taskUnsafe;
-
-            static void PropagateToTaskCompletionSource(
-                Task<(RazorCodeDocument, VersionStamp)> targetTask,
-                TaskCompletionSource<(RazorCodeDocument, VersionStamp)> tcs)
-            {
-                if (targetTask.Status == TaskStatus.RanToCompletion)
-                {
-#pragma warning disable VSTHRD103 // Call async methods when in an async method
-                    tcs.SetResult(targetTask.Result);
-#pragma warning restore VSTHRD103 // Call async methods when in an async method
-                }
-                else if (targetTask.Status == TaskStatus.Canceled)
-                {
-                    tcs.SetCanceled();
-                }
-                else if (targetTask.Status == TaskStatus.Faulted)
-                {
-                    // Faulted tasks area always aggregate exceptions so we need to extract the "true" exception if it's available:
-                    Assumes.NotNull(targetTask.Exception);
-                    var exception = targetTask.Exception.InnerException ?? targetTask.Exception;
-                    tcs.SetException(exception);
-                }
-            }
-        }
-
-        private async Task<(RazorCodeDocument, VersionStamp)> ComputeGeneratedOutputAndVersionAsync(ProjectSnapshot project, IDocumentSnapshot document)
-        {
-            // We only need to produce the generated code if any of our inputs is newer than the
-            // previously cached output.
-            //
-            // First find the versions that are the inputs:
-            // - The project + computed state
-            // - The imports
-            // - This document
-            //
-            // All of these things are cached, so no work is wasted if we do need to generate the code.
-            var configurationVersion = project.State.ConfigurationVersion;
-            var projectWorkspaceStateVersion = project.State.ProjectWorkspaceStateVersion;
-            var documentCollectionVersion = project.State.DocumentCollectionVersion;
-            var imports = await GetImportsAsync(document, project.GetProjectEngine()).ConfigureAwait(false);
-            var documentVersion = await document.GetTextVersionAsync().ConfigureAwait(false);
-
-            // OK now that have the previous output and all of the versions, we can see if anything
-            // has changed that would require regenerating the code.
-            var inputVersion = documentVersion;
-            if (inputVersion.GetNewerVersion(configurationVersion) == configurationVersion)
-            {
-                inputVersion = configurationVersion;
-            }
-
-            if (inputVersion.GetNewerVersion(projectWorkspaceStateVersion) == projectWorkspaceStateVersion)
-            {
-                inputVersion = projectWorkspaceStateVersion;
-            }
-
-            if (inputVersion.GetNewerVersion(documentCollectionVersion) == documentCollectionVersion)
-            {
-                inputVersion = documentCollectionVersion;
-            }
-
-            foreach (var import in imports)
-            {
-                var importVersion = import.Version;
-                if (inputVersion.GetNewerVersion(importVersion) == importVersion)
-                {
-                    inputVersion = importVersion;
-                }
-            }
-
-            if (_older?._taskUnsafeReference != null &&
-                _older._taskUnsafeReference.TryGetTarget(out var taskUnsafe))
-            {
-                var (olderOutput, olderInputVersion) = await taskUnsafe.ConfigureAwait(false);
-                if (inputVersion.GetNewerVersion(olderInputVersion) == olderInputVersion)
-                {
-                    // Nothing has changed, we can use the cached result.
-                    lock (_lock)
-                    {
-                        _taskUnsafeReference = _older._taskUnsafeReference;
-                        _older = null;
-                        return (olderOutput, olderInputVersion);
-                    }
-                }
-            }
-
-            var tagHelpers = await project.GetTagHelpersAsync(CancellationToken.None).ConfigureAwait(false);
-            var codeDocument = await GenerateCodeDocumentAsync(tagHelpers, project.GetProjectEngine(), document, imports).ConfigureAwait(false);
-            return (codeDocument, inputVersion);
-        }
-
-        internal static async Task<RazorCodeDocument> GenerateCodeDocumentAsync(ImmutableArray<TagHelperDescriptor> tagHelpers, RazorProjectEngine projectEngine, IDocumentSnapshot document, ImmutableArray<ImportItem> imports)
-        {
-            // OK we have to generate the code.
-            using var importSources = new PooledArrayBuilder<RazorSourceDocument>(imports.Length);
-            foreach (var item in imports)
-            {
-                var importProjectItem = item.FilePath is null ? null : projectEngine.FileSystem.GetItem(item.FilePath, item.FileKind);
-                var sourceDocument = await GetRazorSourceDocumentAsync(item.Document, importProjectItem).ConfigureAwait(false);
-                importSources.Add(sourceDocument);
-            }
-
-            var projectItem = document.FilePath is null ? null : projectEngine.FileSystem.GetItem(document.FilePath, document.FileKind);
-            var documentSource = await GetRazorSourceDocumentAsync(document, projectItem).ConfigureAwait(false);
-
-            return projectEngine.ProcessDesignTime(documentSource, fileKind: document.FileKind, importSources.DrainToImmutable(), tagHelpers);
-        }
-
-        private static async Task<RazorSourceDocument> GetRazorSourceDocumentAsync(IDocumentSnapshot document, RazorProjectItem? projectItem)
-        {
-            var sourceText = await document.GetTextAsync().ConfigureAwait(false);
-            return RazorSourceDocument.Create(sourceText, RazorSourceDocumentProperties.Create(document.FilePath, projectItem?.RelativePhysicalPath));
-        }
-
-        internal static async Task<ImmutableArray<ImportItem>> GetImportsAsync(IDocumentSnapshot document, RazorProjectEngine projectEngine)
-        {
-            var imports = DocumentState.GetImportsCore(document.Project, projectEngine, document.FilePath.AssumeNotNull(), document.FileKind.AssumeNotNull());
-            using var result = new PooledArrayBuilder<ImportItem>(imports.Length);
-
-            foreach (var snapshot in imports)
-            {
-                var versionStamp = await snapshot.GetTextVersionAsync().ConfigureAwait(false);
-                result.Add(new ImportItem(snapshot.FilePath, versionStamp, snapshot));
-            }
-
-            return result.DrainToImmutable();
-        }
-
-        internal record struct ImportItem(string? FilePath, VersionStamp Version, IDocumentSnapshot Document)
-        {
-            public readonly string? FileKind => Document.FileKind;
-        }
-
-        private class ComputedOutput
-        {
-            private readonly VersionStamp _inputVersion;
-            private readonly WeakReference<RazorCodeDocument> _codeDocumentReference;
-
-            public ComputedOutput(RazorCodeDocument codeDocument, VersionStamp inputVersion)
-            {
-                _codeDocumentReference = new WeakReference<RazorCodeDocument>(codeDocument);
-                _inputVersion = inputVersion;
-            }
-
-            public bool TryGetCachedOutput([NotNullWhen(true)] out RazorCodeDocument? codeDocument, out VersionStamp inputVersion)
-            {
-                // The goal here is to capture a weak reference to the code document so if there's ever a sub-system that's still utilizing it
-                // our computed output maintains its cache.
-
-                if (_codeDocumentReference.TryGetTarget(out codeDocument))
-                {
-                    inputVersion = _inputVersion;
-                    return true;
-                }
-
-                inputVersion = default;
-                return false;
-            }
-        }
+    private static async Task<RazorSourceDocument> GetRazorSourceDocumentAsync(IDocumentSnapshot document, RazorProjectItem? projectItem)
+    {
+        var sourceText = await document.GetTextAsync().ConfigureAwait(false);
+        return RazorSourceDocument.Create(sourceText, RazorSourceDocumentProperties.Create(document.FilePath, projectItem?.RelativePhysicalPath));
     }
 }

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/ProjectSystem/RemoteDocumentSnapshot.cs

@@ -58,8 +58,12 @@ internal class RemoteDocumentSnapshot(TextDocument textDocument, RemoteProjectSn
 
         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);
+        var imports = await DocumentState.GetImportsAsync(this, projectEngine).ConfigureAwait(false);
+
+        // TODO: Get the configuration for forceRuntimeCodeGeneration
+        // var forceRuntimeCodeGeneration = _projectSnapshot.Configuration.LanguageServerFlags?.ForceRuntimeCodeGeneration ?? false;
+
+        _codeDocument = await DocumentState.GenerateCodeDocumentAsync(tagHelpers, projectEngine, this, imports, forceRuntimeCodeGeneration: false).ConfigureAwait(false);
 
         return _codeDocument;
     }

src/Razor/test/Microsoft.AspNetCore.Razor.LanguageServer.Test/Formatting_NetFx/FormattingTestBase.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
 using System;
+using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.IO;
 using System.Linq;
@@ -48,6 +49,13 @@ public class FormattingTestBase : RazorToolingIntegrationTestBase
         bool allowDiagnostics = false,
         RazorLSPOptions? razorLSPOptions = null,
         bool inGlobalNamespace = false)
+    {
+        // Run with and without forceRuntimeCodeGeneration
+        await RunFormattingTestAsync(input, expected, tabSize, insertSpaces, fileKind, tagHelpers, allowDiagnostics, razorLSPOptions, inGlobalNamespace, forceRuntimeCodeGeneration: true);
+        await RunFormattingTestAsync(input, expected, tabSize, insertSpaces, fileKind, tagHelpers, allowDiagnostics, razorLSPOptions, inGlobalNamespace, forceRuntimeCodeGeneration: false);
+    }
+
+    private async Task RunFormattingTestAsync(string input, string expected, int tabSize, bool insertSpaces, string? fileKind, ImmutableArray<TagHelperDescriptor> tagHelpers, bool allowDiagnostics, RazorLSPOptions? razorLSPOptions, bool inGlobalNamespace, bool forceRuntimeCodeGeneration)
     {
         // Arrange
         fileKind ??= FileKinds.Component;
@@ -62,7 +70,7 @@ public class FormattingTestBase : RazorToolingIntegrationTestBase
 
         var path = "file:///path/to/Document." + fileKind;
         var uri = new Uri(path);
-        var (codeDocument, documentSnapshot) = CreateCodeDocumentAndSnapshot(source, uri.AbsolutePath, tagHelpers, fileKind, allowDiagnostics, inGlobalNamespace);
+        var (codeDocument, documentSnapshot) = CreateCodeDocumentAndSnapshot(source, uri.AbsolutePath, tagHelpers, fileKind, allowDiagnostics, inGlobalNamespace, forceRuntimeCodeGeneration);
         var options = new FormattingOptions()
         {
             TabSize = tabSize,
@@ -217,7 +225,7 @@ public class FormattingTestBase : RazorToolingIntegrationTestBase
         return source.WithChanges(changes);
     }
 
-    private static (RazorCodeDocument, IDocumentSnapshot) CreateCodeDocumentAndSnapshot(SourceText text, string path, ImmutableArray<TagHelperDescriptor> tagHelpers = default, string? fileKind = default, bool allowDiagnostics = false, bool inGlobalNamespace = false)
+    private static (RazorCodeDocument, IDocumentSnapshot) CreateCodeDocumentAndSnapshot(SourceText text, string path, ImmutableArray<TagHelperDescriptor> tagHelpers = default, string? fileKind = default, bool allowDiagnostics = false, bool inGlobalNamespace = false, bool forceRuntimeCodeGeneration = false)
     {
         fileKind ??= FileKinds.Component;
         tagHelpers = tagHelpers.NullToEmpty();
@@ -255,12 +263,21 @@ public class FormattingTestBase : RazorToolingIntegrationTestBase
             .Setup(d => d.TargetPath)
             .Returns(importsPath);
 
-        var projectEngine = RazorProjectEngine.Create(builder =>
-        {
-            builder.SetRootNamespace(inGlobalNamespace ? string.Empty : "Test");
-            builder.Features.Add(new DefaultTypeNameFeature());
-            RazorExtensions.Register(builder);
-        });
+        var projectFileSystem = new TestRazorProjectFileSystem([
+            new TestRazorProjectItem(path, fileKind: fileKind),
+            new TestRazorProjectItem(importsPath, fileKind: FileKinds.ComponentImport),
+            ]);
+
+        var projectEngine = RazorProjectEngine.Create(
+            new RazorConfiguration(RazorLanguageVersion.Latest, "TestConfiguration", ImmutableArray<RazorExtension>.Empty, new LanguageServerFlags(forceRuntimeCodeGeneration)),
+            projectFileSystem,
+            builder =>
+            {
+                builder.SetRootNamespace(inGlobalNamespace ? string.Empty : "Test");
+                builder.Features.Add(new DefaultTypeNameFeature());
+                RazorExtensions.Register(builder);
+            });
+
         var codeDocument = projectEngine.ProcessDesignTime(sourceDocument, fileKind, ImmutableArray.Create(importsDocument), tagHelpers);
 
         if (!allowDiagnostics)
@@ -290,9 +307,18 @@ public class FormattingTestBase : RazorToolingIntegrationTestBase
         documentSnapshot
             .Setup(d => d.TargetPath)
             .Returns(path);
+        documentSnapshot
+            .Setup(d => d.Project.Configuration)
+            .Returns(projectEngine.Configuration);
+        documentSnapshot
+            .Setup(d => d.GetTextAsync())
+            .ReturnsAsync(codeDocument.Source.Text);
         documentSnapshot
             .Setup(d => d.Project.GetTagHelpersAsync(It.IsAny<CancellationToken>()))
             .Returns(new ValueTask<ImmutableArray<TagHelperDescriptor>>(tagHelpers));
+        documentSnapshot
+            .Setup(d => d.Project.GetProjectEngine())
+            .Returns(projectEngine);
         documentSnapshot
             .Setup(d => d.FileKind)
             .Returns(fileKind);