Skip to content

fix(amazonq): increase frequency of bearer token sync to server #5636

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
May 12, 2025
Merged

fix(amazonq): increase frequency of bearer token sync to server #5636

Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
e182567
send token to Flare every 10 seconds
samgst-amazon Apr 24, 2025
be8ffb3
feedback
samgst-amazon Apr 25, 2025
800c7ff
Merge branch 'main' into samgst/q-lsp-token-refresh-periodic
samgst-amazon Apr 25, 2025
962a650
Merge branch 'main' into samgst/q-lsp-token-refresh-periodic
samgst-amazon May 2, 2025
5441f3f
Merge branch 'main' into samgst/q-lsp-token-refresh-periodic
samgst-amazon May 12, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 41 additions & 1 deletion ...oftware/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,11 @@

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
Expand All @@ -26,6 +29,9 @@
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,
Expand All @@ -34,7 +40,12 @@
) : AuthCredentialsService,
BearerTokenProviderListener,
ToolkitConnectionManagerListener,
QRegionProfileSelectedListener {
QRegionProfileSelectedListener,
Disposable {

private val scheduler: ScheduledExecutorService = AppExecutorUtil.getAppScheduledExecutorService()

Check warning on line 46 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt

View check run for this annotation

Codecov / codecov/patch

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

Added line #L46 was not covered by tests
private var tokenSyncTask: ScheduledFuture<*>? = null
private val tokenSyncIntervalSeconds = 10L

Check warning on line 48 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt

View check run for this annotation

Codecov / codecov/patch

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

Added line #L48 was not covered by tests

init {
project.messageBus.connect(serverInstance).apply {
Expand All @@ -49,6 +60,26 @@
updateConfiguration()
}
}

// Start periodic token sync
startPeriodicTokenSync()
}

Check warning on line 66 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt

View check run for this annotation

Codecov / codecov/patch

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

Added lines #L65 - L66 were not covered by tests

private fun startPeriodicTokenSync() {
tokenSyncTask = scheduler.scheduleWithFixedDelay(

Check warning on line 69 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt

View check run for this annotation

Codecov / codecov/patch

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

Added line #L69 was not covered by tests
{
try {

Check warning on line 71 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt

View check run for this annotation

Codecov / codecov/patch

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

Added line #L71 was not covered by tests
if (isQConnected(project) && !isQExpired(project)) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

isQExpired

fun isQExpired(project: Project): Boolean {
val manager = ToolkitConnectionManager.getInstance(project)
val qState = manager.connectionStateForFeature(QConnection.getInstance())
val cwState = manager.connectionStateForFeature(CodeWhispererConnection.getInstance())
LOG.debug {
"qConnectionState: $qState; cwConnectionState: $cwState"
}
return qState == BearerTokenAuthState.NEEDS_REFRESH || cwState == BearerTokenAuthState.NEEDS_REFRESH
}

make sure JB itself can automatically switch from NEEDS_REFRESH status to other status.. otherwise if JB stuck in NEEDS_REFRESH, this won't sync

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Otherwise, this PR looks good

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What do you mean by this? Isn't refresh logic handled by calls to reauthConnectionIfNeeded everywhere? If the IDE idles out and the token is no longer syncing because it needs to be refreshed, isn't that expected? We could add something to try and refresh the token here when it's expired, but I think that would increase the frequency of asking users to reauth.

I'm just a little confused about what AI you're implying here.

Copy link
Contributor

@leigaol leigaol Apr 25, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What I meant is: You'd double check if this NEEDS_REFRESH is referring to the automatic bearer token refresh that can be done by the JB IDE itself (calling the getToken API) or it is something that requires a full re-auth which requires user input.

Copy link
Contributor

@leigaol leigaol Apr 25, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

See this code:

aws-toolkit-jetbrains/plugins/core/jetbrains-community/src/software/aws/toolkits/jetbrains/core/credentials/ToolkitAuthManager.kt

Line 321 in 0b592d4

BearerTokenAuthState.NEEDS_REFRESH -> {

When the bearerToken is in the NEEDS_REFRESH state, it can go through tokenProvider.resolveToken() and actually starts to have a new valid bearer token. This is done as long as the user login is valid for about 3 months.

The flow I want to verify is:
Time 0:00: User login, gets a fresh bearer token.
Time 1:00: bearer token expires, BearerTokenAuthState becomes NEEDS_REFRESH state
From 1:00, per 10 seconds, the BearerTokenAuthState should automatically switch from NEEDS_REFRESH state to AUTHORIZED by going through tokenProvider.resolveToken(), even though user does not use the JetBrains IDE. Such that in your PR, the isQExpired can return false and new bearer token can keep being pushed to the Flare.

Otherwise, from time 1:00, BearerTokenAuthState becomes NEEDS_REFRESH state, isQExpired returns true, then no new bearer token will be pushed to Flare, and Flare will make API calls with 4xx.

TLDR: Does this PR work for the case when user login and after one hour of IDE idle (first bearer token expires), Flare can still get the new refresh bearer token and use it to call API?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you. I added some more logic to my implementation to try and address this.

updateTokenFromActiveConnection()

Check warning on line 73 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt

View check run for this annotation

Codecov / codecov/patch

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

Added line #L73 was not covered by tests
}
} catch (e: Exception) {
LOG.warn(e) { "Failed to sync bearer token to Flare" }

Check warning on line 76 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt

View check run for this annotation

Codecov / codecov/patch

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

Added lines #L75 - L76 were not covered by tests
}
},
tokenSyncIntervalSeconds,
tokenSyncIntervalSeconds,
TimeUnit.SECONDS

Check warning on line 81 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt

View check run for this annotation

Codecov / codecov/patch

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

Added lines #L78 - L81 were not covered by tests
)
}

override fun updateTokenCredentials(accessToken: String, encrypted: Boolean): CompletableFuture<ResponseMessage> {
Expand Down Expand Up @@ -134,4 +165,13 @@
server.updateConfiguration(payload)
} ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")))
}

override fun dispose() {
tokenSyncTask?.cancel(false)
tokenSyncTask = null
}

Check warning on line 172 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt

View check run for this annotation

Codecov / codecov/patch

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

Added lines #L171 - L172 were not covered by tests

companion object {
private val LOG = getLogger<DefaultAuthCredentialsService>()

Check warning on line 175 in plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/auth/DefaultAuthCredentialsService.kt

View check run for this annotation

Codecov / codecov/patch

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

Added line #L175 was not covered by tests
}
}
Loading