Skip to content

Commit fe7738d

Browse files
manodnyabmanodnyab
committed
Set up Flare chat connection
1 parent 5456fa0 commit fe7738d

File tree

15 files changed

+452
-15
lines changed

15 files changed

+452
-15
lines changed

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

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.QWebviewPanel
2020
import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQAppInitContext
2121
import software.aws.toolkits.jetbrains.services.amazonq.apps.AppConnection
2222
import software.aws.toolkits.jetbrains.services.amazonq.commands.MessageTypeRegistry
23+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AsyncChatUiListener
2324
import software.aws.toolkits.jetbrains.services.amazonq.messages.AmazonQMessage
2425
import software.aws.toolkits.jetbrains.services.amazonq.messages.MessageConnector
2526
import software.aws.toolkits.jetbrains.services.amazonq.onboarding.OnboardingPageInteraction
@@ -42,7 +43,7 @@ class AmazonQToolWindow private constructor(
4243
private val scope: CoroutineScope,
4344
) : Disposable {
4445
private val appSource = AppSource()
45-
private val browserConnector = BrowserConnector()
46+
private val browserConnector = BrowserConnector(project = project)
4647
private val editorThemeAdapter = EditorThemeAdapter()
4748

4849
private val chatPanel = AmazonQPanel(parent = this)
@@ -55,6 +56,15 @@ class AmazonQToolWindow private constructor(
5556
initConnections()
5657
connectUi()
5758
connectApps()
59+
60+
project.messageBus.connect().subscribe(
61+
AsyncChatUiListener.TOPIC,
62+
object : AsyncChatUiListener {
63+
override fun onChange(message: String) {
64+
chatPanel.browser?.postChat(message)
65+
}
66+
}
67+
)
5868
}
5969

6070
fun disposeAndRecreate() {

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

Lines changed: 94 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import com.intellij.openapi.Disposable
88
import com.intellij.openapi.util.Disposer
99
import com.intellij.ui.jcef.JBCefJSQuery
1010
import org.cef.CefApp
11+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.artifacts.ArtifactHelper
1112
import software.aws.toolkits.jetbrains.services.amazonq.util.HighlightCommand
1213
import software.aws.toolkits.jetbrains.services.amazonq.util.createBrowser
1314
import software.aws.toolkits.jetbrains.settings.MeetQSettings
@@ -46,6 +47,13 @@ class Browser(parent: Disposable) : Disposable {
4647

4748
fun component() = jcefBrowser.component
4849

50+
fun postChat(message: String) {
51+
jcefBrowser
52+
.cefBrowser
53+
.executeJavaScript("window.postMessage($message)", jcefBrowser.cefBrowser.url, 0)
54+
}
55+
56+
// TODO: Remove this once chat has been integrated with agents
4957
fun post(message: String) =
5058
jcefBrowser
5159
.cefBrowser
@@ -82,32 +90,86 @@ class Browser(parent: Disposable) : Disposable {
8290
highlightCommand: HighlightCommand?,
8391
): String {
8492
val postMessageToJavaJsCode = receiveMessageQuery.inject("JSON.stringify(message)")
85-
8693
val jsScripts = """
8794
<script type="text/javascript" src="$WEB_SCRIPT_URI" defer onload="init()"></script>
8895
<script type="text/javascript">
8996
const init = () => {
90-
mynahUI.createMynahUI(
97+
amazonQChat.createChat(
9198
{
9299
postMessage: message => {
93100
$postMessageToJavaJsCode
94101
}
95-
},
96-
${MeetQSettings.getInstance().reinvent2024OnboardingCount < MAX_ONBOARDING_PAGE_COUNT},
97-
${MeetQSettings.getInstance().disclaimerAcknowledged},
98-
$isFeatureDevAvailable, // whether /dev is available
99-
$isCodeTransformAvailable, // whether /transform is available
100-
$isDocAvailable, // whether /doc is available
101-
$isCodeScanAvailable, // whether /scan is available
102-
$isCodeTestAvailable, // whether /test is available
103-
${OBJECT_MAPPER.writeValueAsString(highlightCommand)}
102+
},
103+
{
104+
quickActionCommands: [],
105+
disclaimerAcknowledged: ${MeetQSettings.getInstance().disclaimerAcknowledged}
106+
}
107+
108+
104109
);
105110
}
106111
</script>
107112
""".trimIndent()
108113

114+
addQuickActionCommands(isCodeTransformAvailable, isFeatureDevAvailable, isDocAvailable, isCodeTestAvailable, isCodeScanAvailable, highlightCommand)
109115
return """
110116
<!DOCTYPE html>
117+
<style>
118+
body,
119+
html {
120+
background-color: var(--mynah-color-bg);
121+
color: var(--mynah-color-text-default);
122+
height: 100vh;
123+
width: 100%%;
124+
overflow: hidden;
125+
margin: 0;
126+
padding: 0;
127+
}
128+
.mynah-ui-icon-plus,
129+
.mynah-ui-icon-cancel {
130+
-webkit-mask-size: 155% !important;
131+
mask-size: 155% !important;
132+
mask-position: center;
133+
scale: 60%;
134+
}
135+
.code-snippet-close-button i.mynah-ui-icon-cancel,
136+
.mynah-chat-item-card-related-content-show-more i.mynah-ui-icon-down-open {
137+
-webkit-mask-size: 195.5% !important;
138+
mask-size: 195.5% !important;
139+
mask-position: center;
140+
aspect-ratio: 1/1;
141+
width: 15px;
142+
height: 15px;
143+
scale: 50%
144+
}
145+
.mynah-ui-icon-tabs {
146+
-webkit-mask-size: 102% !important;
147+
mask-size: 102% !important;
148+
mask-position: center;
149+
}
150+
textarea:placeholder-shown {
151+
line-height: 1.5rem;
152+
}
153+
.mynah-ui-spinner-container {
154+
contain: layout !important;
155+
}
156+
.mynah-ui-spinner-container > span.mynah-ui-spinner-logo-part {
157+
position: static !important;
158+
will-change: transform !important;
159+
}
160+
.mynah-ui-spinner-container,
161+
.mynah-ui-spinner-container > span.mynah-ui-spinner-logo-part,
162+
.mynah-ui-spinner-container > span.mynah-ui-spinner-logo-part > .mynah-ui-spinner-logo-mask.text {
163+
border: 0 !important;
164+
outline: none !important;
165+
box-shadow: none !important;
166+
border-radius: 0 !important;
167+
}
168+
.mynah-ui-spinner-container > span.mynah-ui-spinner-logo-part > .mynah-ui-spinner-logo-mask.text {
169+
will-change: transform !important;
170+
transform: translateZ(0) !important;
171+
}
172+
</style>
111173
<html>
112174
<head>
113175
<title>AWS Q</title>
@@ -119,8 +181,28 @@ class Browser(parent: Disposable) : Disposable {
119181
""".trimIndent()
120182
}
121183

184+
fun addQuickActionCommands(
185+
isCodeTransformAvailable: Boolean,
186+
isFeatureDevAvailable: Boolean,
187+
isDocAvailable: Boolean,
188+
isCodeTestAvailable: Boolean,
189+
isCodeScanAvailable: Boolean,
190+
highlightCommand: HighlightCommand?,
191+
) {
192+
// TODO: Remove this once chat has been integrated with agents. This is added temporarily to keep detekt happy.
193+
isCodeScanAvailable
194+
isCodeTestAvailable
195+
isDocAvailable
196+
isFeatureDevAvailable
197+
isCodeTransformAvailable
198+
MAX_ONBOARDING_PAGE_COUNT
199+
OBJECT_MAPPER
200+
highlightCommand
201+
}
202+
122203
companion object {
123-
private const val WEB_SCRIPT_URI = "http://mynah/js/mynah-ui.js"
204+
// TODO: Switch this to respect the overriden paths too
205+
private val WEB_SCRIPT_URI = ArtifactHelper().getLatestLocalLspArtifact().resolve("amazonq-ui.js").toUri()
124206
private const val MAX_ONBOARDING_PAGE_COUNT = 3
125207
private val OBJECT_MAPPER = jacksonObjectMapper()
126208
}

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

Lines changed: 72 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,11 @@
33

44
package software.aws.toolkits.jetbrains.services.amazonq.webview
55

6+
import com.fasterxml.jackson.databind.JsonNode
7+
import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
68
import com.intellij.ide.BrowserUtil
79
import com.intellij.ide.util.RunOnceUtil
10+
import com.intellij.openapi.project.Project
811
import com.intellij.ui.jcef.JBCefJSQuery.Response
912
import kotlinx.coroutines.CompletableDeferred
1013
import kotlinx.coroutines.channels.awaitClose
@@ -17,22 +20,36 @@ import kotlinx.coroutines.flow.merge
1720
import kotlinx.coroutines.flow.onEach
1821
import kotlinx.coroutines.launch
1922
import org.cef.browser.CefBrowser
23+
import org.eclipse.lsp4j.Position
24+
import org.eclipse.lsp4j.Range
2025
import software.aws.toolkits.jetbrains.services.amazonq.apps.AppConnection
2126
import software.aws.toolkits.jetbrains.services.amazonq.commands.MessageSerializer
27+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
28+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager
29+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.ChatCommunicationManager
30+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.getTextDocumentIdentifier
31+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatParams
32+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatPrompt
33+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CursorState
34+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.EncryptedChatParams
35+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SendChatPromptRequest
2236
import software.aws.toolkits.jetbrains.services.amazonq.util.command
2337
import software.aws.toolkits.jetbrains.services.amazonq.util.tabType
2438
import software.aws.toolkits.jetbrains.services.amazonq.webview.theme.AmazonQTheme
2539
import software.aws.toolkits.jetbrains.services.amazonq.webview.theme.ThemeBrowserAdapter
2640
import software.aws.toolkits.jetbrains.settings.MeetQSettings
2741
import software.aws.toolkits.telemetry.MetricResult
2842
import software.aws.toolkits.telemetry.Telemetry
43+
import java.util.concurrent.CompletableFuture
2944
import java.util.function.Function
3045

3146
class BrowserConnector(
3247
private val serializer: MessageSerializer = MessageSerializer.getInstance(),
3348
private val themeBrowserAdapter: ThemeBrowserAdapter = ThemeBrowserAdapter(),
49+
private val project: Project,
3450
) {
3551
var uiReady = CompletableDeferred<Boolean>()
52+
val chatCommunicationManager = ChatCommunicationManager.getInstance(project)
3653

3754
suspend fun connect(
3855
browser: Browser,
@@ -77,7 +94,10 @@ class BrowserConnector(
7794
}
7895
}
7996

80-
val tabType = node.tabType ?: return@onEach
97+
val tabType = node.tabType
98+
if (tabType == null) {
99+
handleFlareChatMessages(browser, node)
100+
}
81101
connections.filter { connection -> connection.app.tabTypes.contains(tabType) }.forEach { connection ->
82102
launch {
83103
val message = serializer.deserialize(node, connection.messageTypeRegistry)
@@ -123,4 +143,55 @@ class BrowserConnector(
123143
browser.receiveMessageQuery.removeHandler(handler)
124144
}
125145
}
146+
147+
private suspend fun handleFlareChatMessages(browser: Browser, node: JsonNode) {
148+
when (node.command) {
149+
"aws/chat/sendChatPrompt" -> {
150+
val requestFromUi = jacksonObjectMapper().readValue(node.toString(), SendChatPromptRequest::class.java)
151+
val chatPrompt = ChatPrompt(
152+
requestFromUi.params.prompt.prompt,
153+
requestFromUi.params.prompt.escapedPrompt,
154+
node.command
155+
)
156+
val textDocumentIdentifier = getTextDocumentIdentifier(project)
157+
val cursorState = CursorState(
158+
Range(
159+
Position(
160+
0,
161+
0
162+
),
163+
Position(
164+
1,
165+
1
166+
)
167+
)
168+
)
169+
val chatParams = ChatParams(
170+
requestFromUi.params.tabId,
171+
chatPrompt,
172+
textDocumentIdentifier,
173+
cursorState
174+
)
175+
val partialResultToken = chatCommunicationManager.addPartialChatMessage(requestFromUi.params.tabId)
176+
177+
var encryptionManager: JwtEncryptionManager? = null
178+
val result = AmazonQLspService.executeIfRunning(project) { server ->
179+
encryptionManager = this.encryptionManager
180+
server.sendChatPrompt(EncryptedChatParams(encryptionManager!!.encrypt(chatParams)))
181+
} ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")))
182+
183+
result.whenComplete {
184+
value, error ->
185+
chatCommunicationManager.removePartialChatMessage(partialResultToken)
186+
val messageToChat = ChatCommunicationManager.convertToJsonToSendToChat(
187+
node.command,
188+
requestFromUi.params.tabId,
189+
encryptionManager?.decrypt(value) ?: "",
190+
isPartialResult = false
191+
)
192+
browser.postChat(messageToChat)
193+
}
194+
}
195+
}
196+
}
126197
}

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,14 +9,17 @@ import org.eclipse.lsp4j.ConfigurationParams
99
import org.eclipse.lsp4j.MessageActionItem
1010
import org.eclipse.lsp4j.MessageParams
1111
import org.eclipse.lsp4j.MessageType
12+
import org.eclipse.lsp4j.ProgressParams
1213
import org.eclipse.lsp4j.PublishDiagnosticsParams
1314
import org.eclipse.lsp4j.ShowMessageRequestParams
1415
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
1516
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
1617
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
18+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.ChatCommunicationManager
1719
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.ConnectionMetadata
1820
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.SsoProfileData
1921
import software.aws.toolkits.jetbrains.settings.CodeWhispererSettings
22+
import software.aws.toolkits.jetbrains.utils.notifyInfo
2023
import java.util.concurrent.CompletableFuture
2124

2225
/**
@@ -93,4 +96,15 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC
9396
}
9497
)
9598
}
99+
100+
override fun notifyProgress(params: ProgressParams?) {
101+
if (params == null) return
102+
val chatCommunicationManager = ChatCommunicationManager.getInstance(project)
103+
try {
104+
chatCommunicationManager.handlePartialResultProgressNotification(project, params)
105+
notifyInfo("hello")
106+
} catch (e: Exception) {
107+
error("cannot handle partial chat")
108+
}
109+
}
96110
}

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import org.eclipse.lsp4j.jsonrpc.services.JsonRequest
99
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
12+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.EncryptedChatParams
1213
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.UpdateCredentialsPayload
1314
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.dependencies.DidChangeDependencyPathsParams
1415
import java.util.concurrent.CompletableFuture
@@ -29,4 +30,7 @@ interface AmazonQLanguageServer : LanguageServer {
2930

3031
@JsonRequest("aws/getConfigurationFromServer")
3132
fun getConfigurationFromServer(params: GetConfigurationFromServerParams): CompletableFuture<LspServerConfigurations>
33+
34+
@JsonRequest("aws/chat/sendChatPrompt")
35+
fun sendChatPrompt(params: EncryptedChatParams): CompletableFuture<String>
3236
}

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,8 @@ class AmazonQLspService(private val project: Project, private val cs: CoroutineS
104104
private var instance: Deferred<AmazonQServerInstance>
105105
val capabilities
106106
get() = instance.getCompleted().initializeResult.getCompleted().capabilities
107+
val encryptionManager
108+
get() = instance.getCompleted().encryptionManager
107109

108110
// dont allow lsp commands if server is restarting
109111
private val mutex = Mutex(false)
@@ -194,7 +196,7 @@ class AmazonQLspService(private val project: Project, private val cs: CoroutineS
194196
}
195197

196198
private class AmazonQServerInstance(private val project: Project, private val cs: CoroutineScope) : Disposable {
197-
private val encryptionManager = JwtEncryptionManager()
199+
val encryptionManager = JwtEncryptionManager()
198200

199201
private val launcher: Launcher<AmazonQLanguageServer>
200202

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,16 @@ class ArtifactHelper(private val lspArtifactsPath: Path = DEFAULT_ARTIFACT_PATH,
7979
.sortedByDescending { (_, semVer) -> semVer }
8080
}
8181

82+
fun getLatestLocalLspArtifact(): Path {
83+
val localFolders = getSubFolders(lspArtifactsPath)
84+
return localFolders.map { localFolder ->
85+
localFolder to SemVer.parseFromText(localFolder.fileName.toString())
86+
}
87+
.sortedByDescending { (_, semVer) -> semVer }
88+
.first()
89+
.first
90+
}
91+
8292
fun getExistingLspArtifacts(versions: List<ManifestManager.Version>, target: ManifestManager.VersionTarget?): Boolean {
8393
if (versions.isEmpty() || target?.contents == null) return false
8494

0 commit comments

Comments
 (0)