Introduce a lifetime service in remote

Author: davidwengier

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/Initialization/ILspLifetimeService.cs

@@ -0,0 +1,12 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using Microsoft.CodeAnalysis.Razor.Remote;
+
+namespace Microsoft.CodeAnalysis.Remote.Razor;
+
+internal interface ILspLifetimeService
+{
+    void OnLspInitialized(RemoteClientLSPInitializationOptions options);
+    void OnLspUninitialized();
+}

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/Initialization/RemoteClientCapabilitiesService.cs

@@ -3,10 +3,21 @@
 
 using System.Composition;
 using Microsoft.CodeAnalysis.Razor.Protocol;
+using Microsoft.CodeAnalysis.Razor.Remote;
 
 namespace Microsoft.CodeAnalysis.Remote.Razor;
 
 [Shared]
 [Export(typeof(IClientCapabilitiesService))]
-[Export(typeof(RemoteClientCapabilitiesService))]
-internal sealed class RemoteClientCapabilitiesService : AbstractClientCapabilitiesService;
+[Export(typeof(ILspLifetimeService))]
+internal sealed class RemoteClientCapabilitiesService : AbstractClientCapabilitiesService, ILspLifetimeService
+{
+    public void OnLspInitialized(RemoteClientLSPInitializationOptions options)
+    {
+        SetCapabilities(options.ClientCapabilities);
+    }
+
+    public void OnLspUninitialized()
+    {
+    }
+}

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/Initialization/RemoteClientInitializationService.cs

@@ -1,10 +1,10 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.Razor.Remote;
-using Microsoft.CodeAnalysis.Remote.Razor.SemanticTokens;
 
 namespace Microsoft.CodeAnalysis.Remote.Razor;
 
@@ -16,9 +16,8 @@ protected override IRemoteClientInitializationService CreateService(in ServiceAr
             => new RemoteClientInitializationService(in args);
     }
 
-    private readonly RemoteClientCapabilitiesService _remoteClientCapabilitiesService = args.ExportProvider.GetExportedValue<RemoteClientCapabilitiesService>();
     private readonly RemoteLanguageServerFeatureOptions _remoteLanguageServerFeatureOptions = args.ExportProvider.GetExportedValue<RemoteLanguageServerFeatureOptions>();
