Skip to content

Commit 88b28a9

Browse files
leigaolleigaol
authored
Merge branch 'feature/q-lsp-chat' into actionCmds
2 parents 77d9cab + 9b4d638 commit 88b28a9

File tree

7 files changed

+136
-82
lines changed

7 files changed

+136
-82
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: 19 additions & 11 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
@@ -139,6 +140,7 @@ class BrowserConnector(
139140
// this is sent when the named agents UI is ready
140141
"ui-is-ready" -> {
141142
uiReady.complete(true)
143+
chatCommunicationManager.setUiReady()
142144
RunOnceUtil.runOnceForApp("AmazonQ-UI-Ready") {
143145
MeetQSettings.getInstance().reinvent2024OnboardingCount += 1
144146
}
@@ -254,7 +256,9 @@ class BrowserConnector(
254256
var encryptionManager: JwtEncryptionManager? = null
255257
val result = AmazonQLspService.executeIfRunning(project) { server ->
256258
encryptionManager = this.encryptionManager
257-
encryptionManager?.encrypt(chatParams)?.let { EncryptedChatParams(it, partialResultToken) }?.let { server.sendChatPrompt(it) }
259+
260+
val encryptedParams = EncryptedChatParams(this.encryptionManager.encrypt(chatParams), partialResultToken)
261+
server.sendChatPrompt(encryptedParams)
258262
} ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")))
259263

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

280282
// We assume there is only one outgoing request per tab because the input is
@@ -323,6 +325,7 @@ class BrowserConnector(
323325
CHAT_READY -> {
324326
handleChatNotification<ChatReadyNotification, Unit>(node) { server, _ ->
325327
uiReady.complete(true)
328+
chatCommunicationManager.setUiReady()
326329
RunOnceUtil.runOnceForApp("AmazonQ-UI-Ready") {
327330
MeetQSettings.getInstance().reinvent2024OnboardingCount += 1
328331
}
@@ -348,7 +351,7 @@ class BrowserConnector(
348351
}
349352
CHAT_OPEN_TAB -> {
350353
val response = serializer.deserializeChatMessages<OpenTabResponse>(node)
351-
ChatCommunicationManager.completeTabOpen(
354+
chatCommunicationManager.completeTabOpen(
352355
response.requestId,
353356
response.params.result.tabId
354357
)
@@ -419,7 +422,7 @@ class BrowserConnector(
419422

420423
GET_SERIALIZED_CHAT_REQUEST_METHOD -> {
421424
val response = serializer.deserializeChatMessages<GetSerializedChatResponse>(node)
422-
ChatCommunicationManager.completeSerializedChatResponse(
425+
chatCommunicationManager.completeSerializedChatResponse(
423426
response.requestId,
424427
response.params.result.content
425428
)
@@ -429,13 +432,18 @@ class BrowserConnector(
429432
handleChatNotification<TabBarActionRequest, TabBarActionParams>(node) {
430433
server, params ->
431434
val result = server.tabBarActions(params)
432-
result.whenComplete { params1, error ->
435+
result.whenComplete { actions, error ->
433436
try {
434437
if (error != null) {
435438
throw error
436439
}
437-
val res = ChatCommunicationManager.convertNotificationToJsonForChat(CHAT_TAB_BAR_ACTIONS, params1)
438-
browser.postChat(res)
440+
441+
browser.postChat(
442+
FlareUiMessage(
443+
command = CHAT_TAB_BAR_ACTIONS,
444+
params = actions ?: emptyMap<Any, Any>()
445+
)
446+
)
439447
} catch (e: Exception) {
440448
LOG.error { "Failed to perform chat tab bar action $e" }
441449
params.tabId?.let {

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

Lines changed: 38 additions & 42 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
@@ -136,20 +137,20 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC
136137
override fun openTab(params: OpenTabParams): CompletableFuture<OpenTabResult> {
137138
val requestId = UUID.randomUUID().toString()
138139
val result = CompletableFuture<OpenTabResult>()
139-
ChatCommunicationManager.pendingTabRequests[requestId] = result
140-
141-
val uiMessage = """
142-
{
143-
"command": "$CHAT_OPEN_TAB",
144-
"params": ${Gson().toJson(params)},
145-
"requestId": "$requestId"
146-
}
147-
""".trimIndent()
148-
AsyncChatUiListener.notifyPartialMessageUpdate(uiMessage)
140+
val chatManager = ChatCommunicationManager.getInstance(project)
141+
chatManager.addTabOpenRequest(requestId, result)
142+
143+
chatManager.notifyUi(
144+
FlareUiMessage(
145+
command = CHAT_OPEN_TAB,
146+
params = params,
147+
requestId = requestId,
148+
)
149+
)
149150

150151
result.orTimeout(30000, TimeUnit.MILLISECONDS)
151152
.whenComplete { _, error ->
152-
ChatCommunicationManager.pendingTabRequests.remove(requestId)
153+
chatManager.removeTabOpenRequest(requestId)
153154
}
154155

155156
return result
@@ -188,21 +189,20 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC
188189
override fun getSerializedChat(params: GetSerializedChatParams): CompletableFuture<GetSerializedChatResult> {
189190
val requestId = UUID.randomUUID().toString()
190191
val result = CompletableFuture<GetSerializedChatResult>()
191-
192-
ChatCommunicationManager.pendingSerializedChatRequests[requestId] = result
193-
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)
192+
val chatManager = ChatCommunicationManager.getInstance(project)
193+
chatManager.addSerializedChatRequest(requestId, result)
194+
195+
chatManager.notifyUi(
196+
FlareUiMessage(
197+
command = GET_SERIALIZED_CHAT_REQUEST_METHOD,
198+
params = params,
199+
requestId = requestId,
200+
)
201+
)
202202

203203
result.orTimeout(30000, TimeUnit.MILLISECONDS)
204204
.whenComplete { _, error ->
205-
ChatCommunicationManager.pendingSerializedChatRequests.remove(requestId)
205+
chatManager.removeSerializedChatRequest(requestId)
206206
}
207207

208208
return result
@@ -257,19 +257,17 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC
257257
try {
258258
chatCommunicationManager.handlePartialResultProgressNotification(project, params)
259259
} catch (e: Exception) {
260-
error("Cannot handle partial chat")
260+
LOG.error(e) { "Cannot handle partial chat" }
261261
}
262262
}
263263

264264
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)
265+
AsyncChatUiListener.notifyPartialMessageUpdate(
266+
FlareUiMessage(
267+
command = CHAT_SEND_UPDATE,
268+
params = params,
269+
)
270+
)
273271

274272
return CompletableFuture.completedFuture(Unit)
275273
}
@@ -343,15 +341,13 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC
343341
)
344342

345343
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)
354-
344+
val chatManager = ChatCommunicationManager.getInstance(project)
345+
chatManager.notifyUi(
346+
FlareUiMessage(
347+
command = CHAT_SEND_CONTEXT_COMMANDS,
348+
params = params ?: error("received empty payload for $CHAT_SEND_CONTEXT_COMMANDS"),
349+
)
350+
)
355351
return CompletableFuture.completedFuture(Unit)
356352
}
357353

‎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
}

0 commit comments

Comments
 (0)