feat(amazonQFlareChat): Set up Flare chat connection

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

@@ -20,6 +20,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.QWebviewPanel
 import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQAppInitContext
 import software.aws.toolkits.jetbrains.services.amazonq.apps.AppConnection
 import software.aws.toolkits.jetbrains.services.amazonq.commands.MessageTypeRegistry
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AsyncChatUiListener
 import software.aws.toolkits.jetbrains.services.amazonq.messages.AmazonQMessage
 import software.aws.toolkits.jetbrains.services.amazonq.messages.MessageConnector
 import software.aws.toolkits.jetbrains.services.amazonq.onboarding.OnboardingPageInteraction
@@ -42,7 +43,7 @@ class AmazonQToolWindow private constructor(
     private val scope: CoroutineScope,
 ) : Disposable {
     private val appSource = AppSource()
-    private val browserConnector = BrowserConnector()
+    private val browserConnector = BrowserConnector(project = project)
     private val editorThemeAdapter = EditorThemeAdapter()
 
     private val chatPanel = AmazonQPanel(parent = this)
@@ -55,6 +56,15 @@ class AmazonQToolWindow private constructor(
         initConnections()
         connectUi()
         connectApps()
+
+        project.messageBus.connect().subscribe(
+            AsyncChatUiListener.TOPIC,
+            object : AsyncChatUiListener {
+                override fun onChange(message: String) {
+                    chatPanel.browser?.postChat(message)
+                }
+            }
+        )
     }
 
     fun disposeAndRecreate() {

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

@@ -8,6 +8,7 @@
 import com.intellij.openapi.util.Disposer
 import com.intellij.ui.jcef.JBCefJSQuery
 import org.cef.CefApp
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.artifacts.ArtifactHelper
 import software.aws.toolkits.jetbrains.services.amazonq.util.HighlightCommand
 import software.aws.toolkits.jetbrains.services.amazonq.util.createBrowser
 import software.aws.toolkits.jetbrains.settings.MeetQSettings
@@ -46,6 +47,13 @@
 
     fun component() = jcefBrowser.component
 
+    fun postChat(message: String) {
+        jcefBrowser
+            .cefBrowser
+            .executeJavaScript("window.postMessage($message)", jcefBrowser.cefBrowser.url, 0)
+    }
+
+    // TODO: Remove this once chat has been integrated with agents
     fun post(message: String) =
         jcefBrowser
             .cefBrowser
@@ -82,32 +90,86 @@
         highlightCommand: HighlightCommand?,
     ): String {
         val postMessageToJavaJsCode = receiveMessageQuery.inject("JSON.stringify(message)")
-
         val jsScripts = """
             <script type="text/javascript" src="$WEB_SCRIPT_URI" defer onload="init()"></script>
             <script type="text/javascript">
                 const init = () => {
-                    mynahUI.createMynahUI(
+                    amazonQChat.createChat(
                         {
                             postMessage: message => {
                                 $postMessageToJavaJsCode
                             }
-                        },
-                        ${MeetQSettings.getInstance().reinvent2024OnboardingCount < MAX_ONBOARDING_PAGE_COUNT},
-                        ${MeetQSettings.getInstance().disclaimerAcknowledged},
-                        $isFeatureDevAvailable, // whether /dev is available
-                        $isCodeTransformAvailable, // whether /transform is available
-                        $isDocAvailable, // whether /doc is available
-                        $isCodeScanAvailable, // whether /scan is available
-                        $isCodeTestAvailable, // whether /test is available
-                        ${OBJECT_MAPPER.writeValueAsString(highlightCommand)}
+                        }, 
+                        {
+                        quickActionCommands: [],
+                        disclaimerAcknowledged: ${MeetQSettings.getInstance().disclaimerAcknowledged}
+                        }
+                       
+                       
                     );
                 }
             </script>        
         """.trimIndent()
 
+        addQuickActionCommands(isCodeTransformAvailable, isFeatureDevAvailable, isDocAvailable, isCodeTestAvailable, isCodeScanAvailable, highlightCommand)
         return """
             <!DOCTYPE html>
+            <style>
+                    body,
+                    html {
+                        background-color: var(--mynah-color-bg);
+                        color: var(--mynah-color-text-default);
+                        height: 100vh;
+                        width: 100%%;
+                        overflow: hidden;
+                        margin: 0;
+                        padding: 0;
+                    }
+                    .mynah-ui-icon-plus,
+                    .mynah-ui-icon-cancel {
+                        -webkit-mask-size: 155% !important;
+                        mask-size: 155% !important;
+                        mask-position: center;
+                        scale: 60%;
+                    }
+                    .code-snippet-close-button i.mynah-ui-icon-cancel,
+                    .mynah-chat-item-card-related-content-show-more i.mynah-ui-icon-down-open {
+                        -webkit-mask-size: 195.5% !important;
+                        mask-size: 195.5% !important;
+                        mask-position: center;
+                        aspect-ratio: 1/1;
+                        width: 15px;
+                        height: 15px;
+                        scale: 50%
+                    }
+                    .mynah-ui-icon-tabs {
+                        -webkit-mask-size: 102% !important;
+                        mask-size: 102% !important;
+                        mask-position: center;
+                    }
+                    textarea:placeholder-shown {
+                        line-height: 1.5rem;
+                    }
+                    .mynah-ui-spinner-container {
+                        contain: layout !important;
+                    }
+                    .mynah-ui-spinner-container > span.mynah-ui-spinner-logo-part {
+                        position: static !important;
+                        will-change: transform !important;
+                    }
+                    .mynah-ui-spinner-container,
+                    .mynah-ui-spinner-container > span.mynah-ui-spinner-logo-part,
+                    .mynah-ui-spinner-container > span.mynah-ui-spinner-logo-part > .mynah-ui-spinner-logo-mask.text {
+                        border: 0 !important;
+                        outline: none !important;
+                        box-shadow: none !important;
+                        border-radius: 0 !important;
+                    }
+                    .mynah-ui-spinner-container > span.mynah-ui-spinner-logo-part > .mynah-ui-spinner-logo-mask.text {
+                        will-change: transform !important;
+                        transform: translateZ(0) !important;
+                    }
+                </style>
             <html>
                 <head>
                     <title>AWS Q</title>
@@ -119,8 +181,28 @@
         """.trimIndent()
     }
 
+    fun addQuickActionCommands(
+        isCodeTransformAvailable: Boolean,
+        isFeatureDevAvailable: Boolean,
+        isDocAvailable: Boolean,
+        isCodeTestAvailable: Boolean,
+        isCodeScanAvailable: Boolean,
+        highlightCommand: HighlightCommand?,
+    ) {
+        // TODO:  Remove this once chat has been integrated with agents. This is added temporarily to keep detekt happy.
+        isCodeScanAvailable
+        isCodeTestAvailable
+        isDocAvailable
+        isFeatureDevAvailable
+        isCodeTransformAvailable
+        MAX_ONBOARDING_PAGE_COUNT
+        OBJECT_MAPPER
+        highlightCommand
+    }
+
     companion object {
-        private const val WEB_SCRIPT_URI = "http://mynah/js/mynah-ui.js"
+        // TODO: Switch this to respect the overriden paths too
+        private val WEB_SCRIPT_URI = ArtifactHelper().getLatestLocalLspArtifact().resolve("amazonq-ui.js").toUri()
         private const val MAX_ONBOARDING_PAGE_COUNT = 3
         private val OBJECT_MAPPER = jacksonObjectMapper()
     }

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

@@ -3,8 +3,11 @@
 
 package software.aws.toolkits.jetbrains.services.amazonq.webview
 
+import com.fasterxml.jackson.databind.JsonNode
+import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
 import com.intellij.ide.BrowserUtil
 import com.intellij.ide.util.RunOnceUtil
+import com.intellij.openapi.project.Project
 import com.intellij.ui.jcef.JBCefJSQuery.Response
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.channels.awaitClose
@@ -17,22 +20,36 @@
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.launch
 import org.cef.browser.CefBrowser
+import org.eclipse.lsp4j.Position
+import org.eclipse.lsp4j.Range
 import software.aws.toolkits.jetbrains.services.amazonq.apps.AppConnection
 import software.aws.toolkits.jetbrains.services.amazonq.commands.MessageSerializer
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.ChatCommunicationManager
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.getTextDocumentIdentifier
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatParams
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ChatPrompt
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.CursorState
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.EncryptedChatParams
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SendChatPromptRequest
 import software.aws.toolkits.jetbrains.services.amazonq.util.command
 import software.aws.toolkits.jetbrains.services.amazonq.util.tabType
 import software.aws.toolkits.jetbrains.services.amazonq.webview.theme.AmazonQTheme
 import software.aws.toolkits.jetbrains.services.amazonq.webview.theme.ThemeBrowserAdapter
 import software.aws.toolkits.jetbrains.settings.MeetQSettings
 import software.aws.toolkits.telemetry.MetricResult
 import software.aws.toolkits.telemetry.Telemetry
+import java.util.concurrent.CompletableFuture
 import java.util.function.Function
 
 class BrowserConnector(
     private val serializer: MessageSerializer = MessageSerializer.getInstance(),
     private val themeBrowserAdapter: ThemeBrowserAdapter = ThemeBrowserAdapter(),
+    private val project: Project,
 ) {
     var uiReady = CompletableDeferred<Boolean>()
+    val chatCommunicationManager = ChatCommunicationManager.getInstance(project)
 
     suspend fun connect(
         browser: Browser,
@@ -77,7 +94,10 @@
                     }
                 }
 
-                val tabType = node.tabType ?: return@onEach
+                val tabType = node.tabType
+                if (tabType == null) {
+                    handleFlareChatMessages(browser, node)
+                }
                 connections.filter { connection -> connection.app.tabTypes.contains(tabType) }.forEach { connection ->
                     launch {
                         val message = serializer.deserialize(node, connection.messageTypeRegistry)
@@ -123,4 +143,55 @@
             browser.receiveMessageQuery.removeHandler(handler)
         }
     }
+
+    private suspend fun handleFlareChatMessages(browser: Browser, node: JsonNode) {
+        when (node.command) {
+            "aws/chat/sendChatPrompt" -> {
+                val requestFromUi = jacksonObjectMapper().readValue(node.toString(), SendChatPromptRequest::class.java)
+                val chatPrompt = ChatPrompt(
+                    requestFromUi.params.prompt.prompt,
+                    requestFromUi.params.prompt.escapedPrompt,
+                    node.command
+                )
+                val textDocumentIdentifier = getTextDocumentIdentifier(project)
+                val cursorState = CursorState(
+                    Range(
+                        Position(
+                            0,
+                            0
+                        ),
+                        Position(
+                            1,
+                            1
+                        )
+                    )
+                )
+                val chatParams = ChatParams(
+                    requestFromUi.params.tabId,
+                    chatPrompt,
+                    textDocumentIdentifier,
+                    cursorState
+                )
+                val partialResultToken = chatCommunicationManager.addPartialChatMessage(requestFromUi.params.tabId)
+
+                var encryptionManager: JwtEncryptionManager? = null
+                val result = AmazonQLspService.executeIfRunning(project) { server ->
+                    encryptionManager = this.encryptionManager
+                    server.sendChatPrompt(EncryptedChatParams(encryptionManager!!.encrypt(chatParams)))
+                } ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")))
+
+                result.whenComplete {
+                        value, error ->
+                    chatCommunicationManager.removePartialChatMessage(partialResultToken)
+                    val messageToChat = ChatCommunicationManager.convertToJsonToSendToChat(
+                        node.command,
+                        requestFromUi.params.tabId,
+                        encryptionManager?.decrypt(value) ?: "",
+                        isPartialResult = false
+                    )
+                    browser.postChat(messageToChat)
+                }
+            }
+        }
+    }
 }

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

