Merge main into feature/q-lsp-chat

Author: aws-toolkit-automation

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>()
+    }
 }