Skip to content

Commit 9981133

Browse files
rlirli
authored
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
1 parent da1446d commit 9981133

File tree

30 files changed

+602
-385
lines changed

30 files changed

+602
-385
lines changed

‎buildSrc/settings.gradle.kts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,5 +43,11 @@ dependencyResolutionManagement {
4343
includeGroupByRegex("org\\.jetbrains\\.intellij\\.platform.*")
4444
}
4545
}
46+
maven {
47+
url = uri("https://s01.oss.sonatype.org/content/repositories/snapshots/")
48+
content {
49+
includeGroupByRegex("org\\.mockito\\.kotlin")
50+
}
51+
}
4652
}
4753
}

‎gradle/libs.versions.toml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ detekt = "1.23.7"
1010
diff-util = "4.12"
1111
intellijExt = "1.1.8"
1212
# match with <root>/settings.gradle.kts
13-
intellijGradle = "2.3.0"
13+
intellijGradle = "2.6.0"
1414
intellijRemoteRobot = "0.11.22"
1515
jackson = "2.17.2"
1616
jacoco = "0.8.12"
@@ -24,7 +24,7 @@ kotlin = "2.1.20"
2424
kotlinCoroutines = "1.8.0"
2525
lsp4j = "0.24.0"
2626
mockito = "5.12.0"
27-
mockitoKotlin = "5.4.0"
27+
mockitoKotlin = "5.4.1-SNAPSHOT"
2828
mockk = "1.13.17"
2929
nimbus-jose-jwt = "9.40"
3030
node-gradle = "7.0.2"

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

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,10 @@ import com.intellij.icons.AllIcons
77
import com.intellij.openapi.actionSystem.ActionUpdateThread
88
import com.intellij.openapi.actionSystem.AnActionEvent
99
import com.intellij.openapi.application.ApplicationManager
10+
import com.intellij.openapi.progress.currentThreadCoroutineScope
1011
import com.intellij.openapi.project.DumbAwareAction
1112
import com.intellij.util.messages.Topic
13+
import kotlinx.coroutines.launch
1214
import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
1315
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.ChatCommunicationManager
1416
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
2224

