fix(amazonq): reduce log spam and frequency of token updates (#5713)

We do not need to update tokens every 10s

Additionally, if we provide the connection metadata in the token update request, then we reduce overall RPC calls

Author: rli

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

@@ -18,6 +18,7 @@ enum class CssVariable(
     TextColorWeak("--mynah-color-text-weak"),
     TextColorLink("--mynah-color-text-link"),
     TextColorInput("--mynah-color-text-input"),
+    TextColorDisabled("--mynah-color-text-disabled"),
 
     Background("--mynah-color-bg"),
     BackgroundAlt("--mynah-color-bg-alt"),

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

@@ -40,6 +40,7 @@ class ThemeBrowserAdapter {
         append(CssVariable.TextColorInput, theme.textFieldForeground)
         append(CssVariable.TextColorLink, theme.linkText)
         append(CssVariable.TextColorWeak, theme.inactiveText)
+        append(CssVariable.TextColorDisabled, theme.inactiveText)
 
         append(CssVariable.Background, theme.background)
         append(CssVariable.BackgroundAlt, theme.background)

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

@@ -29,7 +29,6 @@ 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
 import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.flareChat.AsyncChatUiListener
@@ -45,7 +44,6 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.OpenF
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ShowSaveFileDialogParams
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ShowSaveFileDialogResult
 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.services.amazonq.lsp.util.TelemetryParsingUtil
 import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator
 import software.aws.toolkits.jetbrains.services.telemetry.TelemetryService
@@ -170,19 +168,7 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC
             val connection = ToolkitConnectionManager.getInstance(project)
                 .activeConnectionForFeature(QConnection.getInstance())
 
-            when (connection) {
-                is AwsBearerTokenConnection -> {
-                    ConnectionMetadata(
-                        SsoProfileData(connection.startUrl)
-                    )
-                }
-                else -> {
-                    // If no connection or not a bearer token connection return default builderID start url
-                    ConnectionMetadata(
-                        SsoProfileData(AmazonQLspConstants.AWS_BUILDER_ID_URL)
-                    )
-                }
-            }
+            connection?.let { ConnectionMetadata.fromConnection(it) }
         }
 
     override fun openTab(params: LSPAny): CompletableFuture<LSPAny> {

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

@@ -4,9 +4,10 @@
 package software.aws.toolkits.jetbrains.services.amazonq.lsp.auth
 
 import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage
+import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
 import java.util.concurrent.CompletableFuture
 
 interface AuthCredentialsService {
-    fun updateTokenCredentials(accessToken: String, encrypted: Boolean): CompletableFuture<ResponseMessage>
+    fun updateTokenCredentials(connection: ToolkitConnection, encrypted: Boolean): CompletableFuture<ResponseMessage>
     fun deleteTokenCredentials(): CompletableFuture<Unit>
 }

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

@@ -22,6 +22,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryp
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.LspServerConfigurations
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.UpdateConfigurationParams
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.BearerCredentials
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.ConnectionMetadata
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.UpdateCredentialsPayload
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.UpdateCredentialsPayloadData
 import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile
@@ -46,7 +47,7 @@ class DefaultAuthCredentialsService(
 
     private val scheduler: ScheduledExecutorService = AppExecutorUtil.getAppScheduledExecutorService()
     private var tokenSyncTask: ScheduledFuture<*>? = null
-    private val tokenSyncIntervalSeconds = 900L
+    private val tokenSyncIntervalMinutes = 5L
 
     init {
         project.messageBus.connect(serverInstance).apply {
@@ -98,14 +99,18 @@ class DefaultAuthCredentialsService(
                     LOG.warn(e) { "Failed to sync bearer token to Flare" }
                 }
             },
-            tokenSyncIntervalSeconds,
-            tokenSyncIntervalSeconds,
-            TimeUnit.SECONDS
+            tokenSyncIntervalMinutes,
+            tokenSyncIntervalMinutes,
+            TimeUnit.MINUTES
         )
     }
 
-    override fun updateTokenCredentials(accessToken: String, encrypted: Boolean): CompletableFuture<ResponseMessage> {
-        val payload = createUpdateCredentialsPayload(accessToken, encrypted)
+    override fun updateTokenCredentials(connection: ToolkitConnection, encrypted: Boolean): CompletableFuture<ResponseMessage> {
+        val payload = try {
+            createUpdateCredentialsPayload(connection, encrypted)
+        } catch (e: Exception) {
+            return CompletableFuture.failedFuture(e)
+        }
 
         return AmazonQLspService.executeIfRunning(project) { server ->
             server.updateTokenCredentials(payload)
@@ -142,35 +147,39 @@ class DefaultAuthCredentialsService(
     }
 
     private fun updateTokenFromConnection(connection: ToolkitConnection): CompletableFuture<ResponseMessage> =
-        (connection.getConnectionSettings() as? TokenConnectionSettings)
+        updateTokenCredentials(connection, true)
+
+    override fun invalidate(providerId: String) {
+        deleteTokenCredentials()
+    }
+
+    private fun createUpdateCredentialsPayload(connection: ToolkitConnection, encrypted: Boolean): UpdateCredentialsPayload {
+        val token = (connection.getConnectionSettings() as? TokenConnectionSettings)
             ?.tokenProvider
             ?.delegate
             ?.let { it as? BearerTokenProvider }
             ?.currentToken()
             ?.accessToken
-            ?.let { token -> updateTokenCredentials(token, true) }
-            ?: CompletableFuture.failedFuture(IllegalStateException("Unable to get token from connection"))
+            ?: error("Unable to get token from connection")
 
-    override fun invalidate(providerId: String) {
-        deleteTokenCredentials()
-    }
-
-    private fun createUpdateCredentialsPayload(token: String, encrypted: Boolean): UpdateCredentialsPayload =
-        if (encrypted) {
+        return if (encrypted) {
             UpdateCredentialsPayload(
                 data = encryptionManager.encrypt(
                     UpdateCredentialsPayloadData(
                         BearerCredentials(token)
                     )
                 ),
+                metadata = ConnectionMetadata.fromConnection(connection),
                 encrypted = true
             )
         } else {
             UpdateCredentialsPayload(
                 data = token,
+                metadata = ConnectionMetadata.fromConnection(connection),
                 encrypted = false
             )
         }
+    }
 
     override fun onProfileSelected(project: Project, profile: QRegionProfile?) {
         updateConfiguration()

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/credentials/ConnectionMetadata.kt

@@ -3,9 +3,29 @@
 
 package software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials
 
+import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
+import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspConstants
+
 data class ConnectionMetadata(
     val sso: SsoProfileData,
-)
+) {
+    companion object {
+        fun fromConnection(connection: ToolkitConnection) = when (connection) {
+            is AwsBearerTokenConnection -> {
+                ConnectionMetadata(
+                    SsoProfileData(connection.startUrl)
+                )
+            }
+            else -> {
+                // If no connection or not a bearer token connection return default builderID start url
+                ConnectionMetadata(
+                    SsoProfileData(AmazonQLspConstants.AWS_BUILDER_ID_URL)
+                )
+            }
+        }
+    }
+}
 
 data class SsoProfileData(
     val startUrl: String,

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/credentials/UpdateCredentialsPayload.kt

@@ -5,6 +5,7 @@ package software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentia
 
 data class UpdateCredentialsPayload(
     val data: String,
+    val metadata: ConnectionMetadata,
     val encrypted: Boolean,
 )
 

plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImplTest.kt

@@ -331,7 +331,7 @@ class AmazonQLanguageClientImplTest {
         every { mockConnectionManager.activeConnectionForFeature(QConnection.getInstance()) } returns null
 
         assertThat(sut.getConnectionMetadata().get())
-            .isEqualTo(ConnectionMetadata(SsoProfileData(AmazonQLspConstants.AWS_BUILDER_ID_URL)))
+            .isEqualTo(null)
     }
 
     @Test

plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsServiceTest.kt

@@ -30,6 +30,8 @@ import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.InteractiveBe
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLanguageServer
 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.model.aws.credentials.ConnectionMetadata
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.SsoProfileData
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.UpdateCredentialsPayload
 import software.aws.toolkits.jetbrains.utils.isQConnected
 import software.aws.toolkits.jetbrains.utils.isQExpired
@@ -112,6 +114,7 @@ class DefaultAuthCredentialsServiceTest {
         connectionId: String = "test-connection-id",
     ): AwsBearerTokenConnection = mockk {
         every { id } returns connectionId
+        every { startUrl } returns "startUrl"
         every { getConnectionSettings() } returns createMockTokenSettings(accessToken)
     }
 
@@ -192,17 +195,18 @@ class DefaultAuthCredentialsServiceTest {
 
     @Test
     fun `test updateTokenCredentials unencrypted success`() {
-        sut = DefaultAuthCredentialsService(project, mockEncryptionManager, mockk())
-
-        val token = "unencryptedToken"
         val isEncrypted = false
+        sut = DefaultAuthCredentialsService(project, mockEncryptionManager, mockk())
 
-        sut.updateTokenCredentials(token, isEncrypted)
+        sut.updateTokenCredentials(mockConnection, isEncrypted)
 
         verify(exactly = 1) {
             mockLanguageServer.updateTokenCredentials(
                 UpdateCredentialsPayload(
-                    token,
+                    "test-access-token",
+                    ConnectionMetadata(
+                        SsoProfileData("startUrl")
+                    ),
                     isEncrypted
                 )
             )
@@ -214,17 +218,19 @@ class DefaultAuthCredentialsServiceTest {
         sut = DefaultAuthCredentialsService(project, mockEncryptionManager, mockk())
 
         val encryptedToken = "encryptedToken"
-        val decryptedToken = "decryptedToken"
         val isEncrypted = true
 
         every { mockEncryptionManager.encrypt(any()) } returns encryptedToken
 
-        sut.updateTokenCredentials(decryptedToken, isEncrypted)
+        sut.updateTokenCredentials(mockConnection, isEncrypted)
 
         verify(atLeast = 1) {
             mockLanguageServer.updateTokenCredentials(
                 UpdateCredentialsPayload(
                     encryptedToken,
+                    ConnectionMetadata(
+                        SsoProfileData("startUrl")
+                    ),
                     isEncrypted
                 )
             )