Move capabilities logic back into the endpoint

src/Razor/src/Microsoft.VisualStudio.LanguageServerClient.Razor/Cohost/CohostSemanticTokensRangeEndpoint.cs

@@ -6,6 +6,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor;
+using Microsoft.AspNetCore.Razor.LanguageServer.Common;
 using Microsoft.AspNetCore.Razor.Telemetry;
 using Microsoft.CodeAnalysis.ExternalAccess.Razor.Cohost;
 using Microsoft.CodeAnalysis.Razor.Logging;
@@ -14,19 +15,21 @@
 using Microsoft.Extensions.Logging;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
 using Microsoft.VisualStudio.LanguageServerClient.Razor.Extensions;
+using Newtonsoft.Json;
 
 namespace Microsoft.VisualStudio.LanguageServerClient.Razor.Cohost;
 
 [Shared]
 [CohostEndpoint(Methods.TextDocumentSemanticTokensRangeName)]
+[Export(typeof(IDynamicRegistrationProvider))]
 [ExportCohostStatelessLspService(typeof(CohostSemanticTokensRangeEndpoint))]
 [method: ImportingConstructor]
 internal sealed class CohostSemanticTokensRangeEndpoint(
     IOutOfProcSemanticTokensService semanticTokensInfoService,
     ISemanticTokensLegendService semanticTokensLegendService,
     ITelemetryReporter telemetryReporter,
     IRazorLoggerFactory loggerFactory)