2325
// Notify LSP server about all open tabs being removed
2426
val chatManager = ChatCommunicationManager.getInstance(project)
25-
chatManager.getAllTabIds().forEach { tabId ->
26-
AmazonQLspService.executeIfRunning(project) { server ->
27-
rawEndpoint.notify(CHAT_TAB_REMOVE, mapOf("tabId" to tabId))
27+
currentThreadCoroutineScope().launch {
28+
chatManager.getAllTabIds().forEach { tabId ->
29+
AmazonQLspService.executeAsyncIfRunning(project) {
30+
rawEndpoint.notify(CHAT_TAB_REMOVE, mapOf("tabId" to tabId))
31+
}
2832
}
2933
}
3034

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

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,7 @@ class BrowserConnector(
116116
private val themeBrowserAdapter: ThemeBrowserAdapter = ThemeBrowserAdapter(),
117117
private val project: Project,
118118
) {
119-
var uiReady = CompletableDeferred<Boolean>()
119+
val uiReady = CompletableDeferred<Boolean>()
120120
private val chatCommunicationManager = ChatCommunicationManager.getInstance(project)
121121
private val chatAsyncResultManager = ChatAsyncResultManager.getInstance(project)
122122

@@ -216,7 +216,7 @@ class BrowserConnector(
216216
}
217217
}
218218

219-
private fun handleFlareChatMessages(browser: Browser, node: JsonNode) {
219+
private suspend fun handleFlareChatMessages(browser: Browser, node: JsonNode) {
220220
when (node.command) {
221221
SEND_CHAT_COMMAND_PROMPT -> {
222222
val requestFromUi = serializer.deserializeChatMessages<SendChatPromptRequest>(node)
@@ -238,7 +238,7 @@ class BrowserConnector(
238238
chatCommunicationManager.registerPartialResultToken(partialResultToken)
239239

240240
var encryptionManager: JwtEncryptionManager? = null
241-
val result = AmazonQLspService.executeIfRunning(project) { server ->
241+
val result = AmazonQLspService.executeAsyncIfRunning(project) { server ->
242242
encryptionManager = this.encryptionManager
243243

244244
val encryptedParams = EncryptedChatParams(this.encryptionManager.encrypt(chatParams), partialResultToken)
@@ -258,7 +258,7 @@ class BrowserConnector(
258258
val partialResultToken = chatCommunicationManager.addPartialChatMessage(tabId)
259259
chatCommunicationManager.registerPartialResultToken(partialResultToken)
260260
var encryptionManager: JwtEncryptionManager? = null
261-
val result = AmazonQLspService.executeIfRunning(project) { server ->
261+
val result = AmazonQLspService.executeAsyncIfRunning(project) { server ->
262262
encryptionManager = this.encryptionManager
263263

264264
val encryptedParams = EncryptedQuickActionChatParams(this.encryptionManager.encrypt(quickActionParams), partialResultToken)
@@ -613,7 +613,7 @@ class BrowserConnector(
613613
}
614614
}
615615

616-
private inline fun <reified Request, Response> handleChat(
616+
private suspend inline fun <reified Request, Response> handleChat(
617617
lspMethod: JsonRpcMethod<Request, Response>,
618618
node: JsonNode,
619619
crossinline serverAction: (params: Request, invokeService: () -> CompletableFuture<Response>) -> CompletableFuture<Response>,
@@ -624,7 +624,7 @@ class BrowserConnector(
624624
serializer.deserializeChatMessages<Request>(node.params, lspMethod.params)
625625
}
626626

627-
return AmazonQLspService.executeIfRunning(project) { _ ->
627+
return AmazonQLspService.executeAsyncIfRunning(project) { _ ->
628628
val invokeService = when (lspMethod) {
629629
is JsonRpcNotification<Request> -> {
630630
// notify is Unit
@@ -646,7 +646,7 @@ class BrowserConnector(
646646
} ?: CompletableFuture.failedFuture<Response>(IllegalStateException("LSP Server not running"))
647647
}
648648

649-
private inline fun <reified Request, Response> handleChat(
649+
private suspend inline fun <reified Request, Response> handleChat(
650650
lspMethod: JsonRpcMethod<Request, Response>,
651651
node: JsonNode,
652652
): CompletableFuture<Response> = handleChat(

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -170,7 +170,7 @@ class FeatureDevControllerTest : FeatureDevTestBase() {
170170
whenever(featureDevClient.createTaskAssistConversation()).thenReturn(exampleCreateTaskAssistConversationResponse)
171171
whenever(featureDevClient.sendFeatureDevTelemetryEvent(any())).thenReturn(exampleSendTelemetryEventResponse)
172172
whenever(chatSessionStorage.getSession(any(), any())).thenReturn(spySession)
173-
doNothing().`when`(chatSessionStorage).deleteSession(any())
173+
doNothing().whenever(chatSessionStorage).deleteSession(any())
174174

175175
mockkObject(AmazonqTelemetry)
176176
every { AmazonqTelemetry.endChat(amazonqConversationId = any(), amazonqEndOfTheConversationLatency = any()) } just runs

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

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -187,11 +187,17 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
187187
CodeWhispererInvocationStatus.getInstance().setInvocationStart()
188188
var nextToken: Either<String, Int>? = null
189189
do {
190-
val result = AmazonQLspService.executeIfRunning(requestContext.project) { server ->
190+
val result = AmazonQLspService.executeAsyncIfRunning(requestContext.project) { server ->
191191
val params = createInlineCompletionParams(requestContext.editor, requestContext.triggerTypeInfo, nextToken)
192192
server.inlineCompletionWithReferences(params)
193193
}
194-
val completion = result?.await() ?: break
194+
val completion = result?.await()
195+
if (completion == null) {
196+
// no result / not running
197+
CodeWhispererInvocationStatus.getInstance().finishInvocation()
198+
break
199+
}
200+
195201
nextToken = completion.partialResultToken
196202
val endTime = System.nanoTime()
197203
val latency = TimeUnit.NANOSECONDS.toMillis(endTime - startTime).toDouble()
@@ -426,7 +432,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
426432
CodeWhispererTelemetryService.getInstance().sendUserTriggerDecisionEvent(project, latencyContext, sessionId, recommendationContext)
427433
}
428434

429-
fun getRequestContext(
435+
suspend fun getRequestContext(
430436
triggerTypeInfo: TriggerTypeInfo,
431437
editor: Editor,
432438
project: Project,
@@ -472,11 +478,11 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
472478
)
473479
}
474480

475-
fun getWorkspaceIds(project: Project): CompletableFuture<LspServerConfigurations> {
481+
suspend fun getWorkspaceIds(project: Project): CompletableFuture<LspServerConfigurations> {
476482
val payload = GetConfigurationFromServerParams(
477483
section = "aws.q.workspaceContext"
478484
)
479-
return AmazonQLspService.executeIfRunning(project) { server ->
485+
return AmazonQLspService.executeAsyncIfRunning(project) { server ->
480486
server.getConfigurationFromServer(payload)
481487
} ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")))
482488
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -203,7 +203,7 @@ class CodeWhispererServiceNew(private val cs: CoroutineScope) : Disposable {
203203
var requestCount = 0
204204
var nextToken: Either<String, Int>? = null
205205
do {
206-
val result = AmazonQLspService.executeIfRunning(requestContext.project) { server ->
206+
val result = AmazonQLspService.executeAsyncIfRunning(requestContext.project) { server ->
207207
val params = createInlineCompletionParams(requestContext.editor, requestContext.triggerTypeInfo, nextToken)
208208
server.inlineCompletionWithReferences(params)
209209
}

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

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import com.intellij.openapi.options.BoundConfigurable
1010
import com.intellij.openapi.options.Configurable
1111
import com.intellij.openapi.options.SearchableConfigurable
1212
import com.intellij.openapi.options.ex.Settings
13+
import com.intellij.openapi.progress.currentThreadCoroutineScope
1314
import com.intellij.openapi.project.Project
1415
import com.intellij.openapi.project.ProjectManager
1516
import com.intellij.openapi.ui.emptyText
@@ -22,6 +23,8 @@ import com.intellij.ui.dsl.builder.bindText
2223
import com.intellij.ui.dsl.builder.panel
2324
import com.intellij.util.concurrency.EdtExecutorService
2425
import com.intellij.util.execution.ParametersListUtil
26+
import kotlinx.coroutines.launch
27+
import org.eclipse.lsp4j.DidChangeConfigurationParams
2528
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
2629
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManagerListener
2730
import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
@@ -310,7 +313,12 @@ class CodeWhispererConfigurable(private val project: Project) :
310313
if (project.isDisposed) {
311314
return@forEach
312315
}
313-
AmazonQLspService.didChangeConfiguration(project)
316+
317+
currentThreadCoroutineScope().launch {
318+
AmazonQLspService.executeAsyncIfRunning(project) { server ->
319+
server.workspaceService.didChangeConfiguration(DidChangeConfigurationParams())
320+
}
321+
}
314322
}
315323
}
316324
}

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

Lines changed: 21 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ package software.aws.toolkits.jetbrains.services.codewhisperer.telemetry
66
import com.intellij.openapi.components.Service
77
import com.intellij.openapi.components.service
88
import com.intellij.openapi.project.Project
9+
import kotlinx.coroutines.CoroutineScope
10+
import kotlinx.coroutines.launch
911
import software.aws.toolkits.core.utils.debug
1012
import software.aws.toolkits.core.utils.getLogger
1113
import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
@@ -33,7 +35,7 @@ import java.time.Duration
3335
import java.time.Instant
3436

3537
@Service
36-
class CodeWhispererTelemetryService {
38+
class CodeWhispererTelemetryService(private val cs: CoroutineScope) {
3739
companion object {
3840
fun getInstance(): CodeWhispererTelemetryService = service()
3941
val LOG = getLogger<CodeWhispererTelemetryService>()
@@ -45,22 +47,24 @@ class CodeWhispererTelemetryService {
4547
sessionId: String,
4648
recommendationContext: RecommendationContext,
4749
) {
48-
AmazonQLspService.executeIfRunning(project) { server ->
49-
val params = LogInlineCompletionSessionResultsParams(
50-
sessionId = sessionId,
51-
completionSessionResult = recommendationContext.details.associate {
52-
it.itemId to InlineCompletionStates(
53-
seen = it.hasSeen,
54-
accepted = it.isAccepted,
55-
discarded = it.isDiscarded
56-
)
57-
},
58-
firstCompletionDisplayLatency = latencyContext.perceivedLatency,
59-
totalSessionDisplayTime = CodeWhispererInvocationStatus.getInstance().completionShownTime?.let { Duration.between(it, Instant.now()) }
60-
?.toMillis()?.toDouble(),
61-
typeaheadLength = recommendationContext.userInput.length.toLong()
62-
)
63-
server.logInlineCompletionSessionResults(params)
50+
cs.launch {
51+
AmazonQLspService.executeAsyncIfRunning(project) { server ->
52+
val params = LogInlineCompletionSessionResultsParams(
53+
sessionId = sessionId,
54+
completionSessionResult = recommendationContext.details.associate {
55+
it.itemId to InlineCompletionStates(
56+
seen = it.hasSeen,
57+
accepted = it.isAccepted,
58+
discarded = it.isDiscarded
59+
)
60+
},
61+
firstCompletionDisplayLatency = latencyContext.perceivedLatency,
62+
totalSessionDisplayTime = CodeWhispererInvocationStatus.getInstance().completionShownTime?.let { Duration.between(it, Instant.now()) }
63+
?.toMillis()?.toDouble(),
64+
typeaheadLength = recommendationContext.userInput.length.toLong()
65+
)
66+
server.logInlineCompletionSessionResults(params)
67+
}
6468
}
6569
}
6670

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

Lines changed: 23 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@ package software.aws.toolkits.jetbrains.services.codewhisperer.telemetry
66
import com.intellij.openapi.components.Service
77
import com.intellij.openapi.components.service
88
import com.intellij.openapi.project.Project
9+
import kotlinx.coroutines.CoroutineScope
10+
import kotlinx.coroutines.launch
911
import software.aws.toolkits.core.utils.debug
1012
import software.aws.toolkits.core.utils.getLogger
1113
import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
@@ -29,7 +31,7 @@ import java.time.Duration
2931
import java.time.Instant
3032

3133
@Service
32-
class CodeWhispererTelemetryServiceNew {
34+
class CodeWhispererTelemetryServiceNew(private val cs: CoroutineScope) {
3335

3436
companion object {
3537
fun getInstance(): CodeWhispererTelemetryServiceNew = service()
@@ -109,24 +111,26 @@ class CodeWhispererTelemetryServiceNew {
109111
}
110112

111113
fun sendUserTriggerDecisionEvent(project: Project, latencyContext: LatencyContext) {
112-
AmazonQLspService.executeIfRunning(project) { server ->
113-
CodeWhispererServiceNew.getInstance().getAllPaginationSessions().forEach { jobId, state ->
114-
if (state == null) return@forEach
115-
val params = LogInlineCompletionSessionResultsParams(
116-
sessionId = state.responseContext.sessionId,
117-
completionSessionResult = state.recommendationContext.details.associate {
118-
it.itemId to InlineCompletionStates(
119-
seen = it.hasSeen,
120-
accepted = it.isAccepted,
121-
discarded = it.isDiscarded
122-
)
123-
},
124-
firstCompletionDisplayLatency = latencyContext.perceivedLatency,
125-
totalSessionDisplayTime = CodeWhispererInvocationStatus.getInstance().completionShownTime?.let { Duration.between(it, Instant.now()) }
126-
?.toMillis()?.toDouble(),
127-
typeaheadLength = state.recommendationContext.userInput.length.toLong()
128-
)
129-
server.logInlineCompletionSessionResults(params)
114+
cs.launch {
115+
AmazonQLspService.executeAsyncIfRunning(project) { server ->
116+
CodeWhispererServiceNew.getInstance().getAllPaginationSessions().forEach { jobId, state ->
117+
if (state == null) return@forEach
118+
val params = LogInlineCompletionSessionResultsParams(
119+
sessionId = state.responseContext.sessionId,
120+
completionSessionResult = state.recommendationContext.details.associate {
121+
it.itemId to InlineCompletionStates(
122+
seen = it.hasSeen,
123+
accepted = it.isAccepted,
124+
discarded = it.isDiscarded
125+
)
126+
},
127+
firstCompletionDisplayLatency = latencyContext.perceivedLatency,
128+
totalSessionDisplayTime = CodeWhispererInvocationStatus.getInstance().completionShownTime?.let { Duration.between(it, Instant.now()) }
129+
?.toMillis()?.toDouble(),
130+
typeaheadLength = state.recommendationContext.userInput.length.toLong()
131+
)
132+
server.logInlineCompletionSessionResults(params)
133+
}
130134
}
131135
}
132136
}

0 commit comments

Comments
 (0)