Wait for truffle context init (#1617)

engine/language-server/src/main/scala/org/enso/languageserver/boot/ResourcesInitialization.scala

@@ -38,7 +38,7 @@ object ResourcesInitialization {
     val resources = Seq(
       new DirectoriesInitialization(directoriesConfig),
       new RepoInitialization(eventStream, suggestionsRepo, versionsRepo),
-      new TruffleContextInitialization(truffleContext)
+      new TruffleContextInitialization(eventStream, truffleContext)
     )
     new SequentialResourcesInitialization(resources)
   }

engine/language-server/src/main/scala/org/enso/languageserver/boot/resource/TruffleContextInitialization.scala

@@ -1,4 +1,7 @@
 package org.enso.languageserver.boot.resource
+
+import akka.event.EventStream
+import org.enso.languageserver.event.InitializedEvent
 import org.enso.polyglot.LanguageInfo
 import org.graalvm.polyglot.Context
 import org.slf4j.LoggerFactory
@@ -7,9 +10,13 @@ import scala.concurrent.{ExecutionContext, Future}
 
 /** Initialize the Truffle context.
   *
+  * @param eventStream akka events stream
   * @param truffleContext the Truffle context
   */
-class TruffleContextInitialization(truffleContext: Context)(implicit
+class TruffleContextInitialization(
+  eventStream: EventStream,
+  truffleContext: Context
+)(implicit
   ec: ExecutionContext
 ) extends InitializationComponent {
 
@@ -19,6 +26,7 @@ class TruffleContextInitialization(truffleContext: Context)(implicit
   override def init(): Future[InitializationComponent.Initialized.type] =
     Future {
       truffleContext.initialize(LanguageInfo.ID)
+      eventStream.publish(InitializedEvent.TruffleContextInitialized)
       log.info("Initialized Runtime context.")
       InitializationComponent.Initialized
     }

engine/language-server/src/main/scala/org/enso/languageserver/event/InitializedEvent.scala

@@ -7,6 +7,7 @@ object InitializedEvent {
 
   case object SuggestionsRepoInitialized  extends InitializedEvent
   case object FileVersionsRepoInitialized extends InitializedEvent
+  case object TruffleContextInitialized   extends InitializedEvent
   case object InitializationFinished      extends InitializedEvent
   case object InitializationFailed        extends InitializedEvent
 }

engine/language-server/src/main/scala/org/enso/languageserver/search/SuggestionsHandler.scala

@@ -1,5 +1,7 @@
 package org.enso.languageserver.search
 
+import java.util.UUID
+
 import akka.actor.{Actor, ActorLogging, ActorRef, Props, Stash}
 import akka.pattern.{ask, pipe}
 import org.enso.languageserver.capability.CapabilityProtocol.{
@@ -98,10 +100,8 @@ final class SuggestionsHandler(
     context.system.eventStream.subscribe(self, classOf[FileDeletedEvent])
     context.system.eventStream
       .subscribe(self, InitializedEvent.SuggestionsRepoInitialized.getClass)
-
+    context.system.eventStream
-    runtimeConnector
+      .subscribe(self, InitializedEvent.TruffleContextInitialized.getClass)
-      .ask(Api.Request(Api.GetTypeGraphRequest()))(timeout, self)
-      .pipeTo(self)
 
     config.contentRoots.foreach { case (_, contentRoot) =>
       PackageManager.Default
@@ -119,16 +119,26 @@ final class SuggestionsHandler(
         .renameProject(oldName, newName)
         .map(_ => ProjectNameUpdated(newName))
         .pipeTo(self)
+
     case ProjectNameUpdated(name) =>
       tryInitialize(init.copy(project = Some(name)))
+
     case InitializedEvent.SuggestionsRepoInitialized =>
       tryInitialize(
         init.copy(suggestions =
           Some(InitializedEvent.SuggestionsRepoInitialized)
         )
       )
-    case Api.GetTypeGraphResponse(g) =>
+
+    case InitializedEvent.TruffleContextInitialized =>
+      val requestId = UUID.randomUUID()
+      runtimeConnector
+        .ask(Api.Request(requestId, Api.GetTypeGraphRequest()))(timeout, self)
+        .pipeTo(self)
+
+    case Api.Response(_, Api.GetTypeGraphResponse(g)) =>
       tryInitialize(init.copy(typeGraph = Some(g)))
+
     case _ => stash()
   }
 

engine/language-server/src/test/scala/org/enso/languageserver/search/SuggestionsHandlerSpec.scala

@@ -477,15 +477,6 @@ class SuggestionsHandlerSpec
         Await.ready(repo.insert(Suggestions.atom), Timeout)
         handler ! SearchProtocol.InvalidateSuggestionsDatabase
 
-        connector.expectMsgClass(classOf[Api.Request]) match {
-          case Api.Request(_, Api.GetTypeGraphRequest()) =>
-          case Api.Request(_, msg) =>
-            fail(s"Runtime connector receive unexpected message: $msg")
-        }
-        connector.reply(
-          Api.Response(Api.GetTypeGraphResponse(buildTestTypeGraph))
-        )
-
         connector.expectMsgClass(classOf[Api.Request]) match {
           case Api.Request(_, Api.InvalidateModulesIndexRequest()) =>
           case Api.Request(_, msg) =>
@@ -628,7 +619,12 @@ class SuggestionsHandlerSpec
         )
       )
     handler ! SuggestionsHandler.ProjectNameUpdated("Test")
-    handler ! Api.GetTypeGraphResponse(buildTestTypeGraph)
+    handler ! InitializedEvent.TruffleContextInitialized
+    runtimeConnector.receiveN(1)
+    handler ! Api.Response(
+      UUID.randomUUID(),
+      Api.GetTypeGraphResponse(buildTestTypeGraph)
+    )
     handler
   }
 

engine/language-server/src/test/scala/org/enso/languageserver/websocket/json/BaseServerTest.scala

@@ -16,6 +16,7 @@ import org.enso.languageserver.boot.resource.{
 import org.enso.languageserver.capability.CapabilityRouter
 import org.enso.languageserver.data._
 import org.enso.languageserver.effect.ZioExec
+import org.enso.languageserver.event.InitializedEvent
 import org.enso.languageserver.filemanager.{
   FileManager,
   FileSystem,
@@ -155,8 +156,12 @@ class BaseServerTest extends JsonRpcServerTestKit {
       )
 
     // initialize
+    suggestionsHandler ! InitializedEvent.TruffleContextInitialized
     runtimeConnectorProbe.receiveN(1)
-    suggestionsHandler ! Api.GetTypeGraphResponse(typeGraph)
+    suggestionsHandler ! Api.Response(
+      UUID.randomUUID(),
+      Api.GetTypeGraphResponse(typeGraph)
+    )
     Await.ready(initializationComponent.init(), timeout)
 
     new JsonConnectionControllerFactory(

engine/runtime/src/main/scala/org/enso/interpreter/instrument/command/CommandFactory.scala

@@ -50,13 +50,13 @@ object CommandFactory {
       case payload: Api.ImportSuggestionRequest =>
         new ImportSuggestionCmd(request.requestId, payload)
 
+      case _: Api.GetTypeGraphRequest =>
+        new GetTypeGraphCommand(request.requestId)
+
       case Api.ShutDownRuntimeServer() =>
         throw new IllegalArgumentException(
           "ShutDownRuntimeServer request is not convertible to command object"
         )
-
-      case _: Api.GetTypeGraphRequest =>
-        new GetTypeGraphCommand(request.requestId)
     }
 
 }