From 54fa53bfc215f53bc841f892c77223239337a283 Mon Sep 17 00:00:00 2001 From: Richard Li Date: Tue, 13 May 2025 01:08:05 -0700 Subject: [PATCH 1/3] fix(amazonq): avoid using modeled types in chat message proxy There is a high possibiltiy of drift in our modeling as the UI <-> Server evolves. To avoid a treadmill in making sure these are always up to date, avoid materializing these messages into concrete types --- .../amazonq/commands/MessageSerializer.kt | 3 + .../amazonq/webview/BrowserConnector.kt | 332 +++++++++--------- .../services/amazonq/lsp/AmazonQChatServer.kt | 204 +++++++++++ .../amazonq/lsp/AmazonQLanguageClient.kt | 6 +- .../amazonq/lsp/AmazonQLanguageClientImpl.kt | 5 +- .../amazonq/lsp/AmazonQLanguageServer.kt | 107 ------ .../services/amazonq/lsp/AmazonQLspService.kt | 27 +- .../lsp/auth/DefaultAuthCredentialsService.kt | 2 +- .../lsp/flareChat/ChatCommunicationManager.kt | 13 +- .../amazonq/lsp/model/aws/chat/ChatMessage.kt | 15 +- .../amazonq/lsp/model/aws/chat/TabEvent.kt | 14 +- 11 files changed, 418 insertions(+), 310 deletions(-) create mode 100644 plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQChatServer.kt diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/commands/MessageSerializer.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/commands/MessageSerializer.kt index a83b8ad2978..29b7213a76b 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/commands/MessageSerializer.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/commands/MessageSerializer.kt @@ -40,6 +40,9 @@ class MessageSerializer @VisibleForTesting constructor() { inline fun deserializeChatMessages(value: JsonNode): T = objectMapper.treeToValue(value) + inline fun deserializeChatMessages(value: JsonNode, clazz: Class): T = + objectMapper.treeToValue(value, clazz) + // Provide singleton global access companion object { private val instance = MessageSerializer() 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 4d7b3a34d4e..0f322e0c68d 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 @@ -4,6 +4,8 @@ package software.aws.toolkits.jetbrains.services.amazonq.webview import com.fasterxml.jackson.databind.JsonNode +import com.fasterxml.jackson.databind.node.ObjectNode +import com.fasterxml.jackson.module.kotlin.treeToValue import com.google.gson.Gson import com.intellij.ide.BrowserUtil import com.intellij.ide.util.RunOnceUtil @@ -30,16 +32,17 @@ import software.aws.toolkits.core.utils.getLogger import software.aws.toolkits.core.utils.warn import software.aws.toolkits.jetbrains.services.amazonq.apps.AppConnection import software.aws.toolkits.jetbrains.services.amazonq.commands.MessageSerializer -import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLanguageServer +import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQChatServer import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService +import software.aws.toolkits.jetbrains.services.amazonq.lsp.JsonRpcMethod +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.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 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.AuthFollowUpClickNotification -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ButtonClickNotification -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ButtonClickParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ButtonClickResult import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_BUTTON_CLICK import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_CONVERSATION_CLICK @@ -62,50 +65,23 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_ import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_TAB_BAR_ACTIONS import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_TAB_CHANGE import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_TAB_REMOVE -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatNotification -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.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 -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CreatePromptNotification -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CreatePromptParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.EncryptedChatParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.EncryptedQuickActionChatParams -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.FeedbackNotification -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.FeedbackParams -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.FileClickNotification -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.FileClickParams -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.FollowUpClickNotification -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.FollowUpClickParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.GET_SERIALIZED_CHAT_REQUEST_METHOD import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.GetSerializedChatResponse -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.InfoLinkClickNotification -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.InfoLinkClickParams -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.InsertToCursorPositionNotification -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.InsertToCursorPositionParams -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.LinkClickNotification -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.LinkClickParams -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ListConversationsRequest import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.OPEN_SETTINGS import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.OPEN_WORKSPACE_SETTINGS_KEY import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.OpenSettingsNotification import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.OpenTabResponse +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.OpenTabResult +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.OpenTabResultError +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.OpenTabResultSuccess import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.PROMPT_INPUT_OPTIONS_CHANGE -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.PromptInputOptionChangeNotification -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 -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.TabEventRequest import software.aws.toolkits.jetbrains.services.amazonq.lsp.util.LspEditorUtil import software.aws.toolkits.jetbrains.services.amazonq.lsp.util.LspEditorUtil.toUriString import software.aws.toolkits.jetbrains.services.amazonq.util.command @@ -232,23 +208,19 @@ class BrowserConnector( when (node.command) { SEND_CHAT_COMMAND_PROMPT -> { val requestFromUi = serializer.deserializeChatMessages(node) - val chatPrompt = ChatPrompt( - requestFromUi.params.prompt.prompt, - requestFromUi.params.prompt.escapedPrompt, - node.command - ) val editor = FileEditorManager.getInstance(project).selectedTextEditor val textDocumentIdentifier = editor?.let { TextDocumentIdentifier(toUriString(it.virtualFile)) } val cursorState = editor?.let { LspEditorUtil.getCursorState(it) } - val chatParams = ChatParams( - requestFromUi.params.tabId, - chatPrompt, - textDocumentIdentifier, - cursorState, - context = requestFromUi.params.context + val enrichmentParams = mapOf( + "textDocument" to textDocumentIdentifier, + "cursorState" to cursorState, ) + val serializedEnrichmentParams = serializer.objectMapper.valueToTree(enrichmentParams) + val chatParams: ObjectNode = (node as ObjectNode) + .setAll(serializedEnrichmentParams) + val tabId = requestFromUi.params.tabId val partialResultToken = chatCommunicationManager.addPartialChatMessage(tabId) @@ -256,8 +228,8 @@ class BrowserConnector( val result = AmazonQLspService.executeIfRunning(project) { server -> encryptionManager = this.encryptionManager - val encryptedParams = EncryptedChatParams(this.encryptionManager.encrypt(chatParams), partialResultToken) - server.sendChatPrompt(encryptedParams) + val encryptedParams = EncryptedChatParams(this.encryptionManager.encrypt(chatParams.params), partialResultToken) + rawEndpoint.request(SEND_CHAT_COMMAND_PROMPT, encryptedParams) as CompletableFuture } ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running"))) // We assume there is only one outgoing request per tab because the input is @@ -265,17 +237,18 @@ class BrowserConnector( chatCommunicationManager.setInflightRequestForTab(tabId, result) showResult(result, partialResultToken, tabId, encryptionManager, browser) } + CHAT_QUICK_ACTION -> { val requestFromUi = serializer.deserializeChatMessages(node) val tabId = requestFromUi.params.tabId - val quickActionParams = requestFromUi.params + val quickActionParams = node.params ?: error("empty payload") val partialResultToken = chatCommunicationManager.addPartialChatMessage(tabId) var encryptionManager: JwtEncryptionManager? = null val result = AmazonQLspService.executeIfRunning(project) { server -> encryptionManager = this.encryptionManager val encryptedParams = EncryptedQuickActionChatParams(this.encryptionManager.encrypt(quickActionParams), partialResultToken) - server.sendQuickAction(encryptedParams) + rawEndpoint.request(CHAT_QUICK_ACTION, encryptedParams) as CompletableFuture } ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running"))) // We assume there is only one outgoing request per tab because the input is @@ -284,139 +257,118 @@ class BrowserConnector( showResult(result, partialResultToken, tabId, encryptionManager, browser) } - CHAT_LIST_CONVERSATIONS -> { - val requestFromUi = serializer.deserializeChatMessages(node) - val result = AmazonQLspService.executeIfRunning(project) { server -> - server.listConversations(requestFromUi.params) - } ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running"))) - result.whenComplete { response, _ -> - val uiMessage = """ - { - "command": "$CHAT_LIST_CONVERSATIONS", - "params": ${Gson().toJson(response)} - } - """.trimIndent() - browser.postChat(uiMessage) - } + CHAT_LIST_CONVERSATIONS -> { + handleChat(AmazonQChatServer.listConversations, node) + .whenComplete { response, _ -> + browser.postChat( + FlareUiMessage( + command = CHAT_LIST_CONVERSATIONS, + params = response + ) + ) + } } - CHAT_CONVERSATION_CLICK -> { - val requestFromUi = serializer.deserializeChatMessages(node) - val result = AmazonQLspService.executeIfRunning(project) { server -> - server.conversationClick(requestFromUi.params) - } ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running"))) - result.whenComplete { response, _ -> - val uiMessage = """ - { - "command": "$CHAT_CONVERSATION_CLICK", - "params": ${Gson().toJson(response)} - } - """.trimIndent() - browser.postChat(uiMessage) - } + CHAT_CONVERSATION_CLICK -> { + handleChat(AmazonQChatServer.conversationClick, node) + .whenComplete { response, _ -> + browser.postChat( + FlareUiMessage( + command = CHAT_CONVERSATION_CLICK, + params = response + ) + ) + } } + CHAT_FEEDBACK -> { - handleChatNotification(node) { server, params -> - server.feedback(params) - } + handleChat(AmazonQChatServer.feedback, node) } + CHAT_READY -> { - handleChatNotification(node) { server, _ -> + handleChat(AmazonQChatServer.chatReady, node) { params, invoke -> uiReady.complete(true) chatCommunicationManager.setUiReady() RunOnceUtil.runOnceForApp("AmazonQ-UI-Ready") { MeetQSettings.getInstance().reinvent2024OnboardingCount += 1 } - server.chatReady() + + invoke() } } + CHAT_TAB_ADD -> { - handleChatNotification(node) { server, params -> - server.tabAdd(params) - } + handleChat(AmazonQChatServer.tabAdd, node) } + CHAT_TAB_REMOVE -> { - handleChatNotification(node) { server, params -> + handleChat(AmazonQChatServer.tabRemove, node) { params, invoke -> chatCommunicationManager.removePartialChatMessage(params.tabId) cancelInflightRequests(params.tabId) - server.tabRemove(params) + + invoke() } } + CHAT_TAB_CHANGE -> { - handleChatNotification(node) { server, params -> - server.tabChange(params) - } + handleChat(AmazonQChatServer.tabChange, node) } + CHAT_OPEN_TAB -> { val response = serializer.deserializeChatMessages(node) - chatCommunicationManager.completeTabOpen( - response.requestId, - response.params.result.tabId - ) + val future = chatCommunicationManager.removeTabOpenRequest(response.requestId) ?: return + try { + val id = serializer.deserializeChatMessages(node.params).result.tabId + future.complete(OpenTabResult(id)) + } catch (e: Exception) { + try { + val err = serializer.deserializeChatMessages(node.params) + future.complete(err.error) + } catch (_: Exception) { + future.completeExceptionally(e) + } + } } + CHAT_INSERT_TO_CURSOR -> { - handleChatNotification(node) { server, params -> - server.insertToCursorPosition(params) - } + handleChat(AmazonQChatServer.insertToCursorPosition, node) } + CHAT_LINK_CLICK -> { - handleChatNotification(node) { server, params -> - server.linkClick(params) - } + handleChat(AmazonQChatServer.linkClick, node) } + CHAT_INFO_LINK_CLICK -> { - handleChatNotification(node) { server, params -> - server.infoLinkClick(params) - } + handleChat(AmazonQChatServer.infoLinkClick, node) } + CHAT_SOURCE_LINK_CLICK -> { - handleChatNotification(node) { server, params -> - server.sourceLinkClick(params) - } + handleChat(AmazonQChatServer.sourceLinkClick, node) } + CHAT_FILE_CLICK -> { - handleChatNotification(node) { server, params -> - server.fileClick(params) - } + handleChat(AmazonQChatServer.fileClick, node) } + PROMPT_INPUT_OPTIONS_CHANGE -> { - handleChatNotification(node) { - server, params -> - server.promptInputOptionsChange(params) - } - } - CHAT_PROMPT_OPTION_ACKNOWLEDGED -> { - val acknowledgedMessage = node.get("params").get("messageId") - if (acknowledgedMessage.asText() == "programmerModeCardId") { - MeetQSettings.getInstance().amazonQChatPairProgramming = false - } + handleChat(AmazonQChatServer.promptInputOptionsChange, node) } + CHAT_FOLLOW_UP_CLICK -> { - handleChatNotification(node) { server, params -> - server.followUpClick(params) - } + handleChat(AmazonQChatServer.followUpClick, node) } + CHAT_BUTTON_CLICK -> { - handleChatNotification(node) { server, params -> - server.buttonClick(params) - }.thenApply { response -> + handleChat(AmazonQChatServer.buttonClick, node).thenApply { response -> if (response is ButtonClickResult && !response.success) { LOG.warn { "Failed to execute action associated with button with reason: ${response.failureReason}" } } } } - AUTH_FOLLOW_UP_CLICKED -> { - val message = serializer.deserializeChatMessages(node) - chatCommunicationManager.handleAuthFollowUpClicked( - project, - message.params - ) - } + CHAT_COPY_CODE_TO_CLIPBOARD -> { - handleChatNotification(node) { server, params -> - server.copyCodeToClipboard(params) - } + handleChat(AmazonQChatServer.copyCodeToClipboard, node) } GET_SERIALIZED_CHAT_REQUEST_METHOD -> { @@ -428,36 +380,34 @@ class BrowserConnector( } CHAT_TAB_BAR_ACTIONS -> { - handleChatNotification(node) { - server, params -> - val result = server.tabBarActions(params) - result.whenComplete { actions, error -> - try { - if (error != null) { - throw error - } + handleChat(AmazonQChatServer.tabBarActions, node) { params, invoke -> + invoke() + .whenComplete { actions, error -> + try { + if (error != null) { + throw error + } - browser.postChat( - FlareUiMessage( - command = CHAT_TAB_BAR_ACTIONS, - params = actions ?: emptyMap() + browser.postChat( + FlareUiMessage( + command = CHAT_TAB_BAR_ACTIONS, + params = actions + ) ) - ) - } catch (e: Exception) { - LOG.error { "Failed to perform chat tab bar action $e" } - params.tabId?.let { - browser.postChat(chatCommunicationManager.getErrorUiMessage(it, e, null)) + } catch (e: Exception) { + LOG.error { "Failed to perform chat tab bar action $e" } + params.tabId?.let { + browser.postChat(chatCommunicationManager.getErrorUiMessage(it, e, null)) + } } } - } } } + CHAT_CREATE_PROMPT -> { - handleChatNotification(node) { - server, params -> - server.createPrompt(params) - } + handleChat(AmazonQChatServer.createPrompt, node) } + STOP_CHAT_RESPONSE -> { val stopResponseRequest = serializer.deserializeChatMessages(node) if (!chatCommunicationManager.hasInflightRequest(stopResponseRequest.params.tabId)) { @@ -466,6 +416,22 @@ class BrowserConnector( cancelInflightRequests(stopResponseRequest.params.tabId) chatCommunicationManager.removePartialChatMessage(stopResponseRequest.params.tabId) } + + AUTH_FOLLOW_UP_CLICKED -> { + val message = serializer.deserializeChatMessages(node) + chatCommunicationManager.handleAuthFollowUpClicked( + project, + message.params + ) + } + + CHAT_PROMPT_OPTION_ACKNOWLEDGED -> { + val acknowledgedMessage = node.params?.get("messageId") + if (acknowledgedMessage?.asText() == "programmerModeCardId") { + MeetQSettings.getInstance().amazonQChatPairProgramming = false + } + } + OPEN_SETTINGS -> { val openSettingsNotification = serializer.deserializeChatMessages(node) if (openSettingsNotification.params.settingKey != OPEN_WORKSPACE_SETTINGS_KEY) return @@ -492,15 +458,15 @@ class BrowserConnector( val messageToChat = ChatCommunicationManager.convertToJsonToSendToChat( SEND_CHAT_COMMAND_PROMPT, tabId, - encryptionManager?.decrypt(value).orEmpty(), + value?.let { encryptionManager?.decrypt(it) }.orEmpty(), isPartialResult = false ) browser.postChat(messageToChat) chatCommunicationManager.removeInflightRequestForTab(tabId) - } catch (e: CancellationException) { + } catch (_: CancellationException) { LOG.warn { "Cancelled chat generation" } } catch (e: Exception) { - LOG.error { "Failed to send chat message $e" } + LOG.error(e) { "Failed to send chat message" } browser.postChat(chatCommunicationManager.getErrorUiMessage(tabId, e, partialResultToken)) } } @@ -513,16 +479,50 @@ class BrowserConnector( } } - private inline fun handleChatNotification( + private inline fun handleChat( + lspMethod: JsonRpcMethod, node: JsonNode, - crossinline serverAction: (server: AmazonQLanguageServer, params: R) -> CompletableFuture<*>, - ): CompletableFuture<*> where T : ChatNotification { - val requestFromUi = serializer.deserializeChatMessages(node) - return AmazonQLspService.executeIfRunning(project) { server -> - serverAction(server, requestFromUi.params) - } ?: CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")) + crossinline serverAction: (params: Request, invokeService: () -> CompletableFuture) -> CompletableFuture, + ): CompletableFuture { + val requestFromUi = if (node.params == null) { + Unit as Request + } else { + serializer.deserializeChatMessages(node.params, lspMethod.params) + } + + return AmazonQLspService.executeIfRunning(project) { _ -> + val invokeService = when (lspMethod) { + is JsonRpcNotification -> { + // notify is Unit + { CompletableFuture.completedFuture(rawEndpoint.notify(lspMethod.name, node.params?.let { serializer.objectMapper.treeToValue(it) })) } + } + + is JsonRpcRequest -> { + { + rawEndpoint.request(lspMethod.name, node.params?.let { serializer.objectMapper.treeToValue(it) }).thenApply { + serializer.objectMapper.readValue( + Gson().toJson(it), + lspMethod.response + ) + } + } + } + } as () -> CompletableFuture + serverAction(requestFromUi, invokeService) + } ?: CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")) } + private inline fun handleChat( + lspMethod: JsonRpcMethod, + node: JsonNode, + ): CompletableFuture = handleChat( + lspMethod, + node, + ) { _, invokeService -> invokeService() } + + private val JsonNode.params + get() = get("params") + companion object { private val LOG = getLogger() } diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQChatServer.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQChatServer.kt new file mode 100644 index 00000000000..e65ae09d0a5 --- /dev/null +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQChatServer.kt @@ -0,0 +1,204 @@ +// 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 + +import org.eclipse.lsp4j.jsonrpc.json.JsonRpcMethodProvider +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ButtonClickParams +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ButtonClickResult +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_BUTTON_CLICK +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_CONVERSATION_CLICK +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_COPY_CODE_TO_CLIPBOARD_NOTIFICATION +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_CREATE_PROMPT +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_FEEDBACK +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_FILE_CLICK +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_FOLLOW_UP_CLICK +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_INFO_LINK_CLICK +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_INSERT_TO_CURSOR_NOTIFICATION +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_LINK_CLICK +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_LIST_CONVERSATIONS +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_QUICK_ACTION +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_READY +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_SOURCE_LINK_CLICK +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_TAB_ADD +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_TAB_BAR_ACTIONS +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_TAB_CHANGE +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_TAB_REMOVE +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ConversationClickParams +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CopyCodeToClipboardParams +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CreatePromptParams +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.EncryptedChatParams +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.EncryptedQuickActionChatParams +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.FeedbackParams +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.FileClickParams +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.FollowUpClickParams +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.GET_SERIALIZED_CHAT_REQUEST_METHOD +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.GetSerializedChatParams +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.GetSerializedChatResult +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.InfoLinkClickParams +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.InsertToCursorPositionParams +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.LinkClickParams +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ListConversationsParams +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.PROMPT_INPUT_OPTIONS_CHANGE +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.PromptInputOptionChangeParams +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.SourceLinkClickParams +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.TabBarActionParams +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.TabEventParams +import kotlin.reflect.KProperty +import kotlin.reflect.full.declaredMembers + +sealed interface JsonRpcMethod { + val name: String + val params: Class +} + +data class JsonRpcNotification( + override val name: String, + override val params: Class, +) : JsonRpcMethod + +@Suppress("FunctionNaming") +fun JsonRpcNotification(name: String) = JsonRpcNotification(name, Unit::class.java) + +data class JsonRpcRequest( + override val name: String, + override val params: Class, + val response: Class, +) : JsonRpcMethod + +/** + * Messaging for the Chat feature follows this pattern: + * Mynah-UI <-> Plugin <-> Flare LSP + * + * However, the default scenario is that the plugin only cares about a subset of request/response payload and should otherwise transparently passthrough data. + * To obtain some semblance of type safety, we model the subset of values that are relevant and passthrough the rest. + * + * Generally, methods MUST be modeled here if the response type is needed, or LSP4J will return null + */ +object AmazonQChatServer : JsonRpcMethodProvider { + override fun supportedMethods() = buildMap { + AmazonQChatServer::class.declaredMembers.filter { it is KProperty }.forEach { + val method = it.call(AmazonQChatServer) as JsonRpcMethod<*, *> + + // trick lsp4j into returning the complete message even if we didn't model it completely + val lsp4jMethod = when (method) { + is JsonRpcNotification<*> -> org.eclipse.lsp4j.jsonrpc.json.JsonRpcMethod.notification(method.name, Any::class.java) + is JsonRpcRequest<*, *> -> org.eclipse.lsp4j.jsonrpc.json.JsonRpcMethod.request(method.name, Any::class.java, Any::class.java) + } + + put(method.name, lsp4jMethod) + } + } + + val sendChatPrompt = JsonRpcRequest( + SEND_CHAT_COMMAND_PROMPT, + EncryptedChatParams::class.java, + String::class.java + ) + + val sendQuickAction = JsonRpcRequest( + CHAT_QUICK_ACTION, + EncryptedQuickActionChatParams::class.java, + String::class.java + ) + + val copyCodeToClipboard = JsonRpcNotification( + CHAT_COPY_CODE_TO_CLIPBOARD_NOTIFICATION, + CopyCodeToClipboardParams::class.java, + ) + + val chatReady = JsonRpcNotification( + CHAT_READY, + ) + + val tabAdd = JsonRpcNotification( + CHAT_TAB_ADD, + TabEventParams::class.java + ) + + val tabRemove = JsonRpcNotification( + CHAT_TAB_REMOVE, + TabEventParams::class.java + ) + + val tabChange = JsonRpcNotification( + CHAT_TAB_CHANGE, + TabEventParams::class.java + ) + + val feedback = JsonRpcNotification( + CHAT_FEEDBACK, + FeedbackParams::class.java + ) + + val insertToCursorPosition = JsonRpcNotification( + CHAT_INSERT_TO_CURSOR_NOTIFICATION, + InsertToCursorPositionParams::class.java + ) + + val linkClick = JsonRpcNotification( + CHAT_LINK_CLICK, + LinkClickParams::class.java + ) + + val infoLinkClick = JsonRpcNotification( + CHAT_INFO_LINK_CLICK, + InfoLinkClickParams::class.java + ) + + val sourceLinkClick = JsonRpcNotification( + CHAT_SOURCE_LINK_CLICK, + SourceLinkClickParams::class.java + ) + + val promptInputOptionsChange = JsonRpcNotification( + PROMPT_INPUT_OPTIONS_CHANGE, + PromptInputOptionChangeParams::class.java + ) + + val followUpClick = JsonRpcNotification( + CHAT_FOLLOW_UP_CLICK, + FollowUpClickParams::class.java + ) + + val fileClick = JsonRpcNotification( + CHAT_FILE_CLICK, + FileClickParams::class.java + ) + + val listConversations = JsonRpcRequest( + CHAT_LIST_CONVERSATIONS, + ListConversationsParams::class.java, + Any::class.java + ) + + val conversationClick = JsonRpcRequest( + CHAT_CONVERSATION_CLICK, + ConversationClickParams::class.java, + Any::class.java + ) + + val buttonClick = JsonRpcRequest( + CHAT_BUTTON_CLICK, + ButtonClickParams::class.java, + ButtonClickResult::class.java + ) + + val tabBarActions = JsonRpcRequest( + CHAT_TAB_BAR_ACTIONS, + TabBarActionParams::class.java, + Any::class.java + ) + + val getSerializedActions = JsonRpcRequest( + GET_SERIALIZED_CHAT_REQUEST_METHOD, + GetSerializedChatParams::class.java, + GetSerializedChatResult::class.java + ) + + val createPrompt = JsonRpcNotification( + CHAT_CREATE_PROMPT, + CreatePromptParams::class.java + ) +} diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClient.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClient.kt index 33b51e73a21..7fa611032d4 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClient.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClient.kt @@ -7,13 +7,13 @@ 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_OPEN_TAB 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.GET_SERIALIZED_CHAT_REQUEST_METHOD import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.GetSerializedChatResult import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.OPEN_FILE_DIFF import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.OpenFileDiffParams -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.OpenTabResult import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SHOW_SAVE_FILE_DIALOG_REQUEST_METHOD import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ShowSaveFileDialogParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ShowSaveFileDialogResult @@ -28,8 +28,8 @@ interface AmazonQLanguageClient : LanguageClient { @JsonRequest("aws/credentials/getConnectionMetadata") fun getConnectionMetadata(): CompletableFuture - @JsonRequest("aws/chat/openTab") - fun openTab(params: LSPAny): CompletableFuture + @JsonRequest(CHAT_OPEN_TAB) + fun openTab(params: LSPAny): CompletableFuture @JsonRequest(SHOW_SAVE_FILE_DIALOG_REQUEST_METHOD) fun showSaveFileDialog(params: ShowSaveFileDialogParams): CompletableFuture 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 1042a63bc47..3aa72464a82 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 @@ -42,7 +42,6 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_ import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.GET_SERIALIZED_CHAT_REQUEST_METHOD import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.GetSerializedChatResult import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.OpenFileDiffParams -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.OpenTabResult import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ShowSaveFileDialogParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ShowSaveFileDialogResult import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.ConnectionMetadata @@ -140,9 +139,9 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC } } - override fun openTab(params: LSPAny): CompletableFuture { + override fun openTab(params: LSPAny): CompletableFuture { val requestId = UUID.randomUUID().toString() - val result = CompletableFuture() + val result = CompletableFuture() val chatManager = ChatCommunicationManager.getInstance(project) chatManager.addTabOpenRequest(requestId, result) diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageServer.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageServer.kt index 5f318b1c7dd..6ef5cf818ec 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageServer.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageServer.kt @@ -11,50 +11,6 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.GetConfigu import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.LogInlineCompletionSessionResultsParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.LspServerConfigurations import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.UpdateConfigurationParams -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ButtonClickParams -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ButtonClickResult -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_BUTTON_CLICK -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_CONVERSATION_CLICK -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_COPY_CODE_TO_CLIPBOARD_NOTIFICATION -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_CREATE_PROMPT -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_FEEDBACK -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_FILE_CLICK -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_FOLLOW_UP_CLICK -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_INFO_LINK_CLICK -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_INSERT_TO_CURSOR_NOTIFICATION -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_LINK_CLICK -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_LIST_CONVERSATIONS -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_QUICK_ACTION -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_READY -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_SOURCE_LINK_CLICK -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_TAB_ADD -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_TAB_BAR_ACTIONS -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_TAB_CHANGE -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_TAB_REMOVE -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ConversationClickParams -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ConversationClickResult -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CopyCodeToClipboardParams -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CreatePromptParams -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.EncryptedChatParams -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.EncryptedQuickActionChatParams -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.FeedbackParams -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.FileClickParams -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.FollowUpClickParams -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.GET_SERIALIZED_CHAT_REQUEST_METHOD -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.GetSerializedChatParams -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.GetSerializedChatResult -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.InfoLinkClickParams -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.InsertToCursorPositionParams -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.LinkClickParams -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ListConversationsParams -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ListConversationsResult -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.PROMPT_INPUT_OPTIONS_CHANGE -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.PromptInputOptionChangeParams -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.SourceLinkClickParams -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.TabBarActionParams -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.TabBarActionResult -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.TabEventParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.UpdateCredentialsPayload import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.dependencies.DidChangeDependencyPathsParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument.InlineCompletionListWithReferences @@ -86,67 +42,4 @@ interface AmazonQLanguageServer : LanguageServer { @JsonRequest("aws/updateConfiguration") fun updateConfiguration(params: UpdateConfigurationParams): CompletableFuture - - @JsonRequest(SEND_CHAT_COMMAND_PROMPT) - fun sendChatPrompt(params: EncryptedChatParams): CompletableFuture - - @JsonNotification(CHAT_COPY_CODE_TO_CLIPBOARD_NOTIFICATION) - fun copyCodeToClipboard(params: CopyCodeToClipboardParams): CompletableFuture - - @JsonNotification(CHAT_TAB_ADD) - fun tabAdd(params: TabEventParams): CompletableFuture - - @JsonNotification(CHAT_TAB_REMOVE) - fun tabRemove(params: TabEventParams): CompletableFuture - - @JsonNotification(CHAT_TAB_CHANGE) - fun tabChange(params: TabEventParams): CompletableFuture - - @JsonRequest(CHAT_QUICK_ACTION) - fun sendQuickAction(params: EncryptedQuickActionChatParams): CompletableFuture - - @JsonNotification(CHAT_INSERT_TO_CURSOR_NOTIFICATION) - fun insertToCursorPosition(params: InsertToCursorPositionParams): CompletableFuture - - @JsonNotification(CHAT_FEEDBACK) - fun feedback(params: FeedbackParams): CompletableFuture - - @JsonNotification(CHAT_READY) - fun chatReady(): CompletableFuture - - @JsonNotification(CHAT_LINK_CLICK) - fun linkClick(params: LinkClickParams): CompletableFuture - - @JsonNotification(CHAT_INFO_LINK_CLICK) - fun infoLinkClick(params: InfoLinkClickParams): CompletableFuture - - @JsonNotification(CHAT_SOURCE_LINK_CLICK) - fun sourceLinkClick(params: SourceLinkClickParams): CompletableFuture - - @JsonNotification(PROMPT_INPUT_OPTIONS_CHANGE) - fun promptInputOptionsChange(params: PromptInputOptionChangeParams): CompletableFuture - - @JsonNotification(CHAT_FOLLOW_UP_CLICK) - fun followUpClick(params: FollowUpClickParams): CompletableFuture - - @JsonNotification(CHAT_FILE_CLICK) - fun fileClick(params: FileClickParams): CompletableFuture - - @JsonRequest(CHAT_LIST_CONVERSATIONS) - fun listConversations(params: ListConversationsParams): CompletableFuture - - @JsonRequest(CHAT_CONVERSATION_CLICK) - fun conversationClick(params: ConversationClickParams): CompletableFuture - - @JsonRequest(CHAT_BUTTON_CLICK) - fun buttonClick(params: ButtonClickParams): CompletableFuture - - @JsonRequest(CHAT_TAB_BAR_ACTIONS) - fun tabBarActions(params: TabBarActionParams): CompletableFuture - - @JsonRequest(GET_SERIALIZED_CHAT_REQUEST_METHOD) - fun getSerializedActions(params: GetSerializedChatParams): CompletableFuture - - @JsonNotification(CHAT_CREATE_PROMPT) - fun createPrompt(params: CreatePromptParams): CompletableFuture } diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt index 933027a3b36..09c17796c74 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt @@ -47,6 +47,8 @@ import org.eclipse.lsp4j.TextDocumentClientCapabilities import org.eclipse.lsp4j.WorkspaceClientCapabilities import org.eclipse.lsp4j.jsonrpc.Launcher import org.eclipse.lsp4j.jsonrpc.MessageConsumer +import org.eclipse.lsp4j.jsonrpc.RemoteEndpoint +import org.eclipse.lsp4j.jsonrpc.json.JsonRpcMethod import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage import org.eclipse.lsp4j.launch.LSPLauncher import org.slf4j.event.Level @@ -123,9 +125,13 @@ class AmazonQLspService(private val project: Project, private val cs: CoroutineS private var instance: Deferred val capabilities get() = instance.getCompleted().initializeResult.getCompleted().capabilities + val encryptionManager get() = instance.getCompleted().encryptionManager + val rawEndpoint + get() = instance.getCompleted().rawEndpoint + // dont allow lsp commands if server is restarting private val mutex = Mutex(false) @@ -135,7 +141,7 @@ class AmazonQLspService(private val project: Project, private val cs: CoroutineS var attempts = 0 while (attempts < 3) { try { - return@async withTimeout(30.seconds) { + val result = withTimeout(30.seconds) { val instance = AmazonQServerInstance(project, cs).also { Disposer.register(this@AmazonQLspService, it) } @@ -146,6 +152,9 @@ class AmazonQLspService(private val project: Project, private val cs: CoroutineS _flowInstance.emit(it) } } + + // withTimeout can throw + return@async result } catch (e: Exception) { LOG.warn(e) { "Failed to start LSP server" } } @@ -196,7 +205,7 @@ class AmazonQLspService(private val project: Project, private val cs: CoroutineS return runnable(lsp) } - fun executeSync(runnable: suspend AmazonQLspService.(AmazonQLanguageServer) -> T): T = + private fun executeSync(runnable: suspend AmazonQLspService.(AmazonQLanguageServer) -> T): T = runBlocking(cs.coroutineContext) { execute(runnable) } @@ -205,9 +214,13 @@ class AmazonQLspService(private val project: Project, private val cs: CoroutineS private val LOG = getLogger() fun getInstance(project: Project) = project.service() + @Deprecated("Easy to accidentally freeze EDT") fun executeIfRunning(project: Project, runnable: AmazonQLspService.(AmazonQLanguageServer) -> T): T? = project.serviceIfCreated()?.executeSync(runnable) + suspend fun asyncExecuteIfRunning(project: Project, runnable: suspend AmazonQLspService.(AmazonQLanguageServer) -> T): T? = + project.serviceIfCreated()?.execute(runnable) + fun didChangeConfiguration(project: Project) { executeIfRunning(project) { it.workspaceService.didChangeConfiguration(DidChangeConfigurationParams()) @@ -224,6 +237,9 @@ private class AmazonQServerInstance(private val project: Project, private val cs val languageServer: AmazonQLanguageServer get() = launcher.remoteProxy + val rawEndpoint: RemoteEndpoint + get() = launcher.remoteEndpoint + @Suppress("ForbiddenVoid") private val launcherFuture: Future private val launcherHandler: KillableProcessHandler @@ -322,7 +338,10 @@ private class AmazonQServerInstance(private val project: Project, private val cs launcherHandler.addProcessListener(inputWrapper) launcherHandler.startNotify() - launcher = LSPLauncher.Builder() + launcher = object : LSPLauncher.Builder() { + override fun getSupportedMethods(): Map = + super.getSupportedMethods() + AmazonQChatServer.supportedMethods() + } .wrapMessages { consumer -> MessageConsumer { message -> if (message is ResponseMessage && message.result is AwsExtendedInitializeResult) { @@ -414,5 +433,3 @@ private class AmazonQServerInstance(private val project: Project, private val cs private val LOG = getLogger() } } - -typealias AmazonQInitializeMessageReceivedListener = () -> Unit diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt index 5ee8c258d59..22191c518fa 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt @@ -46,7 +46,7 @@ class DefaultAuthCredentialsService( private val scheduler: ScheduledExecutorService = AppExecutorUtil.getAppScheduledExecutorService() private var tokenSyncTask: ScheduledFuture<*>? = null - private val tokenSyncIntervalSeconds = 10L + private val tokenSyncIntervalSeconds = 900L init { project.messageBus.connect(serverInstance).apply { 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 67dc9fc1b3b..1f00ebc79f8 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 @@ -18,12 +18,12 @@ import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection import software.aws.toolkits.jetbrains.core.credentials.reauthConnectionIfNeeded import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.ProgressNotificationUtils.getObject +import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.LSPAny import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.AuthFollowUpClickedParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.AuthFollowupType import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_ERROR_PARAMS import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ErrorParams import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.GetSerializedChatResult -import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.OpenTabResult import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SEND_CHAT_COMMAND_PROMPT import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener @@ -37,7 +37,7 @@ class ChatCommunicationManager(private val cs: CoroutineScope) { private val chatPartialResultMap = ConcurrentHashMap() private val inflightRequestByTabId = ConcurrentHashMap>() private val pendingSerializedChatRequests = ConcurrentHashMap>() - private val pendingTabRequests = ConcurrentHashMap>() + private val pendingTabRequests = ConcurrentHashMap>() fun setUiReady() { uiReady.complete(true) @@ -85,17 +85,12 @@ class ChatCommunicationManager(private val cs: CoroutineScope) { pendingSerializedChatRequests.remove(requestId) } - fun addTabOpenRequest(requestId: String, result: CompletableFuture) { + fun addTabOpenRequest(requestId: String, result: CompletableFuture) { pendingTabRequests[requestId] = result } - fun completeTabOpen(requestId: String, tabId: String) { - pendingTabRequests.remove(requestId)?.complete(OpenTabResult(tabId)) - } - - fun removeTabOpenRequest(requestId: String) { + fun removeTabOpenRequest(requestId: String) = pendingTabRequests.remove(requestId) - } fun handlePartialResultProgressNotification(project: Project, params: ProgressParams) { val token = ProgressNotificationUtils.getToken(params) diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/ChatMessage.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/ChatMessage.kt index 2729bb3451a..457dedc9ae6 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/ChatMessage.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/ChatMessage.kt @@ -71,18 +71,9 @@ data class Changes( val total: Int? = null, ) -@JsonAdapter(EnumJsonValueAdapter::class) -enum class IconType(@JsonValue val repr: String) { - FILE("file"), - FOLDER("folder"), - CODE_BLOCK("code-block"), - LIST_ADD("list-add"), - MAGIC("magic"), - HELP("help"), - TRASH("trash"), - SEARCH("search"), - CALENDAR("calendar"), -} +// i don't want to model 70+ icon types +// https://github.com/aws/mynah-ui/blob/38608dff905b3790d85c73e2911ec7071c8a8cdf/src/components/icon.ts#L12 +typealias IconType = String @JsonAdapter(EnumJsonValueAdapter::class) enum class Status(@JsonValue val repr: String) { diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/TabEvent.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/TabEvent.kt index d5f3d8d1d31..40daa402bf4 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/TabEvent.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/TabEvent.kt @@ -14,15 +14,21 @@ data class TabEventParams( data class OpenTabResponse( val requestId: String, - val command: String, - val params: OpenTabResponseParams, ) -data class OpenTabResponseParams( - val success: Boolean, +data class OpenTabResultSuccess( val result: OpenTabResult, ) data class OpenTabResult( val tabId: String, ) + +data class OpenTabResultError( + val error: OpenTabResultErrorError, +) + +data class OpenTabResultErrorError( + val type: String, + val message: String, +) From 2244d416db05112fe397fdaed290a9140f50802b Mon Sep 17 00:00:00 2001 From: Richard Li Date: Tue, 13 May 2025 13:08:09 -0700 Subject: [PATCH 2/3] tst --- .../jetbrains/services/amazonq/lsp/AmazonQLspService.kt | 2 +- .../amazonq/lsp/model/aws/chat/ChatMessageTest.kt | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt index 09c17796c74..bc23a28ae03 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt @@ -205,7 +205,7 @@ class AmazonQLspService(private val project: Project, private val cs: CoroutineS return runnable(lsp) } - private fun executeSync(runnable: suspend AmazonQLspService.(AmazonQLanguageServer) -> T): T = + fun executeSync(runnable: suspend AmazonQLspService.(AmazonQLanguageServer) -> T): T = runBlocking(cs.coroutineContext) { execute(runnable) } diff --git a/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/ChatMessageTest.kt b/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/ChatMessageTest.kt index 4c713d98b84..8c8c4ddd1d1 100644 --- a/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/ChatMessageTest.kt +++ b/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/ChatMessageTest.kt @@ -20,14 +20,14 @@ class ChatMessageTest { @Test fun `sanity check`() { val jackson = jacksonObjectMapper() - assertThat(IconType.CODE_BLOCK).satisfiesKt { + assertThat(MessageType.SYSTEM_PROMPT).satisfiesKt { // language=JSON - val expected = """"code-block"""" + val expected = """"system-prompt"""" assertThat(Gson().toJson(it)).isEqualTo(expected) assertThat(jackson.writeValueAsString(it)).isEqualTo(expected) - assertThat(Gson().fromJson(expected, IconType::class.java)).isEqualTo(it) - assertThat(jackson.readValue(expected)).isEqualTo(it) + assertThat(Gson().fromJson(expected, MessageType::class.java)).isEqualTo(it) + assertThat(jackson.readValue(expected)).isEqualTo(it) } } From 043602ddf0b5482767ec33122a0322e7abc37a0b Mon Sep 17 00:00:00 2001 From: Richard Li Date: Tue, 13 May 2025 14:02:00 -0700 Subject: [PATCH 3/3] feedback --- .../jetbrains/services/amazonq/lsp/model/aws/chat/TabEvent.kt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/TabEvent.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/TabEvent.kt index 40daa402bf4..4ab56164be8 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/TabEvent.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/TabEvent.kt @@ -25,10 +25,10 @@ data class OpenTabResult( ) data class OpenTabResultError( - val error: OpenTabResultErrorError, + val error: OpenTabResultErrorData, ) -data class OpenTabResultErrorError( +data class OpenTabResultErrorData( val type: String, val message: String, )