Init semantic token refresh, and queue on request

Author: davidwengier

src/Razor/src/Microsoft.CodeAnalysis.Razor.CohostingShared/SemanticTokens/CohostSemanticTokensInitializeService.cs

@@ -0,0 +1,32 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.ComponentModel.Composition;
+using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.ExternalAccess.Razor.Cohost;
+using Microsoft.CodeAnalysis.Razor.Workspaces;
+
+namespace Microsoft.VisualStudio.Razor.LanguageClient.Cohost;
+
+[Export(typeof(IRazorCohostStartupService))]
+[method: ImportingConstructor]
+internal class CohostSemanticTokensInitializeService(
+    LanguageServerFeatureOptions languageServerFeatureOptions) : IRazorCohostStartupService
+{
+    private readonly LanguageServerFeatureOptions _options = languageServerFeatureOptions;
+
+    public int Order => WellKnownStartupOrder.Default;
+
+    public Task StartupAsync(VSInternalClientCapabilities clientCapabilities, RazorCohostRequestContext requestContext, CancellationToken cancellationToken)
+    {
+        if (_options.UseRazorCohostServer)
+        {
+            // Normally, Roslyn triggers semantic tokens refreshes for additional document changes, but not for Razor documents.
+            // In cohosting, we want it to consider Razor documents too.
+            CodeAnalysis.ExternalAccess.Razor.Cohost.Handlers.SemanticTokensRange.RegisterRefresh(requestContext);
+        }
+
+        return Task.CompletedTask;
+    }
+}

src/Razor/src/Microsoft.CodeAnalysis.Razor.CohostingShared/SemanticTokens/CohostSemanticTokensRangeEndpoint.cs

@@ -39,6 +39,22 @@ internal sealed class CohostSemanticTokensRangeEndpoint(
     protected override RazorTextDocumentIdentifier? GetRazorTextDocumentIdentifier(SemanticTokensRangeParams request)
         => request.TextDocument.ToRazorTextDocumentIdentifier();
 
+    protected override async Task<SemanticTokens?> HandleRequestAsync(SemanticTokensRangeParams request, RazorCohostRequestContext context, TextDocument razorDocument, CancellationToken cancellationToken)
+    {
+        var result = await HandleRequestAsync(request, razorDocument, cancellationToken).ConfigureAwait(false);
+
+        if (result is not null)
+        {
+            // Roslyn uses frozen semantics for semantic tokens, so it could return results from an older project state.
+            // Every time they get a request they queue up a refresh, which will check the project checksums, and if there
+            // hasn't been any changes, will no-op. We call into that same logic here to ensure everything is up to date.
+            // See: https://github.com/dotnet/roslyn/blob/bb57f4643bb3d52eb7626f9863da177d9e219f1e/src/LanguageServer/Protocol/Handler/SemanticTokens/SemanticTokensHelpers.cs#L48-L52
+            await CodeAnalysis.ExternalAccess.Razor.Cohost.Handlers.SemanticTokensRange.TryEnqueueRefreshComputationAsync(context, razorDocument.Project, cancellationToken).ConfigureAwait(false);
+        }
+
+        return result;
+    }
+
     protected override Task<SemanticTokens?> HandleRequestAsync(SemanticTokensRangeParams request, TextDocument razorDocument, CancellationToken cancellationToken)
         => HandleRequestAsync(razorDocument, request.Range.ToLinePositionSpan(), cancellationToken);