Pass CancellationToken to GetGeneratedOutputAsync

This commit is ginormous because I needed to thread cancellation tokens around everywhere.

src/Razor/benchmarks/Microsoft.AspNetCore.Razor.Microbenchmarks/LanguageServer/RazorDocumentMappingBenchmark.cs

@@ -6,6 +6,7 @@
 using System;
 using System.Collections.Generic;
 using System.IO;
+using System.Threading;
 using System.Threading.Tasks;
 using BenchmarkDotNet.Attributes;
 using Microsoft.AspNetCore.Razor.Language;
@@ -50,7 +51,7 @@ public class RazorDocumentMappingBenchmark : RazorLanguageServerBenchmarkBase
 
         DocumentSnapshot = await GetDocumentSnapshotAsync(projectFilePath, _filePath, targetPath);
 
-        var codeDocument = await DocumentSnapshot.GetGeneratedOutputAsync();
+        var codeDocument = await DocumentSnapshot.GetGeneratedOutputAsync(CancellationToken.None);
         CSharpDocument = codeDocument.GetCSharpDocument();
     }
 

src/Razor/benchmarks/Microsoft.AspNetCore.Razor.Microbenchmarks/ProjectSystem/BackgroundCodeGenerationBenchmark.cs

@@ -59,6 +59,6 @@ public class BackgroundCodeGenerationBenchmark : ProjectSnapshotManagerBenchmark
         var project = ProjectManager.GetLoadedProject(e.ProjectKey);
         var document = project.GetDocument(e.DocumentFilePath);
 
-        Tasks.Add(document.GetGeneratedOutputAsync().AsTask());
+        Tasks.Add(document.GetGeneratedOutputAsync(CancellationToken.None).AsTask());
     }
 }

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/CodeActions/CSharp/UnformattedRemappingCSharpCodeActionResolver.cs

@@ -70,7 +70,7 @@ internal sealed class UnformattedRemappingCSharpCodeActionResolver(
             return codeAction;
         }
 
