Initialize feature flags in OOP early

Author: davidwengier

src/Razor/src/Microsoft.CodeAnalysis.Razor.CohostingShared/CohostStartupService.cs

@@ -5,48 +5,60 @@
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.ComponentModel.Composition;
-using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.ExternalAccess.Razor.Cohost;
 using Microsoft.CodeAnalysis.Razor.Logging;
-using Microsoft.CodeAnalysis.Razor.Protocol;
+using Microsoft.CodeAnalysis.Razor.Remote;
 
 namespace Microsoft.VisualStudio.Razor.LanguageClient.Cohost;
 
-[Export(typeof(ICohostStartupService))]
+[Export(typeof(AbstractRazorCohostLifecycleService))]
 [method: ImportingConstructor]
 internal sealed class CohostStartupService(
     [ImportMany] IEnumerable<Lazy<IRazorCohostStartupService>> lazyStartupServices,
-    ILoggerFactory loggerFactory) : ICohostStartupService
+    IRemoteServiceInvoker remoteServiceInvoker,
+    ILoggerFactory loggerFactory) : AbstractRazorCohostLifecycleService
 {
     private readonly ImmutableArray<Lazy<IRazorCohostStartupService>> _lazyStartupServices = [.. lazyStartupServices];
-    private readonly ILogger _logger = loggerFactory.GetOrCreateLogger<CohostStartupService>();
+    private readonly IRemoteServiceInvoker _remoteServiceInvoker = remoteServiceInvoker;
 
-    public async Task StartupAsync(string clientCapabilitiesString, RazorCohostRequestContext requestContext, CancellationToken cancellationToken)
+    public override Task LspServerIntializedAsync(CancellationToken cancellationToken)
     {
-        var clientCapabilities = JsonSerializer.Deserialize<VSInternalClientCapabilities>(clientCapabilitiesString, JsonHelpers.JsonSerializerOptions) ?? new();
+        return _remoteServiceInvoker.InitializeAsync().AsTask();
+    }
+
+    public override async Task RazorActivatedAsync(ClientCapabilities clientCapabilities, RazorCohostRequestContext requestContext, CancellationToken cancellationToken)
+    {
+        // Normally loggers are fields, but this service gets created early so it's better to avoid the work in case Razor
+        // never gets activated.
+        var logger = loggerFactory.GetOrCreateLogger<CohostStartupService>();
 
+        var capabilities = clientCapabilities.ToVSInternalClientCapabilities();
         var providers = _lazyStartupServices.SelectAndOrderByAsArray(p => p.Value, p => p.Order);
 
         foreach (var provider in providers)
         {
             if (cancellationToken.IsCancellationRequested)
             {
-                _logger.LogInformation($"Razor extension startup cancelled.");
+                logger.LogInformation($"Razor extension startup cancelled.");
                 return;
             }
 
             try
             {
-                await provider.StartupAsync(clientCapabilities, requestContext, cancellationToken).ConfigureAwait(false);
+                await provider.StartupAsync(capabilities, requestContext, cancellationToken).ConfigureAwait(false);
             }
             catch (Exception ex) when (ex is not OperationCanceledException)
             {
-                _logger.LogError(ex, $"Error initializing Razor startup service '{provider.GetType().Name}'");
+                logger.LogError(ex, $"Error initializing Razor startup service '{provider.GetType().Name}'");
             }
         }
 
-        _logger.LogInformation($"Razor extension startup finished.");
+        logger.LogInformation($"Razor extension startup finished.");
+    }
+
+    public override void Dispose()
+    {
     }
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Remote/IRemoteServiceInvoker.cs

@@ -11,6 +11,8 @@ namespace Microsoft.CodeAnalysis.Razor.Remote;
 
 internal interface IRemoteServiceInvoker
 {
+    ValueTask InitializeAsync();
+
     ValueTask<TResult?> TryInvokeAsync<TService, TResult>(
         Solution solution,
         Func<TService, RazorPinnedSolutionInfoWrapper, CancellationToken, ValueTask<TResult>> invocation,

src/Razor/src/Microsoft.VisualStudio.LanguageServices.Razor/Remote/RemoteServiceInvoker.cs

@@ -130,7 +130,7 @@ private async static Task<RazorRemoteHostClient> GetJsonClientAsync(IWorkspacePr
             ?? throw new InvalidOperationException($"Couldn't retrieve {nameof(RazorRemoteHostClient)} for JSON serialization.");
     }
 
-    private ValueTask InitializeAsync()
+    public ValueTask InitializeAsync()
     {
         var oopInitialized = _initializeOOPTask is { Status: TaskStatus.RanToCompletion };
         var lspInitialized = _initializeLspTask is { Status: TaskStatus.RanToCompletion };

src/Razor/src/Microsoft.VisualStudioCode.RazorExtension/Services/VSCodeRemoteServiceInvoker.cs

@@ -29,6 +29,8 @@ internal class VSCodeRemoteServiceInvoker(
     private readonly Lock _serviceLock = new();
     private readonly VSCodeBrokeredServiceInterceptor _serviceInterceptor = new();
 
+    public ValueTask InitializeAsync() => ValueTask.CompletedTask; // Initialization in VS Code is handled in VSCodeRemoteServicesInitializer
+
     public async ValueTask<TResult?> TryInvokeAsync<TService, TResult>(
         Solution solution,
         Func<TService, RazorPinnedSolutionInfoWrapper, CancellationToken, ValueTask<TResult>> invocation,

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

@@ -326,6 +326,8 @@ private class RemoteServiceInvoker(TextDocument document, Func<Task>? generateTa
         public bool OOPReturnsNull { get; set; }
         public Func<Task>? GenerateTask { get; set; } = generateTask;
 
+        public ValueTask InitializeAsync() => throw new NotImplementedException();
+
         public async ValueTask<TResult?> TryInvokeAsync<TService, TResult>(Solution solution, Func<TService, RazorPinnedSolutionInfoWrapper, CancellationToken, ValueTask<TResult>> invocation, CancellationToken cancellationToken, [CallerFilePath] string? callerFilePath = null, [CallerMemberName] string? callerMemberName = null) where TService : class
         {
             Assert.Equal(typeof(string), typeof(TResult));

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

@@ -111,6 +111,8 @@ private async Task MakeHtmlRequestAsync<TRequest>(TextDocument document, Uri htm
 
     private class RemoteServiceInvoker : IRemoteServiceInvoker
     {
+        public ValueTask InitializeAsync() => throw new NotImplementedException();
+
         public ValueTask<TResult?> TryInvokeAsync<TService, TResult>(Solution solution, Func<TService, RazorPinnedSolutionInfoWrapper, CancellationToken, ValueTask<TResult>> invocation, CancellationToken cancellationToken, [CallerFilePath] string? callerFilePath = null, [CallerMemberName] string? callerMemberName = null) where TService : class
         {
             return new((TResult?)(object)"");

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

@@ -28,6 +28,8 @@ internal sealed class TestRemoteServiceInvoker(
     private readonly Dictionary<Type, object> _services = [];
     private readonly ReentrantSemaphore _reentrantSemaphore = ReentrantSemaphore.Create(initialCount: 1, joinableTaskContext);
 
+    public ValueTask InitializeAsync() => throw new NotImplementedException();
+
     private async Task<TService> GetOrCreateServiceAsync<TService>()
         where TService : class
     {