Merge main into feature/q-lsp-chat

.changes/next-release/bugfix-1d632e06-24fa-40f1-a8ae-33cca0d86ec4.json

@@ -0,0 +1,4 @@
+{
+  "type" : "bugfix",
+  "description" : "Fix UI freezes that may occur when interacting with large files in the editor"
+}

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

@@ -5,8 +5,6 @@
 
 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.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
@@ -44,28 +45,28 @@
 import java.util.concurrent.CompletableFuture
 import javax.swing.JButton
 
 class AmazonQPanel(val project: Project, private val scope: CoroutineScope) : Disposable {
     private val browser = CompletableFuture<Browser>()
     private val webviewContainer = Wrapper()
     private val appSource = AppSource()
     private val browserConnector = BrowserConnector(project = project)
     private val editorThemeAdapter = EditorThemeAdapter()
     private val appConnections = mutableListOf<AppConnection>()
 
     init {
         project.messageBus.connect().subscribe(
             AsyncChatUiListener.TOPIC,
             object : AsyncChatUiListener {
                 override fun onChange(command: String) {
                     browser.get()?.postChat(command)
                 }
 
                 override fun onChange(command: FlareUiMessage) {
                     browser.get()?.postChat(command)
                 }
             }
         )
     }
 
     val component = panel {
         row {
@@ -78,12 +79,12 @@
             row {
                 cell(
                     JButton("Show Web Debugger").apply {
                         addActionListener {
                             // Code to be executed when the button is clicked
                             // Add your logic here
 
                             browser.get().jcefBrowser.openDevtools()
                         }
                     },
                 )
                     .align(AlignX.CENTER)
@@ -100,107 +101,111 @@
             } else {
                 webviewContainer.add(JBTextArea("JCEF not supported"))
             }
             browser.complete(null)
         } else {
             val loadingPanel = JBLoadingPanel(null, this)
             val wrapper = Wrapper()
             loadingPanel.startLoading()
 
             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()
                         }
                     )
                 }
             }
         }
     }
 
     fun sendMessage(message: AmazonQMessage, tabType: String) {
         appConnections.filter { it.app.tabTypes.contains(tabType) }.forEach {
             scope.launch {
                 it.messagesFromUiToApp.publish(message)
             }
         }
     }
 
     fun sendMessageAppToUi(message: AmazonQMessage, tabType: String) {
         appConnections.filter { it.app.tabTypes.contains(tabType) }.forEach {
             scope.launch {
                 it.messagesFromAppToUi.publish(message)
             }
         }
     }
 
     private fun initConnections() {
         val apps = appSource.getApps(project)
         apps.forEach { app ->
             appConnections += AppConnection(
                 app = app,
                 messagesFromAppToUi = MessageConnector(),
                 messagesFromUiToApp = MessageConnector(),
                 messageTypeRegistry = MessageTypeRegistry(),
             )
         }
     }
 
     private fun connectApps(browser: Browser) {
         val fqnWebviewAdapter = FqnWebviewAdapter(browser.jcefBrowser, browserConnector)
 
         appConnections.forEach { connection ->
             val initContext = AmazonQAppInitContext(
                 project = project,
                 messagesFromAppToUi = connection.messagesFromAppToUi,
                 messagesFromUiToApp = connection.messagesFromUiToApp,
                 messageTypeRegistry = connection.messageTypeRegistry,
                 fqnWebviewAdapter = fqnWebviewAdapter,
             )
             // Connect the app to the UI
             connection.app.init(initContext)
             // Dispose of the app when the tool window is disposed.
             Disposer.register(this, connection.app)
         }
     }
 
     private fun connectUi(browser: Browser) {
         browser.init(
             isCodeTransformAvailable = isCodeTransformAvailable(project),
             isFeatureDevAvailable = isFeatureDevAvailable(project),
             isCodeScanAvailable = isCodeScanAvailable(project),
             isCodeTestAvailable = isCodeTestAvailable(project),
             isDocAvailable = isDocAvailable(project),
             highlightCommand = highlightCommand(),
             activeProfile = QRegionProfileManager.getInstance().takeIf { it.shouldDisplayProfileInfo(project) }?.activeProfile(project)
         )
 
         scope.launch {
             // Pipe messages from the UI to the relevant apps and vice versa
             browserConnector.connect(
                 browser = browser,
                 connections = appConnections,
             )
         }
 
         scope.launch {
             // Update the theme in the UI when the IDE theme changes
             browserConnector.connectTheme(
                 chatBrowser = browser.jcefBrowser.cefBrowser,
                 themeSource = editorThemeAdapter.onThemeChange(),
             )
         }
     }
 
     override fun dispose() {
     }
 }

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