-    : AbstractRazorCohostDocumentRequestHandler<SemanticTokensRangeParams, SemanticTokens?>
+    : AbstractRazorCohostDocumentRequestHandler<SemanticTokensRangeParams, SemanticTokens?>, IDynamicRegistrationProvider
 {
     private readonly IOutOfProcSemanticTokensService _semanticTokensInfoService = semanticTokensInfoService;
     private readonly ISemanticTokensLegendService _semanticTokensLegendService = semanticTokensLegendService;
@@ -36,6 +39,27 @@ internal sealed class CohostSemanticTokensRangeEndpoint(
     protected override bool MutatesSolutionState => false;
     protected override bool RequiresLSPSolution => true;
 
+    public Registration? GetRegistration(VSInternalClientCapabilities clientCapabilities, DocumentFilter[] filter, RazorCohostRequestContext requestContext)
+    {
+        if (clientCapabilities.TextDocument?.SemanticTokens?.DynamicRegistration == true)
+        {
+            var semanticTokensRefreshQueue = requestContext.GetRequiredService<IRazorSemanticTokensRefreshQueue>();
+            var clientCapabilitiesString = JsonConvert.SerializeObject(clientCapabilities);
+            semanticTokensRefreshQueue.Initialize(clientCapabilitiesString);
+
+            return new Registration()
+            {
+                Method = Methods.TextDocumentSemanticTokensRangeName,
+                RegisterOptions = new SemanticTokensRegistrationOptions()
+                {
+                    DocumentSelector = filter,
+                }.EnableSemanticTokens(_semanticTokensLegendService)
+            };
+        }
+
+        return null;
+    }
+
     protected override RazorTextDocumentIdentifier? GetRazorTextDocumentIdentifier(SemanticTokensRangeParams request)
         => request.TextDocument.ToRazorTextDocumentIdentifier();
 

src/Razor/src/Microsoft.VisualStudio.LanguageServerClient.Razor/Cohost/IDynamicRegistrationProvider.cs

@@ -0,0 +1,12 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT license. See License.txt in the project root for license information.
+
+using Microsoft.CodeAnalysis.ExternalAccess.Razor.Cohost;
+using Microsoft.VisualStudio.LanguageServer.Protocol;
+
+namespace Microsoft.VisualStudio.LanguageServerClient.Razor.Cohost;
+
+internal interface IDynamicRegistrationProvider
+{
+    Registration? GetRegistration(VSInternalClientCapabilities clientCapabilities, DocumentFilter[] filter, RazorCohostRequestContext requestContext);
+}

src/Razor/src/Microsoft.VisualStudio.LanguageServerClient.Razor/Cohost/RazorCohostClientCapabilitiesService.cs

@@ -0,0 +1,26 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT license. See License.txt in the project root for license information.
+
+using System;
+using System.Composition;
+using Microsoft.CodeAnalysis.Razor.Workspaces.Protocol;
+using Microsoft.VisualStudio.LanguageServer.Protocol;
+
+namespace Microsoft.VisualStudio.LanguageServerClient.Razor.Cohost;
+
+[Shared]
+[Export(typeof(RazorCohostClientCapabilitiesService))]
+[Export(typeof(IClientCapabilitiesService))]
+internal class RazorCohostClientCapabilitiesService : IClientCapabilitiesService
+{
+    private VSInternalClientCapabilities? _clientCapabilities;
+
+    public bool CanGetClientCapabilities => _clientCapabilities is not null;
+
+    public VSInternalClientCapabilities ClientCapabilities => _clientCapabilities ?? throw new InvalidOperationException("Client capabilities requested before initialized.");
+
+    public void SetCapabilities(VSInternalClientCapabilities clientCapabilities)
+    {
+        _clientCapabilities = clientCapabilities;
+    }
+}

src/Razor/src/Microsoft.VisualStudio.LanguageServerClient.Razor/Cohost/RazorCohostDynamicRegistrationService.cs

@@ -2,16 +2,14 @@
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
 using System;
+using System.Collections.Generic;
 using System.Composition;
 using System.Threading;
 using System.Threading.Tasks;
-using Microsoft.AspNetCore.Razor;
-using Microsoft.AspNetCore.Razor.LanguageServer.Common;
+using Microsoft.AspNetCore.Razor.PooledObjects;
 using Microsoft.CodeAnalysis.ExternalAccess.Razor.Cohost;
 using Microsoft.CodeAnalysis.Razor.Logging;
-using Microsoft.CodeAnalysis.Razor.SemanticTokens;
 using Microsoft.CodeAnalysis.Razor.Workspaces;
-using Microsoft.CodeAnalysis.Razor.Workspaces.Protocol;
 using Microsoft.Extensions.Logging;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
 using Newtonsoft.Json;
@@ -20,61 +18,59 @@ namespace Microsoft.VisualStudio.LanguageServerClient.Razor.Cohost;
 
 [Shared]
 [Export(typeof(IRazorCohostDynamicRegistrationService))]
-[Export(typeof(IClientCapabilitiesService))]
 [method: ImportingConstructor]
-internal class RazorCohostDynamicRegistrationService(LanguageServerFeatureOptions languageServerFeatureOptions, Lazy<ISemanticTokensLegendService> semanticTokensLegendService, IRazorLoggerFactory loggerFactory) : IRazorCohostDynamicRegistrationService, IClientCapabilitiesService
+internal class RazorCohostDynamicRegistrationService(
+    LanguageServerFeatureOptions languageServerFeatureOptions,
+    [ImportMany] IEnumerable<IDynamicRegistrationProvider> registrationProviders,
+    RazorCohostClientCapabilitiesService razorCohostClientCapabilitiesService,
+    IRazorLoggerFactory loggerFactory)
+    : IRazorCohostDynamicRegistrationService
 {
-    private readonly string _id = Guid.NewGuid().ToString();
     private readonly DocumentFilter[] _filter = [new DocumentFilter()
     {
         Language = Constants.RazorLanguageName,
         Pattern = "**/*.{razor,cshtml}"
     }];
 
     private readonly LanguageServerFeatureOptions _languageServerFeatureOptions = languageServerFeatureOptions;
-    private readonly Lazy<ISemanticTokensLegendService> _semanticTokensLegendService = semanticTokensLegendService;
+    private readonly IEnumerable<IDynamicRegistrationProvider> _registrationProviders = registrationProviders;
+    private readonly RazorCohostClientCapabilitiesService _razorCohostClientCapabilitiesService = razorCohostClientCapabilitiesService;
     private readonly ILogger _logger = loggerFactory.CreateLogger<RazorCohostDynamicRegistrationService>();
 
-    private VSInternalClientCapabilities? _clientCapabilities;
-
-    public bool CanGetClientCapabilities => _clientCapabilities is not null;
-
-    public VSInternalClientCapabilities ClientCapabilities => _clientCapabilities.AssumeNotNull();
-
     public async Task RegisterAsync(string clientCapabilitiesString, RazorCohostRequestContext requestContext, CancellationToken cancellationToken)
     {
         if (!_languageServerFeatureOptions.UseRazorCohostServer)
         {
             return;
         }
 
-        var razorCohostClientLanguageServerManager = requestContext.GetRequiredService<IRazorCohostClientLanguageServerManager>();
-        var semanticTokensRefreshQueue = requestContext.GetRequiredService<IRazorSemanticTokensRefreshQueue>();
+        var clientCapabilities = JsonConvert.DeserializeObject<VSInternalClientCapabilities>(clientCapabilitiesString) ?? new();
 
-        _clientCapabilities = JsonConvert.DeserializeObject<VSInternalClientCapabilities>(clientCapabilitiesString) ?? new();
+        _razorCohostClientCapabilitiesService.SetCapabilities(clientCapabilities);
 
-        // TODO: Get the options from the from the endpoints somehow
-        if (_clientCapabilities.TextDocument?.SemanticTokens?.DynamicRegistration == true)
+        _registrationProviders.TryGetCount(out var providerCount);
+        using var registrations = new PooledArrayBuilder<Registration>(providerCount);
+
+        foreach (var provider in _registrationProviders)
         {
-            semanticTokensRefreshQueue.Initialize(clientCapabilitiesString);
+            var registration = provider.GetRegistration(clientCapabilities, _filter, requestContext);
 
-            await razorCohostClientLanguageServerManager.SendRequestAsync(
-                Methods.ClientRegisterCapabilityName,
-                new RegistrationParams()
-                {
-                    Registrations = [
-                        new Registration()
-                        {
-                            Id = _id,
-                            Method = Methods.TextDocumentSemanticTokensRangeName,
-                            RegisterOptions = new SemanticTokensRegistrationOptions()
-                            {
-                                DocumentSelector = _filter,
-                            }.EnableSemanticTokens(_semanticTokensLegendService.Value)
-                        }
-                    ]
-                },
-                cancellationToken).ConfigureAwait(false);
+            if (registration is not null)
+            {
+                // We don't unregister anything, so we don't need to do anything interesting with Ids
+                registration.Id = Guid.NewGuid().ToString();
+                registrations.Add(registration);
+            }
         }
+
+        var razorCohostClientLanguageServerManager = requestContext.GetRequiredService<IRazorCohostClientLanguageServerManager>();
+
+        await razorCohostClientLanguageServerManager.SendRequestAsync(
+            Methods.ClientRegisterCapabilityName,
+            new RegistrationParams()
+            {
+                Registrations = registrations.ToArray()
+            },
+            cancellationToken).ConfigureAwait(false);
     }
 }