Fixed case where dynamic registration would fail if Register is called before Initialize

src/Server/DefaultLanguageServerFacade.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Reactive.Subjects;
 using System.Threading;
 using System.Threading.Tasks;
 using DryIoc;
@@ -25,7 +26,7 @@ internal class DefaultLanguageServerFacade : LanguageProtocolProxy, ILanguageSer
         private readonly Lazy<ISupportedCapabilities> _supportedCapabilities;
         private readonly TextDocumentIdentifiers _textDocumentIdentifiers;
         private readonly IInsanceHasStarted _instancesHasStarted;
-        private readonly TaskCompletionSource<Unit> _hasStarted;
+        private readonly AsyncSubject<System.Reactive.Unit> _hasStarted;
 
         public DefaultLanguageServerFacade(
             IResponseRouter requestRouter,
@@ -54,7 +55,7 @@ IInsanceHasStarted instancesHasStarted
             _supportedCapabilities = supportedCapabilities;
             _textDocumentIdentifiers = textDocumentIdentifiers;
             _instancesHasStarted = instancesHasStarted;
-            _hasStarted = new TaskCompletionSource<Unit>();
+            _hasStarted = new AsyncSubject<System.Reactive.Unit>();
         }
 
         public ITextDocumentLanguageServer TextDocument => _textDocument.Value;
@@ -74,12 +75,13 @@ public IDisposable Register(Action<ILanguageServerRegistry> registryAction)
                 LanguageServerHelpers.InitHandlers(ResolverContext.Resolve<ILanguageServer>(), result);
             }
 
-            return LanguageServerHelpers.RegisterHandlers(_hasStarted.Task, Client, _workDoneManager.Value, _supportedCapabilities.Value, result);
+            return LanguageServerHelpers.RegisterHandlers(_hasStarted, Client, _workDoneManager.Value, _supportedCapabilities.Value, result);
         }
 
         Task IOnLanguageServerStarted.OnStarted(ILanguageServer client, CancellationToken cancellationToken)
         {
-            _hasStarted.TrySetResult(Unit.Value);
+            _hasStarted.OnNext(System.Reactive.Unit.Default);
+            _hasStarted.OnCompleted();
             return Task.CompletedTask;
         }
     }

src/Server/LanguageServer.cs

@@ -332,15 +332,15 @@ async Task<InitializeResult> IRequestHandler<InternalInitializeParams, Initializ
             var windowCapabilities = ClientSettings.Capabilities.Window ??= new WindowClientCapabilities();
             WorkDoneManager.Initialized(windowCapabilities);
 
-            {
+            _disposable.Add(
                 LanguageServerHelpers.RegisterHandlers(
-                    _initializingTask,
+                    _initializeComplete.Select(z => System.Reactive.Unit.Default),
                     Client,
                     WorkDoneManager,
                     _supportedCapabilities,
                     _collection
-                );
-            }
+                )
+            );
 
             await LanguageProtocolEventingHelper.Run(
                 _initializeDelegates,
@@ -508,7 +508,8 @@ public IDisposable Register(Action<ILanguageServerRegistry> registryAction)
                 LanguageServerHelpers.InitHandlers(this, result);
             }
 
-            return LanguageServerHelpers.RegisterHandlers(_initializingTask, Client, WorkDoneManager, _supportedCapabilities, result);
+            return LanguageServerHelpers.RegisterHandlers(_initializeComplete.Select(z => System.Reactive.Unit.Default), Client, WorkDoneManager, _supportedCapabilities, result);
+//            return LanguageServerHelpers.RegisterHandlers(_initializingTask ?? _initializeComplete.ToTask(), Client, WorkDoneManager, _supportedCapabilities, result);
         }
 
         object IServiceProvider.GetService(Type serviceType) => Services.GetService(serviceType);

src/Server/LanguageServerHelpers.cs

@@ -1,6 +1,7 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Reactive;
 using System.Reactive.Disposables;
 using System.Reactive.Linq;
 using System.Reactive.Threading.Tasks;
@@ -51,14 +52,14 @@ internal static void InitHandlers(ILanguageServer client, CompositeDisposable re
         }
 
         internal static IDisposable RegisterHandlers(
-            Task initializeComplete,
+            IObservable<Unit> initializeComplete,
             IClientLanguageServer client,
             IServerWorkDoneManager serverWorkDoneManager,
             ISupportedCapabilities supportedCapabilities,
             IEnumerable<ILspHandlerDescriptor> collection
         )
         {
-            var registrations = new List<Registration>();
+            var descriptors = new List<ILspHandlerDescriptor>();
             foreach (var descriptor in collection)
             {
                 if (descriptor is LspHandlerDescriptor lspHandlerDescriptor &&
@@ -68,51 +69,77 @@ IEnumerable<ILspHandlerDescriptor> collection
                     continue;
                 }
 
-                if (descriptor.HasCapability && supportedCapabilities.AllowsDynamicRegistration(descriptor.CapabilityType))
-                {
-                    if (descriptor.RegistrationOptions is IWorkDoneProgressOptions wdpo)
-                    {
-                        wdpo.WorkDoneProgress = serverWorkDoneManager.IsSupported;
-                    }
-
-                    registrations.Add(
-                        new Registration {
-                            Id = descriptor.Id.ToString(),
-                            Method = descriptor.Method,
-                            RegisterOptions = descriptor.RegistrationOptions
-                        }
-                    );
-                }
+                descriptors.Add(descriptor);
             }
 
-            // Fire and forget
-            DynamicallyRegisterHandlers(client, initializeComplete, registrations.ToArray()).ToObservable().Subscribe();
-
-            return Disposable.Create(
-                () => {
-                    client.UnregisterCapability(
-                        new UnregistrationParams {
-                            Unregisterations = registrations.ToArray()
-                        }
-                    ).ToObservable().Subscribe();
-                }
-            );
+            return DynamicallyRegisterHandlers(client, initializeComplete, serverWorkDoneManager, supportedCapabilities, descriptors);
         }
 
