fixup! fixup! Implement LSP-based code city creation #686

Assets/SEE/Tools/LSP/LSPHandler.cs

@@ -12,6 +12,7 @@
 using OmniSharp.Extensions.LanguageServer.Protocol.General;
 using OmniSharp.Extensions.LanguageServer.Protocol.Models;
 using OmniSharp.Extensions.LanguageServer.Protocol.Server.Capabilities;
+using OmniSharp.Extensions.LanguageServer.Protocol.Window;
 using SEE.UI;
 using SEE.Utils;
 using UnityEngine;
@@ -139,6 +140,7 @@ public async UniTask InitializeAsync()
                 return;
             }
 
+            HashSet<ProgressToken> initialWork = new();
             IDisposable spinner = LoadingSpinner.ShowIndeterminate("Initializing language server...");
             try
             {
@@ -200,18 +202,47 @@ public async UniTask InitializeAsync()
                                                           .WithContentModifiedSupport(false)
                                                           .WithInitializationOptions(Server.InitOptions)
                                                           .DisableProgressTokens()
-                                                          .WithWorkspaceFolder(ProjectPath, "Main"));
+                                                          .WithWorkspaceFolder(ProjectPath, "Main")
+                                                          .OnWorkDoneProgressCreate(HandleInitialWorkDoneProgress));
                 await Client.Initialize(cancellationToken);
-                // FIXME: We need to wait a certain amount until the files are indexed.
-                //        Use progress notifications for this instead of this hack.
-                await UniTask.Delay(TimeSpan.FromSeconds(5), cancellationToken: cancellationToken);
+                do
+                {
+                    // We wait until the initial work is done.
+                    // We detect this by checking if any work progress notifications have been sent,
+                    // and then wait until the progress is done. As soon as there are 500ms without
+                    // any progress, we assume the initial work is done.
+                    await AsyncUtils.RunWithTimeoutAsync(t => UniTask.WaitUntil(() => initialWork.Count == 0,
+                                                                                cancellationToken: t),
+                                                         // We allow more time for the initial work to be done.
+                                                         TimeoutSpan * 5);
+                    await UniTask.Delay(TimeSpan.FromMilliseconds(500), cancellationToken: cancellationToken);
+                } while (initialWork.Count > 0);
                 IsReady = true;
             }
             finally
             {
                 semaphore.Release();
                 spinner.Dispose();
             }
+            return;
+
+            void HandleInitialWorkDoneProgress(WorkDoneProgressCreateParams progressParams)
+            {
+                if (!IsReady && progressParams.Token != null)
+                {
+                    initialWork.Add(progressParams.Token);
+                    MonitorInitialWorkDoneProgress(progressParams.Token).Forget();
+                }
+            }
+
+            async UniTaskVoid MonitorInitialWorkDoneProgress(ProgressToken token)
+            {
+                await foreach (WorkDoneProgress _ in Client.WorkDoneManager.Monitor(token).ToUniTaskAsyncEnumerable()
+                                                           .Where(x => x.Kind == WorkDoneProgressKind.End))
+                {
+                    initialWork.Remove(token);
+                }
+            }
         }
 
         /// <summary>
@@ -433,7 +464,7 @@ public async UniTask ShutdownAsync()
             {
                 if (Client != null)
                 {
-                    await AsyncUtils.RunWithTimeoutAsync(_ => Client.Shutdown().AsUniTask().ContinueWith(() => new object()), TimeoutSpan, throwOnTimeout: false);
+                    await AsyncUtils.RunWithTimeoutAsync(_ => Client.Shutdown().AsUniTask(), TimeoutSpan, throwOnTimeout: false);
                 }
             }
             catch (InvalidParametersException)

Assets/SEE/Utils/AsyncUtils.cs

@@ -62,6 +62,29 @@ public static async UniTask<T> RunWithTimeoutAsync<T>(Func<CancellationToken, Un
             return result;
         }
 
+        /// <summary>
+        /// Runs the given <paramref name="task"/> with a <paramref name="timeout"/>.
+        /// Note that a timeout of <see cref="TimeSpan.Zero"/> will cause no timeout to be applied.
+        ///
+        /// If the task does not complete within the given timeout and <paramref name="throwOnTimeout"/> is set to true,
+        /// a <see cref="TimeoutException"/> is thrown.
+        /// </summary>
+        /// <param name="task">The task to run. It should accept a <see cref="CancellationToken"/> as an argument,
+        /// which will be used to cancel the task if it exceeds the timeout.</param>
+        /// <param name="timeout">The maximum time to wait for the task to complete.</param>
+        /// <param name="throwOnTimeout">Whether to throw a <see cref="TimeoutException"/> if the task times out.</param>
+        /// <exception cref="TimeoutException">Thrown if the task does not complete within the given timeout and
+        /// <paramref name="throwOnTimeout"/> is set to <c>true</c>.</exception>
+        public static async UniTask RunWithTimeoutAsync(Func<CancellationToken, UniTask> task, TimeSpan timeout,
+                                                        bool throwOnTimeout = true)
+        {
+            await RunWithTimeoutAsync<byte>(async token =>
+            {
+                await task(token);
+                return default;
+            }, timeout, throwOnTimeout: throwOnTimeout);
+        }
+
         /// <summary>
         /// Observes the given <paramref name="task"/>'s observable until it completes or <paramref name="timeout"/>
         /// elapses. The results are returned as an asynchronous enumerable that can, for example, be iterated
@@ -139,10 +162,10 @@ public static async UniTask<T> RunOnMainThreadAsync<T>(Func<T> action)
         /// </remarks>
         public static async UniTask RunOnMainThreadAsync(Action action)
         {
-            await RunOnMainThreadAsync(() =>
+            await RunOnMainThreadAsync<byte>(() =>
             {
                 action();
-                return UniTask.CompletedTask;
+                return default;
             });
         }