[LSP] Modify semanticTokens isFinalized logic (#60484)

src/Features/LanguageServer/Protocol/Handler/SemanticTokens/SemanticTokensHelpers.cs

@@ -130,7 +130,7 @@ static SemanticTokensHelpers()
         /// <summary>
         /// Returns the semantic tokens data for a given document with an optional range.
         /// </summary>
-        internal static async Task<(int[], bool isFinalized)> ComputeSemanticTokensDataAsync(
+        internal static async Task<int[]> ComputeSemanticTokensDataAsync(
             Document document,
             Dictionary<string, int> tokenTypesToIndex,
             LSP.Range? range,
@@ -147,10 +147,8 @@ static SemanticTokensHelpers()
 
             // If the full compilation is not yet available, we'll try getting a partial one. It may contain inaccurate
             // results but will speed up how quickly we can respond to the client's request.
-            var frozenDocument = document.WithFrozenPartialSemantics(cancellationToken);
-            var semanticModel = await frozenDocument.GetRequiredSemanticModelAsync(cancellationToken).ConfigureAwait(false);
-            var isFinalized = document.Project.TryGetCompilation(out var compilation) && compilation == semanticModel.Compilation;
-            document = frozenDocument;
+            document = document.WithFrozenPartialSemantics(cancellationToken);
+            options = options with { ForceFrozenPartialSemanticsForCrossProcessOperations = true };
 
             var classifiedSpans = await GetClassifiedSpansForDocumentAsync(
                 document, textSpan, options, includeSyntacticClassifications, cancellationToken).ConfigureAwait(false);
@@ -161,7 +159,7 @@ static SemanticTokensHelpers()
 
             // TO-DO: We should implement support for streaming if LSP adds support for it:
             // https://devdiv.visualstudio.com/DevDiv/_workitems/edit/1276300
-            return (ComputeTokens(text.Lines, updatedClassifiedSpans, tokenTypesToIndex), isFinalized);
+            return ComputeTokens(text.Lines, updatedClassifiedSpans, tokenTypesToIndex);
         }
 
         private static async Task<ClassifiedSpan[]> GetClassifiedSpansForDocumentAsync(

src/Features/LanguageServer/Protocol/Handler/SemanticTokens/SemanticTokensRangeHandler.cs

@@ -3,10 +3,12 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.Collections.Generic;
 using System.Composition;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Classification;
+using Microsoft.CodeAnalysis.Host;
 using Microsoft.CodeAnalysis.Host.Mef;
 using Microsoft.CodeAnalysis.Options;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
@@ -24,6 +26,10 @@ internal class SemanticTokensRangeHandler : AbstractStatelessRequestHandler<LSP.
     {
         private readonly IGlobalOptionService _globalOptions;
 
+        // Lock to guard _projectToCompilation dictionary
+        private readonly object _lock = new();
+        private readonly Dictionary<ProjectId, Task<Compilation?>> _projectIdToCompilation = new();
+
         public override bool MutatesSolutionState => false;
         public override bool RequiresLSPSolution => true;
 
@@ -48,13 +54,20 @@ public SemanticTokensRangeHandler(IGlobalOptionService globalOptions)
             Contract.ThrowIfNull(request.TextDocument, "TextDocument is null.");
             Contract.ThrowIfNull(context.Document, "Document is null.");
 
-            var options = _globalOptions.GetClassificationOptions(context.Document.Project.Language);
+            var project = context.Document.Project;
+            var options = _globalOptions.GetClassificationOptions(project.Language);
+
+            // Razor uses isFinalized to determine whether to cache tokens. We should be able to
+            // remove it altogether once Roslyn implements workspace/semanticTokens/refresh:
+            // https://github.com/dotnet/roslyn/issues/60441
+            var isFinalized = !context.Document.IsRazorDocument() ||
+                await IsDataFinalizedAsync(project, cancellationToken).ConfigureAwait(false);
 
             // The results from the range handler should not be cached since we don't want to cache
             // partial token results. In addition, a range request is only ever called with a whole
             // document request, so caching range results is unnecessary since the whole document
             // handler will cache the results anyway.
-            var (tokensData, isFinalized) = await SemanticTokensHelpers.ComputeSemanticTokensDataAsync(
+            var tokensData = await SemanticTokensHelpers.ComputeSemanticTokensDataAsync(
                 context.Document,
                 SemanticTokensHelpers.TokenTypeToIndex,
                 request.Range,
@@ -64,5 +77,42 @@ public SemanticTokensRangeHandler(IGlobalOptionService globalOptions)
 
             return new RoslynSemanticTokens { Data = tokensData, IsFinalized = isFinalized };
         }
+
+        private async Task<bool> IsDataFinalizedAsync(Project project, CancellationToken cancellationToken)
+        {
+            // We use a combination of IsFullyLoaded + the completed project compilation as the metric
+            // for isFinalized. It may not be completely accurate but this is only a a temporary fix until
+            // workspace/semanticTokens/refresh is implemented.
+            lock (_lock)
+            {
+                if (_projectIdToCompilation.TryGetValue(project.Id, out var compilationTask) && compilationTask.IsCompleted)
+                {
+                    // We don't want to hang on to the compilation since this can be very expensive,
+                    // but we do want to mark the compilation as being successfully retrieved.
+                    if (compilationTask.Result is not null)
+                    {
+                        _projectIdToCompilation[project.Id] = Task.FromResult<Compilation?>(null);
+                    }
+
+                    return true;
+                }
+            }
+
+            var workspaceStatusService = project.Solution.Workspace.Services.GetRequiredService<IWorkspaceStatusService>();
+            var isFullyLoaded = await workspaceStatusService.IsFullyLoadedAsync(cancellationToken).ConfigureAwait(false);
+
+            lock (_lock)
+            {
+                if (isFullyLoaded && !_projectIdToCompilation.ContainsKey(project.Id))
+                {
+                    // If the project's compilation isn't yet available, kick off a task in the background to
+                    // hopefully make it available faster since we'll need it later to compute tokens.
+                    var newCompilationTask = project.GetCompilationAsync(cancellationToken);
+                    _projectIdToCompilation.Add(project.Id, newCompilationTask);
+                }
+            }
+
+            return false;
+        }
     }
 }

src/Features/LanguageServer/ProtocolUnitTests/SemanticTokens/SemanticTokensRangeTests.cs

@@ -53,7 +53,7 @@ static class C { }
             var range = new LSP.Range { Start = new Position(0, 0), End = new Position(2, 0) };
             var options = ClassificationOptions.Default;
 
-            var (results, _) = await SemanticTokensHelpers.ComputeSemanticTokensDataAsync(
+            var results = await SemanticTokensHelpers.ComputeSemanticTokensDataAsync(
                 document, SemanticTokensHelpers.TokenTypeToIndex, range, options, includeSyntacticClassifications: true, CancellationToken.None);
 
             var expectedResults = new LSP.SemanticTokens
@@ -87,7 +87,7 @@ static class C { }
             var document = testLspServer.GetCurrentSolution().Projects.First().Documents.First();
             var range = new LSP.Range { Start = new Position(1, 0), End = new Position(2, 0) };
             var options = ClassificationOptions.Default;
-            var (results, _) = await SemanticTokensHelpers.ComputeSemanticTokensDataAsync(
+            var results = await SemanticTokensHelpers.ComputeSemanticTokensDataAsync(
                 document, SemanticTokensHelpers.TokenTypeToIndex, range, options, includeSyntacticClassifications: true, CancellationToken.None);
 
             var expectedResults = new LSP.SemanticTokens
@@ -122,7 +122,7 @@ public async Task TestGetSemanticTokensRange_MultiLineComment_RazorAsync()
             var document = testLspServer.GetCurrentSolution().Projects.First().Documents.First();
             var range = new LSP.Range { Start = new Position(0, 0), End = new Position(4, 0) };
             var options = ClassificationOptions.Default;
-            var (results, _) = await SemanticTokensHelpers.ComputeSemanticTokensDataAsync(
+            var results = await SemanticTokensHelpers.ComputeSemanticTokensDataAsync(
                 document, SemanticTokensHelpers.TokenTypeToIndex, range, options, includeSyntacticClassifications: true, CancellationToken.None);
 
             var expectedResults = new LSP.SemanticTokens
@@ -197,7 +197,7 @@ void M()
             var document = testLspServer.GetCurrentSolution().Projects.First().Documents.First();
             var range = new LSP.Range { Start = new Position(0, 0), End = new Position(9, 0) };
             var options = ClassificationOptions.Default;
-            var (results, _) = await SemanticTokensHelpers.ComputeSemanticTokensDataAsync(
+            var results = await SemanticTokensHelpers.ComputeSemanticTokensDataAsync(
                 document, SemanticTokensHelpers.TokenTypeToIndex, range, options, includeSyntacticClassifications: true, CancellationToken.None);
 
             var expectedResults = new LSP.SemanticTokens
@@ -250,7 +250,7 @@ void M()
             var document = testLspServer.GetCurrentSolution().Projects.First().Documents.First();
             var range = new LSP.Range { Start = new Position(0, 0), End = new Position(9, 0) };
             var options = ClassificationOptions.Default;
-            var (results, _) = await SemanticTokensHelpers.ComputeSemanticTokensDataAsync(
+            var results = await SemanticTokensHelpers.ComputeSemanticTokensDataAsync(
                 document, SemanticTokensHelpers.TokenTypeToIndex, range, options, includeSyntacticClassifications: true, CancellationToken.None);
 
             var expectedResults = new LSP.SemanticTokens