-    private readonly RemoteSemanticTokensLegendService _remoteSemanticTokensLegendService = args.ExportProvider.GetExportedValue<RemoteSemanticTokensLegendService>();
+    private readonly IEnumerable<ILspLifetimeService> _lspLifetimeServices = args.ExportProvider.GetExportedValues<ILspLifetimeService>();
 
     public ValueTask InitializeAsync(RemoteClientInitializationOptions options, CancellationToken cancellationToken)
         => RunServiceAsync(ct =>
@@ -31,8 +30,11 @@ public ValueTask InitializeAsync(RemoteClientInitializationOptions options, Canc
     public ValueTask InitializeLSPAsync(RemoteClientLSPInitializationOptions options, CancellationToken cancellationToken)
         => RunServiceAsync(ct =>
             {
-                _remoteSemanticTokensLegendService.SetLegend(options.TokenTypes, options.TokenModifiers);
-                _remoteClientCapabilitiesService.SetCapabilities(options.ClientCapabilities);
+                foreach (var service in _lspLifetimeServices)
+                {
+                    service.OnLspInitialized(options);
+                }
+
                 return default;
             },
             cancellationToken);

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/SemanticTokens/RemoteSemanticTokensLegendService.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Composition;
+using Microsoft.CodeAnalysis.Razor.Remote;
 using Microsoft.CodeAnalysis.Razor.SemanticTokens;
 using SemanticTokenModifiers = Microsoft.CodeAnalysis.Razor.SemanticTokens.SemanticTokenModifiers;
 using SemanticTokenTypes = Microsoft.CodeAnalysis.Razor.SemanticTokens.SemanticTokenTypes;
@@ -10,8 +11,8 @@ namespace Microsoft.CodeAnalysis.Remote.Razor.SemanticTokens;
 
 [Shared]
 [Export(typeof(ISemanticTokensLegendService))]
-[Export(typeof(RemoteSemanticTokensLegendService))]
-internal sealed class RemoteSemanticTokensLegendService : ISemanticTokensLegendService
+[Export(typeof(ILspLifetimeService))]
+internal sealed class RemoteSemanticTokensLegendService : ISemanticTokensLegendService, ILspLifetimeService
 {
     private SemanticTokenModifiers _tokenModifiers = null!;
     private SemanticTokenTypes _tokenTypes = null!;
@@ -20,6 +21,15 @@ internal sealed class RemoteSemanticTokensLegendService : ISemanticTokensLegendS
 
     public SemanticTokenTypes TokenTypes => _tokenTypes;
 
+    public void OnLspInitialized(RemoteClientLSPInitializationOptions options)
+    {
+        SetLegend(options.TokenTypes, options.TokenModifiers);
+    }
+
+    public void OnLspUninitialized()
+    {
+    }
+
     public void SetLegend(string[] tokenTypes, string[] tokenModifiers)
     {
         _tokenTypes = new SemanticTokenTypes(tokenTypes);

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

@@ -16,11 +16,11 @@
 using Microsoft.CodeAnalysis.CSharp;
 using Microsoft.CodeAnalysis.Diagnostics;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;
+using Microsoft.CodeAnalysis.Razor.Protocol;
 using Microsoft.CodeAnalysis.Razor.Remote;
 using Microsoft.CodeAnalysis.Razor.Workspaces;
 using Microsoft.CodeAnalysis.Remote.Razor;
 using Microsoft.CodeAnalysis.Remote.Razor.Logging;
-using Microsoft.CodeAnalysis.Remote.Razor.SemanticTokens;
 using Microsoft.CodeAnalysis.Text;
 using Microsoft.NET.Sdk.Razor.SourceGenerators;
 using Microsoft.VisualStudio.Composition;
@@ -43,8 +43,7 @@ public abstract class CohostEndpointTestBase(ITestOutputHelper testOutputHelper)
     private protected TestRemoteServiceInvoker RemoteServiceInvoker => _remoteServiceInvoker.AssumeNotNull();
     private protected IFilePathService FilePathService => _filePathService.AssumeNotNull();
     private protected RemoteLanguageServerFeatureOptions FeatureOptions => OOPExportProvider.GetExportedValue<RemoteLanguageServerFeatureOptions>();
-    private protected RemoteClientCapabilitiesService ClientCapabilitiesService => OOPExportProvider.GetExportedValue<RemoteClientCapabilitiesService>();
-    private protected RemoteSemanticTokensLegendService SemanticTokensLegendService => OOPExportProvider.GetExportedValue<RemoteSemanticTokensLegendService>();
+    private protected RemoteClientCapabilitiesService ClientCapabilitiesService => (RemoteClientCapabilitiesService)OOPExportProvider.GetExportedValue<IClientCapabilitiesService>();
 
     /// <summary>
     /// The export provider for Razor OOP services (not Roslyn)
@@ -145,8 +144,23 @@ private protected void UpdateClientInitializationOptions(Func<RemoteClientInitia
     private protected void UpdateClientLSPInitializationOptions(Func<RemoteClientLSPInitializationOptions, RemoteClientLSPInitializationOptions> mutation)
     {
         _clientLSPInitializationOptions = mutation(_clientLSPInitializationOptions);
-        ClientCapabilitiesService.SetCapabilities(_clientLSPInitializationOptions.ClientCapabilities);
-        SemanticTokensLegendService.SetLegend(_clientLSPInitializationOptions.TokenTypes, _clientLSPInitializationOptions.TokenModifiers);
+
+        var lifetimeServices = OOPExportProvider.GetExportedValues<ILspLifetimeService>();
+        foreach (var service in lifetimeServices)
+        {
+            service.OnLspInitialized(_clientLSPInitializationOptions);
+        }
+    }
+
+    protected override Task DisposeAsync()
+    {
+        var lifetimeServices = OOPExportProvider.GetExportedValues<ILspLifetimeService>();
+        foreach (var service in lifetimeServices)
+        {
+            service.OnLspUninitialized();
+        }
+
+        return Task.CompletedTask;
     }
 
     private protected virtual TestComposition ConfigureRoslynDevenvComposition(TestComposition composition)