From e182567125c5576f6b7ba8c8ed2f0a7c488013cb Mon Sep 17 00:00:00 2001 From: samgst-amazon Date: Thu, 24 Apr 2025 16:38:28 -0700 Subject: [PATCH 1/2] send token to Flare every 10 seconds --- .../lsp/auth/DefaultAuthCredentialsService.kt | 42 ++++++++++++++++++- 1 file changed, 41 insertions(+), 1 deletion(-) 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 index d3a99a1f4fe..5db467149a5 100644 --- 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,8 +5,11 @@ 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 @@ -26,6 +29,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 +40,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 +60,26 @@ class DefaultAuthCredentialsService( updateConfiguration() } } + + // Start periodic token sync + startPeriodicTokenSync() + } + + private fun startPeriodicTokenSync() { + tokenSyncTask = scheduler.scheduleWithFixedDelay( + { + try { + if (isQConnected(project) && !isQExpired(project)) { + 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 { @@ -134,4 +165,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() + } } From be8ffb38178c1fb65145a23ae3ffa42313393684 Mon Sep 17 00:00:00 2001 From: samgst-amazon Date: Fri, 25 Apr 2025 09:51:41 -0700 Subject: [PATCH 2/2] feedback --- .../lsp/auth/DefaultAuthCredentialsService.kt | 26 +++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) 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 index 5db467149a5..5ee8c258d59 100644 --- 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 @@ -14,6 +14,7 @@ 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 @@ -69,8 +70,29 @@ class DefaultAuthCredentialsService( tokenSyncTask = scheduler.scheduleWithFixedDelay( { try { - if (isQConnected(project) && !isQExpired(project)) { - updateTokenFromActiveConnection() + 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" }