Only preinitialize when the UI context is active

Author: davidwengier

src/Tools/ExternalAccess/Razor/Features/Cohost/Constants.cs

@@ -11,6 +11,8 @@ internal static class Constants
 {
     public const string RazorLanguageName = LanguageInfoProvider.RazorLanguageName;
 
-    // The UI context is provided by Razor, so this guid must match the one in https://github.com/dotnet/razor/blob/main/src/Razor/src/Microsoft.VisualStudio.LanguageServices.Razor/RazorConstants.cs
+    // These UI contexts are provided by Razor, so must match https://github.com/dotnet/razor/blob/main/src/Razor/src/Microsoft.VisualStudio.LanguageServices.Razor/RazorConstants.cs
     public static readonly Guid RazorCohostingUIContext = new Guid("6d5b86dc-6b8a-483b-ae30-098a3c7d6774");
+
+    public static readonly Guid RazorCapabilityPresentUIContext = new Guid("2077a158-ee71-484c-be76-350a1d49eaea");
 }

src/Tools/ExternalAccess/Razor/Features/Cohost/RazorStartupServiceFactory.cs

@@ -38,55 +38,73 @@ private class RazorStartupService(
         AbstractRazorCohostLifecycleService? razorCohostLifecycleService) : ILspService, IOnInitialized, IDisposable
     {
         private readonly CancellationTokenSource _disposalTokenSource = new();
-        private IDisposable? _activation;
+        private IDisposable? _cohostActivation;
+        private IDisposable? _razorFilePresentActivation;
 
         public void Dispose()
         {
             razorCohostLifecycleService?.Dispose();
 
-            _activation?.Dispose();
-            _activation = null;
+            _razorFilePresentActivation?.Dispose();
+            _razorFilePresentActivation = null;
+            _cohostActivation?.Dispose();
+            _cohostActivation = null;
             _disposalTokenSource.Cancel();
         }
 
-        public async Task OnInitializedAsync(ClientCapabilities clientCapabilities, RequestContext context, CancellationToken cancellationToken)
+        public Task OnInitializedAsync(ClientCapabilities clientCapabilities, RequestContext context, CancellationToken cancellationToken)
         {
             if (context.ServerKind is not (WellKnownLspServerKinds.AlwaysActiveVSLspServer or WellKnownLspServerKinds.CSharpVisualBasicLspServer))
             {
                 // We have to register this class for Any server, but only want to run in the C# server in VS or VS Code
-                return;
+                return Task.CompletedTask;
             }
 
             if (cohostStartupService is null && razorCohostLifecycleService is null)
             {
-                return;
-            }
-
-            if (razorCohostLifecycleService is not null)
-            {
-                // If we have a cohost lifecycle service, fire pre-initialization, which happens when the LSP server starts up, but before
-                // the UIContext is activated.
-                await razorCohostLifecycleService.LspServerIntializedAsync(cancellationToken).ConfigureAwait(false);
+                return Task.CompletedTask;
             }
 
             if (uIContextActivationService is null)
             {
-                // Outside of VS, we want to initialize immediately.. I think?
+                PreinitializeRazor();
                 InitializeRazor();
             }
             else
             {
-                _activation = uIContextActivationService.ExecuteWhenActivated(Constants.RazorCohostingUIContext, InitializeRazor);
+                // There are two initialization methods for Razor, which looks odd here, but are really controlled by UI contexts.
+                // This method fires for any Roslyn project, but not all Roslyn projects are Razor projects, so the first UI context
+                // triggers where there is a project with a Razor capability present in the solution, and the next is when a Razor file
+                // is opened in the editor. ie these two lines look the same, but really they do different levels of initialization.
+                _razorFilePresentActivation = uIContextActivationService.ExecuteWhenActivated(Constants.RazorCapabilityPresentUIContext, PreinitializeRazor);
+                _cohostActivation = uIContextActivationService.ExecuteWhenActivated(Constants.RazorCohostingUIContext, InitializeRazor);
             }
 
-            return;
+            return Task.CompletedTask;
+
+            void PreinitializeRazor()
+            {
+                this.PreinitializeRazorAsync(_disposalTokenSource.Token).ReportNonFatalErrorAsync();
+            }
 
             void InitializeRazor()
             {
                 this.InitializeRazorAsync(clientCapabilities, context, _disposalTokenSource.Token).ReportNonFatalErrorAsync();
             }
         }
 
+        private async Task PreinitializeRazorAsync(CancellationToken cancellationToken)
+        {
+            if (cancellationToken.IsCancellationRequested) return;
+
+            await TaskScheduler.Default.SwitchTo(alwaysYield: true);
+
+            if (razorCohostLifecycleService is not null)
+            {
+                await razorCohostLifecycleService.LspServerIntializedAsync(cancellationToken).ConfigureAwait(false);
+            }
+        }
+
         private async Task InitializeRazorAsync(ClientCapabilities clientCapabilities, RequestContext context, CancellationToken cancellationToken)
         {
             // The LSP server will dispose us when the server exits, but VS could decide to activate us later.