@@ -9,14 +9,17 @@ import org.eclipse.lsp4j.ConfigurationParams
 import org.eclipse.lsp4j.MessageActionItem
 import org.eclipse.lsp4j.MessageParams
 import org.eclipse.lsp4j.MessageType
+import org.eclipse.lsp4j.ProgressParams
 import org.eclipse.lsp4j.PublishDiagnosticsParams
 import org.eclipse.lsp4j.ShowMessageRequestParams
 import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
 import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.ChatCommunicationManager
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.ConnectionMetadata
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.SsoProfileData
 import software.aws.toolkits.jetbrains.settings.CodeWhispererSettings
+import software.aws.toolkits.jetbrains.utils.notifyInfo
 import java.util.concurrent.CompletableFuture
 
 /**
@@ -93,4 +96,15 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC
             }
         )
     }
+
+    override fun notifyProgress(params: ProgressParams?) {
+        if (params == null) return
+        val chatCommunicationManager = ChatCommunicationManager.getInstance(project)
+        try {
+            chatCommunicationManager.handlePartialResultProgressNotification(project, params)
+            notifyInfo("hello")
+        } catch (e: Exception) {
+            error("cannot handle partial chat")
+        }
+    }
 }

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

@@ -9,6 +9,7 @@ import org.eclipse.lsp4j.jsonrpc.services.JsonRequest
 import org.eclipse.lsp4j.services.LanguageServer
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.GetConfigurationFromServerParams
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.LspServerConfigurations
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.EncryptedChatParams
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.UpdateCredentialsPayload
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.dependencies.DidChangeDependencyPathsParams
 import java.util.concurrent.CompletableFuture
@@ -29,4 +30,7 @@ interface AmazonQLanguageServer : LanguageServer {
 
     @JsonRequest("aws/getConfigurationFromServer")
     fun getConfigurationFromServer(params: GetConfigurationFromServerParams): CompletableFuture<LspServerConfigurations>
+
+    @JsonRequest("aws/chat/sendChatPrompt")
+    fun sendChatPrompt(params: EncryptedChatParams): CompletableFuture<String>
 }

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

@@ -104,6 +104,8 @@ class AmazonQLspService(private val project: Project, private val cs: CoroutineS
     private var instance: Deferred<AmazonQServerInstance>
     val capabilities
         get() = instance.getCompleted().initializeResult.getCompleted().capabilities
+    val encryptionManager
+        get() = instance.getCompleted().encryptionManager
 
     // dont allow lsp commands if server is restarting
     private val mutex = Mutex(false)
@@ -194,7 +196,7 @@ class AmazonQLspService(private val project: Project, private val cs: CoroutineS
 }
 
 private class AmazonQServerInstance(private val project: Project, private val cs: CoroutineScope) : Disposable {
-    private val encryptionManager = JwtEncryptionManager()
+    val encryptionManager = JwtEncryptionManager()
 
     private val launcher: Launcher<AmazonQLanguageServer>
 

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

@@ -79,6 +79,16 @@ class ArtifactHelper(private val lspArtifactsPath: Path = DEFAULT_ARTIFACT_PATH,
             .sortedByDescending { (_, semVer) -> semVer }
     }
 
+    fun getLatestLocalLspArtifact(): Path {
+        val localFolders = getSubFolders(lspArtifactsPath)
+        return localFolders.map { localFolder ->
+            localFolder to SemVer.parseFromText(localFolder.fileName.toString())
+        }
+            .sortedByDescending { (_, semVer) -> semVer }
+            .first()
+            .first
+    }
+
     fun getExistingLspArtifacts(versions: List<ManifestManager.Version>, target: ManifestManager.VersionTarget?): Boolean {
         if (versions.isEmpty() || target?.contents == null) return false