Unitialize Lsp options in OOP when the server is disposed

Author: davidwengier

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

@@ -10,6 +10,7 @@
 using Microsoft.CodeAnalysis.ExternalAccess.Razor.Cohost;
 using Microsoft.CodeAnalysis.Razor.Logging;
 using Microsoft.CodeAnalysis.Razor.Remote;
+using Microsoft.VisualStudio.Threading;
 
 namespace Microsoft.VisualStudio.Razor.LanguageClient.Cohost;
 
@@ -60,5 +61,6 @@ public override async Task RazorActivatedAsync(ClientCapabilities clientCapabili
 
     public override void Dispose()
     {
+        _remoteServiceInvoker.UninitializeLspAsync().Forget();
     }
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Protocol/AbstractClientCapabilitiesService.cs

@@ -13,7 +13,7 @@ internal abstract class AbstractClientCapabilitiesService : IClientCapabilitiesS
 
     public VSInternalClientCapabilities ClientCapabilities => _clientCapabilities ?? throw new InvalidOperationException("Client capabilities requested before initialized.");
 
-    public void SetCapabilities(VSInternalClientCapabilities clientCapabilities)
+    public void SetCapabilities(VSInternalClientCapabilities? clientCapabilities)
     {
         _clientCapabilities = clientCapabilities;
     }

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

@@ -11,4 +11,6 @@ internal interface IRemoteClientInitializationService : IRemoteJsonService
     ValueTask InitializeAsync(RemoteClientInitializationOptions initializationOptions, CancellationToken cancellationToken);
 
     ValueTask InitializeLSPAsync(RemoteClientLSPInitializationOptions lspInitializationOptions, CancellationToken cancellationToken);
+
+    ValueTask UninitializeLspAsync(CancellationToken cancellationToken);
 }

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

@@ -13,6 +13,8 @@ internal interface IRemoteServiceInvoker
 {
     ValueTask InitializeAsync();
 
+    ValueTask UninitializeLspAsync();
+
     ValueTask<TResult?> TryInvokeAsync<TService, TResult>(
         Solution solution,
         Func<TService, RazorPinnedSolutionInfoWrapper, CancellationToken, ValueTask<TResult>> invocation,

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

@@ -19,5 +19,6 @@ public void OnLspInitialized(RemoteClientLSPInitializationOptions options)
 
     public void OnLspUninitialized()
     {
+        SetCapabilities(null);
     }
 }

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

@@ -38,4 +38,16 @@ public ValueTask InitializeLSPAsync(RemoteClientLSPInitializationOptions options
                 return default;
             },
             cancellationToken);
+
+    public ValueTask UninitializeLspAsync(CancellationToken cancellationToken)
+        => RunServiceAsync(ct =>
+            {
+                foreach (var service in _lspLifetimeServices)
+                {
+                    service.OnLspUninitialized();
+                }
+
+                return default;
+            },
+            cancellationToken);
 }

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

@@ -23,16 +23,11 @@ internal sealed class RemoteSemanticTokensLegendService : ISemanticTokensLegendS
 
     public void OnLspInitialized(RemoteClientLSPInitializationOptions options)
     {
-        SetLegend(options.TokenTypes, options.TokenModifiers);
+        _tokenTypes = new SemanticTokenTypes(options.TokenTypes);
+        _tokenModifiers = new SemanticTokenModifiers(options.TokenModifiers);
     }
 
     public void OnLspUninitialized()
     {
     }
-
-    public void SetLegend(string[] tokenTypes, string[] tokenModifiers)
-    {
-        _tokenTypes = new SemanticTokenTypes(tokenTypes);
-        _tokenModifiers = new SemanticTokenModifiers(tokenModifiers);
-    }
 }

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

@@ -48,6 +48,7 @@ internal sealed class RemoteServiceInvoker(
     private readonly object _gate = new();
     private Task? _initializeOOPTask;
     private Task? _initializeLspTask;
+    private CancellationTokenSource? _uninitializeLspTokenSource;
 
     public void Dispose()
     {
@@ -162,6 +163,8 @@ async Task InitializeCoreAsync(bool oopInitialized, bool lspInitialized)
             {
                 lock (_gate)
                 {
+                    _uninitializeLspTokenSource?.Cancel();
+
                     _initializeLspTask ??= InitializeLspAsync(remoteClient);
                 }
 
@@ -216,4 +219,34 @@ Task InitializeLspAsync(RazorRemoteHostClient remoteClient)
             }
         }
     }
+
+    public ValueTask UninitializeLspAsync()
+    {
+        var lspInitialized = _initializeLspTask is { Status: TaskStatus.RanToCompletion };
+
+        lock (_gate)
+        {
+            _uninitializeLspTokenSource?.Cancel();
+
+            _initializeLspTask = null;
+            _uninitializeLspTokenSource = new();
+        }
+
+        return lspInitialized
+            ? new(UninitializeCoreAsync(_uninitializeLspTokenSource.Token))
+            : default;
+
+        async Task UninitializeCoreAsync(CancellationToken cancellationToken)
+        {
+            // Note: IRemoteClientInitializationService is an IRemoteJsonService
+            var remoteClient = await _lazyJsonClient
+                .GetValueAsync(cancellationToken)
+                .ConfigureAwait(false);
+
+            await remoteClient
+                .TryInvokeAsync<IRemoteClientInitializationService>(
+                  (s, ct) => s.UninitializeLspAsync(ct),
+                  cancellationToken);
+        }
+    }
 }

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

@@ -31,6 +31,8 @@ internal class VSCodeRemoteServiceInvoker(
 
     public ValueTask InitializeAsync() => ValueTask.CompletedTask; // Initialization in VS Code is handled in VSCodeRemoteServicesInitializer
 
+    public ValueTask UninitializeLspAsync() => ValueTask.CompletedTask; // When VS Code wants to un-initialize a language server, it just kills the process
+
     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/CohostRoslynRenameTest.cs

@@ -204,6 +204,8 @@ internal void SetInvoker(TestRemoteServiceInvoker remoteServiceInvoker)
 
         public ValueTask InitializeAsync() => throw new NotImplementedException();
 
+        public ValueTask UninitializeLspAsync() => 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
             => _remoteServiceInvoker.AssumeNotNull().TryInvokeAsync(solution, invocation, cancellationToken, callerFilePath, callerMemberName);