fix(amazonq): move workspace lsp messages off blocking thread calls

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

@@ -7,8 +7,10 @@
 import com.intellij.openapi.actionSystem.ActionUpdateThread
 import com.intellij.openapi.actionSystem.AnActionEvent
 import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.progress.currentThreadCoroutineScope
 import com.intellij.openapi.project.DumbAwareAction
 import com.intellij.util.messages.Topic
+import kotlinx.coroutines.launch
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.ChatCommunicationManager
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_TAB_REMOVE
@@ -22,9 +24,11 @@
 
         // Notify LSP server about all open tabs being removed
         val chatManager = ChatCommunicationManager.getInstance(project)
-        chatManager.getAllTabIds().forEach { tabId ->
-            AmazonQLspService.executeIfRunning(project) { server ->
-                rawEndpoint.notify(CHAT_TAB_REMOVE, mapOf("tabId" to tabId))
+        currentThreadCoroutineScope().launch {
+            chatManager.getAllTabIds().forEach { tabId ->
+                AmazonQLspService.executeAsyncIfRunning(project) {
+                    rawEndpoint.notify(CHAT_TAB_REMOVE, mapOf("tabId" to tabId))
+                }
             }
         }
 

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

@@ -112,7 +112,7 @@ class BrowserConnector(
     private val themeBrowserAdapter: ThemeBrowserAdapter = ThemeBrowserAdapter(),
     private val project: Project,
 ) {
-    var uiReady = CompletableDeferred<Boolean>()
+    val uiReady = CompletableDeferred<Boolean>()
     private val chatCommunicationManager = ChatCommunicationManager.getInstance(project)
     private val chatAsyncResultManager = ChatAsyncResultManager.getInstance(project)
 
@@ -212,7 +212,7 @@ class BrowserConnector(
         }
     }
 
-    private fun handleFlareChatMessages(browser: Browser, node: JsonNode) {
+    private suspend fun handleFlareChatMessages(browser: Browser, node: JsonNode) {
         when (node.command) {
             SEND_CHAT_COMMAND_PROMPT -> {
                 val requestFromUi = serializer.deserializeChatMessages<SendChatPromptRequest>(node)
@@ -234,7 +234,7 @@ class BrowserConnector(
                 chatCommunicationManager.registerPartialResultToken(partialResultToken)
 
                 var encryptionManager: JwtEncryptionManager? = null
-                val result = AmazonQLspService.executeIfRunning(project) { server ->
+                val result = AmazonQLspService.executeAsyncIfRunning(project) { server ->
                     encryptionManager = this.encryptionManager
 
                     val encryptedParams = EncryptedChatParams(this.encryptionManager.encrypt(chatParams), partialResultToken)
@@ -254,7 +254,7 @@ class BrowserConnector(
                 val partialResultToken = chatCommunicationManager.addPartialChatMessage(tabId)
                 chatCommunicationManager.registerPartialResultToken(partialResultToken)
                 var encryptionManager: JwtEncryptionManager? = null
-                val result = AmazonQLspService.executeIfRunning(project) { server ->
+                val result = AmazonQLspService.executeAsyncIfRunning(project) { server ->
                     encryptionManager = this.encryptionManager
 
                     val encryptedParams = EncryptedQuickActionChatParams(this.encryptionManager.encrypt(quickActionParams), partialResultToken)
@@ -581,7 +581,7 @@ class BrowserConnector(
         }
     }
 
-    private inline fun <reified Request, Response> handleChat(
+    private suspend inline fun <reified Request, Response> handleChat(
         lspMethod: JsonRpcMethod<Request, Response>,
         node: JsonNode,
         crossinline serverAction: (params: Request, invokeService: () -> CompletableFuture<Response>) -> CompletableFuture<Response>,
@@ -592,7 +592,7 @@ class BrowserConnector(
             serializer.deserializeChatMessages<Request>(node.params, lspMethod.params)
         }
 
-        return AmazonQLspService.executeIfRunning(project) { _ ->
+        return AmazonQLspService.executeAsyncIfRunning(project) { _ ->
             val invokeService = when (lspMethod) {
                 is JsonRpcNotification<Request> -> {
                     // notify is Unit
@@ -614,7 +614,7 @@ class BrowserConnector(
         } ?: CompletableFuture.failedFuture<Response>(IllegalStateException("LSP Server not running"))
     }
 
-    private inline fun <reified Request, Response> handleChat(
+    private suspend inline fun <reified Request, Response> handleChat(
         lspMethod: JsonRpcMethod<Request, Response>,
         node: JsonNode,
     ): CompletableFuture<Response> = handleChat(

plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevControllerTest.kt

@@ -170,7 +170,7 @@ class FeatureDevControllerTest : FeatureDevTestBase() {
         whenever(featureDevClient.createTaskAssistConversation()).thenReturn(exampleCreateTaskAssistConversationResponse)
         whenever(featureDevClient.sendFeatureDevTelemetryEvent(any())).thenReturn(exampleSendTelemetryEventResponse)
         whenever(chatSessionStorage.getSession(any(), any())).thenReturn(spySession)
-        doNothing().`when`(chatSessionStorage).deleteSession(any())
+        doNothing().whenever(chatSessionStorage).deleteSession(any())
 
         mockkObject(AmazonqTelemetry)
         every { AmazonqTelemetry.endChat(amazonqConversationId = any(), amazonqEndOfTheConversationLatency = any()) } just runs

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt

@@ -187,7 +187,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
                 CodeWhispererInvocationStatus.getInstance().setInvocationStart()
                 var nextToken: Either<String, Int>? = null
                 do {
-                    val result = AmazonQLspService.executeIfRunning(requestContext.project) { server ->
+                    val result = AmazonQLspService.executeAsyncIfRunning(requestContext.project) { server ->
                         val params = createInlineCompletionParams(requestContext.editor, requestContext.triggerTypeInfo, nextToken)
                         server.inlineCompletionWithReferences(params)
                     }
@@ -426,7 +426,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
         CodeWhispererTelemetryService.getInstance().sendUserTriggerDecisionEvent(project, latencyContext, sessionId, recommendationContext)
     }
 
-    fun getRequestContext(
+    suspend fun getRequestContext(
         triggerTypeInfo: TriggerTypeInfo,
         editor: Editor,
         project: Project,
@@ -472,11 +472,11 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
         )
     }
 
-    fun getWorkspaceIds(project: Project): CompletableFuture<LspServerConfigurations> {
+    suspend fun getWorkspaceIds(project: Project): CompletableFuture<LspServerConfigurations> {
         val payload = GetConfigurationFromServerParams(
             section = "aws.q.workspaceContext"
         )
-        return AmazonQLspService.executeIfRunning(project) { server ->
+        return AmazonQLspService.executeAsyncIfRunning(project) { server ->
             server.getConfigurationFromServer(payload)
         } ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")))
     }

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererServiceNew.kt

@@ -203,7 +203,7 @@ class CodeWhispererServiceNew(private val cs: CoroutineScope) : Disposable {
             var requestCount = 0
             var nextToken: Either<String, Int>? = null
             do {
-                val result = AmazonQLspService.executeIfRunning(requestContext.project) { server ->
+                val result = AmazonQLspService.executeAsyncIfRunning(requestContext.project) { server ->
                     val params = createInlineCompletionParams(requestContext.editor, requestContext.triggerTypeInfo, nextToken)
                     server.inlineCompletionWithReferences(params)
                 }

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/settings/CodeWhispererConfigurable.kt

@@ -10,6 +10,7 @@
 import com.intellij.openapi.options.Configurable
 import com.intellij.openapi.options.SearchableConfigurable
 import com.intellij.openapi.options.ex.Settings
+import com.intellij.openapi.progress.currentThreadCoroutineScope
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.project.ProjectManager
 import com.intellij.openapi.ui.emptyText
@@ -22,6 +23,8 @@
 import com.intellij.ui.dsl.builder.panel
 import com.intellij.util.concurrency.EdtExecutorService
 import com.intellij.util.execution.ParametersListUtil
+import kotlinx.coroutines.launch
+import org.eclipse.lsp4j.DidChangeConfigurationParams
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManagerListener
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
@@ -310,7 +313,12 @@
                         if (project.isDisposed) {
                             return@forEach
                         }
-                        AmazonQLspService.didChangeConfiguration(project)
+
+                        currentThreadCoroutineScope().launch {
+                            AmazonQLspService.executeAsyncIfRunning(project) { server ->
+                                server.workspaceService.didChangeConfiguration(DidChangeConfigurationParams())
+                            }
+                        }
                     }
                 }
             }

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/telemetry/CodeWhispererTelemetryService.kt

@@ -6,6 +6,8 @@ package software.aws.toolkits.jetbrains.services.codewhisperer.telemetry
 import com.intellij.openapi.components.Service
 import com.intellij.openapi.components.service
 import com.intellij.openapi.project.Project
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.launch
 import software.aws.toolkits.core.utils.debug
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
@@ -33,7 +35,7 @@ import java.time.Duration
 import java.time.Instant
 
 @Service
-class CodeWhispererTelemetryService {
+class CodeWhispererTelemetryService(private val cs: CoroutineScope) {
     companion object {
         fun getInstance(): CodeWhispererTelemetryService = service()
         val LOG = getLogger<CodeWhispererTelemetryService>()
@@ -45,22 +47,24 @@ class CodeWhispererTelemetryService {
         sessionId: String,
         recommendationContext: RecommendationContext,
     ) {
-        AmazonQLspService.executeIfRunning(project) { server ->
-            val params = LogInlineCompletionSessionResultsParams(
-                sessionId = sessionId,
-                completionSessionResult = recommendationContext.details.associate {
-                    it.itemId to InlineCompletionStates(
-                        seen = it.hasSeen,
-                        accepted = it.isAccepted,
-                        discarded = it.isDiscarded
-                    )
-                },
-                firstCompletionDisplayLatency = latencyContext.perceivedLatency,
-                totalSessionDisplayTime = CodeWhispererInvocationStatus.getInstance().completionShownTime?.let { Duration.between(it, Instant.now()) }
-                    ?.toMillis()?.toDouble(),
-                typeaheadLength = recommendationContext.userInput.length.toLong()
-            )
-            server.logInlineCompletionSessionResults(params)
+        cs.launch {
+            AmazonQLspService.executeAsyncIfRunning(project) { server ->
+                val params = LogInlineCompletionSessionResultsParams(
+                    sessionId = sessionId,
+                    completionSessionResult = recommendationContext.details.associate {
+                        it.itemId to InlineCompletionStates(
+                            seen = it.hasSeen,
+                            accepted = it.isAccepted,
+                            discarded = it.isDiscarded
+                        )
+                    },
+                    firstCompletionDisplayLatency = latencyContext.perceivedLatency,
+                    totalSessionDisplayTime = CodeWhispererInvocationStatus.getInstance().completionShownTime?.let { Duration.between(it, Instant.now()) }
+                        ?.toMillis()?.toDouble(),
+                    typeaheadLength = recommendationContext.userInput.length.toLong()
+                )
+                server.logInlineCompletionSessionResults(params)
+            }
         }
     }
 

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/telemetry/CodeWhispererTelemetryServiceNew.kt

@@ -6,6 +6,8 @@ package software.aws.toolkits.jetbrains.services.codewhisperer.telemetry
 import com.intellij.openapi.components.Service
 import com.intellij.openapi.components.service
 import com.intellij.openapi.project.Project
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.launch
 import software.aws.toolkits.core.utils.debug
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
@@ -29,7 +31,7 @@ import java.time.Duration
 import java.time.Instant
 
 @Service
-class CodeWhispererTelemetryServiceNew {
+class CodeWhispererTelemetryServiceNew(private val cs: CoroutineScope) {
 
     companion object {
         fun getInstance(): CodeWhispererTelemetryServiceNew = service()
@@ -109,24 +111,26 @@ class CodeWhispererTelemetryServiceNew {
     }
 
     fun sendUserTriggerDecisionEvent(project: Project, latencyContext: LatencyContext) {
-        AmazonQLspService.executeIfRunning(project) { server ->
-            CodeWhispererServiceNew.getInstance().getAllPaginationSessions().forEach { jobId, state ->
-                if (state == null) return@forEach
-                val params = LogInlineCompletionSessionResultsParams(
-                    sessionId = state.responseContext.sessionId,
-                    completionSessionResult = state.recommendationContext.details.associate {
-                        it.itemId to InlineCompletionStates(
-                            seen = it.hasSeen,
-                            accepted = it.isAccepted,
-                            discarded = it.isDiscarded
-                        )
-                    },
-                    firstCompletionDisplayLatency = latencyContext.perceivedLatency,
-                    totalSessionDisplayTime = CodeWhispererInvocationStatus.getInstance().completionShownTime?.let { Duration.between(it, Instant.now()) }
-                        ?.toMillis()?.toDouble(),
-                    typeaheadLength = state.recommendationContext.userInput.length.toLong()
-                )
-                server.logInlineCompletionSessionResults(params)
+        cs.launch {
+            AmazonQLspService.executeAsyncIfRunning(project) { server ->
+                CodeWhispererServiceNew.getInstance().getAllPaginationSessions().forEach { jobId, state ->
+                    if (state == null) return@forEach
+                    val params = LogInlineCompletionSessionResultsParams(
+                        sessionId = state.responseContext.sessionId,
+                        completionSessionResult = state.recommendationContext.details.associate {
+                            it.itemId to InlineCompletionStates(
+                                seen = it.hasSeen,
+                                accepted = it.isAccepted,
+                                discarded = it.isDiscarded
+                            )
+                        },
+                        firstCompletionDisplayLatency = latencyContext.perceivedLatency,
+                        totalSessionDisplayTime = CodeWhispererInvocationStatus.getInstance().completionShownTime?.let { Duration.between(it, Instant.now()) }
+                            ?.toMillis()?.toDouble(),
+                        typeaheadLength = state.recommendationContext.userInput.length.toLong()
+                    )
+                    server.logInlineCompletionSessionResults(params)
+                }
             }
         }
     }

plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererActionTest.kt

@@ -23,6 +23,7 @@ import org.mockito.kotlin.any
 import org.mockito.kotlin.argumentCaptor
 import org.mockito.kotlin.doNothing
 import org.mockito.kotlin.eq
+import org.mockito.kotlin.whenever
 import software.aws.toolkits.jetbrains.services.amazonq.QConstants
 import software.aws.toolkits.jetbrains.services.codewhisperer.actions.CodeWhispererLearnMoreAction
 import software.aws.toolkits.jetbrains.services.codewhisperer.actions.CodeWhispererShowSettingsAction
@@ -78,7 +79,7 @@ class CodeWhispererActionTest : CodeWhispererTestBase() {
     @Test
     fun `CodeWhispererShowSettingsAction actionPerformed should show settings dialog`() {
         val settingsSpy = spy(ShowSettingsUtil.getInstance())
-        doNothing().`when`(settingsSpy).showSettingsDialog(any(), any<Class<out Configurable>>())
+        doNothing().whenever(settingsSpy).showSettingsDialog(any(), any<Class<out Configurable>>())
         ApplicationManager.getApplication().replaceService(ShowSettingsUtil::class.java, settingsSpy, disposableRule.disposable)
         val action = CodeWhispererShowSettingsAction()
         runInEdtAndWait {

plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererConfigurableTest.kt

@@ -8,6 +8,7 @@ import com.intellij.ui.dsl.builder.components.DslLabel
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.Test
 import org.mockito.kotlin.doNothing
+import org.mockito.kotlin.whenever
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManagerListener
 import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhispererConfigurable
 import software.aws.toolkits.resources.message
@@ -19,9 +20,9 @@ class CodeWhispererConfigurableTest : CodeWhispererTestBase() {
 
     @Test
     fun `test CodeWhisperer configurable`() {
-        doNothing().`when`(codeScanManager).buildCodeScanUI()
-        doNothing().`when`(codeScanManager).showCodeScanUI()
-        doNothing().`when`(codeScanManager).removeCodeScanUI()
+        doNothing().whenever(codeScanManager).buildCodeScanUI()
+        doNothing().whenever(codeScanManager).showCodeScanUI()
+        doNothing().whenever(codeScanManager).removeCodeScanUI()
         val configurable = CodeWhispererConfigurable(projectRule.project)
 
         // A workaround to initialize disposable in the DslConfigurableBase since somehow the disposable is

plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTelemetryTest.kt

@@ -14,6 +14,7 @@ import org.mockito.kotlin.argumentCaptor
 import org.mockito.kotlin.doNothing
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.spy
+import org.mockito.kotlin.whenever
 import software.aws.toolkits.core.telemetry.MetricEvent
 import software.aws.toolkits.core.telemetry.TelemetryBatcher
 import software.aws.toolkits.core.telemetry.TelemetryPublisher
@@ -51,7 +52,7 @@ class CodeWhispererTelemetryTest : CodeWhispererTestBase() {
     @Test
     fun `test toggle autoSuggestion will emit autoSuggestionActivation telemetry (popup)`() {
         val metricCaptor = argumentCaptor<MetricEvent>()
-        doNothing().`when`(batcher).enqueue(metricCaptor.capture())
+        doNothing().whenever(batcher).enqueue(metricCaptor.capture())
 
         Pause().actionPerformed(TestActionEvent { projectRule.project })
         assertEventsContainsFieldsAndCount(