Pass a CancellationToken to GetTextAsync and GetTextVersionAsync

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

@@ -58,7 +58,7 @@ public class RazorCSharpFormattingBenchmark : RazorLanguageServerBenchmarkBase
 
         DocumentUri = new Uri(_filePath);
         DocumentSnapshot = await GetDocumentSnapshotAsync(projectFilePath, _filePath, targetPath);
-        DocumentText = await DocumentSnapshot.GetTextAsync();
+        DocumentText = await DocumentSnapshot.GetTextAsync(CancellationToken.None);
     }
 
     private static void WriteSampleFormattingFile(string filePath, bool preformatted, int blocks)

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

@@ -77,7 +77,7 @@ public class RazorCodeActionsBenchmark : RazorLanguageServerBenchmarkBase
 
         DocumentUri = new Uri(_filePath);
         DocumentSnapshot = await GetDocumentSnapshotAsync(projectFilePath, _filePath, targetPath);
-        DocumentText = await DocumentSnapshot.GetTextAsync();
+        DocumentText = await DocumentSnapshot.GetTextAsync(CancellationToken.None);
 
         RazorCodeActionRange = DocumentText.GetZeroWidthRange(razorCodeActionIndex);
         CSharpCodeActionRange = DocumentText.GetZeroWidthRange(csharpCodeActionIndex);

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

@@ -75,7 +75,7 @@ public class RazorCompletionBenchmark : RazorLanguageServerBenchmarkBase
 
         DocumentUri = new Uri(_filePath);
         DocumentSnapshot = await GetDocumentSnapshotAsync(projectFilePath, _filePath, targetPath);
-        DocumentText = await DocumentSnapshot.GetTextAsync();
+        DocumentText = await DocumentSnapshot.GetTextAsync(CancellationToken.None);
 
         RazorPosition = DocumentText.GetPosition(razorCodeActionIndex);
 

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

@@ -51,7 +51,7 @@ public class RazorSemanticTokensScrollingBenchmark : RazorLanguageServerBenchmar
         var documentSnapshot = await GetDocumentSnapshotAsync(ProjectFilePath, filePath, TargetPath);
         DocumentContext = new DocumentContext(documentUri, documentSnapshot, projectContext: null);
 
-        var text = await DocumentSnapshot.GetTextAsync().ConfigureAwait(false);
+        var text = await DocumentSnapshot.GetTextAsync(CancellationToken.None).ConfigureAwait(false);
         Range = VsLspFactory.CreateRange(
             start: (0, 0),
             end: (text.Lines.Count - 1, 0));

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

@@ -65,7 +65,7 @@ internal sealed class CodeActionEndpoint(
             return null;
         }
 
