Merge remote-tracking branch 'origin/feature/q-lsp-chat' into rli/fix-history-restore

Author: rli

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQPanel.kt

@@ -5,8 +5,6 @@ package software.aws.toolkits.jetbrains.services.amazonq.toolwindow
 
 import com.intellij.idea.AppMode
 import com.intellij.openapi.Disposable
-import com.intellij.openapi.application.ApplicationManager
-import com.intellij.openapi.application.runInEdt
 import com.intellij.openapi.components.service
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Disposer
@@ -19,12 +17,15 @@ import com.intellij.ui.dsl.builder.AlignY
 import com.intellij.ui.dsl.builder.panel
 import com.intellij.ui.jcef.JBCefApp
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.launch
-import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.withContext
+import software.aws.toolkits.jetbrains.core.coroutines.EDT
 import software.aws.toolkits.jetbrains.isDeveloperMode
 import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQAppInitContext
 import software.aws.toolkits.jetbrains.services.amazonq.apps.AppConnection
 import software.aws.toolkits.jetbrains.services.amazonq.commands.MessageTypeRegistry
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.artifacts.ArtifactManager
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AsyncChatUiListener
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.FlareUiMessage
@@ -109,17 +110,21 @@ class AmazonQPanel(val project: Project, private val scope: CoroutineScope) : Di
             webviewContainer.add(wrapper)
             wrapper.setContent(loadingPanel)
 
-            ApplicationManager.getApplication().executeOnPooledThread {
-                val webUri = runBlocking { service<ArtifactManager>().fetchArtifact(project).resolve("amazonq-ui.js").toUri() }
-                loadingPanel.stopLoading()
-                runInEdt {
+            scope.launch {
+                val webUri = service<ArtifactManager>().fetchArtifact(project).resolve("amazonq-ui.js").toUri()
+                // wait for server to be running
+                AmazonQLspService.getInstance(project).instanceFlow.first()
+
+                withContext(EDT) {
                     browser.complete(
-                        Browser(this, webUri, project).also {
+                        Browser(this@AmazonQPanel, webUri, project).also {
                             wrapper.setContent(it.component())
 
                             initConnections()
                             connectUi(it)
                             connectApps(it)
+
+                            loadingPanel.stopLoading()
                         }
                     )
                 }

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/Browser.kt

@@ -11,7 +11,6 @@ import com.intellij.openapi.util.Disposer
 import com.intellij.ui.jcef.JBCefJSQuery
 import org.cef.CefApp
 import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
-import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AwsServerCapabilitiesProvider
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.FlareUiMessage
 import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile
@@ -46,17 +45,16 @@ class Browser(parent: Disposable, private val webUri: URI, val project: Project)
                 "mynah",
                 AssetResourceHandler.AssetResourceHandlerFactory(),
             )
-        AmazonQLspService.getInstance(project).addLspInitializeMessageListener {
-            loadWebView(
-                isCodeTransformAvailable,
-                isFeatureDevAvailable,
-                isDocAvailable,
-                isCodeScanAvailable,
-                isCodeTestAvailable,
-                highlightCommand,
-                activeProfile
-            )
-        }
+
+        loadWebView(
+            isCodeTransformAvailable,
+            isFeatureDevAvailable,
+            isDocAvailable,
+            isCodeScanAvailable,
+            isCodeTestAvailable,
+            highlightCommand,
+            activeProfile
+        )
     }
 
     override fun dispose() {

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt

@@ -25,6 +25,10 @@ import com.intellij.util.net.JdkProxyProvider
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Deferred
 import kotlinx.coroutines.async
+import kotlinx.coroutines.channels.BufferOverflow
+import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.flow.asSharedFlow
+import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.future.asCompletableFuture
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.sync.Mutex
@@ -75,7 +79,6 @@ import java.net.Proxy
 import java.net.URI
 import java.nio.charset.StandardCharsets
 import java.nio.file.Files
-import java.util.Collections
 import java.util.concurrent.Future
 import kotlin.time.Duration.Companion.seconds
 
@@ -114,9 +117,8 @@ internal class LSPProcessListener : ProcessListener {
 
 @Service(Service.Level.PROJECT)
 class AmazonQLspService(private val project: Project, private val cs: CoroutineScope) : Disposable {
-    private val lspInitializedMessageReceivedListener = Collections.synchronizedList(mutableListOf<AmazonQInitializeMessageReceivedListener>())
-    fun addLspInitializeMessageListener(listener: AmazonQInitializeMessageReceivedListener) = lspInitializedMessageReceivedListener.add(listener)
-    fun notifyInitializeMessageReceived() = lspInitializedMessageReceivedListener.forEach { it() }
+    private val _flowInstance = MutableSharedFlow<AmazonQServerInstance>(replay = 1, onBufferOverflow = BufferOverflow.DROP_OLDEST)
+    val instanceFlow = _flowInstance.asSharedFlow().map { it.languageServer }
 
     private var instance: Deferred<AmazonQServerInstance>
     val capabilities
@@ -140,7 +142,9 @@ class AmazonQLspService(private val project: Project, private val cs: CoroutineS
                     // wait for handshake to complete
                     instance.initializeResult.join()
 
-                    instance
+                    instance.also {
+                        _flowInstance.emit(it)
+                    }
                 }
             } catch (e: Exception) {
                 LOG.warn(e) { "Failed to start LSP server" }
@@ -324,7 +328,6 @@ private class AmazonQServerInstance(private val project: Project, private val cs
                     if (message is ResponseMessage && message.result is AwsExtendedInitializeResult) {
                         val result = message.result as AwsExtendedInitializeResult
                         AwsServerCapabilitiesProvider.getInstance(project).setAwsServerCapabilities(result.getAwsServerCapabilities())
-                        AmazonQLspService.getInstance(project).notifyInitializeMessageReceived()
                     }
                     consumer?.consume(message)
                 }