@@ -11,7 +11,6 @@
 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
@@ -24,7 +23,7 @@
 Displays the web view for the Amazon Q tool window
  */
 
 class Browser(parent: Disposable, private val webUri: URI, val project: Project) : Disposable {
 
     val jcefBrowser = createBrowser(parent)
 
@@ -46,17 +45,16 @@
                 "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() {
@@ -65,14 +63,14 @@
 
     fun component() = jcefBrowser.component
 
     fun postChat(command: FlareUiMessage) = postChat(Gson().toJson(command))
 
     @Deprecated("shouldn't need this version")
     fun postChat(message: String) {
         jcefBrowser
             .cefBrowser
             .executeJavaScript("window.postMessage($message)", jcefBrowser.cefBrowser.url, 0)
     }
 
     // Load the chat web app into the jcefBrowser
     private fun loadWebView(
@@ -115,28 +113,28 @@
         activeProfile: QRegionProfile?,
     ): String {
         val postMessageToJavaJsCode = receiveMessageQuery.inject("JSON.stringify(message)")
         val connectorAdapterPath = "http://mynah/js/connectorAdapter.js"
         generateQuickActionConfig()
         // https://github.com/highlightjs/highlight.js/issues/1387
         // language=HTML
         val jsScripts = """
             <script type="text/javascript" charset="UTF-8" src="$connectorAdapterPath"></script>
             <script type="text/javascript" charset="UTF-8" src="$webUri" defer onload="init()"></script>
 
             <script type="text/javascript">
 
                 const init = () => {
                     const hybridChatConnector = connectorAdapter.initiateAdapter(
                      ${MeetQSettings.getInstance().reinvent2024OnboardingCount < MAX_ONBOARDING_PAGE_COUNT},
                      ${MeetQSettings.getInstance().disclaimerAcknowledged},
                      $isFeatureDevAvailable,
                      $isCodeTransformAvailable,
                      $isDocAvailable,
                      $isCodeScanAvailable,
                      $isCodeTestAvailable,
                      {
                             postMessage: message => {
                                 $postMessageToJavaJsCode
                             }
                         },
 
@@ -151,24 +149,24 @@
                         {
                         agenticMode: true,
                         quickActionCommands: commands,
                         disclaimerAcknowledged: ${MeetQSettings.getInstance().disclaimerAcknowledged},
                         pairProgrammingAcknowledged: ${!MeetQSettings.getInstance().amazonQChatPairProgramming}
                         },
                         hybridChatConnector,
                         ${CodeWhispererFeatureConfigService.getInstance().getFeatureConfigJsonString()}                     
                     );
                 }
             </script>        
         """.trimIndent()
 
         addQuickActionCommands(
             isCodeTransformAvailable,
             isFeatureDevAvailable,
             isDocAvailable,
             isCodeTestAvailable,
             isCodeScanAvailable,
             highlightCommand,
             activeProfile
         )
         return """
             <!DOCTYPE html>
@@ -249,20 +247,20 @@
         activeProfile: QRegionProfile?,
     ) {
         // TODO:  Remove this once chat has been integrated with agents. This is added temporarily to keep detekt happy.
         isCodeScanAvailable
         isCodeTestAvailable
         isDocAvailable
         isFeatureDevAvailable
         isCodeTransformAvailable
         MAX_ONBOARDING_PAGE_COUNT
         OBJECT_MAPPER
         highlightCommand
         activeProfile
     }
 
     private fun generateQuickActionConfig() = AwsServerCapabilitiesProvider.getInstance(project).getChatOptions().quickActions.quickActionsCommandGroups
         .let { OBJECT_MAPPER.writeValueAsString(it) }
         ?: "[]"
 
     companion object {
         private const val MAX_ONBOARDING_PAGE_COUNT = 3

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

@@ -12,6 +12,7 @@
 import com.intellij.openapi.options.ShowSettingsUtil
 import com.intellij.openapi.project.Project
 import com.intellij.ui.jcef.JBCefJSQuery.Response
+import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.coroutineScope
@@ -65,7 +66,6 @@
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatParams
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatPrompt
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatReadyNotification
-import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatUiMessageParams
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ConversationClickRequest
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CopyCodeToClipboardNotification
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CopyCodeToClipboardParams
@@ -114,7 +114,6 @@
 import software.aws.toolkits.jetbrains.services.amazonq.webview.theme.ThemeBrowserAdapter
 import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhispererConfigurable
 import software.aws.toolkits.jetbrains.settings.MeetQSettings
-import software.aws.toolkits.resources.AwsCoreBundle
 import software.aws.toolkits.telemetry.MetricResult
 import software.aws.toolkits.telemetry.Telemetry
 import java.util.concurrent.CompletableFuture
@@ -123,10 +122,10 @@
 class BrowserConnector(
     private val serializer: MessageSerializer = MessageSerializer.getInstance(),
     private val themeBrowserAdapter: ThemeBrowserAdapter = ThemeBrowserAdapter(),
     private val project: Project,
 ) {
     var uiReady = CompletableDeferred<Boolean>()
     private val chatCommunicationManager = ChatCommunicationManager.getInstance(project)
 
     suspend fun connect(
         browser: Browser,
@@ -140,7 +139,7 @@
                     // this is sent when the named agents UI is ready
                     "ui-is-ready" -> {
                         uiReady.complete(true)
                         chatCommunicationManager.setUiReady()
                         RunOnceUtil.runOnceForApp("AmazonQ-UI-Ready") {
                             MeetQSettings.getInstance().reinvent2024OnboardingCount += 1
                         }
@@ -172,15 +171,15 @@
                     }
                 }
 
                 val tabType = node.tabType
                 if (tabType == null || tabType == "cwc") {
                     handleFlareChatMessages(browser, node)
                 } else {
                     connections.filter { connection -> connection.app.tabTypes.contains(tabType) }.forEach { connection ->
                         launch {
                             val message = serializer.deserialize(node, connection.messageTypeRegistry)
                             connection.messagesFromUiToApp.publish(message)
                         }
                     }
                 }
             }
@@ -190,12 +189,12 @@
         uiReady.await()
 
         // Chat options including history and quick actions need to be sent in once UI is ready
         val showChatOptions = """{
             "command": "chatOptions",
             "params": ${Gson().toJson(AwsServerCapabilitiesProvider.getInstance(project).getChatOptions())}
             }
         """.trimIndent()
         browser.postChat(showChatOptions)
 
         // Send inbound messages to the browser
         val inboundMessages = connections.map { it.messagesFromAppToUi.flow }.merge()
@@ -232,266 +231,250 @@
     private fun handleFlareChatMessages(browser: Browser, node: JsonNode) {
         when (node.command) {
             SEND_CHAT_COMMAND_PROMPT -> {
                 val requestFromUi = serializer.deserializeChatMessages<SendChatPromptRequest>(node)
                 val chatPrompt = ChatPrompt(
                     requestFromUi.params.prompt.prompt,
                     requestFromUi.params.prompt.escapedPrompt,
                     node.command
                 )
                 val editor = FileEditorManager.getInstance(project).selectedTextEditor
                 val textDocumentIdentifier = editor?.let { TextDocumentIdentifier(toUriString(it.virtualFile)) }
                 val cursorState = editor?.let { LspEditorUtil.getCursorState(it) }
 
                 val chatParams = ChatParams(
                     requestFromUi.params.tabId,
                     chatPrompt,
                     textDocumentIdentifier,
                     cursorState,
                     context = requestFromUi.params.context
                 )
 
                 val tabId = requestFromUi.params.tabId
                 val partialResultToken = chatCommunicationManager.addPartialChatMessage(tabId)
 
                 var encryptionManager: JwtEncryptionManager? = null
                 val result = AmazonQLspService.executeIfRunning(project) { server ->
                     encryptionManager = this.encryptionManager
 
                     val encryptedParams = EncryptedChatParams(this.encryptionManager.encrypt(chatParams), partialResultToken)
                     server.sendChatPrompt(encryptedParams)
                 } ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")))
 
                 // We assume there is only one outgoing request per tab because the input is
                 // blocked when there is an outgoing request
                 chatCommunicationManager.setInflightRequestForTab(tabId, result)
                 showResult(result, partialResultToken, tabId, encryptionManager, browser)
             }
             CHAT_QUICK_ACTION -> {
                 val requestFromUi = serializer.deserializeChatMessages<QuickChatActionRequest>(node)
                 val tabId = requestFromUi.params.tabId
                 val quickActionParams = requestFromUi.params
                 val partialResultToken = chatCommunicationManager.addPartialChatMessage(tabId)
                 var encryptionManager: JwtEncryptionManager? = null
                 val result = AmazonQLspService.executeIfRunning(project) { server ->
                     encryptionManager = this.encryptionManager
 
                     val encryptedParams = EncryptedQuickActionChatParams(this.encryptionManager.encrypt(quickActionParams), partialResultToken)
                     server.sendQuickAction(encryptedParams)
                 } ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")))
 
                 // We assume there is only one outgoing request per tab because the input is
                 // blocked when there is an outgoing request
                 chatCommunicationManager.setInflightRequestForTab(tabId, result)
 
                 showResult(result, partialResultToken, tabId, encryptionManager, browser)
             }
             CHAT_LIST_CONVERSATIONS -> {
                 val requestFromUi = serializer.deserializeChatMessages<ListConversationsRequest>(node)
                 val result = AmazonQLspService.executeIfRunning(project) { server ->
                     server.listConversations(requestFromUi.params)
                 } ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")))
 
                 result.whenComplete { response, _ ->
                     val uiMessage = """
                         {
                             "command": "$CHAT_LIST_CONVERSATIONS",
                             "params": ${Gson().toJson(response)}
                         }
                     """.trimIndent()
                     browser.postChat(uiMessage)
                 }
             }
             CHAT_CONVERSATION_CLICK -> {
                 val requestFromUi = serializer.deserializeChatMessages<ConversationClickRequest>(node)
                 val result = AmazonQLspService.executeIfRunning(project) { server ->
                     server.conversationClick(requestFromUi.params)
                 } ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")))
 
                 result.whenComplete { response, _ ->
                     val uiMessage = """
                         {
                             "command": "$CHAT_CONVERSATION_CLICK",
                             "params": ${Gson().toJson(response)}
                         }
                     """.trimIndent()
                     browser.postChat(uiMessage)
                 }
             }
             CHAT_FEEDBACK -> {
                 handleChatNotification<FeedbackNotification, FeedbackParams>(node) { server, params ->
                     server.feedback(params)
                 }
             }
             CHAT_READY -> {
                 handleChatNotification<ChatReadyNotification, Unit>(node) { server, _ ->
                     uiReady.complete(true)
                     chatCommunicationManager.setUiReady()
                     RunOnceUtil.runOnceForApp("AmazonQ-UI-Ready") {
                         MeetQSettings.getInstance().reinvent2024OnboardingCount += 1
                     }
                     server.chatReady()
                 }
             }
             CHAT_TAB_ADD -> {
                 handleChatNotification<TabEventRequest, TabEventParams>(node) { server, params ->
                     server.tabAdd(params)
                 }
             }
             CHAT_TAB_REMOVE -> {
                 handleChatNotification<TabEventRequest, TabEventParams>(node) { server, params ->
                     chatCommunicationManager.removePartialChatMessage(params.tabId)
                     cancelInflightRequests(params.tabId)
                     server.tabRemove(params)
                 }
             }
             CHAT_TAB_CHANGE -> {
                 handleChatNotification<TabEventRequest, TabEventParams>(node) { server, params ->
                     server.tabChange(params)
                 }
             }
             CHAT_OPEN_TAB -> {
                 val response = serializer.deserializeChatMessages<OpenTabResponse>(node)
                 chatCommunicationManager.completeTabOpen(
                     response.requestId,
                     response.params.result.tabId
                 )
             }
             CHAT_INSERT_TO_CURSOR -> {
                 handleChatNotification<InsertToCursorPositionNotification, InsertToCursorPositionParams>(node) { server, params ->
                     server.insertToCursorPosition(params)
                 }
             }
             CHAT_LINK_CLICK -> {
                 handleChatNotification<LinkClickNotification, LinkClickParams>(node) { server, params ->
                     server.linkClick(params)
                 }
             }
             CHAT_INFO_LINK_CLICK -> {
                 handleChatNotification<InfoLinkClickNotification, InfoLinkClickParams>(node) { server, params ->
                     server.infoLinkClick(params)
                 }
             }
             CHAT_SOURCE_LINK_CLICK -> {
                 handleChatNotification<SourceLinkClickNotification, SourceLinkClickParams>(node) { server, params ->
                     server.sourceLinkClick(params)
                 }
             }
             CHAT_FILE_CLICK -> {
                 handleChatNotification<FileClickNotification, FileClickParams>(node) { server, params ->
                     server.fileClick(params)
                 }
             }
             PROMPT_INPUT_OPTIONS_CHANGE -> {
                 handleChatNotification<PromptInputOptionChangeNotification, PromptInputOptionChangeParams>(node) {
                         server, params ->
                     server.promptInputOptionsChange(params)
                 }
             }
             CHAT_PROMPT_OPTION_ACKNOWLEDGED -> {
                 val acknowledgedMessage = node.get("params").get("messageId")
                 if (acknowledgedMessage.asText() == "programmerModeCardId") {
                     MeetQSettings.getInstance().amazonQChatPairProgramming = false
                 }
             }
             CHAT_FOLLOW_UP_CLICK -> {
                 handleChatNotification<FollowUpClickNotification, FollowUpClickParams>(node) { server, params ->
                     server.followUpClick(params)
                 }
             }
             CHAT_BUTTON_CLICK -> {
                 handleChatNotification<ButtonClickNotification, ButtonClickParams>(node) { server, params ->
                     server.buttonClick(params)
                 }.thenApply { response ->
                     if (response is ButtonClickResult && !response.success) {
                         LOG.warn { "Failed to execute action associated with button with reason: ${response.failureReason}" }
                     }
                 }
             }
             AUTH_FOLLOW_UP_CLICKED -> {
                 val message = serializer.deserializeChatMessages<AuthFollowUpClickNotification>(node)
                 chatCommunicationManager.handleAuthFollowUpClicked(
                     project,
                     message.params
                 )
             }
             CHAT_COPY_CODE_TO_CLIPBOARD -> {
                 handleChatNotification<CopyCodeToClipboardNotification, CopyCodeToClipboardParams>(node) { server, params ->
                     server.copyCodeToClipboard(params)
                 }
             }
 
             GET_SERIALIZED_CHAT_REQUEST_METHOD -> {
                 val response = serializer.deserializeChatMessages<GetSerializedChatResponse>(node)
                 chatCommunicationManager.completeSerializedChatResponse(
                     response.requestId,
                     response.params.result.content
                 )
             }
 
             CHAT_TAB_BAR_ACTIONS -> {
                 handleChatNotification<TabBarActionRequest, TabBarActionParams>(node) {
                         server, params ->
                     val result = server.tabBarActions(params)
                     result.whenComplete { actions, error ->
                         try {
                             if (error != null) {
                                 throw error
                             }
 
                             browser.postChat(
                                 FlareUiMessage(
                                     command = CHAT_TAB_BAR_ACTIONS,
                                     params = actions ?: emptyMap<Any, Any>()
                                 )
                             )
                         } catch (e: Exception) {
                             LOG.error { "Failed to perform chat tab bar action $e" }
                             params.tabId?.let {
                                 browser.postChat(chatCommunicationManager.getErrorUiMessage(it, e, null))
                             }
                         }
                     }
                 }
             }
             CHAT_CREATE_PROMPT -> {
                 handleChatNotification<CreatePromptNotification, CreatePromptParams>(node) {
                         server, params ->
                     server.createPrompt(params)
                 }
             }
             STOP_CHAT_RESPONSE -> {
                 val stopResponseRequest = serializer.deserializeChatMessages<StopResponseMessage>(node)
                 if (!chatCommunicationManager.hasInflightRequest(stopResponseRequest.params.tabId)) {
                     return
                 }
                 cancelInflightRequests(stopResponseRequest.params.tabId)
                 chatCommunicationManager.removePartialChatMessage(stopResponseRequest.params.tabId)
-
-                val paramsJson = Gson().toJson(
-                    // https://github.com/aws/language-servers/blob/1c0d88806087125b6fc561f610cc15e98127c6bf/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts#L403
-                    ChatUiMessageParams(
-                        title = AwsCoreBundle.message("amazonqChat.stopChatResponse"),
-                        body = ""
-                    )
-                )
-
-                val uiMessage = ChatCommunicationManager.convertToJsonToSendToChat(
-                    command = SEND_CHAT_COMMAND_PROMPT,
-                    tabId = stopResponseRequest.params.tabId,
-                    params = paramsJson.toString(),
-                    isPartialResult = false
-                )
-                browser.postChat(uiMessage)
             }
             OPEN_SETTINGS -> {
                 val openSettingsNotification = serializer.deserializeChatMessages<OpenSettingsNotification>(node)
                 if (openSettingsNotification.params.settingKey != OPEN_WORKSPACE_SETTINGS_KEY) return
                 runInEdt {
                     ShowSettingsUtil.getInstance().showSettingsDialog(browser.project, CodeWhispererConfigurable::class.java)
                 }
             }
         }
     }
 
     private fun showResult(
         result: CompletableFuture<String>,
@@ -500,33 +483,35 @@
         encryptionManager: JwtEncryptionManager?,
         browser: Browser,
     ) {
         result.whenComplete { value, error ->
             try {
                 if (error != null) {
                     throw error
                 }
                 chatCommunicationManager.removePartialChatMessage(partialResultToken)
                 val messageToChat = ChatCommunicationManager.convertToJsonToSendToChat(
                     SEND_CHAT_COMMAND_PROMPT,
                     tabId,
                     encryptionManager?.decrypt(value).orEmpty(),
                     isPartialResult = false
                 )
                 browser.postChat(messageToChat)
                 chatCommunicationManager.removeInflightRequestForTab(tabId)
+            } catch (e: CancellationException) {
+                LOG.warn { "Cancelled chat generation" }
             } catch (e: Exception) {
                 LOG.error { "Failed to send chat message $e" }
                 browser.postChat(chatCommunicationManager.getErrorUiMessage(tabId, e, partialResultToken))
             }
         }
     }
 
     private fun cancelInflightRequests(tabId: String) {
         chatCommunicationManager.getInflightRequestForTab(tabId)?.let { request ->
             request.cancel(true)
             chatCommunicationManager.removeInflightRequestForTab(tabId)
         }
     }
 
     private inline fun <reified T, R> handleChatNotification(
         node: JsonNode,
@@ -534,11 +519,11 @@
     ): CompletableFuture<*> where T : ChatNotification<R> {
         val requestFromUi = serializer.deserializeChatMessages<T>(node)
         return AmazonQLspService.executeIfRunning(project) { server ->
             serverAction(server, requestFromUi.params)
         } ?: CompletableFuture.failedFuture<Unit>(IllegalStateException("LSP Server not running"))
     }
 
     companion object {
         private val LOG = getLogger<BrowserConnector>()
     }
 }

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

@@ -25,6 +25,10 @@
 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.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,15 +117,14 @@
 
 @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
         get() = instance.getCompleted().initializeResult.getCompleted().capabilities
     val encryptionManager
         get() = instance.getCompleted().encryptionManager
 
     // dont allow lsp commands if server is restarting
     private val mutex = Mutex(false)
@@ -140,7 +142,9 @@
                     // 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" }
@@ -213,7 +217,7 @@
 }
 
 private class AmazonQServerInstance(private val project: Project, private val cs: CoroutineScope) : Disposable {
     val encryptionManager = JwtEncryptionManager()
 
     private val launcher: Launcher<AmazonQLanguageServer>
 
@@ -268,18 +272,18 @@
 
     init {
         // will cause slow service init, but maybe fine for now. will not block UI since fetch/extract will be under background progress
         val artifact = runBlocking { service<ArtifactManager>().fetchArtifact(project) }.toAbsolutePath()
 
         // more network calls
         // make assumption that all requests will resolve to the same CA
         // also terrible assumption that default endpoint is reachable
         val qUri = URI(QDefaultServiceConfig.ENDPOINT)
         val rtsTrustChain = TrustChainUtil.getTrustChain(qUri)
         val extraCaCerts = Files.createTempFile("q-extra-ca", ".pem").apply {
             writeText(
                 TrustChainUtil.certsToPem(rtsTrustChain)
             )
         }
 
         val node = if (SystemInfo.isWindows) "node.exe" else "node"
         val cmd = GeneralCommandLine(
@@ -287,31 +291,31 @@
             LspSettings.getInstance().getArtifactPath() ?: artifact.resolve("aws-lsp-codewhisperer.js").toString(),
             "--stdio",
             "--set-credentials-encryption-key",
         ).withEnvironment(
             buildMap {
                 put("NODE_EXTRA_CA_CERTS", extraCaCerts.toAbsolutePath().toString())
 
                 val proxy = JdkProxyProvider.getInstance().proxySelector.select(qUri)
                     // log if only socks proxy available
                     .firstOrNull { it.type() == Proxy.Type.HTTP }
 
                 if (proxy != null) {
                     val address = proxy.address()
                     if (address is java.net.InetSocketAddress) {
                         put(
                             "HTTPS_PROXY",
                             URIBuilder("http://${address.hostName}:${address.port}").apply {
                                 val login = HttpConfigurable.getInstance().proxyLogin
                                 if (login != null) {
                                     setUserInfo(login, HttpConfigurable.getInstance().plainProxyPassword)
                                 }
                             }.build().toASCIIString()
                         )
                     }
                 }
             }
         )
             .withParentEnvironmentType(GeneralCommandLine.ParentEnvironmentType.CONSOLE)
 
         launcherHandler = KillableColoredProcessHandler.Silent(cmd)
         val inputWrapper = LSPProcessListener()
@@ -319,15 +323,14 @@
         launcherHandler.startNotify()
 
         launcher = LSPLauncher.Builder<AmazonQLanguageServer>()
             .wrapMessages { consumer ->
                 MessageConsumer { message ->
                     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)
                 }
             }
             .setLocalService(AmazonQLanguageClientImpl(project))
             .setRemoteInterface(AmazonQLanguageServer::class.java)
@@ -337,7 +340,7 @@
 
                 // otherwise Gson treats all numbers as double which causes deser issues
                 it.setObjectToNumberStrategy(ToNumberPolicy.LONG_OR_DOUBLE)
                 it.registerTypeAdapterFactory(AmazonQLspTypeAdapterFactory())
             }.traceMessages(
                 PrintWriter(
                     object : StringWriter() {
@@ -361,7 +364,7 @@
             encryptionManager.writeInitializationPayload(launcherHandler.process.outputStream)
 
             val initializeResult = try {
                 languageServer.initialize(createInitializeParams()).await()
             } catch (e: Exception) {
                 LOG.warn(e) { "LSP initialization failed" }
                 null
@@ -378,7 +381,7 @@
         }
 
         // invokeOnCompletion results in weird lock/timeout error
         initializeResult.asCompletableFuture().handleAsync { lspInitResult, ex ->
             if (ex != null) {
                 return@handleAsync
             }
@@ -386,7 +389,7 @@
             [email protected] {
                 DefaultAuthCredentialsService(project, encryptionManager, this)
                 TextDocumentServiceHandler(project, this)
                 WorkspaceServiceHandler(project, lspInitResult, this)
                 DefaultModuleDependenciesService(project, this)
             }
         }
@@ -412,4 +415,4 @@
     }
 }
 
 typealias AmazonQInitializeMessageReceivedListener = () -> Unit

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

@@ -77,16 +77,18 @@ class TextDocumentServiceHandler(
         }
         AmazonQLspService.executeIfRunning(project) { languageServer ->
             toUriString(file)?.let { uri ->
-                languageServer.textDocumentService.didOpen(
-                    DidOpenTextDocumentParams().apply {
-                        textDocument = TextDocumentItem().apply {
-                            this.uri = uri
-                            text = file.inputStream.readAllBytes().decodeToString()
-                            languageId = file.fileType.name.lowercase()
-                            version = file.modificationStamp.toInt()
+                pluginAwareExecuteOnPooledThread {
+                    languageServer.textDocumentService.didOpen(
+                        DidOpenTextDocumentParams().apply {
+                            textDocument = TextDocumentItem().apply {
+                                this.uri = uri
+                                text = file.inputStream.readAllBytes().decodeToString()
+                                languageId = file.fileType.name.lowercase()
+                                version = file.modificationStamp.toInt()
+                            }
                         }
-                    }
-                )
+                    )
+                }
             }
         }
     }
@@ -95,14 +97,16 @@ class TextDocumentServiceHandler(
         AmazonQLspService.executeIfRunning(project) { languageServer ->
             val file = FileDocumentManager.getInstance().getFile(document) ?: return@executeIfRunning
             toUriString(file)?.let { uri ->
-                languageServer.textDocumentService.didSave(
-                    DidSaveTextDocumentParams().apply {
-                        textDocument = TextDocumentIdentifier().apply {
-                            this.uri = uri
+                pluginAwareExecuteOnPooledThread {
+                    languageServer.textDocumentService.didSave(
+                        DidSaveTextDocumentParams().apply {
+                            textDocument = TextDocumentIdentifier().apply {
+                                this.uri = uri
+                            }
+                            text = document.text
                         }
-                        text = document.text
-                    }
-                )
+                    )
+                }
             }
         }
     }