feedback

Author: samgst-amazon

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

@@ -24,17 +24,21 @@ import org.eclipse.lsp4j.Position
 import org.eclipse.lsp4j.Range
 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.AmazonQLspService
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.ChatCommunicationManager
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.getTextDocumentIdentifier
 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_QUICK_ACTION
+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.CursorState
 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.FollowUpClickNotification
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.FollowUpClickParams
 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.SendChatPromptRequest
@@ -207,10 +211,9 @@ class BrowserConnector(
                 showResult(result, partialResultToken, tabId, encryptionManager, browser)
             }
             CHAT_FOLLOW_UP_CLICK -> {
-                val requestFromUi = serializer.deserializeChatMessages(node, FollowUpClickNotification::class.java)
-                AmazonQLspService.executeIfRunning(project) { server ->
-                    server.followUpClick(requestFromUi.params)
-                } ?: CompletableFuture.failedFuture<Unit>(IllegalStateException("LSP Server not running"))
+                handleChatNotification<FollowUpClickNotification, FollowUpClickParams>(node) { server, params ->
+                    server.followUpClick(params)
+                }
             }
         }
     }
@@ -231,7 +234,17 @@ class BrowserConnector(
                 isPartialResult = false
             )
             browser.postChat(messageToChat)
-                showResult(result, partialResultToken, tabId, encryptionManager, browser)
-            }
+            showResult(result, partialResultToken, tabId, encryptionManager, browser)
         }
+    }
+
+    private inline fun <reified T, R> handleChatNotification(
+        node: JsonNode,
+        crossinline serverAction: (server: AmazonQLanguageServer, params: R) -> CompletableFuture<*>,
+    ): CompletableFuture<*> where T : ChatNotification<R> {
+        val requestFromUi = serializer.deserializeChatMessages<T>(node)
+        return AmazonQLspService.executeIfRunning(project) { server ->
+            serverAction(server, requestFromUi.params)
+        } ?: CompletableFuture.failedFuture<Unit>(IllegalStateException("LSP Server not running"))
+    }
 }

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageServer.kt

@@ -10,11 +10,12 @@ import org.eclipse.lsp4j.services.LanguageServer
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.GetConfigurationFromServerParams
 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.CHAT_FOLLOW_UP_CLICK
 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.EncryptedChatParams
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.EncryptedQuickActionChatParams
-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.FollowUpClickParams
+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.credentials.UpdateCredentialsPayload
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.dependencies.DidChangeDependencyPathsParams
 import java.util.concurrent.CompletableFuture
@@ -45,6 +46,6 @@ interface AmazonQLanguageServer : LanguageServer {
     @JsonRequest(CHAT_QUICK_ACTION)
     fun sendQuickAction(params: EncryptedQuickActionChatParams): CompletableFuture<String>
 
-    @JsonNotification("aws/chat/followUpClick")
+    @JsonNotification(CHAT_FOLLOW_UP_CLICK)
     fun followUpClick(params: FollowUpClickParams): CompletableFuture<Unit>
 }

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/ChatNotification.kt

@@ -0,0 +1,9 @@
+// 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
+
+interface ChatNotification<T> {
+    val command: String
+    val params: T
+}

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/chat/FollowUpClick.kt

@@ -4,9 +4,9 @@
 package software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat
 
 data class FollowUpClickNotification(
-    val command: String,
-    val params: FollowUpClickParams,
-)
+    override val command: String,
+    override val params: FollowUpClickParams,
+) : ChatNotification<FollowUpClickParams>
 
 data class FollowUpClickParams(
     val tabId: String,