-        var razorCodeActionContext = await GenerateRazorCodeActionContextAsync(request, documentContext.Snapshot).ConfigureAwait(false);
+        var razorCodeActionContext = await GenerateRazorCodeActionContextAsync(request, documentContext.Snapshot, cancellationToken).ConfigureAwait(false);
         if (razorCodeActionContext is null)
         {
             return null;
@@ -134,7 +134,10 @@ internal sealed class CodeActionEndpoint(
     }
 
     // internal for testing
-    internal async Task<RazorCodeActionContext?> GenerateRazorCodeActionContextAsync(VSCodeActionParams request, IDocumentSnapshot documentSnapshot)
+    internal async Task<RazorCodeActionContext?> GenerateRazorCodeActionContextAsync(
+        VSCodeActionParams request,
+        IDocumentSnapshot documentSnapshot,
+        CancellationToken cancellationToken)
     {
         var codeDocument = await documentSnapshot.GetGeneratedOutputAsync().ConfigureAwait(false);
         if (codeDocument.IsUnsupported())
@@ -142,7 +145,7 @@ internal sealed class CodeActionEndpoint(
             return null;
         }
 
-        var sourceText = await documentSnapshot.GetTextAsync().ConfigureAwait(false);
+        var sourceText = await documentSnapshot.GetTextAsync(cancellationToken).ConfigureAwait(false);
 
         // VS Provides `CodeActionParams.Context.SelectionRange` in addition to
         // `CodeActionParams.Range`. The `SelectionRange` is relative to where the

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

@@ -9,23 +9,13 @@ using Microsoft.CodeAnalysis.Razor.ProjectSystem;
 
 namespace Microsoft.AspNetCore.Razor.LanguageServer;
 
-internal class DocumentSnapshotTextLoader : TextLoader
+internal class DocumentSnapshotTextLoader(IDocumentSnapshot documentSnapshot) : TextLoader
 {
-    private readonly IDocumentSnapshot _documentSnapshot;
-
-    public DocumentSnapshotTextLoader(IDocumentSnapshot documentSnapshot)
-    {
-        if (documentSnapshot is null)
-        {
-            throw new ArgumentNullException(nameof(documentSnapshot));
-        }
-
-        _documentSnapshot = documentSnapshot;
-    }
+    private readonly IDocumentSnapshot _documentSnapshot = documentSnapshot;
 
     public override async Task<TextAndVersion> LoadTextAndVersionAsync(LoadTextOptions options, CancellationToken cancellationToken)
     {
-        var sourceText = await _documentSnapshot.GetTextAsync().ConfigureAwait(false);
+        var sourceText = await _documentSnapshot.GetTextAsync(cancellationToken).ConfigureAwait(false);
         var textAndVersion = TextAndVersion.Create(sourceText, VersionStamp.Default);
 
         return textAndVersion;

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

@@ -47,7 +47,7 @@ internal sealed class HtmlFormatter(
             return [];
         }
 
-        var sourceText = await documentSnapshot.GetTextAsync().ConfigureAwait(false);
+        var sourceText = await documentSnapshot.GetTextAsync(cancellationToken).ConfigureAwait(false);
         return result.Edits.SelectAsArray(sourceText.GetTextChange);
     }
 
@@ -78,7 +78,7 @@ internal sealed class HtmlFormatter(
             return [];
         }
 
-        var sourceText = await documentSnapshot.GetTextAsync().ConfigureAwait(false);
+        var sourceText = await documentSnapshot.GetTextAsync(cancellationToken).ConfigureAwait(false);
         return result.Edits.SelectAsArray(sourceText.GetTextChange);
     }
 

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

@@ -228,7 +228,7 @@ internal sealed class MapCodeEndpoint(
                     razorNodesToMap.Add(nodeToMap);
                 }
 
-                var sourceText = await documentContext.Snapshot.GetTextAsync().ConfigureAwait(false);
+                var sourceText = await documentContext.Snapshot.GetTextAsync(cancellationToken).ConfigureAwait(false);
 
                 foreach (var nodeToMap in razorNodesToMap)
                 {

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

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

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

@@ -70,9 +70,7 @@ internal class DocumentContext(Uri uri, IDocumentSnapshot snapshot, VSProjectCon
 
         async ValueTask<SourceText> GetSourceTextCoreAsync(CancellationToken cancellationToken)
         {
-            var sourceText = await Snapshot.GetTextAsync().ConfigureAwait(false);
-
-            cancellationToken.ThrowIfCancellationRequested();
+            var sourceText = await Snapshot.GetTextAsync(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

@@ -23,11 +23,11 @@ internal sealed class DocumentSnapshot(ProjectSnapshot project, DocumentState st
     public IProjectSnapshot Project => _project;
     public int Version => _state.Version;
 
-    public ValueTask<SourceText> GetTextAsync()
-        => _state.GetTextAsync();
+    public ValueTask<SourceText> GetTextAsync(CancellationToken cancellationToken)
+        => _state.GetTextAsync(cancellationToken);
 
-    public ValueTask<VersionStamp> GetTextVersionAsync()
-        => _state.GetTextVersionAsync();
+    public ValueTask<VersionStamp> GetTextVersionAsync(CancellationToken cancellationToken)
+        => _state.GetTextVersionAsync(cancellationToken);
 
     public bool TryGetText([NotNullWhen(true)] out SourceText? result)
         => _state.TryGetText(out result);

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

@@ -165,7 +165,7 @@ internal partial class DocumentState
             var projectWorkspaceStateVersion = project.ProjectWorkspaceStateVersion;
             var documentCollectionVersion = project.DocumentCollectionVersion;
             var imports = await GetImportsAsync(document, project.GetProjectEngine()).ConfigureAwait(false);
-            var documentVersion = await document.GetTextVersionAsync().ConfigureAwait(false);
+            var documentVersion = await document.GetTextVersionAsync(CancellationToken.None).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.

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

@@ -90,11 +90,11 @@ internal partial class DocumentState
         }
     }
 
-    public ValueTask<SourceText> GetTextAsync()
+    public ValueTask<SourceText> GetTextAsync(CancellationToken cancellationToken)
     {
         return TryGetText(out var text)
             ? new(text)
-            : GetTextCoreAsync(CancellationToken.None);
+            : GetTextCoreAsync(cancellationToken);
 
         async ValueTask<SourceText> GetTextCoreAsync(CancellationToken cancellationToken)
         {
@@ -104,11 +104,11 @@ internal partial class DocumentState
         }
     }
 
-    public ValueTask<VersionStamp> GetTextVersionAsync()
+    public ValueTask<VersionStamp> GetTextVersionAsync(CancellationToken cancellationToken)
     {
         return TryGetTextVersion(out var version)
             ? new(version)
-            : GetTextVersionCoreAsync(CancellationToken.None);
+            : GetTextVersionCoreAsync(cancellationToken);
 
         async ValueTask<VersionStamp> GetTextVersionCoreAsync(CancellationToken cancellationToken)
         {
@@ -258,7 +258,7 @@ internal partial class DocumentState
 
         foreach (var snapshot in imports)
         {
-            var versionStamp = await snapshot.GetTextVersionAsync().ConfigureAwait(false);
+            var versionStamp = await snapshot.GetTextVersionAsync(CancellationToken.None).ConfigureAwait(false);
             result.Add(new ImportItem(snapshot.FilePath, versionStamp, snapshot));
         }
 
@@ -267,7 +267,7 @@ internal partial class DocumentState
 
     private static async Task<RazorSourceDocument> GetRazorSourceDocumentAsync(IDocumentSnapshot document, RazorProjectItem? projectItem)
     {
-        var sourceText = await document.GetTextAsync().ConfigureAwait(false);
+        var sourceText = await document.GetTextAsync(CancellationToken.None).ConfigureAwait(false);
         return RazorSourceDocument.Create(sourceText, RazorSourceDocumentProperties.Create(document.FilePath, projectItem?.RelativePhysicalPath));
     }
 }

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

@@ -18,8 +18,8 @@ internal interface IDocumentSnapshot
 
     int Version { get; }
 
-    ValueTask<SourceText> GetTextAsync();
-    ValueTask<VersionStamp> GetTextVersionAsync();
+    ValueTask<SourceText> GetTextAsync(CancellationToken cancellationToken);
+    ValueTask<VersionStamp> GetTextVersionAsync(CancellationToken cancellationToken);
     Task<RazorCodeDocument> GetGeneratedOutputAsync(bool forceDesignTimeGeneratedOutput);
 
     /// <summary>

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

@@ -25,7 +25,7 @@ internal sealed class ImportDocumentSnapshot(IProjectSnapshot project, RazorProj
 
     public int Version => 1;
 
-    public ValueTask<SourceText> GetTextAsync()
+    public ValueTask<SourceText> GetTextAsync(CancellationToken cancellationToken)
     {
         return TryGetText(out var text)
             ? new(text)
@@ -44,7 +44,7 @@ internal sealed class ImportDocumentSnapshot(IProjectSnapshot project, RazorProj
     public Task<RazorCodeDocument> GetGeneratedOutputAsync(bool forceDesignTimeGeneratedOutput)
         => throw new NotSupportedException();
 
-    public ValueTask<VersionStamp> GetTextVersionAsync()
+    public ValueTask<VersionStamp> GetTextVersionAsync(CancellationToken cancellationToken)
         => new(VersionStamp.Default);
 
     public bool TryGetText([NotNullWhen(true)] out SourceText? result)

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

@@ -42,18 +42,18 @@ internal sealed class RemoteDocumentSnapshot : IDocumentSnapshot
 
     public int Version => -999; // We don't expect to use this in cohosting, but plenty of existing code logs it's value
 
-    public ValueTask<SourceText> GetTextAsync()
+    public ValueTask<SourceText> GetTextAsync(CancellationToken cancellationToken)
     {
         return TryGetText(out var result)
             ? new(result)
-            : new(TextDocument.GetTextAsync());
+            : new(TextDocument.GetTextAsync(cancellationToken));
     }
 
-    public ValueTask<VersionStamp> GetTextVersionAsync()
+    public ValueTask<VersionStamp> GetTextVersionAsync(CancellationToken cancellationToken)
     {
         return TryGetTextVersion(out var result)
             ? new(result)
-            : new(TextDocument.GetTextVersionAsync());
+            : new(TextDocument.GetTextVersionAsync(cancellationToken));
     }
 
     public bool TryGetText([NotNullWhen(true)] out SourceText? result)

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

@@ -46,7 +46,7 @@ internal class RazorDocumentExcerptService(
             return null;
         }
 
-        var razorDocumentText = await razorDocument.GetTextAsync().ConfigureAwait(false);
+        var razorDocumentText = await razorDocument.GetTextAsync(cancellationToken).ConfigureAwait(false);
         var razorDocumentSpan = razorDocumentText.Lines.GetTextSpan(mappedSpans[0].LinePositionSpan);
 
         var generatedDocument = document;

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

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

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

@@ -337,7 +337,7 @@ $$Path;
             .Setup(x => x.GetGeneratedOutputAsync(It.IsAny<bool>()))
             .ReturnsAsync(codeDocument);
         documentSnapshotMock
-            .Setup(x => x.GetTextAsync())
+            .Setup(x => x.GetTextAsync(It.IsAny<CancellationToken>()))
             .ReturnsAsync(codeDocument.Source.Text);
         documentSnapshotMock
             .Setup(x => x.Project.GetTagHelpersAsync(It.IsAny<CancellationToken>()))

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

@@ -467,7 +467,7 @@ public class TypeAccessibilityCodeActionProviderTest(ITestOutputHelper testOutpu
             .Setup(x => x.GetGeneratedOutputAsync(It.IsAny<bool>()))
             .ReturnsAsync(codeDocument);
         documentSnapshotMock
-            .Setup(x => x.GetTextAsync())
+            .Setup(x => x.GetTextAsync(It.IsAny<CancellationToken>()))
             .ReturnsAsync(codeDocument.Source.Text);
         documentSnapshotMock
             .Setup(x => x.Project.GetTagHelpersAsync(It.IsAny<CancellationToken>()))

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

@@ -471,7 +471,7 @@ public class CodeActionEndpointTest(ITestOutputHelper testOutput) : LanguageServ
         };
 
         // Act
-        var razorCodeActionContext = await codeActionEndpoint.GenerateRazorCodeActionContextAsync(request, documentContext.Snapshot);
+        var razorCodeActionContext = await codeActionEndpoint.GenerateRazorCodeActionContextAsync(request, documentContext.Snapshot, DisposalToken);
 
         // Assert
         Assert.NotNull(razorCodeActionContext);
@@ -499,7 +499,7 @@ public class CodeActionEndpointTest(ITestOutputHelper testOutput) : LanguageServ
         };
 
         // Act
-        var razorCodeActionContext = await codeActionEndpoint.GenerateRazorCodeActionContextAsync(request, documentContext.Snapshot);
+        var razorCodeActionContext = await codeActionEndpoint.GenerateRazorCodeActionContextAsync(request, documentContext.Snapshot, DisposalToken);
 
         // Assert
         Assert.NotNull(razorCodeActionContext);
@@ -524,7 +524,7 @@ public class CodeActionEndpointTest(ITestOutputHelper testOutput) : LanguageServ
             Context = new VSInternalCodeActionContext()
         };
 
-        var context = await codeActionEndpoint.GenerateRazorCodeActionContextAsync(request, documentContext.Snapshot);
+        var context = await codeActionEndpoint.GenerateRazorCodeActionContextAsync(request, documentContext.Snapshot, DisposalToken);
         Assert.NotNull(context);
 
         // Act
@@ -560,7 +560,7 @@ public class CodeActionEndpointTest(ITestOutputHelper testOutput) : LanguageServ
             }
         };
 
-        var context = await codeActionEndpoint.GenerateRazorCodeActionContextAsync(request, documentContext.Snapshot);
+        var context = await codeActionEndpoint.GenerateRazorCodeActionContextAsync(request, documentContext.Snapshot, DisposalToken);
         Assert.NotNull(context);
 
         // Act

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

@@ -155,7 +155,7 @@ public class DefaultHtmlCodeActionProviderTest(ITestOutputHelper testOutput) : L
             .Setup(x => x.GetGeneratedOutputAsync(It.IsAny<bool>()))
             .ReturnsAsync(codeDocument);
         documentSnapshotMock
-            .Setup(x => x.GetTextAsync())
+            .Setup(x => x.GetTextAsync(It.IsAny<CancellationToken>()))
             .ReturnsAsync(codeDocument.Source.Text);
         documentSnapshotMock
             .Setup(x => x.Project.GetTagHelpersAsync(It.IsAny<CancellationToken>()))

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

@@ -468,7 +468,7 @@ public class ComponentAccessibilityCodeActionProviderTest(ITestOutputHelper test
             .Setup(x => x.GetGeneratedOutputAsync(It.IsAny<bool>()))
             .ReturnsAsync(codeDocument);
         documentSnapshotMock
-            .Setup(x => x.GetTextAsync())
+            .Setup(x => x.GetTextAsync(It.IsAny<CancellationToken>()))
             .ReturnsAsync(codeDocument.Source.Text);
         documentSnapshotMock
             .Setup(x => x.Project.GetTagHelpersAsync(It.IsAny<CancellationToken>()))

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

@@ -5,6 +5,7 @@ using System;
 using System.Collections.Immutable;
 using System.Linq;
 using System.Text.Json;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.AspNetCore.Razor.Language.Components;
@@ -401,7 +402,7 @@ public class ExtractToCodeBehindCodeActionProviderTest(ITestOutputHelper testOut
             .Setup(x => x.GetGeneratedOutputAsync(It.IsAny<bool>()))
             .ReturnsAsync(codeDocument);
         documentSnapshotMock
-            .Setup(x => x.GetTextAsync())
+            .Setup(x => x.GetTextAsync(It.IsAny<CancellationToken>()))
             .ReturnsAsync(codeDocument.Source.Text);
 
         return new RazorCodeActionContext(

src/Razor/test/Microsoft.AspNetCore.Razor.LanguageServer.Test/DocumentSnapshotTextLoaderTest.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.Test.Common;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;
@@ -20,7 +21,7 @@ public class DocumentSnapshotTextLoaderTest(ITestOutputHelper testOutput) : Tool
         var expectedSourceText = SourceText.From("Hello World");
         var snapshotMock = new StrictMock<IDocumentSnapshot>();
         snapshotMock
-            .Setup(x => x.GetTextAsync())
+            .Setup(x => x.GetTextAsync(It.IsAny<CancellationToken>()))
             .ReturnsAsync(expectedSourceText);
         var textLoader = new DocumentSnapshotTextLoader(snapshotMock.Object);
 

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

@@ -273,9 +273,9 @@ public class FormattingTestBase : RazorToolingIntegrationTestBase
         var importsPath = new Uri("file:///path/to/_Imports.razor").AbsolutePath;
         var importsSourceText = SourceText.From(DefaultImports);
         var importsDocument = RazorSourceDocument.Create(importsSourceText, RazorSourceDocumentProperties.Create(importsPath, importsPath));
-        var importsSnapshot = new Mock<IDocumentSnapshot>(MockBehavior.Strict);
+        var importsSnapshot = new StrictMock<IDocumentSnapshot>();
         importsSnapshot
-            .Setup(d => d.GetTextAsync())
+            .Setup(d => d.GetTextAsync(It.IsAny<CancellationToken>()))
             .ReturnsAsync(importsSourceText);
         importsSnapshot
             .Setup(d => d.FilePath)
@@ -286,8 +286,7 @@ public class FormattingTestBase : RazorToolingIntegrationTestBase
 
         var projectFileSystem = new TestRazorProjectFileSystem([
             new TestRazorProjectItem(path, fileKind: fileKind),
-            new TestRazorProjectItem(importsPath, fileKind: FileKinds.ComponentImport),
-            ]);
+            new TestRazorProjectItem(importsPath, fileKind: FileKinds.ComponentImport)]);
 
         var projectEngine = RazorProjectEngine.Create(
             new RazorConfiguration(RazorLanguageVersion.Latest, "TestConfiguration", Extensions: [], LanguageServerFlags: new LanguageServerFlags(forceRuntimeCodeGeneration)),
@@ -315,7 +314,7 @@ public class FormattingTestBase : RazorToolingIntegrationTestBase
 
     internal static IDocumentSnapshot CreateDocumentSnapshot(string path, ImmutableArray<TagHelperDescriptor> tagHelpers, string? fileKind, ImmutableArray<RazorSourceDocument> importsDocuments, ImmutableArray<IDocumentSnapshot> imports, RazorProjectEngine projectEngine, RazorCodeDocument codeDocument, bool inGlobalNamespace = false)
     {
-        var documentSnapshot = new Mock<IDocumentSnapshot>(MockBehavior.Strict);
+        var documentSnapshot = new StrictMock<IDocumentSnapshot>();
         documentSnapshot
             .Setup(d => d.GetGeneratedOutputAsync(It.IsAny<bool>()))
             .ReturnsAsync(codeDocument);
@@ -332,7 +331,7 @@ public class FormattingTestBase : RazorToolingIntegrationTestBase
             .Setup(d => d.Project.Configuration)
             .Returns(projectEngine.Configuration);
         documentSnapshot
-            .Setup(d => d.GetTextAsync())
+            .Setup(d => d.GetTextAsync(It.IsAny<CancellationToken>()))
             .ReturnsAsync(codeDocument.Source.Text);
         documentSnapshot
             .Setup(d => d.Project.GetTagHelpersAsync(It.IsAny<CancellationToken>()))

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

@@ -936,7 +936,7 @@ public class HoverServiceTest(ITestOutputHelper testOutput) : TagHelperServiceTe
             .Setup(x => x.GetGeneratedOutputAsync(It.IsAny<bool>()))
             .ReturnsAsync(codeDocument);
         documentSnapshotMock
-            .Setup(x => x.GetTextAsync())
+            .Setup(x => x.GetTextAsync(It.IsAny<CancellationToken>()))
             .ReturnsAsync(codeDocument.Source.Text);
         documentSnapshotMock
             .SetupGet(x => x.FilePath)

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

@@ -958,7 +958,7 @@ public partial class SemanticTokensTest(ITestOutputHelper testOutput) : TagHelpe
             .Setup(x => x.GetGeneratedOutputAsync(It.IsAny<bool>()))
             .ReturnsAsync(document);
         documentSnapshotMock
-            .Setup(x => x.GetTextAsync())
+            .Setup(x => x.GetTextAsync(It.IsAny<CancellationToken>()))
             .ReturnsAsync(document.Source.Text);
         documentSnapshotMock
             .SetupGet(x => x.Version)

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

@@ -75,15 +75,15 @@ internal sealed class TestDocumentSnapshot : IDocumentSnapshot
             : Task.FromResult(_codeDocument);
     }
 
-    public ValueTask<SourceText> GetTextAsync()
+    public ValueTask<SourceText> GetTextAsync(CancellationToken cancellationToken)
     {
         return _codeDocument is null
-            ? RealSnapshot.GetTextAsync()
+            ? RealSnapshot.GetTextAsync(cancellationToken)
             : new(_codeDocument.Source.Text);
     }
 
-    public ValueTask<VersionStamp> GetTextVersionAsync()
-        => RealSnapshot.GetTextVersionAsync();
+    public ValueTask<VersionStamp> GetTextVersionAsync(CancellationToken cancellationToken)
+        => RealSnapshot.GetTextVersionAsync(cancellationToken);
 
     public Task<SyntaxTree> GetCSharpSyntaxTreeAsync(CancellationToken cancellationToken)
     {

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

@@ -4,6 +4,7 @@
 using System;
 using System.IO;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.AspNetCore.Razor.ProjectSystem;
@@ -68,7 +69,7 @@ public abstract class DocumentExcerptServiceTestBase(ITestOutputHelper testOutpu
 
     // Maps a span in the primary buffer to the secondary buffer. This is only valid for C# code
     // that appears in the primary buffer.
-    private static async Task<TextSpan> GetSecondarySpanAsync(IDocumentSnapshot primary, TextSpan primarySpan, Document secondary)
+    private static async Task<TextSpan> GetSecondarySpanAsync(IDocumentSnapshot primary, TextSpan primarySpan, Document secondary, CancellationToken cancellationToken)
     {
         var output = await primary.GetGeneratedOutputAsync();
 
@@ -80,8 +81,8 @@ public abstract class DocumentExcerptServiceTestBase(ITestOutputHelper testOutpu
                 var offset = mapping.GeneratedSpan.AbsoluteIndex - mapping.OriginalSpan.AbsoluteIndex;
                 var secondarySpan = new TextSpan(primarySpan.Start + offset, primarySpan.Length);
                 Assert.Equal(
-                    (await primary.GetTextAsync()).GetSubText(primarySpan).ToString(),
-                    (await secondary.GetTextAsync()).GetSubText(secondarySpan).ToString());
+                    (await primary.GetTextAsync(cancellationToken)).ToString(primarySpan),
+                    (await secondary.GetTextAsync(cancellationToken)).ToString(secondarySpan));
                 return secondarySpan;
             }
         }
@@ -89,19 +90,19 @@ public abstract class DocumentExcerptServiceTestBase(ITestOutputHelper testOutpu
         throw new InvalidOperationException("Could not map the primary span to the generated code.");
     }
 
-    public async Task<(Document generatedDocument, SourceText razorSourceText, TextSpan primarySpan, TextSpan generatedSpan)> InitializeAsync(string razorSource)
+    public async Task<(Document generatedDocument, SourceText razorSourceText, TextSpan primarySpan, TextSpan generatedSpan)> InitializeAsync(string razorSource, CancellationToken cancellationToken)
     {
         var (razorSourceText, primarySpan) = CreateText(razorSource);
         var (primary, generatedDocument) = InitializeDocument(razorSourceText);
-        var generatedSpan = await GetSecondarySpanAsync(primary, primarySpan, generatedDocument);
+        var generatedSpan = await GetSecondarySpanAsync(primary, primarySpan, generatedDocument, cancellationToken);
         return (generatedDocument, razorSourceText, primarySpan, generatedSpan);
     }
 
-    internal async Task<(IDocumentSnapshot primary, Document generatedDocument, TextSpan generatedSpan)> InitializeWithSnapshotAsync(string razorSource)
+    internal async Task<(IDocumentSnapshot primary, Document generatedDocument, TextSpan generatedSpan)> InitializeWithSnapshotAsync(string razorSource, CancellationToken cancellationToken)
     {
         var (razorSourceText, primarySpan) = CreateText(razorSource);
         var (primary, generatedDocument) = InitializeDocument(razorSourceText);
-        var generatedSpan = await GetSecondarySpanAsync(primary, primarySpan, generatedDocument);
+        var generatedSpan = await GetSecondarySpanAsync(primary, primarySpan, generatedDocument, cancellationToken);
         return (primary, generatedDocument, generatedSpan);
     }
 }

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

@@ -30,7 +30,7 @@ public class DocumentStateTest : ToolingTestBase
         var state = DocumentState.Create(_hostDocument, DocumentState.EmptyLoader);
 
         // Assert
-        var text = await state.GetTextAsync();
+        var text = await state.GetTextAsync(DisposalToken);
         Assert.Equal(0, text.Length);
     }
 
@@ -44,7 +44,7 @@ public class DocumentStateTest : ToolingTestBase
         var state = original.WithText(_text, VersionStamp.Create());
 
         // Assert
-        var text = await state.GetTextAsync();
+        var text = await state.GetTextAsync(DisposalToken);
         Assert.Same(_text, text);
     }
 
@@ -58,7 +58,7 @@ public class DocumentStateTest : ToolingTestBase
         var state = original.WithTextLoader(_textLoader);
 
         // Assert
-        var text = await state.GetTextAsync();
+        var text = await state.GetTextAsync(DisposalToken);
         Assert.Same(_text, text);
     }
 
@@ -84,7 +84,7 @@ public class DocumentStateTest : ToolingTestBase
         var original = DocumentState.Create(_hostDocument, DocumentState.EmptyLoader)
             .WithTextLoader(_textLoader);
 
-        await original.GetTextAsync();
+        await original.GetTextAsync(DisposalToken);
 
         // Act
         var state = original.WithConfigurationChange();
@@ -116,7 +116,7 @@ public class DocumentStateTest : ToolingTestBase
         var original = DocumentState.Create(_hostDocument, DocumentState.EmptyLoader)
             .WithTextLoader(_textLoader);
 
-        await original.GetTextAsync();
+        await original.GetTextAsync(DisposalToken);
 
         // Act
         var state = original.WithImportsChange();
@@ -148,7 +148,7 @@ public class DocumentStateTest : ToolingTestBase
         var original = DocumentState.Create(_hostDocument, DocumentState.EmptyLoader)
             .WithTextLoader(_textLoader);
 
-        await original.GetTextAsync();
+        await original.GetTextAsync(DisposalToken);
 
         // Act
         var state = original.WithProjectWorkspaceStateChange();

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

@@ -106,7 +106,7 @@ public class ProjectStateTest : WorkspaceTestBase
         var state = original.WithAddedHostDocument(_documents[0], DocumentState.EmptyLoader);
 
         // Assert
-        var text = await state.Documents[_documents[0].FilePath].GetTextAsync();
+        var text = await state.Documents[_documents[0].FilePath].GetTextAsync(DisposalToken);
         Assert.Equal(0, text.Length);
     }
 
@@ -278,7 +278,7 @@ public class ProjectStateTest : WorkspaceTestBase
         // Assert
         Assert.NotEqual(original.Version, state.Version);
 
-        var text = await state.Documents[_documents[1].FilePath].GetTextAsync();
+        var text = await state.Documents[_documents[1].FilePath].GetTextAsync(DisposalToken);
         Assert.Same(_text, text);
 
         Assert.Equal(original.DocumentCollectionVersion, state.DocumentCollectionVersion);
@@ -298,7 +298,7 @@ public class ProjectStateTest : WorkspaceTestBase
         // Assert
         Assert.NotEqual(original.Version, state.Version);
 
-        var text = await state.Documents[_documents[1].FilePath].GetTextAsync();
+        var text = await state.Documents[_documents[1].FilePath].GetTextAsync(DisposalToken);
         Assert.Same(_text, text);
 
         Assert.Equal(original.DocumentCollectionVersion, state.DocumentCollectionVersion);

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

@@ -35,7 +35,7 @@ public class CSharpDocumentExcerptServiceTest : DocumentExcerptServiceTestBase
 </html>
 ";
 
-        var (generatedDocument, razorSourceText, primarySpan, generatedSpan) = await InitializeAsync(razorSource);
+        var (generatedDocument, razorSourceText, primarySpan, generatedSpan) = await InitializeAsync(razorSource, DisposalToken);
 
 #pragma warning disable CS0618 // Type or member is obsolete
         var excerptService = new CSharpDocumentExcerptService();
@@ -123,7 +123,7 @@ public class CSharpDocumentExcerptServiceTest : DocumentExcerptServiceTestBase
 </html>
 ";
 
-        var (generatedDocument, razorSourceText, primarySpan, generatedSpan) = await InitializeAsync(razorSource);
+        var (generatedDocument, razorSourceText, primarySpan, generatedSpan) = await InitializeAsync(razorSource, DisposalToken);
 
 #pragma warning disable CS0618 // Type or member is obsolete
         var excerptService = new CSharpDocumentExcerptService();
@@ -169,7 +169,7 @@ public class CSharpDocumentExcerptServiceTest : DocumentExcerptServiceTestBase
 </html>
 ";
 
-        var (generatedDocument, razorSourceText, primarySpan, generatedSpan) = await InitializeAsync(razorSource);
+        var (generatedDocument, razorSourceText, primarySpan, generatedSpan) = await InitializeAsync(razorSource, DisposalToken);
 
 #pragma warning disable CS0618 // Type or member is obsolete
         var excerptService = new CSharpDocumentExcerptService();
@@ -216,7 +216,7 @@ public class CSharpDocumentExcerptServiceTest : DocumentExcerptServiceTestBase
 </html>
 ";
 
-        var (generatedDocument, razorSourceText, primarySpan, generatedSpan) = await InitializeAsync(razorSource);
+        var (generatedDocument, razorSourceText, primarySpan, generatedSpan) = await InitializeAsync(razorSource, DisposalToken);
 
 #pragma warning disable CS0618 // Type or member is obsolete
         var excerptService = new CSharpDocumentExcerptService();
@@ -262,7 +262,7 @@ public class CSharpDocumentExcerptServiceTest : DocumentExcerptServiceTestBase
         // Arrange
         var razorSource = @"@{ var [|foo|] = ""Hello, World!""; }";
 
-        var (generatedDocument, razorSourceText, primarySpan, generatedSpan) = await InitializeAsync(razorSource);
+        var (generatedDocument, razorSourceText, primarySpan, generatedSpan) = await InitializeAsync(razorSource, DisposalToken);
 
 #pragma warning disable CS0618 // Type or member is obsolete
         var excerptService = new CSharpDocumentExcerptService();

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

@@ -28,7 +28,7 @@ public class RazorDocumentExcerptServiceTest(ITestOutputHelper testOutput) : Doc
 </html>
 ";
 
-        var (primary, secondary, secondarySpan) = await InitializeWithSnapshotAsync(razorSource);
+        var (primary, secondary, secondarySpan) = await InitializeWithSnapshotAsync(razorSource, DisposalToken);
 
         var service = CreateExcerptService(primary);
 
@@ -106,7 +106,7 @@ public class RazorDocumentExcerptServiceTest(ITestOutputHelper testOutput) : Doc
 </html>
 ";
 
-        var (primary, secondary, secondarySpan) = await InitializeWithSnapshotAsync(razorSource);
+        var (primary, secondary, secondarySpan) = await InitializeWithSnapshotAsync(razorSource, DisposalToken);
 
         var service = CreateExcerptService(primary);
 
@@ -159,7 +159,7 @@ public class RazorDocumentExcerptServiceTest(ITestOutputHelper testOutput) : Doc
 </html>
 ";
 
-        var (primary, secondary, secondarySpan) = await InitializeWithSnapshotAsync(razorSource);
+        var (primary, secondary, secondarySpan) = await InitializeWithSnapshotAsync(razorSource, DisposalToken);
 
         var service = CreateExcerptService(primary);
 
@@ -266,7 +266,7 @@ public class RazorDocumentExcerptServiceTest(ITestOutputHelper testOutput) : Doc
             </html>
             """;
 
-        var (primary, secondary, secondarySpan) = await InitializeWithSnapshotAsync(razorSource);
+        var (primary, secondary, secondarySpan) = await InitializeWithSnapshotAsync(razorSource, DisposalToken);
 
         var service = CreateExcerptService(primary);
 
@@ -370,7 +370,7 @@ public class RazorDocumentExcerptServiceTest(ITestOutputHelper testOutput) : Doc
             </html>
             """;
 
-        var (primary, secondary, secondarySpan) = await InitializeWithSnapshotAsync(razorSource);
+        var (primary, secondary, secondarySpan) = await InitializeWithSnapshotAsync(razorSource, DisposalToken);
 
         var service = CreateExcerptService(primary);
 
@@ -444,7 +444,7 @@ public class RazorDocumentExcerptServiceTest(ITestOutputHelper testOutput) : Doc
 </html>
 ";
 
-        var (primary, secondary, secondarySpan) = await InitializeWithSnapshotAsync(razorSource);
+        var (primary, secondary, secondarySpan) = await InitializeWithSnapshotAsync(razorSource, DisposalToken);
 
         var service = CreateExcerptService(primary);
 
@@ -553,7 +553,7 @@ public class RazorDocumentExcerptServiceTest(ITestOutputHelper testOutput) : Doc
         // Arrange
         var razorSource = @"@{ var [|foo|] = ""Hello, World!""; }";
 
-        var (primary, secondary, secondarySpan) = await InitializeWithSnapshotAsync(razorSource);
+        var (primary, secondary, secondarySpan) = await InitializeWithSnapshotAsync(razorSource, DisposalToken);
 
         var service = CreateExcerptService(primary);
 

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

@@ -43,7 +43,7 @@ public class RazorSpanMappingServiceTest(ITestOutputHelper testOutput) : Workspa
         var span = new TextSpan(generated.GeneratedCode.IndexOf(symbol, StringComparison.Ordinal), symbol.Length);
 
         // Act
-        var result = RazorSpanMappingService.TryGetMappedSpans(span, await document.GetTextAsync(), generated, out var mappedLinePositionSpan, out var mappedSpan);
+        var result = RazorSpanMappingService.TryGetMappedSpans(span, await document.GetTextAsync(DisposalToken), generated, out var mappedLinePositionSpan, out var mappedSpan);
 
         // Assert
         Assert.True(result);
@@ -76,7 +76,7 @@ public class RazorSpanMappingServiceTest(ITestOutputHelper testOutput) : Workspa
         var span = new TextSpan(generated.GeneratedCode.IndexOf(symbol, generated.GeneratedCode.IndexOf(symbol, StringComparison.Ordinal) + symbol.Length, StringComparison.Ordinal), symbol.Length);
 
         // Act
-        var result = RazorSpanMappingService.TryGetMappedSpans(span, await document.GetTextAsync(), generated, out var mappedLinePositionSpan, out var mappedSpan);
+        var result = RazorSpanMappingService.TryGetMappedSpans(span, await document.GetTextAsync(DisposalToken), generated, out var mappedLinePositionSpan, out var mappedSpan);
 
         // Assert
         Assert.True(result);
@@ -108,7 +108,7 @@ public class RazorSpanMappingServiceTest(ITestOutputHelper testOutput) : Workspa
         var span = new TextSpan(generated.GeneratedCode.IndexOf(symbol, StringComparison.Ordinal), symbol.Length);
 
         // Act
-        var result = RazorSpanMappingService.TryGetMappedSpans(span, await document.GetTextAsync(), generated, out var mappedLinePositionSpan, out var mappedSpan);
+        var result = RazorSpanMappingService.TryGetMappedSpans(span, await document.GetTextAsync(DisposalToken), generated, out var mappedLinePositionSpan, out var mappedSpan);
 
         // Assert
         Assert.True(result);
@@ -140,7 +140,7 @@ public class RazorSpanMappingServiceTest(ITestOutputHelper testOutput) : Workspa
         var span = new TextSpan(generated.GeneratedCode.IndexOf(symbol, StringComparison.Ordinal), symbol.Length);
 
         // Act
-        var result = RazorSpanMappingService.TryGetMappedSpans(span, await document.GetTextAsync(), generated, out _, out _);
+        var result = RazorSpanMappingService.TryGetMappedSpans(span, await document.GetTextAsync(DisposalToken), generated, out _, out _);
 
         // Assert
         Assert.False(result);

src/Razor/test/Microsoft.VisualStudio.LanguageServices.Razor.Test/ProjectSystem/ProjectSnapshotManagerTest.cs

@@ -230,7 +230,7 @@ public class ProjectSnapshotManagerTest : VisualStudioWorkspaceTestBase
         var document = project.GetDocument(documentFilePath);
         Assert.NotNull(document);
 
-        var text = await document.GetTextAsync();
+        var text = await document.GetTextAsync(DisposalToken);
         Assert.Equal(0, text.Length);
     }
 
@@ -257,7 +257,7 @@ public class ProjectSnapshotManagerTest : VisualStudioWorkspaceTestBase
         var document = project.GetDocument(documentFilePath);
         Assert.NotNull(document);
 
-        var actual = await document.GetTextAsync();
+        var actual = await document.GetTextAsync(DisposalToken);
         Assert.Same(expected, actual);
     }
 
@@ -465,7 +465,7 @@ public class ProjectSnapshotManagerTest : VisualStudioWorkspaceTestBase
         var project = _projectManager.GetLoadedProject(s_hostProject.Key);
         var document = project.GetDocument(s_documents[0].FilePath);
         Assert.NotNull(document);
-        var text = await document.GetTextAsync();
+        var text = await document.GetTextAsync(DisposalToken);
         Assert.Same(_sourceText, text);
 
         Assert.True(_projectManager.IsDocumentOpen(s_documents[0].FilePath));
@@ -504,7 +504,7 @@ public class ProjectSnapshotManagerTest : VisualStudioWorkspaceTestBase
         var project = _projectManager.GetLoadedProject(s_hostProject.Key);
         var document = project.GetDocument(s_documents[0].FilePath);
         Assert.NotNull(document);
-        var text = await document.GetTextAsync();
+        var text = await document.GetTextAsync(DisposalToken);
         Assert.Same(expected, text);
         Assert.False(_projectManager.IsDocumentOpen(s_documents[0].FilePath));
         Assert.Equal(3, document.Version);
@@ -538,7 +538,7 @@ public class ProjectSnapshotManagerTest : VisualStudioWorkspaceTestBase
         var project = _projectManager.GetLoadedProject(s_hostProject.Key);
         var document = project.GetDocument(s_documents[0].FilePath);
         Assert.NotNull(document);
-        var text = await document.GetTextAsync();
+        var text = await document.GetTextAsync(DisposalToken);
         Assert.Same(expected, text);
     }
 
@@ -570,7 +570,7 @@ public class ProjectSnapshotManagerTest : VisualStudioWorkspaceTestBase
         var project = _projectManager.GetLoadedProject(s_hostProject.Key);
         var document = project.GetDocument(s_documents[0].FilePath);
         Assert.NotNull(document);
-        var text = await document.GetTextAsync();
+        var text = await document.GetTextAsync(DisposalToken);
         Assert.Same(expected, text);
         Assert.Equal(3, document.Version);
     }
@@ -604,7 +604,7 @@ public class ProjectSnapshotManagerTest : VisualStudioWorkspaceTestBase
         var project = _projectManager.GetLoadedProject(s_hostProject.Key);
         var document = project.GetDocument(s_documents[0].FilePath);
         Assert.NotNull(document);
-        var text = await document.GetTextAsync();
+        var text = await document.GetTextAsync(DisposalToken);
         Assert.Same(expected, text);
         Assert.Equal(3, document.Version);
     }