Skip to content

Commit 9979e79

Browse files
rlirli
authored
refactor: avoid manually crafting JSON blobs for postMessage(...) (#5693)
We are calling `Gson().toJson(...)` all over the place due to poor modeling of the Flare UI contract
1 parent 7bc1baa commit 9979e79

File tree

7 files changed

+78
-57
lines changed

7 files changed

+78
-57
lines changed

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

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.apps.AppConnection
2727
import software.aws.toolkits.jetbrains.services.amazonq.commands.MessageTypeRegistry
2828
import software.aws.toolkits.jetbrains.services.amazonq.lsp.artifacts.ArtifactManager
2929
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AsyncChatUiListener
30+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.FlareUiMessage
3031
import software.aws.toolkits.jetbrains.services.amazonq.messages.AmazonQMessage
3132
import software.aws.toolkits.jetbrains.services.amazonq.messages.MessageConnector
3233
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
@@ -55,8 +56,12 @@ class AmazonQPanel(val project: Project, private val scope: CoroutineScope) : Di
5556
project.messageBus.connect().subscribe(
5657
AsyncChatUiListener.TOPIC,
5758
object : AsyncChatUiListener {
58-
override fun onChange(message: String) {
59-
browser.get()?.postChat(message)
59+
override fun onChange(command: String) {
60+
browser.get()?.postChat(command)
61+
}
62+
63+
override fun onChange(command: FlareUiMessage) {
64+
browser.get()?.postChat(command)
6065
}
6166
}
6267
)

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
11
// Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
22
// SPDX-License-Identifier: Apache-2.0
3-
3+
@file:Suppress("BannedImports")
44
package software.aws.toolkits.jetbrains.services.amazonq.webview
55

66
import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
7+
import com.google.gson.Gson
78
import com.intellij.openapi.Disposable
89
import com.intellij.openapi.project.Project
910
import com.intellij.openapi.util.Disposer
@@ -12,6 +13,7 @@ import org.cef.CefApp
1213
import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
1314
import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
1415
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AwsServerCapabilitiesProvider
16+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.FlareUiMessage
1517
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile
1618
import software.aws.toolkits.jetbrains.services.amazonq.util.HighlightCommand
1719
import software.aws.toolkits.jetbrains.services.amazonq.util.createBrowser
@@ -63,6 +65,9 @@ class Browser(parent: Disposable, private val webUri: URI, val project: Project)
6365

6466
fun component() = jcefBrowser.component
6567

68+
fun postChat(command: FlareUiMessage) = postChat(Gson().toJson(command))
69+
70+
@Deprecated("shouldn't need this version")
6671
fun postChat(message: String) {
6772
jcefBrowser
6873
.cefBrowser

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

Lines changed: 15 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
3434
import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager
3535
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AwsServerCapabilitiesProvider
3636
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.ChatCommunicationManager
37+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.FlareUiMessage
3738
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.AUTH_FOLLOW_UP_CLICKED
3839
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.AuthFollowUpClickNotification
3940
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ButtonClickNotification
@@ -254,7 +255,9 @@ class BrowserConnector(
254255
var encryptionManager: JwtEncryptionManager? = null
255256
val result = AmazonQLspService.executeIfRunning(project) { server ->
256257
encryptionManager = this.encryptionManager
257-
encryptionManager?.encrypt(chatParams)?.let { EncryptedChatParams(it, partialResultToken) }?.let { server.sendChatPrompt(it) }
258+
259+
val encryptedParams = EncryptedChatParams(this.encryptionManager.encrypt(chatParams), partialResultToken)
260+
server.sendChatPrompt(encryptedParams)
258261
} ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")))
259262

260263
// We assume there is only one outgoing request per tab because the input is
@@ -270,11 +273,9 @@ class BrowserConnector(
270273
var encryptionManager: JwtEncryptionManager? = null
271274
val result = AmazonQLspService.executeIfRunning(project) { server ->
272275
encryptionManager = this.encryptionManager
273-
encryptionManager?.encrypt(quickActionParams)?.let {
274-
EncryptedQuickActionChatParams(it, partialResultToken)
275-
}?.let {
276-
server.sendQuickAction(it)
277-
}
276+
277+
val encryptedParams = EncryptedQuickActionChatParams(this.encryptionManager.encrypt(quickActionParams), partialResultToken)
278+
server.sendQuickAction(encryptedParams)
278279
} ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")))
279280

280281
// We assume there is only one outgoing request per tab because the input is
@@ -429,13 +430,18 @@ class BrowserConnector(
429430
handleChatNotification<TabBarActionRequest, TabBarActionParams>(node) {
430431
server, params ->
431432
val result = server.tabBarActions(params)
432-
result.whenComplete { params1, error ->
433+
result.whenComplete { actions, error ->
433434
try {
434435
if (error != null) {
435436
throw error
436437
}
437-
val res = ChatCommunicationManager.convertNotificationToJsonForChat(CHAT_TAB_BAR_ACTIONS, params1)
438-
browser.postChat(res)
438+
439+
browser.postChat(
440+
FlareUiMessage(
441+
command = CHAT_TAB_BAR_ACTIONS,
442+
params = actions ?: emptyMap<Any, Any>()
443+
)
444+
)
439445
} catch (e: Exception) {
440446
LOG.error { "Failed to perform chat tab bar action $e" }
441447
params.tabId?.let {

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

Lines changed: 29 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33
@file:Suppress("BannedImports")
44
package software.aws.toolkits.jetbrains.services.amazonq.lsp
55

6-
import com.google.gson.Gson
76
import com.intellij.diff.DiffContentFactory
87
import com.intellij.diff.DiffManager
98
import com.intellij.diff.DiffManagerEx
@@ -26,13 +25,15 @@ import org.eclipse.lsp4j.PublishDiagnosticsParams
2625
import org.eclipse.lsp4j.ShowDocumentParams
2726
import org.eclipse.lsp4j.ShowDocumentResult
2827
import org.eclipse.lsp4j.ShowMessageRequestParams
28+
import software.aws.toolkits.core.utils.error
2929
import software.aws.toolkits.core.utils.getLogger
3030
import software.aws.toolkits.core.utils.warn
3131
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
3232
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
3333
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
3434
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AsyncChatUiListener
3535
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.ChatCommunicationManager
36+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.FlareUiMessage
3637
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.LSPAny
3738
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_OPEN_TAB
3839
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CHAT_SEND_CONTEXT_COMMANDS
@@ -138,14 +139,13 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC
138139
val result = CompletableFuture<OpenTabResult>()
139140
ChatCommunicationManager.pendingTabRequests[requestId] = result
140141

141-
val uiMessage = """
142-
{
143-
"command": "$CHAT_OPEN_TAB",
144-
"params": ${Gson().toJson(params)},
145-
"requestId": "$requestId"
146-
}
147-
""".trimIndent()
148-
AsyncChatUiListener.notifyPartialMessageUpdate(uiMessage)
142+
AsyncChatUiListener.notifyPartialMessageUpdate(
143+
FlareUiMessage(
144+
command = CHAT_OPEN_TAB,
145+
params = params,
146+
requestId = requestId,
147+
)
148+
)
149149

150150
result.orTimeout(30000, TimeUnit.MILLISECONDS)
151151
.whenComplete { _, error ->
@@ -191,14 +191,13 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC
191191

192192
ChatCommunicationManager.pendingSerializedChatRequests[requestId] = result
193193

194-
val uiMessage = """
195-
{
196-
"command": "$GET_SERIALIZED_CHAT_REQUEST_METHOD",
197-
"params": ${Gson().toJson(params)},
198-
"requestId": "$requestId"
199-
}
200-
""".trimIndent()
201-
AsyncChatUiListener.notifyPartialMessageUpdate(uiMessage)
194+
AsyncChatUiListener.notifyPartialMessageUpdate(
195+
FlareUiMessage(
196+
command = GET_SERIALIZED_CHAT_REQUEST_METHOD,
197+
params = params,
198+
requestId = requestId,
199+
)
200+
)
202201

203202
result.orTimeout(30000, TimeUnit.MILLISECONDS)
204203
.whenComplete { _, error ->
@@ -257,19 +256,17 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC
257256
try {
258257
chatCommunicationManager.handlePartialResultProgressNotification(project, params)
259258
} catch (e: Exception) {
260-
error("Cannot handle partial chat")
259+
LOG.error(e) { "Cannot handle partial chat" }
261260
}
262261
}
263262

264263
override fun sendChatUpdate(params: ChatUpdateParams): CompletableFuture<Unit> {
265-
val uiMessage = """
266-
{
267-
"command":"$CHAT_SEND_UPDATE",
268-
"params":${Gson().toJson(params)}
269-
}
270-
""".trimIndent()
271-
272-
AsyncChatUiListener.notifyPartialMessageUpdate(uiMessage)
264+
AsyncChatUiListener.notifyPartialMessageUpdate(
265+
FlareUiMessage(
266+
command = CHAT_SEND_UPDATE,
267+
params = params,
268+
)
269+
)
273270

274271
return CompletableFuture.completedFuture(Unit)
275272
}
@@ -343,14 +340,12 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC
343340
)
344341

345342
override fun sendContextCommands(params: LSPAny): CompletableFuture<Unit> {
346-
val showContextCommands = """
347-
{
348-
"command":"$CHAT_SEND_CONTEXT_COMMANDS",
349-
"params": ${Gson().toJson(params)}
350-
}
351-
""".trimIndent()
352-
353-
AsyncChatUiListener.notifyPartialMessageUpdate(showContextCommands)
343+
AsyncChatUiListener.notifyPartialMessageUpdate(
344+
FlareUiMessage(
345+
command = CHAT_SEND_CONTEXT_COMMANDS,
346+
params = params ?: error("received empty payload for $CHAT_SEND_CONTEXT_COMMANDS"),
347+
)
348+
)
354349

355350
return CompletableFuture.completedFuture(Unit)
356351
}

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

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,14 +8,22 @@ import com.intellij.util.messages.Topic
88
import java.util.EventListener
99

1010
interface AsyncChatUiListener : EventListener {
11-
fun onChange(message: String) {}
11+
@Deprecated("shouldn't need this version")
12+
fun onChange(command: String) {}
13+
14+
fun onChange(command: FlareUiMessage) {}
1215

1316
companion object {
1417
@Topic.AppLevel
1518
val TOPIC = Topic.create("Partial chat message provider", AsyncChatUiListener::class.java)
1619

17-
fun notifyPartialMessageUpdate(message: String) {
18-
ApplicationManager.getApplication().messageBus.syncPublisher(TOPIC).onChange(message)
20+
fun notifyPartialMessageUpdate(command: FlareUiMessage) {
21+
ApplicationManager.getApplication().messageBus.syncPublisher(TOPIC).onChange(command)
22+
}
23+
24+
@Deprecated("shouldn't need this version")
25+
fun notifyPartialMessageUpdate(command: String) {
26+
ApplicationManager.getApplication().messageBus.syncPublisher(TOPIC).onChange(command)
1927
}
2028
}
2129
}

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

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -154,13 +154,5 @@ class ChatCommunicationManager {
154154
fun completeTabOpen(requestId: String, tabId: String) {
155155
pendingTabRequests.remove(requestId)?.complete(OpenTabResult(tabId))
156156
}
157-
158-
inline fun <reified T> convertNotificationToJsonForChat(command: String, params: T? = null) =
159-
"""
160-
{
161-
"command":"$command",
162-
"params": ${if (params != null) Gson().toJson(params) else "{}"}
163-
}
164-
""".trimIndent()
165157
}
166158
}

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat
5+
6+
data class FlareUiMessage(
7+
val command: String,
8+
val params: Any,
9+
val requestId: String? = null,
10+
)

0 commit comments

Comments
 (0)