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 b7d76719348..ebe5858d97e 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 @@ -18,7 +18,6 @@ import kotlinx.coroutines.flow.distinctUntilChanged import kotlinx.coroutines.flow.launchIn import kotlinx.coroutines.flow.merge import kotlinx.coroutines.flow.onEach -import kotlinx.coroutines.future.await import kotlinx.coroutines.launch import org.cef.browser.CefBrowser import org.eclipse.lsp4j.Position @@ -32,6 +31,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.ChatCommunicationManager.Companion.convertToJsonToSendToChat 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 @@ -61,6 +61,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatN import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatPrompt import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatReadyNotification +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatUiMessageParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ConversationClickRequest import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CopyCodeToClipboardNotification import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CopyCodeToClipboardParams @@ -90,9 +91,11 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.Promp import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.PromptInputOptionChangeParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.QuickChatActionRequest import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SEND_CHAT_COMMAND_PROMPT +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.STOP_CHAT_RESPONSE import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SendChatPromptRequest import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SourceLinkClickNotification import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SourceLinkClickParams +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.StopResponseMessage import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.TabBarActionParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.TabBarActionRequest import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.TabEventParams @@ -102,6 +105,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.util.tabType import software.aws.toolkits.jetbrains.services.amazonq.webview.theme.AmazonQTheme import software.aws.toolkits.jetbrains.services.amazonq.webview.theme.ThemeBrowserAdapter import software.aws.toolkits.jetbrains.settings.MeetQSettings +import software.aws.toolkits.resources.AwsCoreBundle import software.aws.toolkits.telemetry.MetricResult import software.aws.toolkits.telemetry.Telemetry import java.util.concurrent.CompletableFuture @@ -245,6 +249,10 @@ class BrowserConnector( encryptionManager = this.encryptionManager encryptionManager?.encrypt(chatParams)?.let { EncryptedChatParams(it, partialResultToken) }?.let { server.sendChatPrompt(it) } } ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running"))) + + // We assume there is only one outgoing request per tab because the input is + // blocked when there is an outgoing request + chatCommunicationManager.setInflightRequestForTab(tabId, result) showResult(result, partialResultToken, tabId, encryptionManager, browser) } CHAT_QUICK_ACTION -> { @@ -262,6 +270,10 @@ class BrowserConnector( } } ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running"))) + // We assume there is only one outgoing request per tab because the input is + // blocked when there is an outgoing request + chatCommunicationManager.setInflightRequestForTab(tabId, result) + showResult(result, partialResultToken, tabId, encryptionManager, browser) } CHAT_LIST_CONVERSATIONS -> { @@ -318,6 +330,7 @@ class BrowserConnector( CHAT_TAB_REMOVE -> { handleChatNotification(node) { server, params -> chatCommunicationManager.removePartialChatMessage(params.tabId) + cancelInflightRequests(params.tabId) server.tabRemove(params) } } @@ -414,6 +427,30 @@ class BrowserConnector( server.createPrompt(params) } } + STOP_CHAT_RESPONSE -> { + val stopResponseRequest = serializer.deserializeChatMessages(node) + if (!chatCommunicationManager.hasInflightRequest(stopResponseRequest.params.tabId)) { + return + } + cancelInflightRequests(stopResponseRequest.params.tabId) + chatCommunicationManager.removePartialChatMessage(stopResponseRequest.params.tabId) + + val paramsJson = Gson().toJson( + // https://github.com/aws/language-servers/blob/1c0d88806087125b6fc561f610cc15e98127c6bf/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts#L403 + ChatUiMessageParams( + title = AwsCoreBundle.message("amazonqChat.stopChatResponse"), + body = "" + ) + ) + + val uiMessage = convertToJsonToSendToChat( + command = SEND_CHAT_COMMAND_PROMPT, + tabId = stopResponseRequest.params.tabId, + params = paramsJson.toString(), + isPartialResult = false + ) + browser.postChat(uiMessage) + } } } @@ -433,6 +470,14 @@ class BrowserConnector( isPartialResult = false ) browser.postChat(messageToChat) + chatCommunicationManager.removeInflightRequestForTab(tabId) + } + } + + private fun cancelInflightRequests(tabId: String) { + chatCommunicationManager.getInflightRequestForTab(tabId)?.let { request -> + request.cancel(true) + chatCommunicationManager.removeInflightRequestForTab(tabId) } } 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 f148bc2ffe7..90a321b24c5 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 @@ -20,8 +20,21 @@ import java.util.concurrent.ConcurrentHashMap @Service(Service.Level.PROJECT) class ChatCommunicationManager { private val chatPartialResultMap = ConcurrentHashMap() - private fun getPartialChatMessage(partialResultToken: String): String = - chatPartialResultMap.getValue(partialResultToken) + private fun getPartialChatMessage(partialResultToken: String): String? = + chatPartialResultMap.getOrDefault(partialResultToken, null) + + private val inflightRequestByTabId = ConcurrentHashMap>() + + fun setInflightRequestForTab(tabId: String, result: CompletableFuture) { + inflightRequestByTabId[tabId] = result + } + fun removeInflightRequestForTab(tabId: String) { + inflightRequestByTabId.remove(tabId) + } + + fun getInflightRequestForTab(tabId: String): CompletableFuture? = inflightRequestByTabId[tabId] + + fun hasInflightRequest(tabId: String): Boolean = inflightRequestByTabId.containsKey(tabId) fun addPartialChatMessage(tabId: String): String { val partialResultToken: String = UUID.randomUUID().toString() @@ -35,7 +48,7 @@ class ChatCommunicationManager { fun handlePartialResultProgressNotification(project: Project, params: ProgressParams) { val token = ProgressNotificationUtils.getToken(params) val tabId = getPartialChatMessage(token) - if (tabId == null || tabId.isEmpty()) { + if (tabId.isNullOrEmpty()) { return } if (params.value.isLeft || params.value.right == null) { @@ -47,7 +60,6 @@ class ChatCommunicationManager { val encryptedPartialChatResult = getObject(params, String::class.java) if (encryptedPartialChatResult != null) { val partialChatResult = AmazonQLspService.getInstance(project).encryptionManager.decrypt(encryptedPartialChatResult) - val uiMessage = convertToJsonToSendToChat( command = SEND_CHAT_COMMAND_PROMPT, tabId = tabId, diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/ChatUiMessageParams.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/ChatUiMessageParams.kt new file mode 100644 index 00000000000..c6057bab9fe --- /dev/null +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/ChatUiMessageParams.kt @@ -0,0 +1,14 @@ +// 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.model.aws.chat +import java.util.UUID + +data class ChatUiMessageParams( + val title: String, + val additionalMessages: List = emptyList(), + val messageId: String = UUID.randomUUID().toString(), + val buttons: List = emptyList(), + val codeReference: List = emptyList(), + val body: String = "", +) diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/FlareChatCommands.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/FlareChatCommands.kt index 02934b87db1..614385a7abc 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/FlareChatCommands.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/FlareChatCommands.kt @@ -31,3 +31,4 @@ const val CHAT_SEND_UPDATE = "aws/chat/sendChatUpdate" const val CHAT_CREATE_PROMPT = "aws/chat/createPrompt" const val SHOW_SAVE_FILE_DIALOG_REQUEST_METHOD = "aws/showSaveFileDialog" const val GET_SERIALIZED_CHAT_REQUEST_METHOD = "aws/chat/getSerializedChat" +const val STOP_CHAT_RESPONSE = "stopChatResponse" diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/StopResponseParams.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/StopResponseParams.kt new file mode 100644 index 00000000000..5ddd9d4a430 --- /dev/null +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/StopResponseParams.kt @@ -0,0 +1,12 @@ +// 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.model.aws.chat + +data class StopResponseMessage( + val params: StopResponseParams, +) + +data class StopResponseParams( + val tabId: String, +) diff --git a/plugins/core/resources/resources/software/aws/toolkits/resources/MessagesBundle.properties b/plugins/core/resources/resources/software/aws/toolkits/resources/MessagesBundle.properties index d21089142ba..74c794110a4 100644 --- a/plugins/core/resources/resources/software/aws/toolkits/resources/MessagesBundle.properties +++ b/plugins/core/resources/resources/software/aws/toolkits/resources/MessagesBundle.properties @@ -38,6 +38,7 @@ action.aws.toolkit.toolwindow.newConnection.text=Add Another Connection... action.dynamic.open.text=Open Resource... action.q.openchat.text=Open Chat Panel amazonqChat.project_context.index_in_progress=By the way, I'm still indexing this project for full context from your workspace. I may have a better response in a few minutes when it's complete if you'd like to try again then. +amazonqChat.stopChatResponse=You stopped your current work, please provide additional examples or ask another question. amazonqDoc.answer.codeResult=You can accept the changes to your files, or describe any additional changes you'd like me to make. amazonqDoc.answer.readmeCreated=I've created a README for your code. amazonqDoc.answer.readmeUpdated=I've updated your README.