Skip to content

Commit 34efb26

Browse files
rlirli
authored
Merge branch 'feature/q-lsp-chat' into rli/fix-history-restore
2 parents 4d99a80 + 9b4d638 commit 34efb26

File tree

3 files changed

+59
-26
lines changed

3 files changed

+59
-26
lines changed

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -140,6 +140,7 @@ class BrowserConnector(
140140
// this is sent when the named agents UI is ready
141141
"ui-is-ready" -> {
142142
uiReady.complete(true)
143+
chatCommunicationManager.setUiReady()
143144
RunOnceUtil.runOnceForApp("AmazonQ-UI-Ready") {
144145
MeetQSettings.getInstance().reinvent2024OnboardingCount += 1
145146
}
@@ -324,6 +325,7 @@ class BrowserConnector(
324325
CHAT_READY -> {
325326
handleChatNotification<ChatReadyNotification, Unit>(node) { server, _ ->
326327
uiReady.complete(true)
328+
chatCommunicationManager.setUiReady()
327329
RunOnceUtil.runOnceForApp("AmazonQ-UI-Ready") {
328330
MeetQSettings.getInstance().reinvent2024OnboardingCount += 1
329331
}
@@ -349,7 +351,7 @@ class BrowserConnector(
349351
}
350352
CHAT_OPEN_TAB -> {
351353
val response = serializer.deserializeChatMessages<OpenTabResponse>(node)
352-
ChatCommunicationManager.completeTabOpen(
354+
chatCommunicationManager.completeTabOpen(
353355
response.requestId,
354356
response.params.result.tabId
355357
)
@@ -420,7 +422,7 @@ class BrowserConnector(
420422

421423
GET_SERIALIZED_CHAT_REQUEST_METHOD -> {
422424
val response = serializer.deserializeChatMessages<GetSerializedChatResponse>(node)
423-
ChatCommunicationManager.completeSerializedChatResponse(
425+
chatCommunicationManager.completeSerializedChatResponse(
424426
response.requestId,
425427
response.params.result.content
426428
)

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

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -134,9 +134,10 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC
134134
override fun openTab(params: LSPAny): CompletableFuture<OpenTabResult> {
135135
val requestId = UUID.randomUUID().toString()
136136
val result = CompletableFuture<OpenTabResult>()
137-
ChatCommunicationManager.pendingTabRequests[requestId] = result
137+
val chatManager = ChatCommunicationManager.getInstance(project)
138+
chatManager.addTabOpenRequest(requestId, result)
138139

139-
AsyncChatUiListener.notifyPartialMessageUpdate(
140+
chatManager.notifyUi(
140141
FlareUiMessage(
141142
command = CHAT_OPEN_TAB,
142143
params = params,
@@ -146,7 +147,7 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC
146147

147148
result.orTimeout(30000, TimeUnit.MILLISECONDS)
148149
.whenComplete { _, error ->
149-
ChatCommunicationManager.pendingTabRequests.remove(requestId)
150+
chatManager.removeTabOpenRequest(requestId)
150151
}
151152

152153
return result
@@ -185,10 +186,10 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC
185186
override fun getSerializedChat(params: LSPAny): CompletableFuture<GetSerializedChatResult> {
186187
val requestId = UUID.randomUUID().toString()
187188
val result = CompletableFuture<GetSerializedChatResult>()
189+
val chatManager = ChatCommunicationManager.getInstance(project)
190+
chatManager.addSerializedChatRequest(requestId, result)
188191

189-
ChatCommunicationManager.pendingSerializedChatRequests[requestId] = result
190-
191-
AsyncChatUiListener.notifyPartialMessageUpdate(
192+
chatManager.notifyUi(
192193
FlareUiMessage(
193194
command = GET_SERIALIZED_CHAT_REQUEST_METHOD,
194195
params = params,
@@ -198,7 +199,7 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC
198199

199200
result.orTimeout(30000, TimeUnit.MILLISECONDS)
200201
.whenComplete { _, error ->
201-
ChatCommunicationManager.pendingSerializedChatRequests.remove(requestId)
202+
chatManager.removeSerializedChatRequest(requestId)
202203
}
203204

204205
return result
@@ -337,13 +338,13 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC
337338
)
338339

339340
override fun sendContextCommands(params: LSPAny): CompletableFuture<Unit> {
340-
AsyncChatUiListener.notifyPartialMessageUpdate(
341+
val chatManager = ChatCommunicationManager.getInstance(project)
342+
chatManager.notifyUi(
341343
FlareUiMessage(
342344
command = CHAT_SEND_CONTEXT_COMMANDS,
343345
params = params,
344346
)
345347
)
346-
347348
return CompletableFuture.completedFuture(Unit)
348349
}
349350

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

Lines changed: 45 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,9 @@ import com.google.gson.Gson
77
import com.intellij.openapi.components.Service
88
import com.intellij.openapi.components.service
99
import com.intellij.openapi.project.Project
10+
import kotlinx.coroutines.CompletableDeferred
11+
import kotlinx.coroutines.CoroutineScope
12+
import kotlinx.coroutines.launch
1013
import org.eclipse.lsp4j.ProgressParams
1114
import software.aws.toolkits.core.utils.getLogger
1215
import software.aws.toolkits.core.utils.warn
@@ -29,12 +32,23 @@ import java.util.concurrent.CompletableFuture
2932
import java.util.concurrent.ConcurrentHashMap
3033

3134
@Service(Service.Level.PROJECT)
32-
class ChatCommunicationManager {
35+
class ChatCommunicationManager(private val cs: CoroutineScope) {
36+
val uiReady = CompletableDeferred<Boolean>()
3337
private val chatPartialResultMap = ConcurrentHashMap<String, String>()
34-
private fun getPartialChatMessage(partialResultToken: String): String? =
35-
chatPartialResultMap.getOrDefault(partialResultToken, null)
36-
3738
private val inflightRequestByTabId = ConcurrentHashMap<String, CompletableFuture<String>>()
39+
private val pendingSerializedChatRequests = ConcurrentHashMap<String, CompletableFuture<GetSerializedChatResult>>()
40+
private val pendingTabRequests = ConcurrentHashMap<String, CompletableFuture<OpenTabResult>>()
41+
42+
fun setUiReady() {
43+
uiReady.complete(true)
44+
}
45+
46+
fun notifyUi(uiMessage: FlareUiMessage) {
47+
cs.launch {
48+
uiReady.await()
49+
AsyncChatUiListener.notifyPartialMessageUpdate(uiMessage)
50+
}
51+
}
3852

3953
fun setInflightRequestForTab(tabId: String, result: CompletableFuture<String>) {
4054
inflightRequestByTabId[tabId] = result
@@ -53,9 +67,36 @@ class ChatCommunicationManager {
5367
return partialResultToken
5468
}
5569

70+
private fun getPartialChatMessage(partialResultToken: String): String? =
71+
chatPartialResultMap.getOrDefault(partialResultToken, null)
72+
5673
fun removePartialChatMessage(partialResultToken: String) =
5774
chatPartialResultMap.remove(partialResultToken)
5875

76+
fun addSerializedChatRequest(requestId: String, result: CompletableFuture<GetSerializedChatResult>) {
77+
pendingSerializedChatRequests[requestId] = result
78+
}
79+
80+
fun completeSerializedChatResponse(requestId: String, content: String) {
81+
pendingSerializedChatRequests.remove(requestId)?.complete(GetSerializedChatResult((content)))
82+
}
83+
84+
fun removeSerializedChatRequest(requestId: String) {
85+
pendingSerializedChatRequests.remove(requestId)
86+
}
87+
88+
fun addTabOpenRequest(requestId: String, result: CompletableFuture<OpenTabResult>) {
89+
pendingTabRequests[requestId] = result
90+
}
91+
92+
fun completeTabOpen(requestId: String, tabId: String) {
93+
pendingTabRequests.remove(requestId)?.complete(OpenTabResult(tabId))
94+
}
95+
96+
fun removeTabOpenRequest(requestId: String) {
97+
pendingTabRequests.remove(requestId)
98+
}
99+
59100
fun handlePartialResultProgressNotification(project: Project, params: ProgressParams) {
60101
val token = ProgressNotificationUtils.getToken(params)
61102
val tabId = getPartialChatMessage(token)
@@ -134,11 +175,6 @@ class ChatCommunicationManager {
134175

135176
private val LOG = getLogger<ChatCommunicationManager>()
136177

137-
val pendingSerializedChatRequests = ConcurrentHashMap<String, CompletableFuture<GetSerializedChatResult>>()
138-
fun completeSerializedChatResponse(requestId: String, content: String) {
139-
pendingSerializedChatRequests.remove(requestId)?.complete(GetSerializedChatResult((content)))
140-
}
141-
142178
fun convertToJsonToSendToChat(command: String, tabId: String, params: String, isPartialResult: Boolean): String =
143179
"""
144180
{
@@ -148,11 +184,5 @@ class ChatCommunicationManager {
148184
"isPartialResult": $isPartialResult
149185
}
150186
""".trimIndent()
151-
152-
val pendingTabRequests = ConcurrentHashMap<String, CompletableFuture<OpenTabResult>>()
153-
154-
fun completeTabOpen(requestId: String, tabId: String) {
155-
pendingTabRequests.remove(requestId)?.complete(OpenTabResult(tabId))
156-
}
157187
}
158188
}

0 commit comments

Comments
 (0)