-        internal static async Task DynamicallyRegisterHandlers(IClientLanguageServer client, Task initializeComplete, Registration[] registrations)
+        internal static IDisposable DynamicallyRegisterHandlers(
+            IClientLanguageServer client,
+            IObservable<Unit> initializeComplete,
+            IServerWorkDoneManager serverWorkDoneManager,
+            ISupportedCapabilities supportedCapabilities,
+            IReadOnlyList<ILspHandlerDescriptor> descriptors
+        )
         {
-            if (registrations.Length == 0)
-                return; // No dynamic registrations supported by client.
+            if (descriptors.Count == 0)
+                return Disposable.Empty; // No dynamic registrations supported by client.
+
+            var disposable = new CompositeDisposable();
 
-            var @params = new RegistrationParams { Registrations = registrations };
+            var result = initializeComplete
+                        .LastOrDefaultAsync()
+                        .Select(
+                             _ => {
+                                 var registrations = new List<Registration>();
+                                 foreach (var descriptor in descriptors)
+                                 {
+                                     if (descriptor.HasCapability && supportedCapabilities.AllowsDynamicRegistration(descriptor.CapabilityType))
+                                     {
+                                         if (descriptor.RegistrationOptions is IWorkDoneProgressOptions wdpo)
+                                         {
+                                             wdpo.WorkDoneProgress = serverWorkDoneManager.IsSupported;
+                                         }
 
-            await initializeComplete;
+                                         registrations.Add(
+                                             new Registration {
+                                                 Id = descriptor.Id.ToString(),
+                                                 Method = descriptor.Method,
+                                                 RegisterOptions = descriptor.RegistrationOptions
+                                             }
+                                         );
+                                     }
+                                 }
 
-            await client.RegisterCapability(@params);
+                                 return registrations;
+                             }
+                         )
+                        .SelectMany(
+                             registrations => Observable.FromAsync(ct => client.RegisterCapability(new RegistrationParams { Registrations = registrations }, ct)), (a, b) => a
+                         )
+                        .Aggregate((z, b) => z)
+                        .Subscribe(
+                             registrations => {
+                                 disposable.Add(
+                                     Disposable.Create(
+                                         () => {
+                                             client.UnregisterCapability(
+                                                 new UnregistrationParams {
+                                                     Unregisterations = registrations
+                                                 }
+                                             ).ToObservable().Subscribe();
+                                         }
+                                     )
+                                 );
+                             }
+                         );
+            disposable.Add(result);
+            return disposable;
         }
 
         internal static IDisposable RegisterHandlers(
-            Task initializeComplete,
+            IObservable<Unit> initializeComplete,
             IClientLanguageServer client,
             IServerWorkDoneManager serverWorkDoneManager,
             ISupportedCapabilities supportedCapabilities,
@@ -145,4 +172,4 @@ IDisposable handlerDisposable
             return cd;
         }
     }
-}
+}

test/Lsp.Tests/Integration/InitializationTests.cs

@@ -1,4 +1,6 @@
 using System.Collections.Generic;
+using System.Reactive.Linq;
+using System.Reactive.Threading.Tasks;
 using System.Threading;
 using System.Threading.Tasks;
 using FluentAssertions;
@@ -42,11 +44,34 @@ public async Task Facades_should_be_resolvable()
             server.Services.GetService<ILanguageServerFacade>().Should().NotBeNull();
             // This ensures that the facade made it.
             var response = await client.RequestCodeLens(new CodeLensParams(), CancellationToken);
+            response.Should().NotBeNull();
+        }
+
+        [Fact]
+        public async Task Should_Be_Able_To_Register_Before_Initialize()
+        {
+            var (client, server) = Create(options => options.EnableDynamicRegistration().EnableAllCapabilities(), options => {});
+
+            server.Register(
+                r => {
+                    r.AddHandler<CodeLensHandlerA>();
+                }
+            );
+
+            await Observable.FromAsync(client.Initialize)
+                            .ForkJoin(Observable.FromAsync(server.Initialize), (a, b) => ( client, server ))
+                            .ToTask(CancellationToken);
+
+            client.Services.GetService<ILanguageClientFacade>().Should().NotBeNull();
+            server.Services.GetService<ILanguageServerFacade>().Should().NotBeNull();
+            // This ensures that the facade made it.
+            var response = await client.RequestCodeLens(new CodeLensParams(), CancellationToken);
+            response.Should().NotBeNull();
         }
 
         class CodeLensHandlerA : CodeLensHandler
         {
-            public CodeLensHandlerA(ILanguageServerFacade languageServerFacade) : base(new CodeLensRegistrationOptions())
+            public CodeLensHandlerA(ILanguageServerFacade languageServerFacade) : base(new CodeLensRegistrationOptions() {  })
             {
                 languageServerFacade.Should().NotBeNull();
             }

test/Lsp.Tests/Integration/WorkspaceFolderTests.cs

@@ -26,6 +26,8 @@ public class WorkspaceFolderTests : LanguageProtocolTestBase
     {
         public WorkspaceFolderTests(ITestOutputHelper outputHelper) : base(
             new JsonRpcTestOptions().ConfigureForXUnit(outputHelper, LogEventLevel.Verbose)
+                                    .WithTimeout(TimeSpan.FromMilliseconds(1200))
+                                    .WithWaitTime(TimeSpan.FromMilliseconds(600))
         )
         {
         }