fix(amazonq): fix chat refresh button (#5701)

Previously, the browser waits for the server initialization to succeed before unlocking the loading process. When reloading the webview, since the server has already been initialized, it does not receive an initialization event and is stuck loading forever.

Fix by moving to a hot flow that always returns the latest instance of AmazonQLanguageServer. While this could be any value, the idea is to bring us closer to the state where we can tie the lifetime of ChatCommunicationManager to the webview instance so that we can delete AsyncChatUiListener.

Additionally this should also address the rare race condition where the server finishes startup before the webview finishes drawing.

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)
                 }