Skip to content

Commit b85d204

Browse files
manodnyabmanodnyab
authored
feat(amazonq): Add quick actions to Flare chat (#5561)
* Set up Flare chat connection * Partial chat results * feedback * detekt * syntax error * detekt * detekt * Add quick actions to chat * adding commands from initialize result * added trace logging * Feedback * detekt
1 parent f948277 commit b85d204

File tree

11 files changed

+240
-24
lines changed

11 files changed

+240
-24
lines changed

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/commands/MessageSerializer.kt

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.util.command
2020

2121
class MessageSerializer @VisibleForTesting constructor() {
2222

23-
private val objectMapper = jacksonObjectMapper()
23+
val objectMapper = jacksonObjectMapper()
2424
.registerModule(JavaTimeModule())
2525
.enable(MapperFeature.ACCEPT_CASE_INSENSITIVE_ENUMS)
2626
.disable(SerializationFeature.WRITE_DATES_AS_TIMESTAMPS)
@@ -37,8 +37,8 @@ class MessageSerializer @VisibleForTesting constructor() {
3737

3838
fun serialize(value: Any): String = objectMapper.writeValueAsString(value)
3939

40-
fun <T> deserializeChatMessages(value: JsonNode, clazz: Class<T>): T =
41-
objectMapper.treeToValue(value, clazz)
40+
inline fun <reified T> deserializeChatMessages(value: JsonNode): T =
41+
objectMapper.treeToValue<T>(value)
4242

4343
// Provide singleton global access
4444
companion object {

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQPanel.kt

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import com.intellij.idea.AppMode
77
import com.intellij.openapi.Disposable
88
import com.intellij.openapi.application.ApplicationManager
99
import com.intellij.openapi.application.runInEdt
10+
import com.intellij.openapi.project.Project
1011
import com.intellij.openapi.util.Disposer
1112
import com.intellij.ui.components.JBLoadingPanel
1213
import com.intellij.ui.components.JBPanelWithEmptyText
@@ -24,7 +25,7 @@ import java.awt.event.ActionListener
2425
import java.util.concurrent.CompletableFuture
2526
import javax.swing.JButton
2627

27-
class AmazonQPanel(private val parent: Disposable) {
28+
class AmazonQPanel(private val parent: Disposable, val project: Project) {
2829
private val webviewContainer = Wrapper()
2930
val browser = CompletableFuture<Browser>()
3031

@@ -91,7 +92,7 @@ class AmazonQPanel(private val parent: Disposable) {
9192
loadingPanel.stopLoading()
9293
runInEdt {
9394
browser.complete(
94-
Browser(parent, webUri).also {
95+
Browser(parent, webUri, project).also {
9596
wrapper.setContent(it.component())
9697
}
9798
)

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/toolwindow/AmazonQToolWindow.kt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ class AmazonQToolWindow private constructor(
4949
private val browserConnector = BrowserConnector(project = project)
5050
private val editorThemeAdapter = EditorThemeAdapter()
5151

52-
private val chatPanel = AmazonQPanel(parent = this)
52+
private val chatPanel = AmazonQPanel(parent = this, project)
5353

5454
val component: JComponent = chatPanel.component
5555

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

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,12 @@ package software.aws.toolkits.jetbrains.services.amazonq.webview
55

66
import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
77
import com.intellij.openapi.Disposable
8+
import com.intellij.openapi.project.Project
89
import com.intellij.openapi.util.Disposer
910
import com.intellij.ui.jcef.JBCefJSQuery
1011
import org.cef.CefApp
12+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
13+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AwsServerCapabilitiesProvider
1114
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile
1215
import software.aws.toolkits.jetbrains.services.amazonq.util.HighlightCommand
1316
import software.aws.toolkits.jetbrains.services.amazonq.util.createBrowser
@@ -17,7 +20,8 @@ import java.net.URI
1720
/*
1821
Displays the web view for the Amazon Q tool window
1922
*/
20-
class Browser(parent: Disposable, private val webUri: URI) : Disposable {
23+
24+
class Browser(parent: Disposable, private val webUri: URI, val project: Project) : Disposable {
2125

2226
val jcefBrowser = createBrowser(parent)
2327

@@ -39,8 +43,17 @@ class Browser(parent: Disposable, private val webUri: URI) : Disposable {
3943
"mynah",
4044
AssetResourceHandler.AssetResourceHandlerFactory(),
4145
)
42-
43-
loadWebView(isCodeTransformAvailable, isFeatureDevAvailable, isDocAvailable, isCodeScanAvailable, isCodeTestAvailable, highlightCommand, activeProfile)
46+
AmazonQLspService.getInstance(project).addLspInitializeMessageListener {
47+
loadWebView(
48+
isCodeTransformAvailable,
49+
isFeatureDevAvailable,
50+
isDocAvailable,
51+
isCodeScanAvailable,
52+
isCodeTestAvailable,
53+
highlightCommand,
54+
activeProfile
55+
)
56+
}
4457
}
4558

4659
override fun dispose() {
@@ -101,6 +114,7 @@ class Browser(parent: Disposable, private val webUri: URI) : Disposable {
101114
highlightCommand: HighlightCommand?,
102115
activeProfile: QRegionProfile?,
103116
): String {
117+
val quickActionConfig = generateQuickActionConfig()
104118
val postMessageToJavaJsCode = receiveMessageQuery.inject("JSON.stringify(message)")
105119
val jsScripts = """
106120
<script type="text/javascript" src="$webUri" defer onload="init()"></script>
@@ -113,7 +127,7 @@ class Browser(parent: Disposable, private val webUri: URI) : Disposable {
113127
}
114128
},
115129
{
116-
quickActionCommands: [],
130+
quickActionCommands: $quickActionConfig,
117131
disclaimerAcknowledged: ${MeetQSettings.getInstance().disclaimerAcknowledged}
118132
}
119133
);
@@ -220,6 +234,10 @@ class Browser(parent: Disposable, private val webUri: URI) : Disposable {
220234
activeProfile
221235
}
222236

237+
private fun generateQuickActionConfig() = AwsServerCapabilitiesProvider.getInstance(project).getChatOptions().quickActions.quickActionsCommandGroups
238+
.let { OBJECT_MAPPER.writeValueAsString(it) }
239+
?: "[]"
240+
223241
companion object {
224242
private const val MAX_ONBOARDING_PAGE_COUNT = 3
225243
private val OBJECT_MAPPER = jacksonObjectMapper()

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

Lines changed: 44 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import kotlinx.coroutines.flow.distinctUntilChanged
1717
import kotlinx.coroutines.flow.launchIn
1818
import kotlinx.coroutines.flow.merge
1919
import kotlinx.coroutines.flow.onEach
20+
import kotlinx.coroutines.future.await
2021
import kotlinx.coroutines.launch
2122
import org.cef.browser.CefBrowser
2223
import org.eclipse.lsp4j.Position
@@ -27,10 +28,13 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
2728
import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager
2829
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.ChatCommunicationManager
2930
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.getTextDocumentIdentifier
31+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_QUICK_ACTION
3032
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatParams
3133
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatPrompt
3234
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CursorState
3335
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.EncryptedChatParams
36+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.EncryptedQuickActionChatParams
37+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.QuickChatActionRequest
3438
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SEND_CHAT_COMMAND_PROMPT
3539
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SendChatPromptRequest
3640
import software.aws.toolkits.jetbrains.services.amazonq.util.command
@@ -147,7 +151,7 @@ class BrowserConnector(
147151
private fun handleFlareChatMessages(browser: Browser, node: JsonNode) {
148152
when (node.command) {
149153
SEND_CHAT_COMMAND_PROMPT -> {
150-
val requestFromUi = serializer.deserializeChatMessages(node, SendChatPromptRequest::class.java)
154+
val requestFromUi = serializer.deserializeChatMessages<SendChatPromptRequest>(node)
151155
val chatPrompt = ChatPrompt(
152156
requestFromUi.params.prompt.prompt,
153157
requestFromUi.params.prompt.escapedPrompt,
@@ -167,32 +171,59 @@ class BrowserConnector(
167171
)
168172
)
169173

170-
val partialResultToken = chatCommunicationManager.addPartialChatMessage(requestFromUi.params.tabId)
171174
val chatParams = ChatParams(
172175
requestFromUi.params.tabId,
173176
chatPrompt,
174177
textDocumentIdentifier,
175178
cursorState
176179
)
177180

181+
val tabId = requestFromUi.params.tabId
182+
val partialResultToken = chatCommunicationManager.addPartialChatMessage(tabId)
183+
178184
var encryptionManager: JwtEncryptionManager? = null
179185
val result = AmazonQLspService.executeIfRunning(project) { server ->
180186
encryptionManager = this.encryptionManager
181187
encryptionManager?.encrypt(chatParams)?.let { EncryptedChatParams(it, partialResultToken) }?.let { server.sendChatPrompt(it) }
182188
} ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")))
189+
showResult(result, partialResultToken, tabId, encryptionManager, browser)
190+
}
191+
CHAT_QUICK_ACTION -> {
192+
val requestFromUi = serializer.deserializeChatMessages<QuickChatActionRequest>(node)
193+
val tabId = requestFromUi.params.tabId
194+
val quickActionParams = requestFromUi.params
195+
val partialResultToken = chatCommunicationManager.addPartialChatMessage(tabId)
196+
var encryptionManager: JwtEncryptionManager? = null
197+
val result = AmazonQLspService.executeIfRunning(project) { server ->
198+
encryptionManager = this.encryptionManager
199+
encryptionManager?.encrypt(quickActionParams)?.let {
200+
EncryptedQuickActionChatParams(it, partialResultToken)
201+
}?.let {
202+
server.sendQuickAction(it)
203+
}
204+
} ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")))
183205

184-
result.whenComplete {
185-
value, error ->
186-
chatCommunicationManager.removePartialChatMessage(partialResultToken)
187-
val messageToChat = ChatCommunicationManager.convertToJsonToSendToChat(
188-
node.command,
189-
requestFromUi.params.tabId,
190-
encryptionManager?.decrypt(value).orEmpty(),
191-
isPartialResult = false
192-
)
193-
browser.postChat(messageToChat)
194-
}
206+
showResult(result, partialResultToken, tabId, encryptionManager, browser)
195207
}
196208
}
197209
}
210+
211+
private fun showResult(
212+
result: CompletableFuture<String>,
213+
partialResultToken: String,
214+
tabId: String,
215+
encryptionManager: JwtEncryptionManager?,
216+
browser: Browser,
217+
) {
218+
result.whenComplete { value, error ->
219+
chatCommunicationManager.removePartialChatMessage(partialResultToken)
220+
val messageToChat = ChatCommunicationManager.convertToJsonToSendToChat(
221+
SEND_CHAT_COMMAND_PROMPT,
222+
tabId,
223+
encryptionManager?.decrypt(value).orEmpty(),
224+
isPartialResult = false
225+
)
226+
browser.postChat(messageToChat)
227+
}
228+
}
198229
}

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

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,10 @@ import org.eclipse.lsp4j.services.LanguageServer
1010
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.GetConfigurationFromServerParams
1111
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.LspServerConfigurations
1212
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.UpdateConfigurationParams
13+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_QUICK_ACTION
1314
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.EncryptedChatParams
15+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.EncryptedQuickActionChatParams
16+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SEND_CHAT_COMMAND_PROMPT
1417
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.UpdateCredentialsPayload
1518
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.dependencies.DidChangeDependencyPathsParams
1619
import java.util.concurrent.CompletableFuture
@@ -35,6 +38,9 @@ interface AmazonQLanguageServer : LanguageServer {
3538
@JsonRequest("aws/updateConfiguration")
3639
fun updateConfiguration(params: UpdateConfigurationParams): CompletableFuture<LspServerConfigurations>
3740

38-
@JsonRequest("aws/chat/sendChatPrompt")
41+
@JsonRequest(SEND_CHAT_COMMAND_PROMPT)
3942
fun sendChatPrompt(params: EncryptedChatParams): CompletableFuture<String>
43+
44+
@JsonRequest(CHAT_QUICK_ACTION)
45+
fun sendQuickAction(params: EncryptedQuickActionChatParams): CompletableFuture<String>
4046
}

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

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import com.intellij.openapi.project.Project
1919
import com.intellij.openapi.util.Disposer
2020
import com.intellij.openapi.util.Key
2121
import com.intellij.openapi.util.SystemInfo
22+
import com.intellij.util.animation.consumer
2223
import com.intellij.util.io.await
2324
import kotlinx.coroutines.CoroutineScope
2425
import kotlinx.coroutines.Deferred
@@ -40,6 +41,10 @@ import org.eclipse.lsp4j.SynchronizationCapabilities
4041
import org.eclipse.lsp4j.TextDocumentClientCapabilities
4142
import org.eclipse.lsp4j.WorkspaceClientCapabilities
4243
import org.eclipse.lsp4j.jsonrpc.Launcher
44+
import org.eclipse.lsp4j.jsonrpc.Launcher.Builder
45+
import org.eclipse.lsp4j.jsonrpc.MessageConsumer
46+
import org.eclipse.lsp4j.jsonrpc.messages.Message
47+
import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage
4348
import org.eclipse.lsp4j.launch.LSPLauncher
4449
import org.slf4j.event.Level
4550
import software.aws.toolkits.core.utils.getLogger
@@ -50,6 +55,9 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.artifacts.ArtifactMa
5055
import software.aws.toolkits.jetbrains.services.amazonq.lsp.auth.DefaultAuthCredentialsService
5156
import software.aws.toolkits.jetbrains.services.amazonq.lsp.dependencies.DefaultModuleDependenciesService
5257
import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager
58+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AmazonQLspTypeAdapterFactory
59+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AwsExtendedInitializeResult
60+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AwsServerCapabilitiesProvider
5361
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.createExtendedClientMetadata
5462
import software.aws.toolkits.jetbrains.services.amazonq.lsp.textdocument.TextDocumentServiceHandler
5563
import software.aws.toolkits.jetbrains.services.amazonq.lsp.util.WorkspaceFolderUtil.createWorkspaceFolders
@@ -63,6 +71,7 @@ import java.io.PipedOutputStream
6371
import java.io.PrintWriter
6472
import java.io.StringWriter
6573
import java.nio.charset.StandardCharsets
74+
import java.util.Collections
6675
import java.util.concurrent.Future
6776
import kotlin.time.Duration.Companion.seconds
6877

@@ -101,6 +110,10 @@ internal class LSPProcessListener : ProcessListener {
101110

102111
@Service(Service.Level.PROJECT)
103112
class AmazonQLspService(private val project: Project, private val cs: CoroutineScope) : Disposable {
113+
private val lspInitializedMessageReceivedListener = Collections.synchronizedList(mutableListOf<AmazonQInitializeMessageReceivedListener>())
114+
fun addLspInitializeMessageListener(listener: AmazonQInitializeMessageReceivedListener) = lspInitializedMessageReceivedListener.add(listener)
115+
fun notifyInitializeMessageReceived() = lspInitializedMessageReceivedListener.forEach { it() }
116+
104117
private var instance: Deferred<AmazonQServerInstance>
105118
val capabilities
106119
get() = instance.getCompleted().initializeResult.getCompleted().capabilities
@@ -266,6 +279,17 @@ private class AmazonQServerInstance(private val project: Project, private val cs
266279
launcherHandler.startNotify()
267280

268281
launcher = LSPLauncher.Builder<AmazonQLanguageServer>()
282+
.wrapMessages { consumer ->
283+
MessageConsumer {
284+
message ->
285+
if (message is ResponseMessage && message.result is AwsExtendedInitializeResult) {
286+
val result = message.result as AwsExtendedInitializeResult
287+
AwsServerCapabilitiesProvider.getInstance(project).setAwsServerCapabilities(result.getAwsServerCapabilities())
288+
AmazonQLspService.getInstance(project).notifyInitializeMessageReceived()
289+
}
290+
consumer?.consume(message)
291+
}
292+
}
269293
.setLocalService(AmazonQLanguageClientImpl(project))
270294
.setRemoteInterface(AmazonQLanguageServer::class.java)
271295
.configureGson {
@@ -274,6 +298,7 @@ private class AmazonQServerInstance(private val project: Project, private val cs
274298

275299
// otherwise Gson treats all numbers as double which causes deser issues
276300
it.setObjectToNumberStrategy(ToNumberPolicy.LONG_OR_DOUBLE)
301+
it.registerTypeAdapterFactory(AmazonQLspTypeAdapterFactory())
277302
}.traceMessages(
278303
PrintWriter(
279304
object : StringWriter() {
@@ -348,7 +373,21 @@ private class AmazonQServerInstance(private val project: Project, private val cs
348373
}
349374
}
350375

376+
class MessageTracer {
377+
private val traceLogger = LOG.atLevel(if (isDeveloperMode()) Level.INFO else Level.DEBUG)
378+
379+
fun trace(direction: String, message: Message) {
380+
traceLogger.log {
381+
buildString {
382+
append("$direction: ")
383+
append(message.toString())
384+
}
385+
}
386+
}
387+
}
351388
companion object {
352389
private val LOG = getLogger<AmazonQServerInstance>()
353390
}
354391
}
392+
393+
typealias AmazonQInitializeMessageReceivedListener = () -> Unit

‎plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/flareChat/AmazonQLspTypeAdapterFactory.kt

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
@file:Suppress("BannedImports")
4+
package software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat
5+
6+
import com.google.gson.Gson
7+
import com.google.gson.TypeAdapter
8+
import com.google.gson.TypeAdapterFactory
9+
import com.google.gson.reflect.TypeToken
10+
import com.google.gson.stream.JsonReader
11+
import com.google.gson.stream.JsonWriter
12+
import org.eclipse.lsp4j.InitializeResult
13+
import java.io.IOException
14+
15+
class AmazonQLspTypeAdapterFactory : TypeAdapterFactory {
16+
override fun <T : Any?> create(gson: Gson, type: TypeToken<T>): TypeAdapter<T>? {
17+
if (type.rawType === InitializeResult::class.java) {
18+
val delegate: TypeAdapter<InitializeResult?> = gson.getDelegateAdapter(this, type) as TypeAdapter<InitializeResult?>
19+
20+
return object : TypeAdapter<InitializeResult>() {
21+
@Throws(IOException::class)
22+
override fun write(out: JsonWriter, value: InitializeResult?) {
23+
delegate.write(out, value)
24+
}
25+
26+
@Throws(IOException::class)
27+
override fun read(`in`: JsonReader): InitializeResult =
28+
gson.fromJson(`in`, AwsExtendedInitializeResult::class.java)
29+
} as TypeAdapter<T>
30+
}
31+
return null
32+
}
33+
}
34+
35+
class AwsExtendedInitializeResult(awsServerCapabilities: AwsServerCapabilities? = null) : InitializeResult() {
36+
private var awsServerCapabilities: AwsServerCapabilities? = null
37+
38+
fun getAwsServerCapabilities(): AwsServerCapabilities? = awsServerCapabilities
39+
40+
fun setAwsServerCapabilities(awsServerCapabilities: AwsServerCapabilities?) {
41+
this.awsServerCapabilities = awsServerCapabilities
42+
}
43+
}

0 commit comments

Comments
 (0)