Merge pull request #4759 from lkk214/feat/Improve-init-and-loading

feat: Improve ProtocolClient initialization and GUI loading

Author: sestinj

extensions/intellij/src/main/kotlin/com/github/continuedev/continueintellijextension/services/ContinuePluginService.kt

@@ -13,12 +13,16 @@ import com.intellij.openapi.project.DumbAware
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.cancel
+import kotlin.properties.Delegates
 
 @Service(Service.Level.PROJECT)
 class ContinuePluginService : Disposable, DumbAware {
     private val coroutineScope = CoroutineScope(Dispatchers.Main)
     var continuePluginWindow: ContinuePluginToolWindowFactory.ContinuePluginWindow? = null
-    var ideProtocolClient: IdeProtocolClient? = null
+    var listener: (() -> Unit)? = null
+    var ideProtocolClient: IdeProtocolClient? by Delegates.observable(null) { _, _, _ ->
+        synchronized(this) { listener?.also { listener = null }?.invoke() }
+    }
     var coreMessengerManager: CoreMessengerManager? = null
     val coreMessenger: CoreMessenger?
         get() = coreMessengerManager?.coreMessenger
@@ -42,6 +46,25 @@ class ContinuePluginService : Disposable, DumbAware {
         continuePluginWindow?.browser?.sendToWebview(messageType, data, messageId)
     }
 
+    /**
+     * Add a listener for protocolClient initialization.
+     * Currently, only one needs to be processed. If there are more than one,
+     * we can use an array to add listeners to ensure that the message is processed.
+     */
+    fun onProtocolClientInitialized(listener: () -> Unit) {
+        if (ideProtocolClient == null) {
+            synchronized(this) {
+                if (ideProtocolClient == null) {
+                    this.listener = listener
+                } else {
+                    listener()
+                }
+            }
+        } else {
+            listener()
+        }
+    }
+
     fun updateLastFileSaveTimestamp() {
         ideProtocolClient?.updateLastFileSaveTimestamp()
     }

extensions/intellij/src/main/kotlin/com/github/continuedev/continueintellijextension/toolWindow/ContinueBrowser.kt

@@ -1,7 +1,6 @@
 package com.github.continuedev.continueintellijextension.toolWindow
 
 import com.github.continuedev.continueintellijextension.activities.ContinuePluginDisposable
-import com.github.continuedev.continueintellijextension.constants.MessageTypes
 import com.github.continuedev.continueintellijextension.constants.MessageTypes.Companion.PASS_THROUGH_TO_CORE
 import com.github.continuedev.continueintellijextension.factories.CustomSchemeHandlerFactory
 import com.github.continuedev.continueintellijextension.services.ContinueExtensionSettings
@@ -10,11 +9,11 @@ import com.github.continuedev.continueintellijextension.utils.uuid
 import com.google.gson.Gson
 import com.google.gson.JsonObject
 import com.google.gson.JsonParser
-import com.intellij.openapi.components.ServiceManager
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Disposer
 import com.intellij.ui.jcef.*
 import com.intellij.ui.jcef.JBCefClient.Properties
+import com.intellij.util.application
 import org.cef.CefApp
 import org.cef.browser.CefBrowser
 import org.cef.handler.CefLoadHandlerAdapter
@@ -30,8 +29,10 @@ class ContinueBrowser(val project: Project, url: String) {
 
     val browser: JBCefBrowser
 
+    val continuePluginService: ContinuePluginService = project.getService(ContinuePluginService::class.java)
+
     init {
-        val isOSREnabled = ServiceManager.getService(ContinueExtensionSettings::class.java).continueState.enableOSR
+        val isOSREnabled = application.getService(ContinueExtensionSettings::class.java).continueState.enableOSR
 
         this.browser = JBCefBrowser.createBuilder().setOffScreenRendering(isOSREnabled).build().apply {
             // To avoid using System.setProperty to affect other plugins,
@@ -42,7 +43,6 @@ class ContinueBrowser(val project: Project, url: String) {
         }
 
         registerAppSchemeHandler()
-        browser.loadURL(url)
         Disposer.register(ContinuePluginDisposable.getInstance(project), browser)
 
         // Listen for events sent from browser
@@ -55,11 +55,6 @@ class ContinueBrowser(val project: Project, url: String) {
             val data = json.get("data")
             val messageId = json.get("messageId")?.asString
 
-            val continuePluginService = ServiceManager.getService(
-                project,
-                ContinuePluginService::class.java
-            )
-
             val respond = fun(data: Any?) {
                 sendToWebview(messageType, data, messageId ?: uuid())
             }
@@ -101,6 +96,13 @@ class ContinueBrowser(val project: Project, url: String) {
             }
         }, browser.cefBrowser)
 
+        // Load the url only after the protocolClient is initialized,
+        // otherwise some messages will be lost, which are some configurations when the page is loaded.
+        // Moreover, we should add LoadHandler before loading the url.
+        continuePluginService.onProtocolClientInitialized {
+            browser.loadURL(url)
+        }
+
     }
 
     fun executeJavaScript(browser: CefBrowser?, myJSQueryOpenInBrowser: JBCefJSQuery) {