diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQPanel.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQPanel.kt index 4ea60b26c6e..e3f15935479 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQPanel.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQPanel.kt @@ -32,6 +32,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.isQSupportedInThisVersio 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.ChatCommunicationManager 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 @@ -66,6 +67,7 @@ class AmazonQPanel(val project: Project, private val scope: CoroutineScope) : Di private val appConnections = mutableListOf() init { + // will be removed in next iteration. project.messageBus.connect().subscribe( AsyncChatUiListener.TOPIC, object : AsyncChatUiListener { @@ -135,6 +137,11 @@ class AmazonQPanel(val project: Project, private val scope: CoroutineScope) : Di Browser(this@AmazonQPanel, webUri, project).also { browserInstance -> wrapper.setContent(browserInstance.component()) + // Register direct callback instead of using message bus + ChatCommunicationManager.getInstance(project).setChatUpdateCallback { message -> + browserInstance.postChat(message) + } + // Add DropTarget to the browser component // JCEF does not propagate OS-level dragenter, dragOver and drop into DOM. // As an alternative, enabling the native drag in JCEF, diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/BrowserConnector.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/BrowserConnector.kt index 8008c552278..71efd5845c2 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/BrowserConnector.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/BrowserConnector.kt @@ -43,7 +43,6 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.JsonRpcNotification import software.aws.toolkits.jetbrains.services.amazonq.lsp.JsonRpcRequest 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.ChatAsyncResultManager 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.chat.AUTH_FOLLOW_UP_CLICKED @@ -123,7 +122,6 @@ class BrowserConnector( ) { val uiReady = CompletableDeferred() private val chatCommunicationManager = ChatCommunicationManager.getInstance(project) - private val chatAsyncResultManager = ChatAsyncResultManager.getInstance(project) suspend fun connect( browser: Browser, @@ -240,7 +238,6 @@ class BrowserConnector( val tabId = requestFromUi.params.tabId val partialResultToken = chatCommunicationManager.addPartialChatMessage(tabId) - chatCommunicationManager.registerPartialResultToken(partialResultToken) var encryptionManager: JwtEncryptionManager? = null val result = AmazonQLspService.executeAsyncIfRunning(project) { server -> @@ -261,7 +258,6 @@ class BrowserConnector( val tabId = requestFromUi.params.tabId val quickActionParams = node.params ?: error("empty payload") val partialResultToken = chatCommunicationManager.addPartialChatMessage(tabId) - chatCommunicationManager.registerPartialResultToken(partialResultToken) var encryptionManager: JwtEncryptionManager? = null val result = AmazonQLspService.executeAsyncIfRunning(project) { server -> encryptionManager = this.encryptionManager @@ -595,17 +591,6 @@ class BrowserConnector( chatCommunicationManager.removeInflightRequestForTab(tabId) } catch (e: CancellationException) { LOG.warn { "Cancelled chat generation" } - try { - chatAsyncResultManager.createRequestId(partialResultToken) - chatAsyncResultManager.getResult(partialResultToken) - handleCancellation(tabId, browser) - } catch (ex: Exception) { - LOG.warn(ex) { "An error occurred while processing cancellation" } - } finally { - chatAsyncResultManager.removeRequestId(partialResultToken) - chatCommunicationManager.removePartialResultLock(partialResultToken) - chatCommunicationManager.removeFinalResultProcessed(partialResultToken) - } } catch (e: Exception) { LOG.warn(e) { "Failed to send chat message" } browser.postChat(chatCommunicationManager.getErrorUiMessage(tabId, e, partialResultToken)) @@ -613,12 +598,6 @@ class BrowserConnector( } } - private fun handleCancellation(tabId: String, browser: Browser) { - // Send a message to hide the stop button without showing an error - val cancelMessage = chatCommunicationManager.getCancellationUiMessage(tabId) - browser.postChat(cancelMessage) - } - private fun updateQuickActionsInBrowser(browser: Browser) { val isFeatureDevAvailable = isFeatureDevAvailable(project) val isCodeTransformAvailable = isCodeTransformAvailable(project) diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/commands/ActionRegistrar.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/commands/ActionRegistrar.kt index fe93c71f505..284bfc30c91 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/commands/ActionRegistrar.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/commands/ActionRegistrar.kt @@ -8,7 +8,7 @@ import com.intellij.openapi.project.Project import kotlinx.coroutines.flow.MutableSharedFlow import kotlinx.coroutines.flow.asSharedFlow import kotlinx.coroutines.runBlocking -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.chat.GENERIC_COMMAND import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.GenericCommandParams @@ -40,7 +40,7 @@ class ActionRegistrar { val params = SendToPromptParams(selection = codeSelection, triggerType = TriggerType.CONTEXT_MENU) uiMessage = FlareUiMessage(command = SEND_TO_PROMPT, params = params) } - AsyncChatUiListener.notifyPartialMessageUpdate(project, uiMessage) + ChatCommunicationManager.getInstance(project).notifyUi(uiMessage) } } } diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/commands/codescan/actions/ExplainCodeIssueAction.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/commands/codescan/actions/ExplainCodeIssueAction.kt index 94f572984a9..33a8bcf857c 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/commands/codescan/actions/ExplainCodeIssueAction.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/commands/codescan/actions/ExplainCodeIssueAction.kt @@ -11,7 +11,7 @@ import com.intellij.openapi.actionSystem.DataKey import com.intellij.openapi.application.ApplicationManager import com.intellij.openapi.project.DumbAware import kotlinx.coroutines.runBlocking -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.chat.ChatPrompt import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SEND_TO_PROMPT @@ -58,7 +58,8 @@ class ExplainCodeIssueAction : AnAction(), DumbAware { ) val uiMessage = FlareUiMessage(SEND_TO_PROMPT, params) - AsyncChatUiListener.notifyPartialMessageUpdate(project, uiMessage) + ChatCommunicationManager.getInstance(project).notifyUi(uiMessage) +// AsyncChatUiListener.notifyPartialMessageUpdate(project, uiMessage) } } } diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt index 47d490cb93d..34a87e04737 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt @@ -44,7 +44,6 @@ import software.aws.toolkits.core.utils.info import software.aws.toolkits.core.utils.warn 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 @@ -399,8 +398,7 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC } override fun sendChatUpdate(params: LSPAny) { - AsyncChatUiListener.notifyPartialMessageUpdate( - project, + chatManager.notifyUi( FlareUiMessage( command = CHAT_SEND_UPDATE, params = params, diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/AsyncChatUiListener.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/AsyncChatUiListener.kt index 21d812cf075..6b56199aa21 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/AsyncChatUiListener.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/AsyncChatUiListener.kt @@ -23,6 +23,7 @@ interface AsyncChatUiListener : EventListener { project.messageBus.syncPublisher(TOPIC).onChange(command) } + // will be removed in next iteration. @Deprecated("shouldn't need this version") fun notifyPartialMessageUpdate(project: Project, command: String) { project.messageBus.syncPublisher(TOPIC).onChange(command) diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/ChatAsyncResultManager.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/ChatAsyncResultManager.kt deleted file mode 100644 index aa729fbf12f..00000000000 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/ChatAsyncResultManager.kt +++ /dev/null @@ -1,74 +0,0 @@ -// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. -// SPDX-License-Identifier: Apache-2.0 - -package software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat - -import com.intellij.openapi.components.Service -import com.intellij.openapi.components.service -import com.intellij.openapi.project.Project -import java.util.concurrent.CompletableFuture -import java.util.concurrent.ConcurrentHashMap -import java.util.concurrent.TimeUnit -import java.util.concurrent.TimeoutException - -/** - * Manages asynchronous results for chat operations, particularly handling the coordination - * between partial results and final results during cancellation. - */ -@Service(Service.Level.PROJECT) -class ChatAsyncResultManager { - private val results = ConcurrentHashMap>() - private val completedResults = ConcurrentHashMap() - private val timeout = 30L - private val timeUnit = TimeUnit.SECONDS - - fun createRequestId(requestId: String) { - if (!completedResults.containsKey(requestId)) { - results[requestId] = CompletableFuture() - } - } - - fun removeRequestId(requestId: String) { - val future = results.remove(requestId) - if (future != null && !future.isDone) { - future.cancel(true) - } - completedResults.remove(requestId) - } - - fun setResult(requestId: String, result: Any) { - val future = results[requestId] - if (future != null) { - future.complete(result) - results.remove(requestId) - } - completedResults[requestId] = result - } - - fun getResult(requestId: String): Any? = - getResult(requestId, timeout, timeUnit) - - private fun getResult(requestId: String, timeout: Long, unit: TimeUnit): Any? { - val completedResult = completedResults[requestId] - if (completedResult != null) { - return completedResult - } - - val future = results[requestId] ?: throw IllegalArgumentException("Request ID not found: $requestId") - - try { - val result = future.get(timeout, unit) - completedResults[requestId] = result - results.remove(requestId) - return result - } catch (e: TimeoutException) { - future.cancel(true) - results.remove(requestId) - throw TimeoutException("Operation timed out for requestId: $requestId") - } - } - - companion object { - fun getInstance(project: Project) = project.service() - } -} diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/ChatCommunicationManager.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/ChatCommunicationManager.kt index 57f7aef7b45..a4481219033 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/ChatCommunicationManager.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/ChatCommunicationManager.kt @@ -43,18 +43,22 @@ class ChatCommunicationManager(private val project: Project, private val cs: Cor private val inflightRequestByTabId = ConcurrentHashMap>() private val pendingSerializedChatRequests = ConcurrentHashMap>() private val pendingTabRequests = ConcurrentHashMap>() - private val partialResultLocks = ConcurrentHashMap() - private val finalResultProcessed = ConcurrentHashMap() private val openTabs = mutableSetOf() fun setUiReady() { uiReady.complete(true) } + private var chatUpdateCallback: ((FlareUiMessage) -> Unit)? = null + + fun setChatUpdateCallback(callback: (FlareUiMessage) -> Unit) { + chatUpdateCallback = callback + } + fun notifyUi(uiMessage: FlareUiMessage) { cs.launch { uiReady.await() - AsyncChatUiListener.notifyPartialMessageUpdate(project, uiMessage) + chatUpdateCallback?.invoke(uiMessage) } } @@ -118,20 +122,6 @@ class ChatCommunicationManager(private val project: Project, private val cs: Cor fun removeTabOpenRequest(requestId: String) = pendingTabRequests.remove(requestId) - fun removePartialResultLock(token: String) { - partialResultLocks.remove(token) - } - - fun removeFinalResultProcessed(token: String) { - finalResultProcessed.remove(token) - } - - fun registerPartialResultToken(partialResultToken: String) { - val lock = Any() - partialResultLocks[partialResultToken] = lock - finalResultProcessed[partialResultToken] = false - } - fun handlePartialResultProgressNotification(project: Project, params: ProgressParams) { val token = ProgressNotificationUtils.getToken(params) val tabId = getPartialChatMessage(token) @@ -147,49 +137,18 @@ class ChatCommunicationManager(private val project: Project, private val cs: Cor val encryptedPartialChatResult = getObject(params, String::class.java) if (encryptedPartialChatResult != null) { val partialChatResult = AmazonQLspService.getInstance(project).encryptionManager.decrypt(encryptedPartialChatResult) - - // Special case: check for stop message before proceeding - val partialResultMap = tryOrNull { + val partialResultMap: Any = tryOrNull { Gson().fromJson(partialChatResult, Map::class.java) - } + } ?: partialChatResult - if (partialResultMap != null) { - @Suppress("UNCHECKED_CAST") - val additionalMessages = partialResultMap["additionalMessages"] as? List> - if (additionalMessages != null) { - for (message in additionalMessages) { - val messageId = message["messageId"] as? String - if (messageId != null && messageId.startsWith("stopped")) { - // Process stop messages immediately - val uiMessage = convertToJsonToSendToChat( - command = SEND_CHAT_COMMAND_PROMPT, - tabId = tabId, - params = partialChatResult, - isPartialResult = true - ) - AsyncChatUiListener.notifyPartialMessageUpdate(project, uiMessage) - finalResultProcessed[token] = true - ChatAsyncResultManager.getInstance(project).setResult(token, partialResultMap) - return - } - } - } - } - - // Normal processing for non-stop messages - val lock = partialResultLocks[token] ?: return - synchronized(lock) { - if (finalResultProcessed[token] == true || partialResultLocks[token] == null) { - return@synchronized - } - val uiMessage = convertToJsonToSendToChat( + notifyUi( + FlareUiMessage( command = SEND_CHAT_COMMAND_PROMPT, tabId = tabId, - params = partialChatResult, + params = partialResultMap, isPartialResult = true ) - AsyncChatUiListener.notifyPartialMessageUpdate(project, uiMessage) - } + ) } } @@ -219,21 +178,6 @@ class ChatCommunicationManager(private val project: Project, private val cs: Cor return uiMessage } - fun getCancellationUiMessage(tabId: String): String { - // Create a minimal error params with empty error message to hide the stop button - // without showing an actual error message to the user - val errorParams = Gson().toJson(ErrorParams(tabId, null, "", "")).toString() - - return """ - { - "command":"$CHAT_ERROR_PARAMS", - "tabId": "$tabId", - "params": $errorParams, - "isPartialResult": false - } - """.trimIndent() - } - fun handleAuthFollowUpClicked(project: Project, params: AuthFollowUpClickedParams) { val incomingType = params.authFollowupType val connectionManager = ToolkitConnectionManager.getInstance(project) diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/FlareUiMessage.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/FlareUiMessage.kt index 7bc5fe2461f..7faaee02f76 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/FlareUiMessage.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/FlareUiMessage.kt @@ -7,4 +7,6 @@ data class FlareUiMessage( val command: String, val params: Any, val requestId: String? = null, + val tabId: String? = null, + val isPartialResult: Boolean? = false, )