Skip to content

fix(amazonq): ChatCommunicationManager blocks messages to UI until ready #5697

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 15 commits into from
May 9, 2025
Merged
Show file tree
Hide file tree
Changes from 6 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ 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.artifacts.ArtifactManager
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AsyncChatUiListener
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.FlareUiMessage
import software.aws.toolkits.jetbrains.services.amazonq.messages.AmazonQMessage
import software.aws.toolkits.jetbrains.services.amazonq.messages.MessageConnector
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
Expand Down Expand Up @@ -55,9 +56,15 @@ class AmazonQPanel(val project: Project, private val scope: CoroutineScope) : Di
project.messageBus.connect().subscribe(
AsyncChatUiListener.TOPIC,
object : AsyncChatUiListener {
override fun onChange(message: String) {
override fun onChange(command: String) {
runInEdt {
browser.get()?.postChat(message)
browser.get()?.postChat(command)
}
}

override fun onChange(command: FlareUiMessage) {
runInEdt {
browser.get()?.postChat(command)
}
}
}
Expand Down
Original file line number Diff line number Diff line change
@@ -1,16 +1,19 @@
// Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

@file:Suppress("BannedImports")
package software.aws.toolkits.jetbrains.services.amazonq.webview

import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
import com.google.gson.Gson
import com.intellij.openapi.Disposable
import com.intellij.openapi.project.Project
import com.intellij.openapi.util.Disposer
import com.intellij.ui.jcef.JBCefJSQuery
import org.cef.CefApp
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
import software.aws.toolkits.jetbrains.services.amazonq.util.HighlightCommand
import software.aws.toolkits.jetbrains.services.amazonq.util.createBrowser
Expand Down Expand Up @@ -62,6 +65,9 @@ class Browser(parent: Disposable, private val webUri: URI, val project: Project)

fun component() = jcefBrowser.component

fun postChat(command: FlareUiMessage) = postChat(Gson().toJson(command))

@Deprecated("shouldn't need this version")
fun postChat(message: String) {
jcefBrowser
.cefBrowser
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AwsServerCapabilitiesProvider
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.ChatCommunicationManager
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.FlareUiMessage
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.getTextDocumentIdentifier
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ButtonClickNotification
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ButtonClickParams
Expand Down Expand Up @@ -256,7 +257,9 @@ class BrowserConnector(
var encryptionManager: JwtEncryptionManager? = null
val result = AmazonQLspService.executeIfRunning(project) { server ->
encryptionManager = this.encryptionManager
encryptionManager?.encrypt(chatParams)?.let { EncryptedChatParams(it, partialResultToken) }?.let { server.sendChatPrompt(it) }

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
Expand All @@ -272,11 +275,9 @@ class BrowserConnector(
var encryptionManager: JwtEncryptionManager? = null
val result = AmazonQLspService.executeIfRunning(project) { server ->
encryptionManager = this.encryptionManager
encryptionManager?.encrypt(quickActionParams)?.let {
EncryptedQuickActionChatParams(it, partialResultToken)
}?.let {
server.sendQuickAction(it)
}

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
Expand Down Expand Up @@ -325,6 +326,7 @@ class BrowserConnector(
CHAT_READY -> {
handleChatNotification<ChatReadyNotification, Unit>(node) { server, _ ->
uiReady.complete(true)
chatCommunicationManager.setUiReady(true)
RunOnceUtil.runOnceForApp("AmazonQ-UI-Ready") {
MeetQSettings.getInstance().reinvent2024OnboardingCount += 1
}
Expand Down Expand Up @@ -424,13 +426,18 @@ class BrowserConnector(
handleChatNotification<TabBarActionRequest, TabBarActionParams>(node) {
server, params ->
val result = server.tabBarActions(params)
result.whenComplete { params1, error ->
result.whenComplete { actions, error ->
try {
if (error != null) {
throw error
}
val res = ChatCommunicationManager.convertNotificationToJsonForChat(CHAT_TAB_BAR_ACTIONS, params1)
browser.postChat(res)

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 {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@ package software.aws.toolkits.jetbrains.services.amazonq.lsp
import org.eclipse.lsp4j.jsonrpc.services.JsonNotification
import org.eclipse.lsp4j.jsonrpc.services.JsonRequest
import org.eclipse.lsp4j.services.LanguageClient
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.LSPAny
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_SEND_CONTEXT_COMMANDS
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_SEND_UPDATE
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatUpdateParams
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.GET_SERIALIZED_CHAT_REQUEST_METHOD
Expand Down Expand Up @@ -43,4 +45,7 @@ interface AmazonQLanguageClient : LanguageClient {

@JsonNotification(OPEN_FILE_DIFF)
fun openFileDiff(params: OpenFileDiffParams): CompletableFuture<Unit>

@JsonNotification(CHAT_SEND_CONTEXT_COMMANDS)
fun sendContextCommands(params: LSPAny): CompletableFuture<Unit>
}
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
@file:Suppress("BannedImports")
package software.aws.toolkits.jetbrains.services.amazonq.lsp

import com.google.gson.Gson
import com.intellij.diff.DiffContentFactory
import com.intellij.diff.DiffManager
import com.intellij.diff.requests.SimpleDiffRequest
Expand All @@ -24,13 +23,16 @@
import org.eclipse.lsp4j.ShowDocumentParams
import org.eclipse.lsp4j.ShowDocumentResult
import org.eclipse.lsp4j.ShowMessageRequestParams
import software.aws.toolkits.core.utils.error
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.core.utils.warn
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AsyncChatUiListener
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.ChatCommunicationManager
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.FlareUiMessage
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.LSPAny
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_OPEN_TAB
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_SEND_UPDATE
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatUpdateParams
Expand Down Expand Up @@ -133,14 +135,13 @@
val result = CompletableFuture<OpenTabResult>()
ChatCommunicationManager.pendingTabRequests[requestId] = result

val uiMessage = """
{
"command": "$CHAT_OPEN_TAB",
"params": ${Gson().toJson(params)},
"requestId": "$requestId"
}
""".trimIndent()
AsyncChatUiListener.notifyPartialMessageUpdate(uiMessage)
AsyncChatUiListener.notifyPartialMessageUpdate(
FlareUiMessage(
command = CHAT_OPEN_TAB,
params = params,
requestId = requestId,
)
)

result.orTimeout(30000, TimeUnit.MILLISECONDS)
.whenComplete { _, error ->
Expand Down Expand Up @@ -186,14 +187,13 @@

ChatCommunicationManager.pendingSerializedChatRequests[requestId] = result

val uiMessage = """
{
"command": "$GET_SERIALIZED_CHAT_REQUEST_METHOD",
"params": ${Gson().toJson(params)},
"requestId": "$requestId"
}
""".trimIndent()
AsyncChatUiListener.notifyPartialMessageUpdate(uiMessage)
AsyncChatUiListener.notifyPartialMessageUpdate(
FlareUiMessage(
command = GET_SERIALIZED_CHAT_REQUEST_METHOD,
params = params,
requestId = requestId,
)
)

result.orTimeout(30000, TimeUnit.MILLISECONDS)
.whenComplete { _, error ->
Expand Down Expand Up @@ -252,19 +252,17 @@
try {
chatCommunicationManager.handlePartialResultProgressNotification(project, params)
} catch (e: Exception) {
error("Cannot handle partial chat")
LOG.error(e) { "Cannot handle partial chat" }
}
}

override fun sendChatUpdate(params: ChatUpdateParams): CompletableFuture<Unit> {
val uiMessage = """
{
"command":"$CHAT_SEND_UPDATE",
"params":${Gson().toJson(params)}
}
""".trimIndent()

AsyncChatUiListener.notifyPartialMessageUpdate(uiMessage)
AsyncChatUiListener.notifyPartialMessageUpdate(
FlareUiMessage(
command = CHAT_SEND_UPDATE,
params = params,
)
)

return CompletableFuture.completedFuture(Unit)
}
Expand Down Expand Up @@ -308,6 +306,11 @@
ApplicationManager.getApplication()::invokeLater
)

override fun sendContextCommands(params: LSPAny): CompletableFuture<Unit> {
ChatCommunicationManager.getInstance(project).queueContextCommands(params)

Check warning on line 310 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt

View check run for this annotation

Codecov / codecov/patch

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt#L310

Added line #L310 was not covered by tests
return CompletableFuture.completedFuture(Unit)
}

companion object {
private val LOG = getLogger<AmazonQLanguageClientImpl>()
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
@file:Suppress("BannedImports")
package software.aws.toolkits.jetbrains.services.amazonq.lsp

import com.google.gson.Gson
import com.google.gson.ToNumberPolicy
import com.intellij.execution.configurations.GeneralCommandLine
import com.intellij.execution.impl.ExecutionManagerImpl
Expand Down Expand Up @@ -44,7 +43,6 @@ import org.eclipse.lsp4j.WorkspaceClientCapabilities
import org.eclipse.lsp4j.jsonrpc.Launcher
import org.eclipse.lsp4j.jsonrpc.Launcher.Builder
import org.eclipse.lsp4j.jsonrpc.MessageConsumer
import org.eclipse.lsp4j.jsonrpc.messages.NotificationMessage
import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage
import org.eclipse.lsp4j.launch.LSPLauncher
import org.slf4j.event.Level
Expand All @@ -57,7 +55,6 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.auth.DefaultAuthCred
import software.aws.toolkits.jetbrains.services.amazonq.lsp.dependencies.DefaultModuleDependenciesService
import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AmazonQLspTypeAdapterFactory
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AsyncChatUiListener
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AwsExtendedInitializeResult
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AwsServerCapabilitiesProvider
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.createExtendedClientMetadata
Expand Down Expand Up @@ -282,23 +279,12 @@ private class AmazonQServerInstance(private val project: Project, private val cs

launcher = LSPLauncher.Builder<AmazonQLanguageServer>()
.wrapMessages { consumer ->
MessageConsumer {
message ->
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()
}
if (message is NotificationMessage && message.method == "aws/chat/sendContextCommands") {
val showContextCommands = """
{
"command":"aws/chat/sendContextCommands",
"params": ${Gson().toJson(message.params)}
}
""".trimIndent()

AsyncChatUiListener.notifyPartialMessageUpdate(showContextCommands)
}
consumer?.consume(message)
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -8,14 +8,22 @@ import com.intellij.util.messages.Topic
import java.util.EventListener

interface AsyncChatUiListener : EventListener {
fun onChange(message: String) {}
@Deprecated("shouldn't need this version")
fun onChange(command: String) {}

fun onChange(command: FlareUiMessage) {}

companion object {
@Topic.AppLevel
val TOPIC = Topic.create("Partial chat message provider", AsyncChatUiListener::class.java)

fun notifyPartialMessageUpdate(message: String) {
ApplicationManager.getApplication().messageBus.syncPublisher(TOPIC).onChange(message)
fun notifyPartialMessageUpdate(command: FlareUiMessage) {
ApplicationManager.getApplication().messageBus.syncPublisher(TOPIC).onChange(command)
}

@Deprecated("shouldn't need this version")
fun notifyPartialMessageUpdate(command: String) {
ApplicationManager.getApplication().messageBus.syncPublisher(TOPIC).onChange(command)
}
}
}
Loading
Loading