-        var codeDocument = await documentContext.Snapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
+        var codeDocument = await documentContext.Snapshot.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
         if (codeDocument.IsUnsupported())
         {
             return codeAction;

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

@@ -139,13 +139,13 @@ internal sealed class CodeActionEndpoint(
         IDocumentSnapshot documentSnapshot,
         CancellationToken cancellationToken)
     {
-        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
+        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
         if (codeDocument.IsUnsupported())
         {
             return null;
         }
 
-        var sourceText = await documentSnapshot.GetTextAsync(cancellationToken).ConfigureAwait(false);
+        var sourceText = codeDocument.Source.Text;
 
         // VS Provides `CodeActionParams.Context.SelectionRange` in addition to
         // `CodeActionParams.Range`. The `SelectionRange` is relative to where the

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/CodeActions/Html/DefaultHtmlCodeActionProvider.cs

@@ -49,7 +49,7 @@ internal sealed class DefaultHtmlCodeActionProvider(IEditMappingService editMapp
 
         if (codeAction.Edit.TryGetTextDocumentEdits(out var documentEdits))
         {
-            var codeDocument = await documentSnapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
+            var codeDocument = await documentSnapshot.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
             var htmlSourceText = codeDocument.GetHtmlSourceText();
 
             foreach (var edit in documentEdits)

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

@@ -45,7 +45,7 @@ internal sealed class AddUsingsCodeActionResolver(IDocumentContextFactory docume
 
         var documentSnapshot = documentContext.Snapshot;
 
-        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
+        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
         if (codeDocument.IsUnsupported())
         {
             return null;

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Diagnostics/DocumentPullDiagnosticsEndpoint.cs

@@ -80,7 +80,7 @@ internal class DocumentPullDiagnosticsEndpoint : IRazorRequestHandler<VSInternal
 
         var documentSnapshot = documentContext.Snapshot;
 
-        var razorDiagnostics = await GetRazorDiagnosticsAsync(documentSnapshot).ConfigureAwait(false);
+        var razorDiagnostics = await GetRazorDiagnosticsAsync(documentSnapshot, cancellationToken).ConfigureAwait(false);
 
         await ReportRZ10012TelemetryAsync(documentContext, razorDiagnostics, cancellationToken).ConfigureAwait(false);
 
@@ -106,7 +106,9 @@ internal class DocumentPullDiagnosticsEndpoint : IRazorRequestHandler<VSInternal
             {
                 if (report.Diagnostics is not null)
                 {
-                    var mappedDiagnostics = await _translateDiagnosticsService.TranslateAsync(RazorLanguageKind.CSharp, report.Diagnostics, documentSnapshot).ConfigureAwait(false);
+                    var mappedDiagnostics = await _translateDiagnosticsService
+                        .TranslateAsync(RazorLanguageKind.CSharp, report.Diagnostics, documentSnapshot, cancellationToken)
+                        .ConfigureAwait(false);
                     report.Diagnostics = mappedDiagnostics;
                 }
 
@@ -120,7 +122,9 @@ internal class DocumentPullDiagnosticsEndpoint : IRazorRequestHandler<VSInternal
             {
                 if (report.Diagnostics is not null)
                 {
-                    var mappedDiagnostics = await _translateDiagnosticsService.TranslateAsync(RazorLanguageKind.Html, report.Diagnostics, documentSnapshot).ConfigureAwait(false);
+                    var mappedDiagnostics = await _translateDiagnosticsService
+                        .TranslateAsync(RazorLanguageKind.Html, report.Diagnostics, documentSnapshot, cancellationToken)
+                        .ConfigureAwait(false);
                     report.Diagnostics = mappedDiagnostics;
                 }
 
@@ -131,9 +135,9 @@ internal class DocumentPullDiagnosticsEndpoint : IRazorRequestHandler<VSInternal
         return allDiagnostics.ToArray();
     }
 
-    private static async Task<VSInternalDiagnosticReport[]?> GetRazorDiagnosticsAsync(IDocumentSnapshot documentSnapshot)
+    private static async Task<VSInternalDiagnosticReport[]?> GetRazorDiagnosticsAsync(IDocumentSnapshot documentSnapshot, CancellationToken cancellationToken)
     {
-        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
+        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
         var sourceText = codeDocument.Source.Text;
         var csharpDocument = codeDocument.GetCSharpDocument();
         var diagnostics = csharpDocument.Diagnostics;

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Diagnostics/RazorDiagnosticsPublisher.cs

@@ -113,10 +113,10 @@ internal partial class RazorDiagnosticsPublisher : IDocumentProcessedListener, I
         }
     }
 
-    private async Task PublishDiagnosticsAsync(IDocumentSnapshot document, CancellationToken token)
+    private async Task PublishDiagnosticsAsync(IDocumentSnapshot document, CancellationToken cancellationToken)
     {
-        var result = await document.GetGeneratedOutputAsync().ConfigureAwait(false);
-        var csharpDiagnostics = await GetCSharpDiagnosticsAsync(document, token).ConfigureAwait(false);
+        var result = await document.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
+        var csharpDiagnostics = await GetCSharpDiagnosticsAsync(document, cancellationToken).ConfigureAwait(false);
         var razorDiagnostics = result.GetCSharpDocument().Diagnostics;
 
         lock (_publishedDiagnostics)
@@ -188,7 +188,7 @@ internal partial class RazorDiagnosticsPublisher : IDocumentProcessedListener, I
                     if (_documentContextFactory.Value.TryCreate(delegatedParams.TextDocument.Uri, projectContext: null, out var documentContext))
                     {
                         return await _translateDiagnosticsService.Value
-                            .TranslateAsync(RazorLanguageKind.CSharp, fullDiagnostics.Items, documentContext.Snapshot)
+                            .TranslateAsync(RazorLanguageKind.CSharp, fullDiagnostics.Items, documentContext.Snapshot, cancellationToken)
                             .ConfigureAwait(false);
                     }
                 }

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

@@ -1,6 +1,7 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.CodeAnalysis.Razor.Formatting;
@@ -10,9 +11,9 @@ namespace Microsoft.AspNetCore.Razor.LanguageServer.Formatting;
 
 internal sealed class LspFormattingCodeDocumentProvider : IFormattingCodeDocumentProvider
 {
-    public ValueTask<RazorCodeDocument> GetCodeDocumentAsync(IDocumentSnapshot snapshot)
+    public ValueTask<RazorCodeDocument> GetCodeDocumentAsync(IDocumentSnapshot snapshot, CancellationToken cancellationToken)
     {
         // Formatting always uses design time
-        return snapshot.GetGeneratedOutputAsync(forceDesignTimeGeneratedOutput: true);
+        return snapshot.GetGeneratedOutputAsync(forceDesignTimeGeneratedOutput: true, cancellationToken);
     }
 }

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/MapCode/MapCodeEndpoint.cs

@@ -105,7 +105,7 @@ internal sealed class MapCodeEndpoint(
 
                 // We create a new Razor file based on each content in each mapping order to get the syntax tree that we'll later use to map.
                 var newSnapshot = snapshot.WithText(SourceText.From(content));
-                var codeToMap = await newSnapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
+                var codeToMap = await newSnapshot.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
 
                 var mappingSuccess = await TryMapCodeAsync(
                     codeToMap, mapping.FocusLocations, changes, mapCodeCorrelationId, documentContext, cancellationToken).ConfigureAwait(false);

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Mapping/RazorLanguageQueryEndpoint.cs

@@ -42,7 +42,7 @@ internal sealed class RazorLanguageQueryEndpoint(IDocumentMappingService documen
         var documentSnapshot = documentContext.Snapshot;
         var documentVersion = documentContext.Snapshot.Version;
 
-        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
+        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
         var sourceText = codeDocument.Source.Text;
         var hostDocumentIndex = sourceText.GetPosition(request.Position);
         var responsePosition = request.Position;

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/OpenDocumentGenerator.cs

@@ -75,7 +75,7 @@ internal partial class OpenDocumentGenerator : IRazorStartupService, IDisposable
                 return;
             }
 
-            var codeDocument = await document.GetGeneratedOutputAsync().ConfigureAwait(false);
+            var codeDocument = await document.GetGeneratedOutputAsync(token).ConfigureAwait(false);
 
             foreach (var listener in _listeners)
             {

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Diagnostics/RazorTranslateDiagnosticsService.cs

@@ -6,6 +6,7 @@ using System.Collections.Frozen;
 using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.AspNetCore.Razor.Language.Syntax;
@@ -45,18 +46,27 @@ internal class RazorTranslateDiagnosticsService(IDocumentMappingService document
     ]).ToFrozenSet();
 
     /// <summary>
-    /// Translates code diagnostics from one representation into another.
+    ///  Translates code diagnostics from one representation into another.
     /// </summary>
-    /// <param name="diagnosticKind">The `RazorLanguageKind` of the `Diagnostic` objects included in `diagnostics`.</param>
-    /// <param name="diagnostics">An array of `Diagnostic` objects to translate.</param>
-    /// <param name="documentSnapshot">The `DocumentContext` for the code document associated with the diagnostics.</param>
+    /// <param name="diagnosticKind">
+    ///  The <see cref="RazorLanguageKind"/> of the <see cref="Diagnostic"/> objects
+    ///  included in <paramref name="diagnostics"/>.
+    /// </param>
+    /// <param name="diagnostics">
+    ///  An array of <see cref="Diagnostic"/> objects to translate.
+    /// </param>
+    /// <param name="documentSnapshot">
+    ///  The <see cref="IDocumentSnapshot"/> for the code document associated with the diagnostics.
+    /// </param>
+    /// <param name="cancellationToken">A token that can be checked to cancel work.</param>
     /// <returns>An array of translated diagnostics</returns>
     internal async Task<LspDiagnostic[]> TranslateAsync(
         RazorLanguageKind diagnosticKind,
         LspDiagnostic[] diagnostics,
-        IDocumentSnapshot documentSnapshot)
+        IDocumentSnapshot documentSnapshot,
+        CancellationToken cancellationToken)
     {
-        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
+        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
         if (codeDocument.IsUnsupported() != false)
         {
             _logger.LogInformation($"Unsupported code document.");

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Formatting/FormattingContext.cs

@@ -7,6 +7,7 @@ using System.Collections.Immutable;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.AspNetCore.Razor.Language.Syntax;
@@ -224,11 +225,13 @@ internal sealed class FormattingContext
         return false;
     }
 
-    public async Task<FormattingContext> WithTextAsync(SourceText changedText)
+    public async Task<FormattingContext> WithTextAsync(SourceText changedText, CancellationToken cancellationToken)
     {
         var changedSnapshot = OriginalSnapshot.WithText(changedText);
 
-        var codeDocument = await _codeDocumentProvider.GetCodeDocumentAsync(changedSnapshot).ConfigureAwait(false);
+        var codeDocument = await _codeDocumentProvider
+            .GetCodeDocumentAsync(changedSnapshot, cancellationToken)
+            .ConfigureAwait(false);
 
         DEBUG_ValidateComponents(CodeDocument, codeDocument);
 

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Formatting/IFormattingCodeDocumentProvider.cs

@@ -1,6 +1,7 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;
@@ -9,5 +10,5 @@ namespace Microsoft.CodeAnalysis.Razor.Formatting;
 
 internal interface IFormattingCodeDocumentProvider
 {
-    ValueTask<RazorCodeDocument> GetCodeDocumentAsync(IDocumentSnapshot snapshot);
+    ValueTask<RazorCodeDocument> GetCodeDocumentAsync(IDocumentSnapshot snapshot, CancellationToken cancellationToken);
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Formatting/Passes/CSharpFormattingPass.cs

@@ -37,7 +37,7 @@ internal sealed class CSharpFormattingPass(
         if (changes.Length > 0)
         {
             changedText = changedText.WithChanges(changes);
-            changedContext = await context.WithTextAsync(changedText).ConfigureAwait(false);
+            changedContext = await context.WithTextAsync(changedText, cancellationToken).ConfigureAwait(false);
         }
 
         cancellationToken.ThrowIfCancellationRequested();
@@ -48,7 +48,7 @@ internal sealed class CSharpFormattingPass(
         if (csharpChanges.Length > 0)
         {
             changedText = changedText.WithChanges(csharpChanges);
-            changedContext = await changedContext.WithTextAsync(changedText).ConfigureAwait(false);
+            changedContext = await changedContext.WithTextAsync(changedText, cancellationToken).ConfigureAwait(false);
 
             _logger.LogTestOnly($"After FormatCSharpAsync:\r\n{changedText}");
         }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Formatting/Passes/CSharpOnTypeFormattingPass.cs

@@ -109,7 +109,7 @@ internal sealed class CSharpOnTypeFormattingPass(
         var formattedText = ApplyChangesAndTrackChange(originalText, filteredChanges, out _, out var spanAfterFormatting);
         _logger.LogTestOnly($"After C# changes:\r\n{formattedText}");
 
-        var changedContext = await context.WithTextAsync(formattedText).ConfigureAwait(false);
+        var changedContext = await context.WithTextAsync(formattedText, cancellationToken).ConfigureAwait(false);
         var linePositionSpanAfterFormatting = formattedText.GetLinePositionSpan(spanAfterFormatting);
 
         cancellationToken.ThrowIfCancellationRequested();
@@ -119,7 +119,7 @@ internal sealed class CSharpOnTypeFormattingPass(
         var cleanedText = formattedText.WithChanges(cleanupChanges);
         _logger.LogTestOnly($"After CleanupDocument:\r\n{cleanedText}");
 
-        changedContext = await changedContext.WithTextAsync(cleanedText).ConfigureAwait(false);
+        changedContext = await changedContext.WithTextAsync(cleanedText, cancellationToken).ConfigureAwait(false);
 
         cancellationToken.ThrowIfCancellationRequested();
 

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Formatting/Passes/FormattingDiagnosticValidationPass.cs

@@ -27,7 +27,7 @@ internal sealed class FormattingDiagnosticValidationPass(ILoggerFactory loggerFa
 
         var text = context.SourceText;
         var changedText = text.WithChanges(changes);
-        var changedContext = await context.WithTextAsync(changedText).ConfigureAwait(false);
+        var changedContext = await context.WithTextAsync(changedText, cancellationToken).ConfigureAwait(false);
         var changedDiagnostics = changedContext.CodeDocument.GetSyntaxTree().Diagnostics;
 
         // We want to ensure diagnostics didn't change, but since we're formatting things, its expected

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Formatting/Passes/HtmlFormattingPassBase.cs

@@ -34,7 +34,7 @@ internal abstract class HtmlFormattingPassBase(ILogger logger) : IFormattingPass
 
             changedText = originalText.WithChanges(filteredChanges);
             // Create a new formatting context for the changed razor document.
-            changedContext = await context.WithTextAsync(changedText).ConfigureAwait(false);
+            changedContext = await context.WithTextAsync(changedText, cancellationToken).ConfigureAwait(false);
 
             _logger.LogTestOnly($"After normalizedEdits:\r\n{changedText}");
         }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Formatting/Passes/RazorFormattingPass.cs

@@ -30,7 +30,7 @@ internal sealed class RazorFormattingPass : IFormattingPass
         if (changes.Length > 0)
         {
             changedText = changedText.WithChanges(changes);
-            changedContext = await context.WithTextAsync(changedText).ConfigureAwait(false);
+            changedContext = await context.WithTextAsync(changedText, cancellationToken).ConfigureAwait(false);
 
             cancellationToken.ThrowIfCancellationRequested();
         }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Formatting/RazorFormattingService.cs

@@ -67,7 +67,9 @@ internal class RazorFormattingService : IRazorFormattingService
         RazorFormattingOptions options,
         CancellationToken cancellationToken)
     {
-        var codeDocument = await _codeDocumentProvider.GetCodeDocumentAsync(documentContext.Snapshot).ConfigureAwait(false);
+        var codeDocument = await _codeDocumentProvider
+            .GetCodeDocumentAsync(documentContext.Snapshot, cancellationToken)
+            .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:
@@ -223,8 +225,8 @@ internal class RazorFormattingService : IRazorFormattingService
         // Code actions were computed on the regular document, which with FUSE could be a runtime document. We have to make
         // sure for code actions specifically we are formatting that same document, or TextChange spans may not line up
         var codeDocument = isCodeActionFormattingRequest
-            ? await documentSnapshot.GetGeneratedOutputAsync(forceDesignTimeGeneratedOutput: false).ConfigureAwait(false)
-            : await _codeDocumentProvider.GetCodeDocumentAsync(documentSnapshot).ConfigureAwait(false);
+            ? await documentSnapshot.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false)
+            : await _codeDocumentProvider.GetCodeDocumentAsync(documentSnapshot, cancellationToken).ConfigureAwait(false);
 
         var context = FormattingContext.CreateForOnTypeFormatting(
             documentSnapshot,

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/GoToDefinition/AbstractRazorComponentDefinitionService.cs

@@ -44,7 +44,7 @@ internal abstract class AbstractRazorComponentDefinitionService(
             return null;
         }
 
-        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
+        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
 
         if (!RazorComponentDefinitionHelpers.TryGetBoundTagHelpers(codeDocument, positionInfo.HostDocumentIndex, ignoreAttributes, _logger, out var boundTagHelper, out var boundAttribute))
         {
@@ -52,7 +52,10 @@ internal abstract class AbstractRazorComponentDefinitionService(
             return null;
         }
 
-        var componentDocument = await _componentSearchEngine.TryLocateComponentAsync(boundTagHelper, solutionQueryOperations).ConfigureAwait(false);
+        var componentDocument = await _componentSearchEngine
+            .TryLocateComponentAsync(boundTagHelper, solutionQueryOperations, cancellationToken)
+            .ConfigureAwait(false);
+
         if (componentDocument is null)
         {
             _logger.LogInformation($"Could not locate component document.");

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/GoToDefinition/RazorComponentDefinitionHelpers.cs

@@ -150,7 +150,7 @@ internal static class RazorComponentDefinitionHelpers
 
         var csharpSyntaxTree = await documentSnapshot.GetCSharpSyntaxTreeAsync(cancellationToken).ConfigureAwait(false);
         var root = await csharpSyntaxTree.GetRootAsync(cancellationToken).ConfigureAwait(false);
-        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
+        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
 
         // Since we know how the compiler generates the C# source we can be a little specific here, and avoid
         // long tree walks. If the compiler ever changes how they generate their code, the tests for this will break

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/IRazorComponentSearchEngine.cs

@@ -1,6 +1,7 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;
@@ -9,5 +10,8 @@ namespace Microsoft.CodeAnalysis.Razor.Workspaces;
 
 internal interface IRazorComponentSearchEngine
 {
-    Task<IDocumentSnapshot?> TryLocateComponentAsync(TagHelperDescriptor tagHelper, ISolutionQueryOperations solutionQueryOperations);
+    Task<IDocumentSnapshot?> TryLocateComponentAsync(
+        TagHelperDescriptor tagHelper,
+        ISolutionQueryOperations solutionQueryOperations,
+        CancellationToken cancellationToken);
 }

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

@@ -51,9 +51,9 @@ internal class DocumentContext(Uri uri, IDocumentSnapshot snapshot, VSProjectCon
 
         async ValueTask<RazorCodeDocument> GetCodeDocumentCoreAsync(CancellationToken cancellationToken)
         {
-            var codeDocument = await Snapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
-
-            cancellationToken.ThrowIfCancellationRequested();
+            var codeDocument = await Snapshot
+                .GetGeneratedOutputAsync(cancellationToken)
+                .ConfigureAwait(false);
 
             // Interlock to ensure that we only ever return one instance of RazorCodeDocument.
             // In race scenarios, when more than one RazorCodeDocument is produced, we want to

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

@@ -36,11 +36,9 @@ internal sealed class DocumentSnapshot(ProjectSnapshot project, DocumentState st
 
     public bool TryGetGeneratedOutput([NotNullWhen(true)] out RazorCodeDocument? result)
     {
-        if (_state.IsGeneratedOutputResultAvailable)
+        if (_state.TryGetGeneratedOutputAndVersion(out var outputAndVersion))
         {
-#pragma warning disable VSTHRD002 // Avoid problematic synchronous waits
-            result = _state.GetGeneratedOutputAndVersionAsync(_project, this).Result.output;
-#pragma warning restore VSTHRD002 // Avoid problematic synchronous waits
+            result = outputAndVersion.output;
             return true;
         }
 
@@ -61,27 +59,32 @@ internal sealed class DocumentSnapshot(ProjectSnapshot project, DocumentState st
 
         async Task<SyntaxTree> GetCSharpSyntaxTreeCoreAsync(CancellationToken cancellationToken)
         {
-            var codeDocument = await GetGeneratedOutputAsync(forceDesignTimeGeneratedOutput: false).ConfigureAwait(false);
+            var codeDocument = await GetGeneratedOutputAsync(forceDesignTimeGeneratedOutput: false, cancellationToken).ConfigureAwait(false);
             return codeDocument.GetCSharpSyntaxTree(cancellationToken);
         }
     }
 
-    public async ValueTask<RazorCodeDocument> GetGeneratedOutputAsync(bool forceDesignTimeGeneratedOutput)
+    public async ValueTask<RazorCodeDocument> GetGeneratedOutputAsync(bool forceDesignTimeGeneratedOutput, CancellationToken cancellationToken)
     {
         if (forceDesignTimeGeneratedOutput)
         {
-            return await GetDesignTimeGeneratedOutputAsync().ConfigureAwait(false);
+            return await GetDesignTimeGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
         }
 
-        var (output, _) = await _state.GetGeneratedOutputAndVersionAsync(_project, this).ConfigureAwait(false);
+        var (output, _) = await _state
+            .GetGeneratedOutputAndVersionAsync(_project, this, cancellationToken)
+            .ConfigureAwait(false);
+
         return output;
     }
 
-    private async Task<RazorCodeDocument> GetDesignTimeGeneratedOutputAsync()
+    private async Task<RazorCodeDocument> GetDesignTimeGeneratedOutputAsync(CancellationToken cancellationToken)
     {
-        var tagHelpers = await Project.GetTagHelpersAsync(CancellationToken.None).ConfigureAwait(false);
+        var tagHelpers = await Project.GetTagHelpersAsync(cancellationToken).ConfigureAwait(false);
         var projectEngine = Project.GetProjectEngine();
-        var imports = await DocumentState.GetImportsAsync(this, projectEngine).ConfigureAwait(false);
-        return await DocumentState.GenerateCodeDocumentAsync(this, projectEngine, imports, tagHelpers, forceRuntimeCodeGeneration: false).ConfigureAwait(false);
+        var imports = await DocumentState.GetImportsAsync(this, projectEngine, cancellationToken).ConfigureAwait(false);
+        return await DocumentState
+            .GenerateCodeDocumentAsync(this, projectEngine, imports, tagHelpers, forceRuntimeCodeGeneration: false, cancellationToken)
+            .ConfigureAwait(false);
     }
 }

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

@@ -46,20 +46,39 @@ internal partial class DocumentState
             }
         }
 
-        public async Task<(RazorCodeDocument, VersionStamp)> GetGeneratedOutputAndVersionAsync(ProjectSnapshot project, IDocumentSnapshot document)
+        public bool TryGetGeneratedOutputAndVersion(out (RazorCodeDocument Output, VersionStamp InputVersion) result)
+        {
+            if (_computedOutput?.TryGetCachedOutput(out var output, out var version) == true)
+            {
+                result = (output, version);
+                return true;
+            }
+
+            result = default;
+            return false;
+        }
+
+        public async Task<(RazorCodeDocument, VersionStamp)> GetGeneratedOutputAndVersionAsync(
+            ProjectSnapshot project,
+            IDocumentSnapshot document,
+            CancellationToken cancellationToken)
         {
             if (_computedOutput?.TryGetCachedOutput(out var cachedCodeDocument, out var cachedInputVersion) == true)
             {
                 return (cachedCodeDocument, cachedInputVersion);
             }
 
-            var (codeDocument, inputVersion) = await GetMemoizedGeneratedOutputAndVersionAsync(project, document).ConfigureAwait(false);
+            var (codeDocument, inputVersion) = await GetMemoizedGeneratedOutputAndVersionAsync(project, document, cancellationToken)
+                .ConfigureAwait(false);
 
             _computedOutput = new ComputedOutput(codeDocument, inputVersion);
             return (codeDocument, inputVersion);
         }
 
-        private Task<(RazorCodeDocument, VersionStamp)> GetMemoizedGeneratedOutputAndVersionAsync(ProjectSnapshot project, IDocumentSnapshot document)
+        private Task<(RazorCodeDocument, VersionStamp)> GetMemoizedGeneratedOutputAndVersionAsync(
+            ProjectSnapshot project,
+            IDocumentSnapshot document,
+            CancellationToken cancellationToken)
         {
             if (project is null)
             {
@@ -99,7 +118,7 @@ internal partial class DocumentState
                 }
 
                 // Typically in VS scenarios this will run synchronously because all resources are readily available.
-                var outputTask = ComputeGeneratedOutputAndVersionAsync(project, document);
+                var outputTask = ComputeGeneratedOutputAndVersionAsync(project, document, cancellationToken);
                 if (outputTask.IsCompleted)
                 {
                     // Compiling ran synchronously, lets just immediately propagate to the TCS
@@ -150,7 +169,10 @@ internal partial class DocumentState
             }
         }
 
-        private async Task<(RazorCodeDocument, VersionStamp)> ComputeGeneratedOutputAndVersionAsync(ProjectSnapshot project, IDocumentSnapshot document)
+        private async Task<(RazorCodeDocument, VersionStamp)> ComputeGeneratedOutputAndVersionAsync(
+            ProjectSnapshot project,
+            IDocumentSnapshot document,
+            CancellationToken cancellationToken)
         {
             // We only need to produce the generated code if any of our inputs is newer than the
             // previously cached output.
@@ -164,8 +186,8 @@ internal partial class DocumentState
             var configurationVersion = project.ConfigurationVersion;
             var projectWorkspaceStateVersion = project.ProjectWorkspaceStateVersion;
             var documentCollectionVersion = project.DocumentCollectionVersion;
-            var imports = await GetImportsAsync(document, project.GetProjectEngine()).ConfigureAwait(false);
-            var documentVersion = await document.GetTextVersionAsync(CancellationToken.None).ConfigureAwait(false);
+            var imports = await GetImportsAsync(document, project.GetProjectEngine(), cancellationToken).ConfigureAwait(false);
+            var documentVersion = await document.GetTextVersionAsync(cancellationToken).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.
@@ -210,9 +232,9 @@ internal partial class DocumentState
                 }
             }
 
-            var tagHelpers = await project.GetTagHelpersAsync(CancellationToken.None).ConfigureAwait(false);
+            var tagHelpers = await project.GetTagHelpersAsync(cancellationToken).ConfigureAwait(false);
             var forceRuntimeCodeGeneration = project.Configuration.LanguageServerFlags?.ForceRuntimeCodeGeneration ?? false;
-            var codeDocument = await GenerateCodeDocumentAsync(document, project.GetProjectEngine(), imports, tagHelpers, forceRuntimeCodeGeneration).ConfigureAwait(false);
+            var codeDocument = await GenerateCodeDocumentAsync(document, project.GetProjectEngine(), imports, tagHelpers, forceRuntimeCodeGeneration, cancellationToken).ConfigureAwait(false);
             return (codeDocument, inputVersion);
         }
 

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

@@ -69,9 +69,17 @@ internal partial class DocumentState
     private ComputedStateTracker ComputedState
         => _computedState ??= InterlockedOperations.Initialize(ref _computedState, new ComputedStateTracker());
 
-    public Task<(RazorCodeDocument output, VersionStamp inputVersion)> GetGeneratedOutputAndVersionAsync(ProjectSnapshot project, DocumentSnapshot document)
+    public bool TryGetGeneratedOutputAndVersion(out (RazorCodeDocument output, VersionStamp inputVersion) result)
     {
-        return ComputedState.GetGeneratedOutputAndVersionAsync(project, document);
+        return ComputedState.TryGetGeneratedOutputAndVersion(out result);
+    }
+
+    public Task<(RazorCodeDocument output, VersionStamp inputVersion)> GetGeneratedOutputAndVersionAsync(
+        ProjectSnapshot project,
+        DocumentSnapshot document,
+        CancellationToken cancellationToken)
+    {
+        return ComputedState.GetGeneratedOutputAndVersionAsync(project, document, cancellationToken);
     }
 
     public ValueTask<TextAndVersion> GetTextAndVersionAsync(CancellationToken cancellationToken)
@@ -229,19 +237,25 @@ internal partial class DocumentState
         return imports.DrainToImmutable();
     }
 
-    internal static async Task<RazorCodeDocument> GenerateCodeDocumentAsync(IDocumentSnapshot document, RazorProjectEngine projectEngine, ImmutableArray<ImportItem> imports, ImmutableArray<TagHelperDescriptor> tagHelpers, bool forceRuntimeCodeGeneration)
+    internal static async Task<RazorCodeDocument> GenerateCodeDocumentAsync(
+        IDocumentSnapshot document,
+        RazorProjectEngine projectEngine,
+        ImmutableArray<ImportItem> imports,
+        ImmutableArray<TagHelperDescriptor> tagHelpers,
+        bool forceRuntimeCodeGeneration,
+        CancellationToken cancellationToken)
     {
         // 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);
+            var sourceDocument = await GetRazorSourceDocumentAsync(item.Document, importProjectItem, cancellationToken).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);
+        var documentSource = await GetRazorSourceDocumentAsync(document, projectItem, cancellationToken).ConfigureAwait(false);
 
         if (forceRuntimeCodeGeneration)
         {
@@ -251,23 +265,26 @@ internal partial class DocumentState
         return projectEngine.ProcessDesignTime(documentSource, fileKind: document.FileKind, importSources.DrainToImmutable(), tagHelpers);
     }
 
-    internal static async Task<ImmutableArray<ImportItem>> GetImportsAsync(IDocumentSnapshot document, RazorProjectEngine projectEngine)
+    internal static async Task<ImmutableArray<ImportItem>> GetImportsAsync(IDocumentSnapshot document, RazorProjectEngine projectEngine, CancellationToken cancellationToken)
     {
         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)
         {
-            var versionStamp = await snapshot.GetTextVersionAsync(CancellationToken.None).ConfigureAwait(false);
+            var versionStamp = await snapshot.GetTextVersionAsync(cancellationToken).ConfigureAwait(false);
             result.Add(new ImportItem(snapshot.FilePath, versionStamp, snapshot));
         }
 
         return result.DrainToImmutable();
     }
 
-    private static async Task<RazorSourceDocument> GetRazorSourceDocumentAsync(IDocumentSnapshot document, RazorProjectItem? projectItem)
+    private static async Task<RazorSourceDocument> GetRazorSourceDocumentAsync(
+        IDocumentSnapshot document,
+        RazorProjectItem? projectItem,
+        CancellationToken cancellationToken)
     {
-        var sourceText = await document.GetTextAsync(CancellationToken.None).ConfigureAwait(false);
+        var sourceText = await document.GetTextAsync(cancellationToken).ConfigureAwait(false);
         return RazorSourceDocument.Create(sourceText, RazorSourceDocumentProperties.Create(document.FilePath, projectItem?.RelativePhysicalPath));
     }
 }

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

@@ -1,9 +1,6 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
-#nullable disable
-
-using System;
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
@@ -12,30 +9,19 @@ using Microsoft.CodeAnalysis.Text;
 
 namespace Microsoft.CodeAnalysis.Razor.ProjectSystem;
 
-internal class GeneratedDocumentTextLoader : TextLoader
+internal class GeneratedDocumentTextLoader(IDocumentSnapshot document, string filePath) : TextLoader
 {
-    private readonly IDocumentSnapshot _document;
-    private readonly string _filePath;
-    private readonly VersionStamp _version;
-
-    public GeneratedDocumentTextLoader(IDocumentSnapshot document, string filePath)
-    {
-        if (document is null)
-        {
-            throw new ArgumentNullException(nameof(document));
-        }
-
-        _document = document;
-        _filePath = filePath;
-        _version = VersionStamp.Create();
-    }
+    private readonly IDocumentSnapshot _document = document;
+    private readonly string _filePath = filePath;
+    private readonly VersionStamp _version = VersionStamp.Create();
 
     public override async Task<TextAndVersion> LoadTextAndVersionAsync(LoadTextOptions options, CancellationToken cancellationToken)
     {
-        var output = await _document.GetGeneratedOutputAsync().ConfigureAwait(false);
+        var output = await _document.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
 
-        // Providing an encoding here is important for debuggability. Without this edit-and-continue
-        // won't work for projects with Razor files.
-        return TextAndVersion.Create(SourceText.From(output.GetCSharpDocument().GeneratedCode, Encoding.UTF8), _version, _filePath);
+        // Providing an encoding here is important for debuggability.
+        // Without this, edit-and-continue won't work for projects with Razor files.
+        var csharpSourceText = SourceText.From(output.GetCSharpDocument().GeneratedCode, Encoding.UTF8);
+        return TextAndVersion.Create(csharpSourceText, _version, _filePath);
     }
 }

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

@@ -20,7 +20,9 @@ internal interface IDocumentSnapshot
 
     ValueTask<SourceText> GetTextAsync(CancellationToken cancellationToken);
     ValueTask<VersionStamp> GetTextVersionAsync(CancellationToken cancellationToken);
-    ValueTask<RazorCodeDocument> GetGeneratedOutputAsync(bool forceDesignTimeGeneratedOutput);
+    ValueTask<RazorCodeDocument> GetGeneratedOutputAsync(
+        bool forceDesignTimeGeneratedOutput,
+        CancellationToken cancellationToken);
 
     /// <summary>
     /// Gets the Roslyn syntax tree for the generated C# for this Razor document

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

@@ -11,7 +11,9 @@ namespace Microsoft.CodeAnalysis.Razor.ProjectSystem;
 
 internal static class IDocumentSnapshotExtensions
 {
-    public static async Task<TagHelperDescriptor?> TryGetTagHelperDescriptorAsync(this IDocumentSnapshot documentSnapshot, CancellationToken cancellationToken)
+    public static async Task<TagHelperDescriptor?> TryGetTagHelperDescriptorAsync(
+        this IDocumentSnapshot documentSnapshot,
+        CancellationToken cancellationToken)
     {
         // No point doing anything if its not a component
         if (documentSnapshot.FileKind != FileKinds.Component)
@@ -19,7 +21,7 @@ internal static class IDocumentSnapshotExtensions
             return null;
         }
 
-        var razorCodeDocument = await documentSnapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
+        var razorCodeDocument = await documentSnapshot.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
         if (razorCodeDocument is null)
         {
             return null;
@@ -53,8 +55,10 @@ internal static class IDocumentSnapshotExtensions
         return fileName.AsSpan().Equals(path.Span, FilePathComparison.Instance);
     }
 
-    public static ValueTask<RazorCodeDocument> GetGeneratedOutputAsync(this IDocumentSnapshot documentSnapshot)
+    public static ValueTask<RazorCodeDocument> GetGeneratedOutputAsync(
+        this IDocumentSnapshot documentSnapshot,
+        CancellationToken cancellationToken)
     {
-        return documentSnapshot.GetGeneratedOutputAsync(forceDesignTimeGeneratedOutput: false);
+        return documentSnapshot.GetGeneratedOutputAsync(forceDesignTimeGeneratedOutput: false, cancellationToken);
     }
 }

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

@@ -41,7 +41,9 @@ internal sealed class ImportDocumentSnapshot(IProjectSnapshot project, RazorProj
         }
     }
 
-    public ValueTask<RazorCodeDocument> GetGeneratedOutputAsync(bool forceDesignTimeGeneratedOutput)
+    public ValueTask<RazorCodeDocument> GetGeneratedOutputAsync(
+        bool forceDesignTimeGeneratedOutput,
+        CancellationToken cancellationToken)
         => throw new NotSupportedException();
 
     public ValueTask<VersionStamp> GetTextVersionAsync(CancellationToken cancellationToken)

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/RazorComponentSearchEngine.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
 using System;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.CodeAnalysis.Razor.Logging;
@@ -13,18 +14,35 @@ internal class RazorComponentSearchEngine(ILoggerFactory loggerFactory) : IRazor
 {
     private readonly ILogger _logger = loggerFactory.GetOrCreateLogger<RazorComponentSearchEngine>();
 
-    /// <summary>Search for a component in a project based on its tag name and fully qualified name.</summary>
+    /// <summary>
+    ///  Search for a component in a project based on its tag name and fully qualified name.
+    /// </summary>
+    /// <param name="tagHelper">
+    ///  A <see cref="TagHelperDescriptor"/> to find the corresponding Razor component for.
+    /// </param>
+    /// <param name="solutionQueryOperations">
+    ///  An <see cref="ISolutionQueryOperations"/> to enumerate project snapshots.
+    /// </param>
+    /// <param name="cancellationToken">
+    ///  A token that is checked to cancel work.
+    /// </param>
+    /// <returns>
+    ///  The corresponding <see cref="IDocumentSnapshot"/> if found, <see langword="null"/> otherwise.
+    /// </returns>
     /// <remarks>
-    /// This method makes several assumptions about the nature of components. First, it assumes that a component
-    /// a given name "Name" will be located in a file "Name.razor". Second, it assumes that the namespace the
-    /// component is present in has the same name as the assembly its corresponding tag helper is loaded from.
-    /// Implicitly, this method inherits any assumptions made by TrySplitNamespaceAndType.
+    ///  This method makes several assumptions about the nature of components. First,
+    ///  it assumes that a component a given name "Name" will be located in a file
+    ///  "Name.razor". Second, it assumes that the namespace the component is present in
+    ///  has the same name as the assembly its corresponding tag helper is loaded from.
+    ///  Implicitly, this method inherits any assumptions made by TrySplitNamespaceAndType.
     /// </remarks>
-    /// <param name="tagHelper">A TagHelperDescriptor to find the corresponding Razor component for.</param>
-    /// <param name="solutionQueryOperations">An <see cref="ISolutionQueryOperations"/> to enumerate project snapshots.</param>
-    /// <returns>The corresponding DocumentSnapshot if found, null otherwise.</returns>
-    /// <exception cref="ArgumentNullException">Thrown if <paramref name="tagHelper"/> is null.</exception>
-    public async Task<IDocumentSnapshot?> TryLocateComponentAsync(TagHelperDescriptor tagHelper, ISolutionQueryOperations solutionQueryOperations)
+    /// <exception cref="ArgumentNullException">
+    ///  Thrown if <paramref name="tagHelper"/> is <see langword="null"/>.
+    /// </exception>
+    public async Task<IDocumentSnapshot?> TryLocateComponentAsync(
+        TagHelperDescriptor tagHelper,
+        ISolutionQueryOperations solutionQueryOperations,
+        CancellationToken cancellationToken)
     {
         var typeName = tagHelper.GetTypeNameIdentifier();
         var namespaceName = tagHelper.GetTypeNamespace();
@@ -52,7 +70,7 @@ internal class RazorComponentSearchEngine(ILoggerFactory loggerFactory) : IRazor
                     continue;
                 }
 
-                var razorCodeDocument = await document.GetGeneratedOutputAsync().ConfigureAwait(false);
+                var razorCodeDocument = await document.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
                 if (razorCodeDocument is null)
                 {
                     continue;

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Rename/RenameService.cs

@@ -58,7 +58,9 @@ internal class RenameService(
             return null;
         }
 
-        var originComponentDocumentSnapshot = await _componentSearchEngine.TryLocateComponentAsync(originTagHelpers.First(), solutionQueryOperations).ConfigureAwait(false);
+        var originComponentDocumentSnapshot = await _componentSearchEngine
+            .TryLocateComponentAsync(originTagHelpers.First(), solutionQueryOperations, cancellationToken)
+            .ConfigureAwait(false);
         if (originComponentDocumentSnapshot is null)
         {
             return null;
@@ -80,7 +82,7 @@ internal class RenameService(
 
         foreach (var documentSnapshot in documentSnapshots)
         {
-            await AddEditsForCodeDocumentAsync(documentChanges, originTagHelpers, newName, documentSnapshot).ConfigureAwait(false);
+            await AddEditsForCodeDocumentAsync(documentChanges, originTagHelpers, newName, documentSnapshot, cancellationToken).ConfigureAwait(false);
         }
 
         foreach (var documentChange in documentChanges)
@@ -165,14 +167,15 @@ internal class RenameService(
         List<SumType<TextDocumentEdit, CreateFile, RenameFile, DeleteFile>> documentChanges,
         ImmutableArray<TagHelperDescriptor> originTagHelpers,
         string newName,
-        IDocumentSnapshot documentSnapshot)
+        IDocumentSnapshot documentSnapshot,
+        CancellationToken cancellationToken)
     {
         if (!FileKinds.IsComponent(documentSnapshot.FileKind))
         {
             return;
         }
 
-        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
+        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
         if (codeDocument.IsUnsupported())
         {
             return;

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/AutoInsert/RemoteAutoInsertService.cs

@@ -137,7 +137,9 @@ internal sealed class RemoteAutoInsertService(in ServiceArgs args)
             return Response.NoFurtherHandling;
         }
 
-        var generatedDocument = await remoteDocumentContext.Snapshot.GetGeneratedDocumentAsync().ConfigureAwait(false);
+        var generatedDocument = await remoteDocumentContext.Snapshot
+            .GetGeneratedDocumentAsync(cancellationToken)
+            .ConfigureAwait(false);
 
         var autoInsertResponseItem = await OnAutoInsert.GetOnAutoInsertResponseAsync(
             generatedDocument,

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/Diagnostics/RemoteDiagnosticsService.cs

@@ -49,8 +49,8 @@ internal sealed class RemoteDiagnosticsService(in ServiceArgs args) : RazorDocum
 
         return [
             .. RazorDiagnosticConverter.Convert(razorDiagnostics, codeDocument.Source.Text, context.Snapshot),
-            .. await _translateDiagnosticsService.TranslateAsync(RazorLanguageKind.CSharp, csharpDiagnostics, context.Snapshot),
-            .. await _translateDiagnosticsService.TranslateAsync(RazorLanguageKind.Html, htmlDiagnostics, context.Snapshot)
+            .. await _translateDiagnosticsService.TranslateAsync(RazorLanguageKind.CSharp, csharpDiagnostics, context.Snapshot, cancellationToken),
+            .. await _translateDiagnosticsService.TranslateAsync(RazorLanguageKind.Html, htmlDiagnostics, context.Snapshot, cancellationToken)
         ];
     }
 }

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/DocumentHighlight/RemoteDocumentHighlightService.cs

@@ -64,7 +64,9 @@ internal sealed partial class RemoteDocumentHighlightService(in ServiceArgs args
         var csharpDocument = codeDocument.GetCSharpDocument();
         if (DocumentMappingService.TryMapToGeneratedDocumentPosition(csharpDocument, index, out var mappedPosition, out _))
         {
-            var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync().ConfigureAwait(false);
+            var generatedDocument = await context.Snapshot
+                .GetGeneratedDocumentAsync(cancellationToken)
+                .ConfigureAwait(false);
 
             var highlights = await DocumentHighlights.GetHighlightsAsync(generatedDocument, mappedPosition, cancellationToken).ConfigureAwait(false);
 

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/DocumentMapping/RemoteDocumentMappingService.cs

@@ -54,8 +54,9 @@ internal sealed class RemoteDocumentMappingService(
 
         var razorDocumentSnapshot = _snapshotManager.GetSnapshot(razorDocument);
 
-        var razorCodeDocument = await razorDocumentSnapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
-        cancellationToken.ThrowIfCancellationRequested();
+        var razorCodeDocument = await razorDocumentSnapshot
+            .GetGeneratedOutputAsync(cancellationToken)
+            .ConfigureAwait(false);
 
         if (razorCodeDocument is null)
         {

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/DocumentSymbols/RemoteDocumentSymbolService.cs

@@ -33,7 +33,10 @@ internal sealed partial class RemoteDocumentSymbolService(in ServiceArgs args) :
 
     private async ValueTask<SumType<DocumentSymbol[], SymbolInformation[]>?> GetDocumentSymbolsAsync(RemoteDocumentContext context, bool useHierarchicalSymbols, CancellationToken cancellationToken)
     {
-        var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync().ConfigureAwait(false);
+        var generatedDocument = await context.Snapshot
+            .GetGeneratedDocumentAsync(cancellationToken)
+            .ConfigureAwait(false);
+
         var csharpSymbols = await ExternalHandlers.DocumentSymbols.GetDocumentSymbolsAsync(generatedDocument, useHierarchicalSymbols, cancellationToken).ConfigureAwait(false);
 
         var codeDocument = await context.GetCodeDocumentAsync(cancellationToken).ConfigureAwait(false);

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/FoldingRanges/RemoteFoldingRangeService.cs

@@ -41,7 +41,9 @@ internal sealed class RemoteFoldingRangeService(in ServiceArgs args) : RazorDocu
         ImmutableArray<RemoteFoldingRange> htmlRanges,
         CancellationToken cancellationToken)
     {
-        var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync().ConfigureAwait(false);
+        var generatedDocument = await context.Snapshot
+            .GetGeneratedDocumentAsync(cancellationToken)
+            .ConfigureAwait(false);
 
         var csharpRanges = await ExternalHandlers.FoldingRanges.GetFoldingRangesAsync(generatedDocument, cancellationToken).ConfigureAwait(false);
 

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/Formatting/RemoteFormattingCodeDocumentProvider.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
 using System.Composition;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.CodeAnalysis.Razor.Formatting;
@@ -12,9 +13,9 @@ namespace Microsoft.CodeAnalysis.Remote.Razor.Formatting;
 [Export(typeof(IFormattingCodeDocumentProvider)), Shared]
 internal sealed class RemoteFormattingCodeDocumentProvider : IFormattingCodeDocumentProvider
 {
-    public ValueTask<RazorCodeDocument> GetCodeDocumentAsync(IDocumentSnapshot snapshot)
+    public ValueTask<RazorCodeDocument> GetCodeDocumentAsync(IDocumentSnapshot snapshot, CancellationToken cancellationToken)
     {
         // Formatting always uses design time
-        return snapshot.GetGeneratedOutputAsync(forceDesignTimeGeneratedOutput: true);
+        return snapshot.GetGeneratedOutputAsync(forceDesignTimeGeneratedOutput: true, cancellationToken);
     }
 }

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/GoToDefinition/RemoteGoToDefinitionService.cs

@@ -83,7 +83,9 @@ internal sealed class RemoteGoToDefinitionService(in ServiceArgs args) : RazorDo
         }
 
         // Finally, call into C#.
-        var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync().ConfigureAwait(false);
+        var generatedDocument = await context.Snapshot
+            .GetGeneratedDocumentAsync(cancellationToken)
+            .ConfigureAwait(false);
 
         var locations = await ExternalHandlers.GoToDefinition
             .GetDefinitionsAsync(

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/GoToImplementation/RemoteGoToImplementationService.cs

@@ -73,7 +73,9 @@ internal sealed class RemoteGoToImplementationService(in ServiceArgs args) : Raz
         }
 
         // Finally, call into C#.
-        var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync().ConfigureAwait(false);
+        var generatedDocument = await context.Snapshot
+            .GetGeneratedDocumentAsync(cancellationToken)
+            .ConfigureAwait(false);
 
         var locations = await ExternalHandlers.GoToImplementation
             .FindImplementationsAsync(

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/InlayHints/RemoteInlayHintService.cs

@@ -61,7 +61,9 @@ internal sealed partial class RemoteInlayHintService(in ServiceArgs args) : Razo
             return null;
         }
 
-        var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync().ConfigureAwait(false);
+        var generatedDocument = await context.Snapshot
+            .GetGeneratedDocumentAsync(cancellationToken)
+            .ConfigureAwait(false);
 
         var textDocument = inlayHintParams.TextDocument.WithUri(generatedDocument.CreateUri());
         var range = projectedLinePositionSpan.ToRange();
@@ -115,7 +117,9 @@ internal sealed partial class RemoteInlayHintService(in ServiceArgs args) : Razo
 
     private async ValueTask<InlayHint> ResolveInlayHintAsync(RemoteDocumentContext context, InlayHint inlayHint, CancellationToken cancellationToken)
     {
-        var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync().ConfigureAwait(false);
+        var generatedDocument = await context.Snapshot
+            .GetGeneratedDocumentAsync(cancellationToken)
+            .ConfigureAwait(false);
 
         return await InlayHints.ResolveInlayHintAsync(generatedDocument, inlayHint, cancellationToken).ConfigureAwait(false);
     }

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

@@ -62,16 +62,18 @@ internal sealed class RemoteDocumentSnapshot : IDocumentSnapshot
     public bool TryGetTextVersion(out VersionStamp result)
         => TextDocument.TryGetTextVersion(out result);
 
-    public ValueTask<RazorCodeDocument> GetGeneratedOutputAsync(bool forceDesignTimeGeneratedOutput)
+    public ValueTask<RazorCodeDocument> GetGeneratedOutputAsync(
+        bool forceDesignTimeGeneratedOutput,
+        CancellationToken cancellationToken)
     {
         // TODO: We don't need to worry about locking if we get called from the didOpen/didChange LSP requests, as CLaSP
         //       takes care of that for us, and blocks requests until those are complete. If that doesn't end up happening,
         //       then a locking mechanism here would prevent concurrent compilations.
         return TryGetGeneratedOutput(out var codeDocument)
             ? new(codeDocument)
-            : new(GetGeneratedOutputCoreAsync());
+            : new(GetGeneratedOutputCoreAsync(cancellationToken));
 
-        async Task<RazorCodeDocument> GetGeneratedOutputCoreAsync()
+        async Task<RazorCodeDocument> GetGeneratedOutputCoreAsync(CancellationToken cancellationToken)
         {
             // The non-cohosted DocumentSnapshot implementation uses DocumentState to get the generated output, and we could do that too
             // but most of that code is optimized around caching pre-computed results when things change that don't affect the compilation.
@@ -81,15 +83,15 @@ internal sealed class RemoteDocumentSnapshot : IDocumentSnapshot
             // 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 projectEngine = await ProjectSnapshot.GetProjectEngine_CohostOnlyAsync(CancellationToken.None).ConfigureAwait(false);
-            var tagHelpers = await ProjectSnapshot.GetTagHelpersAsync(CancellationToken.None).ConfigureAwait(false);
-            var imports = await DocumentState.GetImportsAsync(this, projectEngine).ConfigureAwait(false);
+            var projectEngine = await ProjectSnapshot.GetProjectEngine_CohostOnlyAsync(cancellationToken).ConfigureAwait(false);
+            var tagHelpers = await ProjectSnapshot.GetTagHelpersAsync(cancellationToken).ConfigureAwait(false);
+            var imports = await DocumentState.GetImportsAsync(this, projectEngine, cancellationToken).ConfigureAwait(false);
 
             // TODO: Get the configuration for forceRuntimeCodeGeneration
             // var forceRuntimeCodeGeneration = _projectSnapshot.Configuration.LanguageServerFlags?.ForceRuntimeCodeGeneration ?? false;
 
             codeDocument = await DocumentState
-                .GenerateCodeDocumentAsync(this, projectEngine, imports, tagHelpers, forceRuntimeCodeGeneration: false)
+                .GenerateCodeDocumentAsync(this, projectEngine, imports, tagHelpers, forceRuntimeCodeGeneration: false, cancellationToken)
                 .ConfigureAwait(false);
 
             return _codeDocument ??= InterlockedOperations.Initialize(ref _codeDocument, codeDocument);
@@ -114,15 +116,15 @@ internal sealed class RemoteDocumentSnapshot : IDocumentSnapshot
         return result is not null;
     }
 
-    public ValueTask<Document> GetGeneratedDocumentAsync()
+    public ValueTask<Document> GetGeneratedDocumentAsync(CancellationToken cancellationToken)
     {
         return TryGetGeneratedDocument(out var generatedDocument)
             ? new(generatedDocument)
-            : GetGeneratedDocumentCoreAsync();
+            : GetGeneratedDocumentCoreAsync(cancellationToken);
 
-        async ValueTask<Document> GetGeneratedDocumentCoreAsync()
+        async ValueTask<Document> GetGeneratedDocumentCoreAsync(CancellationToken cancellationToken)
         {
-            var generatedDocument = await HACK_GenerateDocumentAsync().ConfigureAwait(false);
+            var generatedDocument = await HACK_GenerateDocumentAsync(cancellationToken).ConfigureAwait(false);
             return _generatedDocument ??= InterlockedOperations.Initialize(ref _generatedDocument, generatedDocument);
         }
     }
@@ -133,7 +135,7 @@ internal sealed class RemoteDocumentSnapshot : IDocumentSnapshot
         return generatedDocument is not null;
     }
 
-    private async Task<Document> HACK_GenerateDocumentAsync()
+    private async Task<Document> HACK_GenerateDocumentAsync(CancellationToken cancellationToken)
     {
         // TODO: A real implementation needs to get the SourceGeneratedDocument from the solution
 
@@ -144,7 +146,7 @@ internal sealed class RemoteDocumentSnapshot : IDocumentSnapshot
         var generatedDocumentId = solution.GetDocumentIdsWithFilePath(generatedFilePath).First(d => d.ProjectId == projectId);
         var generatedDocument = solution.GetRequiredDocument(generatedDocumentId);
 
-        var codeDocument = await this.GetGeneratedOutputAsync().ConfigureAwait(false);
+        var codeDocument = await this.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
         var csharpSourceText = codeDocument.GetCSharpSourceText();
 
         // HACK: We're not in the same solution fork as the LSP server that provides content for this document
@@ -163,7 +165,7 @@ internal sealed class RemoteDocumentSnapshot : IDocumentSnapshot
 
         async ValueTask<SyntaxTree> GetCSharpSyntaxTreeCoreAsync(Document? document, CancellationToken cancellationToken)
         {
-            document ??= await GetGeneratedDocumentAsync();
+            document ??= await GetGeneratedDocumentAsync(cancellationToken).ConfigureAwait(false);
 
             var tree = await document.GetSyntaxTreeAsync(cancellationToken).ConfigureAwait(false);
             return tree.AssumeNotNull();

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/Rename/RemoteRenameService.cs

@@ -52,7 +52,9 @@ internal sealed class RemoteRenameService(in ServiceArgs args) : RazorDocumentSe
             return NoFurtherHandling;
         }
 
-        var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync().ConfigureAwait(false);
+        var generatedDocument = await context.Snapshot
+            .GetGeneratedDocumentAsync(cancellationToken)
+            .ConfigureAwait(false);
 
         var razorEdit = await _renameService
             .TryGetRazorRenameEditsAsync(context, positionInfo, newName, context.GetSolutionQueryOperations(), cancellationToken)

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/SemanticTokens/RemoteCSharpSemanticTokensProvider.cs

@@ -31,13 +31,17 @@ internal class RemoteCSharpSemanticTokensProvider(IFilePathService filePathServi
         // We have a razor document, lets find the generated C# document
         Debug.Assert(documentContext is RemoteDocumentContext, "This method only works on document snapshots created in the OOP process");
         var snapshot = (RemoteDocumentSnapshot)documentContext.Snapshot;
-        var generatedDocument = await snapshot.GetGeneratedDocumentAsync().ConfigureAwait(false);
+        var generatedDocument = await snapshot
+            .GetGeneratedDocumentAsync(cancellationToken)
+            .ConfigureAwait(false);
 
-        var data = await SemanticTokensRange.GetSemanticTokensAsync(
-            generatedDocument,
-            csharpRanges,
-            supportsVisualStudioExtensions: true,
-            cancellationToken).ConfigureAwait(false);
+        var data = await SemanticTokensRange
+            .GetSemanticTokensAsync(
+                generatedDocument,
+                csharpRanges,
+                supportsVisualStudioExtensions: true,
+                cancellationToken)
+            .ConfigureAwait(false);
 
         return data;
     }

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/SignatureHelp/RemoteSignatureHelpService.cs

@@ -35,7 +35,9 @@ internal sealed class RemoteSignatureHelpService(in ServiceArgs args) : RazorDoc
         var linePosition = new LinePosition(position.Line, position.Character);
         var absoluteIndex = codeDocument.Source.Text.GetRequiredAbsoluteIndex(linePosition);
 
-        var generatedDocument = await context.Snapshot.GetGeneratedDocumentAsync().ConfigureAwait(false);
+        var generatedDocument = await context.Snapshot
+            .GetGeneratedDocumentAsync(cancellationToken)
+            .ConfigureAwait(false);
 
         if (DocumentMappingService.TryMapToGeneratedDocumentPosition(codeDocument.GetCSharpDocument(), absoluteIndex, out var mappedPosition, out _))
         {

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/SpellCheck/RemoteCSharpSpellCheckRangeProvider.cs

@@ -21,7 +21,7 @@ internal sealed class RemoteCSharpSpellCheckRangeProvider() : ICSharpSpellCheckR
         // We have a razor document, lets find the generated C# document
         Debug.Assert(documentContext is RemoteDocumentContext, "This method only works on document snapshots created in the OOP process");
         var snapshot = (RemoteDocumentSnapshot)documentContext.Snapshot;
-        var generatedDocument = await snapshot.GetGeneratedDocumentAsync().ConfigureAwait(false);
+        var generatedDocument = await snapshot.GetGeneratedDocumentAsync(cancellationToken).ConfigureAwait(false);
 
         var csharpRanges = await ExternalAccess.Razor.Cohost.Handlers.SpellCheck.GetSpellCheckSpansAsync(generatedDocument, cancellationToken).ConfigureAwait(false);
 

src/Razor/src/Microsoft.VisualStudio.LanguageServices.Razor/Documents/RazorCodeDocumentProvidingSnapshotChangeTrigger.cs

@@ -5,6 +5,7 @@ using System;
 using System.Collections.Generic;
 using System.ComponentModel.Composition;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.AspNetCore.Razor.ProjectSystem;
@@ -61,7 +62,7 @@ internal class RazorCodeDocumentProvidingSnapshotChangeTrigger : IRazorStartupSe
         }
     }
 
-    public async Task<RazorCodeDocument?> GetRazorCodeDocumentAsync(string filePath)
+    public async Task<RazorCodeDocument?> GetRazorCodeDocumentAsync(string filePath, CancellationToken cancellationToken)
     {
         if (!_documentProjectMap.TryGetValue(filePath, out var projectKey))
         {
@@ -79,8 +80,6 @@ internal class RazorCodeDocumentProvidingSnapshotChangeTrigger : IRazorStartupSe
             return null;
         }
 
-        var razorDocument = await document.GetGeneratedOutputAsync().ConfigureAwait(false);
-
-        return razorDocument;
+        return await document.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
     }
 }

src/Razor/src/Microsoft.VisualStudio.LanguageServices.Razor/DynamicFiles/BackgroundDocumentGenerator.cs

@@ -71,9 +71,9 @@ internal partial class BackgroundDocumentGenerator : IRazorStartupService, IDisp
     protected Task WaitUntilCurrentBatchCompletesAsync()
         => _workQueue.WaitUntilCurrentBatchCompletesAsync();
 
-    protected virtual async Task ProcessDocumentAsync(IProjectSnapshot project, IDocumentSnapshot document)
+    protected virtual async Task ProcessDocumentAsync(IProjectSnapshot project, IDocumentSnapshot document, CancellationToken cancellationToken)
     {
-        await document.GetGeneratedOutputAsync().ConfigureAwait(false);
+        await document.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
 
         UpdateFileInfo(project, document);
     }
@@ -116,7 +116,7 @@ internal partial class BackgroundDocumentGenerator : IRazorStartupService, IDisp
 
             try
             {
-                await ProcessDocumentAsync(project, document).ConfigureAwait(false);
+                await ProcessDocumentAsync(project, document, token).ConfigureAwait(false);
             }
             catch (UnauthorizedAccessException)
             {

src/Razor/src/Microsoft.VisualStudio.LanguageServices.Razor/DynamicFiles/RazorDocumentExcerptService.cs

@@ -56,7 +56,7 @@ internal class RazorDocumentExcerptService(
 
         // Then we'll classify the spans based on the primary document, since that's the coordinate
         // space that our output mappings use.
-        var output = await _document.GetGeneratedOutputAsync().ConfigureAwait(false);
+        var output = await _document.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
         var mappings = output.GetCSharpDocument().SourceMappings;
         var classifiedSpans = await ClassifyPreviewAsync(
             excerptSpan,

src/Razor/src/Microsoft.VisualStudio.LanguageServices.Razor/DynamicFiles/RazorSpanMappingService.cs

@@ -42,8 +42,8 @@ internal class RazorSpanMappingService(IDocumentSnapshot document) : IRazorSpanM
             return ImmutableArray<RazorMappedSpanResult>.Empty;
         }
 
-        var source = await _document.GetTextAsync(cancellationToken).ConfigureAwait(false);
-        var output = await _document.GetGeneratedOutputAsync().ConfigureAwait(false);
+        var output = await _document.GetGeneratedOutputAsync(cancellationToken).ConfigureAwait(false);
+        var source = output.Source.Text;
 
         var csharpDocument = output.GetCSharpDocument();
         var filePath = output.Source.FilePath.AssumeNotNull();

src/Razor/src/Microsoft.VisualStudio.RazorExtension/SyntaxVisualizer/SourceMappingTagger.cs

@@ -3,7 +3,7 @@
 
 using System;
 using System.Collections.Generic;
-using System.Linq;
+using System.Threading;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.VisualStudio.Razor.Documents;
 using Microsoft.VisualStudio.Shell;
@@ -39,21 +39,22 @@ internal sealed class SourceMappingTagger : ITagger<SourceMappingTag>
     {
         if (!Enabled || spans.Count == 0)
         {
-            return Enumerable.Empty<ITagSpan<SourceMappingTag>>();
+            return [];
         }
 
         var snapshot = spans[0].Snapshot;
 
         if (!_textDocumentFactoryService.TryGetTextDocument(_buffer, out var textDocument))
         {
-            return Enumerable.Empty<ITagSpan<SourceMappingTag>>();
+            return [];
         }
 
-        var codeDocument = ThreadHelper.JoinableTaskFactory.Run(() => _sourceMappingProjectChangeTrigger.Value.GetRazorCodeDocumentAsync(textDocument.FilePath));
+        var codeDocument = ThreadHelper.JoinableTaskFactory.Run(
+            () => _sourceMappingProjectChangeTrigger.Value.GetRazorCodeDocumentAsync(textDocument.FilePath, CancellationToken.None));
 
         if (codeDocument is null)
         {
-            return Enumerable.Empty<ITagSpan<SourceMappingTag>>();
+            return [];
         }
 
         return GetTagsWorker(codeDocument, snapshot);

src/Razor/src/Microsoft.VisualStudio.RazorExtension/SyntaxVisualizer/SyntaxVisualizerControl.xaml.cs

@@ -5,6 +5,7 @@ using System;
 using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
 using System.IO;
+using System.Threading;
 using System.Windows;
 using System.Windows.Controls;
 using Microsoft.AspNetCore.Razor.Language;
@@ -539,7 +540,8 @@ internal partial class SyntaxVisualizerControl : UserControl, IVsRunningDocTable
 
         var filePath = hostDocumentUri.GetAbsoluteOrUNCPath().Replace('/', '\\');
 
-        var codeDocument = _joinableTaskFactory.Run(() => _codeDocumentProvider.GetRazorCodeDocumentAsync(filePath));
+        var codeDocument = _joinableTaskFactory.Run(
+            () => _codeDocumentProvider.GetRazorCodeDocumentAsync(filePath, CancellationToken.None));
         if (codeDocument is null)
         {
             return null;

src/Razor/test/Microsoft.AspNetCore.Razor.LanguageServer.Test/CodeActions/CSharp/DefaultCSharpCodeActionProviderTest.cs

@@ -334,7 +334,7 @@ $$Path;
 
         var documentSnapshotMock = new StrictMock<IDocumentSnapshot>();
         documentSnapshotMock
-            .Setup(x => x.GetGeneratedOutputAsync(It.IsAny<bool>()))
+            .Setup(x => x.GetGeneratedOutputAsync(It.IsAny<bool>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(codeDocument);
         documentSnapshotMock
             .Setup(x => x.GetTextAsync(It.IsAny<CancellationToken>()))

src/Razor/test/Microsoft.AspNetCore.Razor.LanguageServer.Test/CodeActions/CSharp/TypeAccessibilityCodeActionProviderTest.cs

@@ -464,7 +464,7 @@ public class TypeAccessibilityCodeActionProviderTest(ITestOutputHelper testOutpu
 
         var documentSnapshotMock = new StrictMock<IDocumentSnapshot>();
         documentSnapshotMock
-            .Setup(x => x.GetGeneratedOutputAsync(It.IsAny<bool>()))
+            .Setup(x => x.GetGeneratedOutputAsync(It.IsAny<bool>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(codeDocument);
         documentSnapshotMock
             .Setup(x => x.GetTextAsync(It.IsAny<CancellationToken>()))

src/Razor/test/Microsoft.AspNetCore.Razor.LanguageServer.Test/CodeActions/Html/DefaultHtmlCodeActionProviderTest.cs

@@ -152,7 +152,7 @@ public class DefaultHtmlCodeActionProviderTest(ITestOutputHelper testOutput) : L
 
         var documentSnapshotMock = new StrictMock<IDocumentSnapshot>();
         documentSnapshotMock
-            .Setup(x => x.GetGeneratedOutputAsync(It.IsAny<bool>()))
+            .Setup(x => x.GetGeneratedOutputAsync(It.IsAny<bool>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(codeDocument);
         documentSnapshotMock
             .Setup(x => x.GetTextAsync(It.IsAny<CancellationToken>()))

src/Razor/test/Microsoft.AspNetCore.Razor.LanguageServer.Test/CodeActions/Razor/ComponentAccessibilityCodeActionProviderTest.cs

@@ -465,7 +465,7 @@ public class ComponentAccessibilityCodeActionProviderTest(ITestOutputHelper test
 
         var documentSnapshotMock = new StrictMock<IDocumentSnapshot>();
         documentSnapshotMock
-            .Setup(x => x.GetGeneratedOutputAsync(It.IsAny<bool>()))
+            .Setup(x => x.GetGeneratedOutputAsync(It.IsAny<bool>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(codeDocument);
         documentSnapshotMock
             .Setup(x => x.GetTextAsync(It.IsAny<CancellationToken>()))

src/Razor/test/Microsoft.AspNetCore.Razor.LanguageServer.Test/CodeActions/Razor/ExtractToCodeBehindCodeActionProviderTest.cs

@@ -399,7 +399,7 @@ public class ExtractToCodeBehindCodeActionProviderTest(ITestOutputHelper testOut
 
         var documentSnapshotMock = new StrictMock<IDocumentSnapshot>();
         documentSnapshotMock
-            .Setup(x => x.GetGeneratedOutputAsync(It.IsAny<bool>()))
+            .Setup(x => x.GetGeneratedOutputAsync(It.IsAny<bool>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(codeDocument);
         documentSnapshotMock
             .Setup(x => x.GetTextAsync(It.IsAny<CancellationToken>()))

src/Razor/test/Microsoft.AspNetCore.Razor.LanguageServer.Test/CodeDocumentReferenceHolderTest.cs

@@ -44,7 +44,7 @@ public class CodeDocumentReferenceHolderTest(ITestOutputHelper testOutput) : Lan
     {
         // Arrange
         var documentSnapshot = await CreateDocumentSnapshotAsync();
-        var codeDocumentReference = await ProcessDocumentAndRetrieveOutputAsync(documentSnapshot);
+        var codeDocumentReference = await ProcessDocumentAndRetrieveOutputAsync(documentSnapshot, DisposalToken);
 
         // Act
         PerformFullGC();
@@ -70,8 +70,8 @@ public class CodeDocumentReferenceHolderTest(ITestOutputHelper testOutput) : Lan
 
         Assert.NotNull(unrelatedDocumentSnapshot);
 
-        var mainCodeDocumentReference = await ProcessDocumentAndRetrieveOutputAsync(documentSnapshot);
-        var unrelatedCodeDocumentReference = await ProcessDocumentAndRetrieveOutputAsync(unrelatedDocumentSnapshot);
+        var mainCodeDocumentReference = await ProcessDocumentAndRetrieveOutputAsync(documentSnapshot, DisposalToken);
+        var unrelatedCodeDocumentReference = await ProcessDocumentAndRetrieveOutputAsync(unrelatedDocumentSnapshot, DisposalToken);
 
         // Act
         await _projectManager.UpdateAsync(updater =>
@@ -91,7 +91,7 @@ public class CodeDocumentReferenceHolderTest(ITestOutputHelper testOutput) : Lan
     {
         // Arrange
         var documentSnapshot = await CreateDocumentSnapshotAsync();
-        var codeDocumentReference = await ProcessDocumentAndRetrieveOutputAsync(documentSnapshot);
+        var codeDocumentReference = await ProcessDocumentAndRetrieveOutputAsync(documentSnapshot, DisposalToken);
 
         // Act
 
@@ -111,7 +111,7 @@ public class CodeDocumentReferenceHolderTest(ITestOutputHelper testOutput) : Lan
     {
         // Arrange
         var documentSnapshot = await CreateDocumentSnapshotAsync();
-        var codeDocumentReference = await ProcessDocumentAndRetrieveOutputAsync(documentSnapshot);
+        var codeDocumentReference = await ProcessDocumentAndRetrieveOutputAsync(documentSnapshot, DisposalToken);
 
         // Act
         await _projectManager.UpdateAsync(updater =>
@@ -130,7 +130,7 @@ public class CodeDocumentReferenceHolderTest(ITestOutputHelper testOutput) : Lan
     {
         // Arrange
         var documentSnapshot = await CreateDocumentSnapshotAsync();
-        var codeDocumentReference = await ProcessDocumentAndRetrieveOutputAsync(documentSnapshot);
+        var codeDocumentReference = await ProcessDocumentAndRetrieveOutputAsync(documentSnapshot, DisposalToken);
 
         // Act
         await _projectManager.UpdateAsync(updater =>
@@ -149,7 +149,7 @@ public class CodeDocumentReferenceHolderTest(ITestOutputHelper testOutput) : Lan
     {
         // Arrange
         var documentSnapshot = await CreateDocumentSnapshotAsync();
-        var codeDocumentReference = await ProcessDocumentAndRetrieveOutputAsync(documentSnapshot);
+        var codeDocumentReference = await ProcessDocumentAndRetrieveOutputAsync(documentSnapshot, DisposalToken);
 
         // Act
         await _projectManager.UpdateAsync(updater =>
@@ -176,9 +176,9 @@ public class CodeDocumentReferenceHolderTest(ITestOutputHelper testOutput) : Lan
     }
 
     [MethodImpl(MethodImplOptions.NoInlining)]
-    private async Task<WeakReference<RazorCodeDocument>> ProcessDocumentAndRetrieveOutputAsync(IDocumentSnapshot documentSnapshot)
+    private async Task<WeakReference<RazorCodeDocument>> ProcessDocumentAndRetrieveOutputAsync(IDocumentSnapshot documentSnapshot, CancellationToken cancellationToken)
     {
-        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync();
+        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync(cancellationToken);
 
         _referenceHolder.DocumentProcessed(codeDocument, documentSnapshot);
 

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

@@ -99,9 +99,9 @@ public class FormattingContentValidationPassTest(ITestOutputHelper testOutput) :
         var projectEngine = RazorProjectEngine.Create(builder => builder.SetRootNamespace("Test"));
         var codeDocument = projectEngine.ProcessDesignTime(sourceDocument, fileKind, importSources: default, tagHelpers);
 
-        var documentSnapshot = new Mock<IDocumentSnapshot>(MockBehavior.Strict);
+        var documentSnapshot = new StrictMock<IDocumentSnapshot>();
         documentSnapshot
-            .Setup(d => d.GetGeneratedOutputAsync(It.IsAny<bool>()))
+            .Setup(d => d.GetGeneratedOutputAsync(It.IsAny<bool>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(codeDocument);
         documentSnapshot
             .Setup(d => d.TargetPath)

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

@@ -316,7 +316,7 @@ public class FormattingTestBase : RazorToolingIntegrationTestBase
     {
         var documentSnapshot = new StrictMock<IDocumentSnapshot>();
         documentSnapshot
-            .Setup(d => d.GetGeneratedOutputAsync(It.IsAny<bool>()))
+            .Setup(d => d.GetGeneratedOutputAsync(It.IsAny<bool>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(codeDocument);
         documentSnapshot
             .Setup(d => d.FilePath)

src/Razor/test/Microsoft.AspNetCore.Razor.LanguageServer.Test/Hover/HoverServiceTest.cs

@@ -933,7 +933,7 @@ public class HoverServiceTest(ITestOutputHelper testOutput) : TagHelperServiceTe
 
         var documentSnapshotMock = new StrictMock<IDocumentSnapshot>();
         documentSnapshotMock
-            .Setup(x => x.GetGeneratedOutputAsync(It.IsAny<bool>()))
+            .Setup(x => x.GetGeneratedOutputAsync(It.IsAny<bool>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(codeDocument);
         documentSnapshotMock
             .Setup(x => x.GetTextAsync(It.IsAny<CancellationToken>()))

src/Razor/test/Microsoft.AspNetCore.Razor.LanguageServer.Test/RazorComponentSearchEngineTest.cs

@@ -101,8 +101,8 @@ public class RazorComponentSearchEngineTest(ITestOutputHelper testOutput) : Lang
         var searchEngine = new RazorComponentSearchEngine(LoggerFactory);
 
         // Act
-        var documentSnapshot1 = await searchEngine.TryLocateComponentAsync(tagHelperDescriptor1, _projectManager.GetQueryOperations());
-        var documentSnapshot2 = await searchEngine.TryLocateComponentAsync(tagHelperDescriptor2, _projectManager.GetQueryOperations());
+        var documentSnapshot1 = await searchEngine.TryLocateComponentAsync(tagHelperDescriptor1, _projectManager.GetQueryOperations(), DisposalToken);
+        var documentSnapshot2 = await searchEngine.TryLocateComponentAsync(tagHelperDescriptor2, _projectManager.GetQueryOperations(), DisposalToken);
 
         // Assert
         Assert.NotNull(documentSnapshot1);
@@ -120,8 +120,8 @@ public class RazorComponentSearchEngineTest(ITestOutputHelper testOutput) : Lang
         var searchEngine = new RazorComponentSearchEngine(LoggerFactory);
 
         // Act
-        var documentSnapshot1 = await searchEngine.TryLocateComponentAsync(tagHelperDescriptor1, _projectManager.GetQueryOperations());
-        var documentSnapshot2 = await searchEngine.TryLocateComponentAsync(tagHelperDescriptor2, _projectManager.GetQueryOperations());
+        var documentSnapshot1 = await searchEngine.TryLocateComponentAsync(tagHelperDescriptor1, _projectManager.GetQueryOperations(), DisposalToken);
+        var documentSnapshot2 = await searchEngine.TryLocateComponentAsync(tagHelperDescriptor2, _projectManager.GetQueryOperations(), DisposalToken);
 
         // Assert
         Assert.NotNull(documentSnapshot1);
@@ -138,7 +138,7 @@ public class RazorComponentSearchEngineTest(ITestOutputHelper testOutput) : Lang
         var searchEngine = new RazorComponentSearchEngine(LoggerFactory);
 
         // Act
-        var documentSnapshot = await searchEngine.TryLocateComponentAsync(tagHelperDescriptor, _projectManager.GetQueryOperations());
+        var documentSnapshot = await searchEngine.TryLocateComponentAsync(tagHelperDescriptor, _projectManager.GetQueryOperations(), DisposalToken);
 
         // Assert
         Assert.NotNull(documentSnapshot);
@@ -153,7 +153,7 @@ public class RazorComponentSearchEngineTest(ITestOutputHelper testOutput) : Lang
         var searchEngine = new RazorComponentSearchEngine(LoggerFactory);
 
         // Act
-        var documentSnapshot = await searchEngine.TryLocateComponentAsync(tagHelperDescriptor, _projectManager.GetQueryOperations());
+        var documentSnapshot = await searchEngine.TryLocateComponentAsync(tagHelperDescriptor, _projectManager.GetQueryOperations(), DisposalToken);
 
         // Assert
         Assert.Null(documentSnapshot);
@@ -167,7 +167,7 @@ public class RazorComponentSearchEngineTest(ITestOutputHelper testOutput) : Lang
         var searchEngine = new RazorComponentSearchEngine(LoggerFactory);
 
         // Act
-        var documentSnapshot = await searchEngine.TryLocateComponentAsync(tagHelperDescriptor, _projectManager.GetQueryOperations());
+        var documentSnapshot = await searchEngine.TryLocateComponentAsync(tagHelperDescriptor, _projectManager.GetQueryOperations(), DisposalToken);
 
         // Assert
         Assert.Null(documentSnapshot);
@@ -181,7 +181,7 @@ public class RazorComponentSearchEngineTest(ITestOutputHelper testOutput) : Lang
         var searchEngine = new RazorComponentSearchEngine(LoggerFactory);
 
         // Act
-        var documentSnapshot = await searchEngine.TryLocateComponentAsync(tagHelperDescriptor, _projectManager.GetQueryOperations());
+        var documentSnapshot = await searchEngine.TryLocateComponentAsync(tagHelperDescriptor, _projectManager.GetQueryOperations(), DisposalToken);
 
         // Assert
         Assert.Null(documentSnapshot);
@@ -195,7 +195,7 @@ public class RazorComponentSearchEngineTest(ITestOutputHelper testOutput) : Lang
         var searchEngine = new RazorComponentSearchEngine(LoggerFactory);
 
         // Act
-        var documentSnapshot = await searchEngine.TryLocateComponentAsync(tagHelperDescriptor, _projectManager.GetQueryOperations());
+        var documentSnapshot = await searchEngine.TryLocateComponentAsync(tagHelperDescriptor, _projectManager.GetQueryOperations(), DisposalToken);
 
         // Assert
         Assert.NotNull(documentSnapshot);

src/Razor/test/Microsoft.AspNetCore.Razor.LanguageServer.Test/Semantic/SemanticTokensTest.cs

@@ -955,7 +955,7 @@ public partial class SemanticTokensTest(ITestOutputHelper testOutput) : TagHelpe
             .SetupGet(x => x.Project)
             .Returns(projectSnapshot.Object);
         documentSnapshotMock
-            .Setup(x => x.GetGeneratedOutputAsync(It.IsAny<bool>()))
+            .Setup(x => x.GetGeneratedOutputAsync(It.IsAny<bool>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(document);
         documentSnapshotMock
             .Setup(x => x.GetTextAsync(It.IsAny<CancellationToken>()))

src/Razor/test/Microsoft.AspNetCore.Razor.Test.Common.Tooling/ProjectSystem/TestDocumentSnapshot.cs

@@ -67,10 +67,12 @@ internal sealed class TestDocumentSnapshot : IDocumentSnapshot
     public IProjectSnapshot Project => RealSnapshot.Project;
     public int Version => RealSnapshot.Version;
 
-    public ValueTask<RazorCodeDocument> GetGeneratedOutputAsync(bool forceDesignTimeGeneratedOutput)
+    public ValueTask<RazorCodeDocument> GetGeneratedOutputAsync(
+        bool forceDesignTimeGeneratedOutput,
+        CancellationToken cancellationToken)
     {
         return _codeDocument is null
-            ? RealSnapshot.GetGeneratedOutputAsync(forceDesignTimeGeneratedOutput)
+            ? RealSnapshot.GetGeneratedOutputAsync(forceDesignTimeGeneratedOutput, cancellationToken)
             : new(_codeDocument);
     }
 

src/Razor/test/Microsoft.AspNetCore.Razor.Test.Common.Tooling/Workspaces/DocumentExcerptServiceTestBase.cs

@@ -71,7 +71,7 @@ public abstract class DocumentExcerptServiceTestBase(ITestOutputHelper testOutpu
     // that appears in the primary buffer.
     private static async Task<TextSpan> GetSecondarySpanAsync(IDocumentSnapshot primary, TextSpan primarySpan, Document secondary, CancellationToken cancellationToken)
     {
-        var output = await primary.GetGeneratedOutputAsync();
+        var output = await primary.GetGeneratedOutputAsync(cancellationToken);
 
         foreach (var mapping in output.GetCSharpDocument().SourceMappings)
         {

src/Razor/test/Microsoft.CodeAnalysis.Razor.Workspaces.Test/ProjectSystem/DefaultDocumentSnapshotTest.cs

@@ -55,7 +55,7 @@ public class DefaultDocumentSnapshotTest : WorkspaceTestBase
     public async Task GCCollect_OutputIsNoLongerCached()
     {
         // Arrange
-        await Task.Run(async () => { await _legacyDocument.GetGeneratedOutputAsync(); });
+        await Task.Run(async () => { await _legacyDocument.GetGeneratedOutputAsync(DisposalToken); });
 
         // Act
 
@@ -70,7 +70,7 @@ public class DefaultDocumentSnapshotTest : WorkspaceTestBase
     public async Task RegeneratingWithReference_CachesOutput()
     {
         // Arrange
-        var output = await _legacyDocument.GetGeneratedOutputAsync();
+        var output = await _legacyDocument.GetGeneratedOutputAsync(DisposalToken);
 
         // Mostly doing this to ensure "var output" doesn't get optimized out
         Assert.NotNull(output);
@@ -86,7 +86,7 @@ public class DefaultDocumentSnapshotTest : WorkspaceTestBase
     public async Task GetGeneratedOutputAsync_CshtmlComponent_ContainsComponentImports()
     {
         // Act
-        var codeDocument = await _componentCshtmlDocument.GetGeneratedOutputAsync();
+        var codeDocument = await _componentCshtmlDocument.GetGeneratedOutputAsync(DisposalToken);
 
         // Assert
         Assert.Contains("using global::Microsoft.AspNetCore.Components", codeDocument.GetCSharpSourceText().ToString(), StringComparison.Ordinal);
@@ -96,7 +96,7 @@ public class DefaultDocumentSnapshotTest : WorkspaceTestBase
     public async Task GetGeneratedOutputAsync_Component()
     {
         // Act
-        var codeDocument = await _componentDocument.GetGeneratedOutputAsync();
+        var codeDocument = await _componentDocument.GetGeneratedOutputAsync(DisposalToken);
 
         // Assert
         Assert.Contains("ComponentBase", codeDocument.GetCSharpSourceText().ToString(), StringComparison.Ordinal);
@@ -106,7 +106,7 @@ public class DefaultDocumentSnapshotTest : WorkspaceTestBase
     public async Task GetGeneratedOutputAsync_NestedComponentDocument_SetsCorrectNamespaceAndClassName()
     {
         // Act
-        var codeDocument = await _nestedComponentDocument.GetGeneratedOutputAsync();
+        var codeDocument = await _nestedComponentDocument.GetGeneratedOutputAsync(DisposalToken);
 
         // Assert
         Assert.Contains("ComponentBase", codeDocument.GetCSharpSourceText().ToString(), StringComparison.Ordinal);
@@ -120,7 +120,7 @@ public class DefaultDocumentSnapshotTest : WorkspaceTestBase
     public async Task GetGeneratedOutputAsync_Legacy()
     {
         // Act
-        var codeDocument = await _legacyDocument.GetGeneratedOutputAsync();
+        var codeDocument = await _legacyDocument.GetGeneratedOutputAsync(DisposalToken);
 
         // Assert
         Assert.Contains("Template", codeDocument.GetCSharpSourceText().ToString(), StringComparison.Ordinal);

src/Razor/test/Microsoft.CodeAnalysis.Razor.Workspaces.Test/ProjectSystem/ProjectStateGeneratedOutputTest.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
 using System.Collections.Immutable;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.AspNetCore.Razor.ProjectSystem;
@@ -45,13 +46,13 @@ public class ProjectStateGeneratedOutputTest : WorkspaceTestBase
             ProjectState.Create(ProjectEngineFactoryProvider, _hostProject, ProjectWorkspaceState.Default)
             .WithAddedHostDocument(_hostDocument, DocumentState.EmptyLoader);
 
-        var (originalOutput, originalInputVersion) = await GetOutputAsync(original, _hostDocument);
+        var (originalOutput, originalInputVersion) = await GetOutputAsync(original, _hostDocument, DisposalToken);
 
         // Act
         var state = original.WithAddedHostDocument(TestProjectData.AnotherProjectFile1, DocumentState.EmptyLoader);
 
         // Assert
-        var (actualOutput, actualInputVersion) = await GetOutputAsync(state, _hostDocument);
+        var (actualOutput, actualInputVersion) = await GetOutputAsync(state, _hostDocument, DisposalToken);
         Assert.Same(originalOutput, actualOutput);
         Assert.Equal(originalInputVersion, actualInputVersion);
     }
@@ -64,13 +65,13 @@ public class ProjectStateGeneratedOutputTest : WorkspaceTestBase
             ProjectState.Create(ProjectEngineFactoryProvider, _hostProject, ProjectWorkspaceState.Default)
             .WithAddedHostDocument(_hostDocument, DocumentState.EmptyLoader);
 
-        var (originalOutput, originalInputVersion) = await GetOutputAsync(original, _hostDocument);
+        var (originalOutput, originalInputVersion) = await GetOutputAsync(original, _hostDocument, DisposalToken);
 
         // Act
         var state = original.WithAddedHostDocument(TestProjectData.SomeProjectImportFile, DocumentState.EmptyLoader);
 
         // Assert
-        var (actualOutput, actualInputVersion) = await GetOutputAsync(state, _hostDocument);
+        var (actualOutput, actualInputVersion) = await GetOutputAsync(state, _hostDocument, DisposalToken);
         Assert.NotSame(originalOutput, actualOutput);
         Assert.NotEqual(originalInputVersion, actualInputVersion);
         Assert.Equal(state.DocumentCollectionVersion, actualInputVersion);
@@ -85,14 +86,14 @@ public class ProjectStateGeneratedOutputTest : WorkspaceTestBase
             .WithAddedHostDocument(_hostDocument, DocumentState.EmptyLoader)
             .WithAddedHostDocument(TestProjectData.SomeProjectImportFile, DocumentState.EmptyLoader);
 
-        var (originalOutput, originalInputVersion) = await GetOutputAsync(original, _hostDocument);
+        var (originalOutput, originalInputVersion) = await GetOutputAsync(original, _hostDocument, DisposalToken);
 
         // Act
         var version = VersionStamp.Create();
         var state = original.WithChangedHostDocument(_hostDocument, TestMocks.CreateTextLoader("@using System", version));
 
         // Assert
-        var (actualOutput, actualInputVersion) = await GetOutputAsync(state, _hostDocument);
+        var (actualOutput, actualInputVersion) = await GetOutputAsync(state, _hostDocument, DisposalToken);
         Assert.NotSame(originalOutput, actualOutput);
         Assert.NotEqual(originalInputVersion, actualInputVersion);
         Assert.Equal(version, actualInputVersion);
@@ -107,14 +108,14 @@ public class ProjectStateGeneratedOutputTest : WorkspaceTestBase
             .WithAddedHostDocument(_hostDocument, DocumentState.EmptyLoader)
             .WithAddedHostDocument(TestProjectData.SomeProjectImportFile, DocumentState.EmptyLoader);
 
-        var (originalOutput, originalInputVersion) = await GetOutputAsync(original, _hostDocument);
+        var (originalOutput, originalInputVersion) = await GetOutputAsync(original, _hostDocument, DisposalToken);
 
         // Act
         var version = VersionStamp.Create();
         var state = original.WithChangedHostDocument(TestProjectData.SomeProjectImportFile, TestMocks.CreateTextLoader("@using System", version));
 
         // Assert
-        var (actualOutput, actualInputVersion) = await GetOutputAsync(state, _hostDocument);
+        var (actualOutput, actualInputVersion) = await GetOutputAsync(state, _hostDocument, DisposalToken);
         Assert.NotSame(originalOutput, actualOutput);
         Assert.NotEqual(originalInputVersion, actualInputVersion);
         Assert.Equal(version, actualInputVersion);
@@ -129,13 +130,13 @@ public class ProjectStateGeneratedOutputTest : WorkspaceTestBase
             .WithAddedHostDocument(_hostDocument, DocumentState.EmptyLoader)
             .WithAddedHostDocument(TestProjectData.SomeProjectImportFile, DocumentState.EmptyLoader);
 
-        var (originalOutput, originalInputVersion) = await GetOutputAsync(original, _hostDocument);
+        var (originalOutput, originalInputVersion) = await GetOutputAsync(original, _hostDocument, DisposalToken);
 
         // Act
         var state = original.WithRemovedHostDocument(TestProjectData.SomeProjectImportFile);
 
         // Assert
-        var (actualOutput, actualInputVersion) = await GetOutputAsync(state, _hostDocument);
+        var (actualOutput, actualInputVersion) = await GetOutputAsync(state, _hostDocument, DisposalToken);
         Assert.NotSame(originalOutput, actualOutput);
         Assert.NotEqual(originalInputVersion, actualInputVersion);
         Assert.Equal(state.DocumentCollectionVersion, actualInputVersion);
@@ -149,14 +150,14 @@ public class ProjectStateGeneratedOutputTest : WorkspaceTestBase
             ProjectState.Create(ProjectEngineFactoryProvider, _hostProject, ProjectWorkspaceState.Default)
             .WithAddedHostDocument(_hostDocument, DocumentState.EmptyLoader);
 
-        var (originalOutput, originalInputVersion) = await GetOutputAsync(original, _hostDocument);
+        var (originalOutput, originalInputVersion) = await GetOutputAsync(original, _hostDocument, DisposalToken);
         var changed = ProjectWorkspaceState.Default;
 
         // Act
         var state = original.WithProjectWorkspaceState(changed);
 
         // Assert
-        var (actualOutput, actualInputVersion) = await GetOutputAsync(state, _hostDocument);
+        var (actualOutput, actualInputVersion) = await GetOutputAsync(state, _hostDocument, DisposalToken);
         Assert.Same(originalOutput, actualOutput);
         Assert.Equal(originalInputVersion, actualInputVersion);
     }
@@ -170,14 +171,14 @@ public class ProjectStateGeneratedOutputTest : WorkspaceTestBase
             ProjectState.Create(ProjectEngineFactoryProvider, _hostProject, ProjectWorkspaceState.Default)
             .WithAddedHostDocument(_hostDocument, DocumentState.EmptyLoader);
 
-        var (originalOutput, originalInputVersion) = await GetOutputAsync(original, _hostDocument);
+        var (originalOutput, originalInputVersion) = await GetOutputAsync(original, _hostDocument, DisposalToken);
         var changed = ProjectWorkspaceState.Create(_someTagHelpers);
 
         // Act
         var state = original.WithProjectWorkspaceState(changed);
 
         // Assert
-        var (actualOutput, actualInputVersion) = await GetOutputAsync(state, _hostDocument);
+        var (actualOutput, actualInputVersion) = await GetOutputAsync(state, _hostDocument, DisposalToken);
         Assert.NotSame(originalOutput, actualOutput);
         Assert.NotEqual(originalInputVersion, actualInputVersion);
         Assert.Equal(state.ProjectWorkspaceStateVersion, actualInputVersion);
@@ -195,13 +196,13 @@ public class ProjectStateGeneratedOutputTest : WorkspaceTestBase
             .WithAddedHostDocument(_hostDocument, TestMocks.CreateTextLoader("@DateTime.Now", VersionStamp.Default));
         var changedWorkspaceState = ProjectWorkspaceState.Create(_someTagHelpers, LanguageVersion.CSharp8);
 
-        var (originalOutput, originalInputVersion) = await GetOutputAsync(original, _hostDocument);
+        var (originalOutput, originalInputVersion) = await GetOutputAsync(original, _hostDocument, DisposalToken);
 
         // Act
         var state = original.WithProjectWorkspaceState(changedWorkspaceState);
 
         // Assert
-        var (actualOutput, actualInputVersion) = await GetOutputAsync(state, _hostDocument);
+        var (actualOutput, actualInputVersion) = await GetOutputAsync(state, _hostDocument, DisposalToken);
         Assert.NotSame(originalOutput, actualOutput);
         Assert.NotEqual(originalInputVersion, actualInputVersion);
         Assert.Equal(state.ProjectWorkspaceStateVersion, actualInputVersion);
@@ -215,29 +216,28 @@ public class ProjectStateGeneratedOutputTest : WorkspaceTestBase
             ProjectState.Create(ProjectEngineFactoryProvider, _hostProject, ProjectWorkspaceState.Default)
             .WithAddedHostDocument(_hostDocument, DocumentState.EmptyLoader);
 
-        var (originalOutput, originalInputVersion) = await GetOutputAsync(original, _hostDocument);
+        var (originalOutput, originalInputVersion) = await GetOutputAsync(original, _hostDocument, DisposalToken);
 
         // Act
         var state = original.WithHostProject(_hostProjectWithConfigurationChange);
 
         // Assert
-        var (actualOutput, actualInputVersion) = await GetOutputAsync(state, _hostDocument);
+        var (actualOutput, actualInputVersion) = await GetOutputAsync(state, _hostDocument, DisposalToken);
         Assert.NotSame(originalOutput, actualOutput);
         Assert.NotEqual(originalInputVersion, actualInputVersion);
         Assert.NotEqual(state.ProjectWorkspaceStateVersion, actualInputVersion);
     }
 
-    private static Task<(RazorCodeDocument, VersionStamp)> GetOutputAsync(ProjectState project, HostDocument hostDocument)
+    private static Task<(RazorCodeDocument, VersionStamp)> GetOutputAsync(ProjectState project, HostDocument hostDocument, CancellationToken cancellationToken)
     {
         var document = project.Documents[hostDocument.FilePath];
-        return GetOutputAsync(project, document);
+        return GetOutputAsync(project, document, cancellationToken);
     }
 
-    private static Task<(RazorCodeDocument, VersionStamp)> GetOutputAsync(ProjectState project, DocumentState document)
+    private static Task<(RazorCodeDocument, VersionStamp)> GetOutputAsync(ProjectState project, DocumentState document, CancellationToken cancellationToken)
     {
-
         var projectSnapshot = new ProjectSnapshot(project);
         var documentSnapshot = new DocumentSnapshot(projectSnapshot, document);
-        return document.GetGeneratedOutputAndVersionAsync(projectSnapshot, documentSnapshot);
+        return document.GetGeneratedOutputAndVersionAsync(projectSnapshot, documentSnapshot, cancellationToken);
     }
 }

src/Razor/test/Microsoft.VisualStudio.LanguageServices.Razor.Test/Cohost/CohostEndpointTestBase.cs

@@ -4,6 +4,7 @@
 using System;
 using System.IO;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 using Basic.Reference.Assemblies;
 using Microsoft.AspNetCore.Razor;
@@ -71,7 +72,11 @@ public abstract class CohostEndpointTestBase(ITestOutputHelper testOutputHelper)
         FeatureOptions.SetOptions(_clientInitializationOptions);
     }
 
-    protected Task<TextDocument> CreateProjectAndRazorDocumentAsync(string contents, string? fileKind = null, (string fileName, string contents)[]? additionalFiles = null, bool createSeparateRemoteAndLocalWorkspaces = false)
+    protected Task<TextDocument> CreateProjectAndRazorDocumentAsync(
+        string contents,
+        string? fileKind = null,
+        (string fileName, string contents)[]? additionalFiles = null,
+        bool createSeparateRemoteAndLocalWorkspaces = false)
     {
         // Using IsLegacy means null == component, so easier for test authors
         var isComponent = !FileKinds.IsLegacy(fileKind);
@@ -95,7 +100,14 @@ public abstract class CohostEndpointTestBase(ITestOutputHelper testOutputHelper)
             // actual solution syncing set up for testing, and don't really use a service broker, but since we also would
             // expect to never make changes to a workspace, it should be fine to simply create duplicated solutions as part
             // of test setup.
-            return CreateLocalProjectAndRazorDocumentAsync(remoteDocument.Project.Solution, projectId, projectName, documentId, documentFilePath, contents, additionalFiles);
+            return CreateLocalProjectAndRazorDocumentAsync(
+                remoteDocument.Project.Solution,
+                projectId,
+                projectName,
+                documentId,
+                documentFilePath,
+                contents,
+                additionalFiles);
         }
 
         // If we're just creating one workspace, then its the remote one and we just return the remote document
@@ -104,7 +116,14 @@ public abstract class CohostEndpointTestBase(ITestOutputHelper testOutputHelper)
         return Task.FromResult(remoteDocument);
     }
 
-    private async Task<TextDocument> CreateLocalProjectAndRazorDocumentAsync(Solution remoteSolution, ProjectId projectId, string projectName, DocumentId documentId, string documentFilePath, string contents, (string fileName, string contents)[]? additionalFiles)
+    private async Task<TextDocument> CreateLocalProjectAndRazorDocumentAsync(
+        Solution remoteSolution,
+        ProjectId projectId,
+        string projectName,
+        DocumentId documentId,
+        string documentFilePath,
+        string contents,
+        (string fileName, string contents)[]? additionalFiles)
     {
         var exportProvider = TestComposition.Roslyn.ExportProviderFactory.CreateExportProvider();
         AddDisposable(exportProvider);
@@ -120,7 +139,7 @@ public abstract class CohostEndpointTestBase(ITestOutputHelper testOutputHelper)
         var snapshotManager = _exportProvider.AssumeNotNull().GetExportedValue<RemoteSnapshotManager>();
         var snapshot = snapshotManager.GetSnapshot(razorDocument);
         // Compile the Razor file
-        var codeDocument = await snapshot.GetGeneratedOutputAsync(false);
+        var codeDocument = await snapshot.GetGeneratedOutputAsync(forceDesignTimeGeneratedOutput: false, DisposalToken);
         // Update the generated doc contents
         var generatedDocumentIds = solution.GetDocumentIdsWithFilePath(documentFilePath + CSharpVirtualDocumentSuffix);
         solution = solution.WithDocumentText(generatedDocumentIds, codeDocument.GetCSharpSourceText());

src/Razor/test/Microsoft.VisualStudio.LanguageServices.Razor.Test/Cohost/RazorComponentDefinitionServiceTest.cs

@@ -83,7 +83,7 @@ public class RazorComponentDefinitionServiceTest(ITestOutputHelper testOutputHel
         var documentMappingService = OOPExportProvider.GetExportedValue<IDocumentMappingService>();
 
         var documentSnapshot = snapshotManager.GetSnapshot(document);
-        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync();
+        var codeDocument = await documentSnapshot.GetGeneratedOutputAsync(DisposalToken);
         var positionInfo = documentMappingService.GetPositionInfo(codeDocument, input.Position);
 
         var location = await service.GetDefinitionAsync(

src/Razor/test/Microsoft.VisualStudio.LanguageServices.Razor.Test/DocumentGenerator/BackgroundDocumentGeneratorTest.cs

@@ -408,12 +408,12 @@ public class BackgroundDocumentGeneratorTest(ITestOutputHelper testOutput) : Vis
             base.Enqueue(project, document);
         }
 
-        protected override Task ProcessDocumentAsync(IProjectSnapshot project, IDocumentSnapshot document)
+        protected override Task ProcessDocumentAsync(IProjectSnapshot project, IDocumentSnapshot document, CancellationToken cancellationToken)
         {
             var key = GetKey(project, document);
             PendingWork.Remove(key);
 
-            var task = base.ProcessDocumentAsync(project, document);
+            var task = base.ProcessDocumentAsync(project, document, cancellationToken);
 
             CompletedWork.Add(key);
 

src/Razor/test/Microsoft.VisualStudio.LanguageServices.Razor.Test/LanguageClient/DynamicFiles/RazorSpanMappingServiceTest.cs

@@ -36,7 +36,7 @@ public class RazorSpanMappingServiceTest(ITestOutputHelper testOutput) : Workspa
         var document = project.GetDocument(_hostDocument.FilePath);
         Assert.NotNull(document);
 
-        var output = await document.GetGeneratedOutputAsync();
+        var output = await document.GetGeneratedOutputAsync(DisposalToken);
         var generated = output.GetCSharpDocument();
 
         var symbol = "SomeProperty";
@@ -68,7 +68,7 @@ public class RazorSpanMappingServiceTest(ITestOutputHelper testOutput) : Workspa
         var document = project.GetDocument(_hostDocument.FilePath);
         Assert.NotNull(document);
 
-        var output = await document.GetGeneratedOutputAsync();
+        var output = await document.GetGeneratedOutputAsync(DisposalToken);
         var generated = output.GetCSharpDocument();
 
         var symbol = "SomeProperty";
@@ -101,7 +101,7 @@ public class RazorSpanMappingServiceTest(ITestOutputHelper testOutput) : Workspa
         var document = project.GetDocument(_hostDocument.FilePath);
         Assert.NotNull(document);
 
-        var output = await document.GetGeneratedOutputAsync();
+        var output = await document.GetGeneratedOutputAsync(DisposalToken);
         var generated = output.GetCSharpDocument();
 
         var symbol = "SomeProperty";
@@ -133,7 +133,7 @@ public class RazorSpanMappingServiceTest(ITestOutputHelper testOutput) : Workspa
         var document = project.GetDocument(_hostDocument.FilePath);
         Assert.NotNull(document);
 
-        var output = await document.GetGeneratedOutputAsync();
+        var output = await document.GetGeneratedOutputAsync(DisposalToken);
         var generated = output.GetCSharpDocument();
 
         var symbol = "ExecuteAsync";

src/Razor/test/Microsoft.VisualStudio.LanguageServices.Razor.Test/LanguageClient/RazorDocumentOptionsServiceTest.cs

@@ -101,6 +101,7 @@ public class RazorDocumentOptionsServiceTest(ITestOutputHelper testOutput) : Wor
             .WithAddedHostDocument(hostDocument, TestMocks.CreateTextLoader(sourceText, VersionStamp.Create())));
 
         var documentSnapshot = project.GetDocument(hostDocument.FilePath);
+        Assert.NotNull(documentSnapshot);
 
         var solution = Workspace.CurrentSolution.AddProject(ProjectInfo.Create(
             ProjectId.CreateNewId(Path.GetFileNameWithoutExtension(hostDocument.FilePath)),