Unsubscribe from UIContext events when the LSP server exits

src/EditorFeatures/Core/Shared/Utilities/IUIContextActivationService.cs

@@ -11,5 +11,5 @@ internal interface IUIContextActivationService
     /// <summary>
     /// Executes the specified action when the UIContext first becomes active, or immediately if it is already active
     /// </summary>
-    void ExecuteWhenActivated(Guid uiContext, Action action);
+    IDisposable ExecuteWhenActivated(Guid uiContext, Action action);
 }

src/Tools/ExternalAccess/Razor/Cohost/RazorDynamicRegistrationServiceFactory.cs

@@ -36,9 +36,12 @@ private class RazorDynamicRegistrationService(
         IClientLanguageServerManager? clientLanguageServerManager) : ILspService, IOnInitialized, IDisposable
     {
         private readonly CancellationTokenSource _disposalTokenSource = new();
+        private IDisposable? _activation;
 
         public void Dispose()
         {
+            _activation?.Dispose();
+            _activation = null;
             _disposalTokenSource.Cancel();
         }
 
@@ -56,7 +59,7 @@ public Task OnInitializedAsync(ClientCapabilities clientCapabilities, RequestCon
             }
             else
             {
-                uIContextActivationService.ExecuteWhenActivated(Constants.RazorCohostingUIContext, InitializeRazor);
+                _activation = uIContextActivationService.ExecuteWhenActivated(Constants.RazorCohostingUIContext, InitializeRazor);
             }
 
             return Task.CompletedTask;

src/VisualStudio/Core/Def/Implementation/VisualStudioUIContextActivationService.cs

@@ -7,6 +7,7 @@
 using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
 using Microsoft.CodeAnalysis.Host.Mef;
 using Microsoft.VisualStudio.Shell;
+using Roslyn.Utilities;
 
 namespace Microsoft.VisualStudio.LanguageServices.Implementation;
 
@@ -15,9 +16,60 @@ namespace Microsoft.VisualStudio.LanguageServices.Implementation;
 [method: Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
 internal sealed class VisualStudioUIContextActivationService() : IUIContextActivationService
 {
-    public void ExecuteWhenActivated(Guid uiContext, Action action)
+    public IDisposable ExecuteWhenActivated(Guid uiContext, Action action)
     {
         var context = UIContext.FromUIContextGuid(uiContext);
-        context.WhenActivated(action);
+        if (context.IsActive)
+        {
+            action();
+            return EmptyDisposable.Instance;
+        }
+        else
+        {
+            return new WhenActivatedHandler(context, action);
+        }
+    }
+
+    private sealed class EmptyDisposable : IDisposable
+    {
+        public static EmptyDisposable Instance = new();
+
+        public void Dispose()
+        {
+        }
+    }
+
+    private sealed class WhenActivatedHandler : IDisposable
+    {
+        private readonly Action _action;
+        private UIContext? _context;
+
+        public WhenActivatedHandler(UIContext context, Action action)
+        {
+            _context = context;
+            _action = action;
+            _context.UIContextChanged += OnContextChanged;
+        }
+
+        public void Dispose()
+        {
+            if (_context is not null)
+            {
+                _context.UIContextChanged -= OnContextChanged;
+            }
+
+            _context = null;
+        }
+
+        private void OnContextChanged(object sender, UIContextChangedEventArgs e)
+        {
+            Contract.ThrowIfNull(_context);
+
+            if (e.Activated)
+            {
+                _action();
+                Dispose();
+            }
+        }
     }
 }