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

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

we are inappropriately freezing EDT while passing workspace context to the LSP

* logging

* wip

* tst

* tst

* coroutines/verbosity

* c

* tst

* tst

* tst

* aaaaaaaaaaa

* tst

* revert manifest

* tst

* tst

* lint

Author: rli

buildSrc/settings.gradle.kts

@@ -43,5 +43,11 @@ dependencyResolutionManagement {
                 includeGroupByRegex("org\\.jetbrains\\.intellij\\.platform.*")
             }
         }
+        maven {
+            url = uri("https://s01.oss.sonatype.org/content/repositories/snapshots/")
+            content {
+                includeGroupByRegex("org\\.mockito\\.kotlin")
+            }
+        }
     }
 }

gradle/libs.versions.toml

@@ -10,7 +10,7 @@ detekt = "1.23.7"
 diff-util = "4.12"
 intellijExt = "1.1.8"
 # match with <root>/settings.gradle.kts
-intellijGradle = "2.3.0"
+intellijGradle = "2.6.0"
 intellijRemoteRobot = "0.11.22"
 jackson = "2.17.2"
 jacoco = "0.8.12"
@@ -24,7 +24,7 @@ kotlin = "2.1.20"
 kotlinCoroutines = "1.8.0"
 lsp4j = "0.24.0"
 mockito = "5.12.0"
-mockitoKotlin = "5.4.0"
+mockitoKotlin = "5.4.1-SNAPSHOT"
 mockk = "1.13.17"
 nimbus-jose-jwt = "9.40"
 node-gradle = "7.0.2"

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

@@ -7,8 +7,10 @@ import com.intellij.icons.AllIcons
 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 @@ class QRefreshPanelAction : DumbAwareAction(AmazonQBundle.message("amazonq.refre
 
         // 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

@@ -116,7 +116,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)
 
@@ -216,7 +216,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)
@@ -238,7 +238,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)
@@ -258,7 +258,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)
@@ -613,7 +613,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>,
@@ -624,7 +624,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
@@ -646,7 +646,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,11 +187,17 @@ 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)
                     }
-                    val completion = result?.await() ?: break
+                    val completion = result?.await()
+                    if (completion == null) {
+                        // no result / not running
+                        CodeWhispererInvocationStatus.getInstance().finishInvocation()
+                        break
+                    }
+
                     nextToken = completion.partialResultToken
                     val endTime = System.nanoTime()
                     val latency = TimeUnit.NANOSECONDS.toMillis(endTime - startTime).toDouble()
@@ -426,7 +432,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 +478,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.BoundConfigurable
 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.bindText
 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 @@ class CodeWhispererConfigurable(private val project: Project) :
                         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)
+                }
             }
         }
     }