Merge branch 'feature/q-lsp-chat' into manodnyb/handleCodeScanContextCommand

manodnyab · web-flow · commit 00fe2811f01b · 2025-05-13T13:54:16.000-07:00
diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/theme/AmazonQTheme.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/theme/AmazonQTheme.kt
@@ -37,8 +37,6 @@ data class AmazonQTheme(
     val warning: Color,
     val error: Color,
 
-    val cardBackground: Color,
-
     val editorFont: Font,
     val editorBackground: Color,
     val editorForeground: Color,
@@ -49,5 +47,5 @@ data class AmazonQTheme(
     val editorKeyword: Color,
     val editorString: Color,
     val editorProperty: Color,
-
+    val editorClassName: Color,
 )
diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/theme/CssVariable.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/theme/CssVariable.kt
@@ -13,6 +13,7 @@ enum class CssVariable(
     FontFamily("--mynah-font-family"),
 
     TextColorDefault("--mynah-color-text-default"),
+    TextColorAlt("--mynah-color-text-alternate"),
     TextColorStrong("--mynah-color-text-strong"),
     TextColorWeak("--mynah-color-text-weak"),
     TextColorLink("--mynah-color-text-link"),
@@ -25,7 +26,7 @@ enum class CssVariable(
     ColorDeep("--mynah-color-deep"),
     ColorDeepReverse("--mynah-color-deep-reverse"),
     BorderDefault("--mynah-color-border-default"),
-    InputBackground("--mynah-color-input-bg"),
+    InputBackground("--mynah-input-bg"),
 
     SyntaxBackground("--mynah-color-syntax-bg"),
     SyntaxVariable("--mynah-color-syntax-variable"),
@@ -36,6 +37,9 @@ enum class CssVariable(
     SyntaxProperty("--mynah-color-syntax-property"),
     SyntaxComment("--mynah-color-syntax-comment"),
     SyntaxCode("--mynah-color-syntax-code"),
+    SyntaxKeyword("--mynah-color-syntax-keyword"),
+    SyntaxString("--mynah-color-syntax-string"),
+    SyntaxClassName("--mynah-color-syntax-class-name"),
     SyntaxCodeFontFamily("--mynah-syntax-code-font-family"),
     SyntaxCodeFontSize("--mynah-syntax-code-font-size"),
 
@@ -50,10 +54,9 @@ enum class CssVariable(
     SecondaryButtonBackground("--mynah-color-alternate"),
     SecondaryButtonForeground("--mynah-color-alternate-reverse"),
 
-    CodeText("--mynah-color-code-text"),
-
     MainBackground("--mynah-color-main"),
     MainForeground("--mynah-color-main-reverse"),
 
     CardBackground("--mynah-card-bg"),
+    CardBackgroundAlt("--mynah-card-bg-alternate"),
 }
diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/theme/EditorThemeAdapter.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/theme/EditorThemeAdapter.kt
@@ -113,18 +113,17 @@ class EditorThemeAdapter {
                 warning = themeColor("Component.warningFocusColor", default = 0xE2A53A),
                 error = themeColor("ProgressBar.failedColor", default = 0xD64F4F, darkDefault = 0xE74848),
 
-                cardBackground = cardBackground,
-
                 editorFont = currentScheme.getFont(EditorFontType.PLAIN),
-                editorBackground = chatBackground,
-                editorForeground = text,
+                editorBackground = currentScheme.defaultBackground,
+                editorForeground = currentScheme.defaultForeground,
                 editorVariable = currentScheme.foregroundColor(DefaultLanguageHighlighterColors.LOCAL_VARIABLE),
                 editorOperator = currentScheme.foregroundColor(DefaultLanguageHighlighterColors.OPERATION_SIGN),
                 editorFunction = currentScheme.foregroundColor(DefaultLanguageHighlighterColors.FUNCTION_DECLARATION),
                 editorComment = currentScheme.foregroundColor(DefaultLanguageHighlighterColors.LINE_COMMENT),
                 editorKeyword = currentScheme.foregroundColor(DefaultLanguageHighlighterColors.KEYWORD),
                 editorString = currentScheme.foregroundColor(DefaultLanguageHighlighterColors.STRING),
                 editorProperty = currentScheme.foregroundColor(DefaultLanguageHighlighterColors.INSTANCE_FIELD),
+                editorClassName = currentScheme.foregroundColor(DefaultLanguageHighlighterColors.CLASS_NAME),
             )
         }
 
diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/theme/ThemeBrowserAdapter.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/webview/theme/ThemeBrowserAdapter.kt
@@ -35,14 +35,16 @@ class ThemeBrowserAdapter {
         append(CssVariable.FontFamily, theme.font.toCssFontFamily())
 
         append(CssVariable.TextColorDefault, theme.defaultText)
+        append(CssVariable.TextColorAlt, theme.defaultText)
         append(CssVariable.TextColorStrong, theme.textFieldForeground)
         append(CssVariable.TextColorInput, theme.textFieldForeground)
         append(CssVariable.TextColorLink, theme.linkText)
         append(CssVariable.TextColorWeak, theme.inactiveText)
 
         append(CssVariable.Background, theme.background)
         append(CssVariable.BackgroundAlt, theme.background)
-        append(CssVariable.CardBackground, theme.cardBackground)
+        append(CssVariable.CardBackground, theme.editorBackground)
+        append(CssVariable.CardBackgroundAlt, theme.editorBackground)
         append(CssVariable.BorderDefault, theme.border)
         append(CssVariable.TabActive, theme.activeTab)
 
@@ -63,6 +65,7 @@ class ThemeBrowserAdapter {
 
         append(CssVariable.SyntaxCodeFontFamily, theme.editorFont.toCssFontFamily("monospace"))
         append(CssVariable.SyntaxCodeFontSize, theme.editorFont.toCssSize())
+        append(CssVariable.SyntaxCode, theme.editorForeground)
         append(CssVariable.SyntaxBackground, theme.editorBackground)
         append(CssVariable.SyntaxVariable, theme.editorVariable)
         append(CssVariable.SyntaxOperator, theme.editorOperator)
@@ -71,7 +74,9 @@ class ThemeBrowserAdapter {
         append(CssVariable.SyntaxAttributeValue, theme.editorKeyword)
         append(CssVariable.SyntaxAttribute, theme.editorString)
         append(CssVariable.SyntaxProperty, theme.editorProperty)
-        append(CssVariable.SyntaxCode, theme.editorForeground)
+        append(CssVariable.SyntaxKeyword, theme.editorKeyword)
+        append(CssVariable.SyntaxString, theme.editorString)
+        append(CssVariable.SyntaxClassName, theme.editorClassName)
 
         append(CssVariable.MainBackground, theme.buttonBackground)
         append(CssVariable.MainForeground, theme.buttonForeground)
diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/controller/CodeTestChatController.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqCodeTest/controller/CodeTestChatController.kt
@@ -56,6 +56,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQAppInitConte
 import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthController
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument.InlineCompletionReference
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument.InlineCompletionReferencePosition
+import software.aws.toolkits.jetbrains.services.amazonq.messages.AmazonQMessage
 import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
 import software.aws.toolkits.jetbrains.services.amazonq.project.RelevantDocument
 import software.aws.toolkits.jetbrains.services.amazonqCodeTest.CodeWhispererUTGChatManager
@@ -110,6 +111,12 @@ import java.util.UUID
 import software.amazon.awssdk.services.codewhispererstreaming.model.Position as StreamingPosition
 import software.amazon.awssdk.services.codewhispererstreaming.model.Range as StreamingRange
 
+data class TestCommandMessage(
+    val sender: String = "codetest",
+    val command: String = "test",
+    val type: String = "addAnswer",
+) : AmazonQMessage
+
 class CodeTestChatController(
     private val context: AmazonQAppInitContext,
     private val chatSessionStorage: ChatSessionStorage,
diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/commands/ActionRegistrar.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/commands/ActionRegistrar.kt
@@ -1,8 +1,9 @@
 // Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
-
+@file:Suppress("BannedImports")
 package software.aws.toolkits.jetbrains.services.cwc.commands
 
+import com.google.gson.Gson
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.project.Project
 import kotlinx.coroutines.flow.MutableSharedFlow
@@ -15,6 +16,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.Gener
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.SendToPromptParams
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.TriggerType
 import software.aws.toolkits.jetbrains.services.amazonq.messages.AmazonQMessage
+import software.aws.toolkits.jetbrains.services.amazonqCodeTest.controller.TestCommandMessage
 import software.aws.toolkits.jetbrains.services.cwc.editor.context.ActiveFileContextExtractor
 import software.aws.toolkits.jetbrains.services.cwc.editor.context.ExtractionTriggerType
 
@@ -26,8 +28,7 @@ class ActionRegistrar {
 
     fun reportMessageClick(command: EditorContextCommand, project: Project) {
         if (command == EditorContextCommand.GenerateUnitTests) {
-            // pre-existing old chat code path
-            _messages.tryEmit(ContextMenuActionMessage(command, project))
+            AsyncChatUiListener.notifyPartialMessageUpdate(Gson().toJson(TestCommandMessage()))
         } else {
             // new agentic chat route
             ApplicationManager.getApplication().executeOnPooledThread {
diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/ChatController.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/ChatController.kt
@@ -38,7 +38,6 @@ import software.aws.toolkits.jetbrains.core.coroutines.EDT
 import software.aws.toolkits.jetbrains.services.amazonq.apps.AmazonQAppInitContext
 import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthController
 import software.aws.toolkits.jetbrains.services.amazonq.auth.AuthNeededState
-import software.aws.toolkits.jetbrains.services.amazonq.messages.AmazonQMessage
 import software.aws.toolkits.jetbrains.services.amazonq.messages.MessagePublisher
 import software.aws.toolkits.jetbrains.services.amazonq.onboarding.OnboardingPageInteraction
 import software.aws.toolkits.jetbrains.services.amazonq.onboarding.OnboardingPageInteractionType
@@ -52,7 +51,6 @@ import software.aws.toolkits.jetbrains.services.cwc.clients.chat.model.TriggerTy
 import software.aws.toolkits.jetbrains.services.cwc.clients.chat.v1.ChatSessionFactoryV1
 import software.aws.toolkits.jetbrains.services.cwc.commands.CodeScanIssueActionMessage
 import software.aws.toolkits.jetbrains.services.cwc.commands.ContextMenuActionMessage
-import software.aws.toolkits.jetbrains.services.cwc.commands.EditorContextCommand
 import software.aws.toolkits.jetbrains.services.cwc.controller.chat.StaticPrompt
 import software.aws.toolkits.jetbrains.services.cwc.controller.chat.StaticTextResponse
 import software.aws.toolkits.jetbrains.services.cwc.controller.chat.messenger.ChatPromptHandler
@@ -76,12 +74,6 @@ import software.aws.toolkits.jetbrains.services.cwc.storage.ChatSessionStorage
 import software.aws.toolkits.telemetry.CwsprChatCommandType
 import java.util.UUID
 
-data class TestCommandMessage(
-    val sender: String = "codetest",
-    val command: String = "test",
-    val type: String = "addAnswer",
-) : AmazonQMessage
-
 class ChatController private constructor(
     private val context: AmazonQAppInitContext,
     private val chatSessionStorage: ChatSessionStorage,
@@ -285,16 +277,7 @@ class ChatController private constructor(
 
     // JB specific (not in vscode)
     override suspend fun processContextMenuCommand(message: ContextMenuActionMessage) {
-        // Extract context
-        if (message.project != context.project) {
-            return
-        }
-
-        if (message.command == EditorContextCommand.GenerateUnitTests) {
-            // Publish an event to "codetest" tab with command as "test" and type as "addAnswer"
-            val messageToPublish = TestCommandMessage()
-            context.messagesFromAppToUi.publish(messageToPublish)
-        }
+        // No-op since context commands are handled elsewhere. This function will be deprecated once we remove this class
     }
 
     private suspend fun processPromptActions(
diff --git a/plugins/amazonq/mynah-ui/src/mynah-ui/ui/quickActions/handler.ts b/plugins/amazonq/mynah-ui/src/mynah-ui/ui/quickActions/handler.ts
@@ -331,7 +331,7 @@ private handleDocCommand(chatPrompt: ChatPrompt, tabID: string, taskName: string
         this.showScanInTab(affectedTabId)
     }
 
-    private handleCodeTestCommand(chatPrompt: ChatPrompt, tabID: string, eventId: string | undefined) {
+    private handleCodeTestCommand(chatPrompt: ChatPrompt, tabID: string | undefined, eventId: string | undefined) {
         if (!this.isCodeTestEnabled) {
             return
         }
@@ -343,10 +343,17 @@ private handleDocCommand(chatPrompt: ChatPrompt, tabID: string, taskName: string
             this.connector.startTestGen(testTabId, realPromptText)
             return
         }
+        /**
+         * right click -> generate test has no tab id
+         * we have to manually create one if a testgen tab
+         * wasn't previously created
+         */
+        if (!tabID) {
+            tabID = this.mynahUI?.updateStore('', {})
+        }
         const affectedTabId: string | undefined = this.addTab(tabID)
 
         // if there is no test tab, open a new one
-
         if (affectedTabId === undefined) {
             this.mynahUI?.notify({
                 content: uiComponentsTexts.noMoreTabsTooltip,
diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt
@@ -7,7 +7,6 @@ import com.intellij.diff.DiffContentFactory
 import com.intellij.diff.DiffManager
 import com.intellij.diff.DiffManagerEx
 import com.intellij.diff.requests.SimpleDiffRequest
-import com.intellij.notification.NotificationType
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.fileChooser.FileChooserFactory
 import com.intellij.openapi.fileChooser.FileSaverDescriptor
@@ -25,8 +24,10 @@ import org.eclipse.lsp4j.PublishDiagnosticsParams
 import org.eclipse.lsp4j.ShowDocumentParams
 import org.eclipse.lsp4j.ShowDocumentResult
 import org.eclipse.lsp4j.ShowMessageRequestParams
+import org.slf4j.event.Level
 import software.aws.toolkits.core.utils.error
 import software.aws.toolkits.core.utils.getLogger
+import software.aws.toolkits.core.utils.info
 import software.aws.toolkits.core.utils.warn
 import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
@@ -70,11 +71,19 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC
 
     override fun showMessage(messageParams: MessageParams) {
         val type = when (messageParams.type) {
-            MessageType.Error -> NotificationType.ERROR
-            MessageType.Warning -> NotificationType.WARNING
-            MessageType.Info, MessageType.Log -> NotificationType.INFORMATION
+            MessageType.Error -> Level.ERROR
+            MessageType.Warning -> Level.WARN
+            MessageType.Info, MessageType.Log -> Level.INFO
         }
-        println("$type: ${messageParams.message}")
+
+        if (type == Level.ERROR &&
+            messageParams.message.lineSequence().firstOrNull()?.contains("NOTE: The AWS SDK for JavaScript (v2) is in maintenance mode.") == true
+        ) {
+            LOG.info { "Suppressed Flare AWS JS SDK v2 EoL error message" }
+            return
+        }
+
+        LOG.atLevel(type).log(messageParams.message)
     }
 
     override fun showMessageRequest(requestParams: ShowMessageRequestParams): CompletableFuture<MessageActionItem?>? {
diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt
@@ -5,12 +5,16 @@ package software.aws.toolkits.jetbrains.services.amazonq.lsp.auth
 
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.project.Project
+import com.intellij.util.concurrency.AppExecutorUtil
 import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage
 import software.aws.toolkits.core.TokenConnectionSettings
+import software.aws.toolkits.core.utils.getLogger
+import software.aws.toolkits.core.utils.warn
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManagerListener
 import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
+import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenAuthState
 import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProvider
 import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.BearerTokenProviderListener
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
@@ -26,6 +30,9 @@ import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSe
 import software.aws.toolkits.jetbrains.utils.isQConnected
 import software.aws.toolkits.jetbrains.utils.isQExpired
 import java.util.concurrent.CompletableFuture
+import java.util.concurrent.ScheduledExecutorService
+import java.util.concurrent.ScheduledFuture
+import java.util.concurrent.TimeUnit
 
 class DefaultAuthCredentialsService(
     private val project: Project,
@@ -34,7 +41,12 @@ class DefaultAuthCredentialsService(
 ) : AuthCredentialsService,
     BearerTokenProviderListener,
     ToolkitConnectionManagerListener,
-    QRegionProfileSelectedListener {
+    QRegionProfileSelectedListener,
+    Disposable {
+
+    private val scheduler: ScheduledExecutorService = AppExecutorUtil.getAppScheduledExecutorService()
+    private var tokenSyncTask: ScheduledFuture<*>? = null
+    private val tokenSyncIntervalSeconds = 10L
 
     init {
         project.messageBus.connect(serverInstance).apply {
@@ -49,6 +61,47 @@ class DefaultAuthCredentialsService(
                     updateConfiguration()
                 }
         }
+
+        // Start periodic token sync
+        startPeriodicTokenSync()
+    }
+
+    private fun startPeriodicTokenSync() {
+        tokenSyncTask = scheduler.scheduleWithFixedDelay(
+            {
+                try {
+                    if (isQConnected(project)) {
+                        if (isQExpired(project)) {
+                            val manager = ToolkitConnectionManager.getInstance(project)
+                            val connection = manager.activeConnectionForFeature(QConnection.getInstance()) ?: return@scheduleWithFixedDelay
+
+                            // Try to refresh the token if it's in NEEDS_REFRESH state
+                            val tokenProvider = (connection.getConnectionSettings() as? TokenConnectionSettings)
+                                ?.tokenProvider
+                                ?.delegate
+                                ?.let { it as? BearerTokenProvider } ?: return@scheduleWithFixedDelay
+
+                            if (tokenProvider.state() == BearerTokenAuthState.NEEDS_REFRESH) {
+                                try {
+                                    tokenProvider.resolveToken()
+                                    // Now that the token is refreshed, update it in Flare
+                                    updateTokenFromActiveConnection()
+                                } catch (e: Exception) {
+                                    LOG.warn(e) { "Failed to refresh bearer token" }
+                                }
+                            }
+                        } else {
+                            updateTokenFromActiveConnection()
+                        }
+                    }
+                } catch (e: Exception) {
+                    LOG.warn(e) { "Failed to sync bearer token to Flare" }
+                }
+            },
+            tokenSyncIntervalSeconds,
+            tokenSyncIntervalSeconds,
+            TimeUnit.SECONDS
+        )
     }
 
     override fun updateTokenCredentials(accessToken: String, encrypted: Boolean): CompletableFuture<ResponseMessage> {
@@ -134,4 +187,13 @@ class DefaultAuthCredentialsService(
             server.updateConfiguration(payload)
         } ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")))
     }
+
+    override fun dispose() {
+        tokenSyncTask?.cancel(false)
+        tokenSyncTask = null
+    }
+
+    companion object {
+        private val LOG = getLogger<DefaultAuthCredentialsService>()
+    }
 }
