Skip to content

Commit d0aed14

Browse files
authored
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
1 parent 440fd77 commit d0aed14

File tree

9 files changed

+65
-40
lines changed

9 files changed

+65
-40
lines changed

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ enum class CssVariable(
1818
TextColorWeak("--mynah-color-text-weak"),
1919
TextColorLink("--mynah-color-text-link"),
2020
TextColorInput("--mynah-color-text-input"),
21+
TextColorDisabled("--mynah-color-text-disabled"),
2122

2223
Background("--mynah-color-bg"),
2324
BackgroundAlt("--mynah-color-bg-alt"),

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ class ThemeBrowserAdapter {
4040
append(CssVariable.TextColorInput, theme.textFieldForeground)
4141
append(CssVariable.TextColorLink, theme.linkText)
4242
append(CssVariable.TextColorWeak, theme.inactiveText)
43+
append(CssVariable.TextColorDisabled, theme.inactiveText)
4344

4445
append(CssVariable.Background, theme.background)
4546
append(CssVariable.BackgroundAlt, theme.background)

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

Lines changed: 1 addition & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,6 @@ import software.aws.toolkits.core.utils.error
2929
import software.aws.toolkits.core.utils.getLogger
3030
import software.aws.toolkits.core.utils.info
3131
import software.aws.toolkits.core.utils.warn
32-
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
3332
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
3433
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
3534
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
4544
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ShowSaveFileDialogParams
4645
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.chat.ShowSaveFileDialogResult
4746
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.ConnectionMetadata
48-
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.SsoProfileData
4947
import software.aws.toolkits.jetbrains.services.amazonq.lsp.util.TelemetryParsingUtil
5048
import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator
5149
import software.aws.toolkits.jetbrains.services.telemetry.TelemetryService
@@ -170,19 +168,7 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC
170168
val connection = ToolkitConnectionManager.getInstance(project)
171169
.activeConnectionForFeature(QConnection.getInstance())
172170

173-
when (connection) {
174-
is AwsBearerTokenConnection -> {
175-
ConnectionMetadata(
176-
SsoProfileData(connection.startUrl)
177-
)
178-
}
179-
else -> {
180-
// If no connection or not a bearer token connection return default builderID start url
181-
ConnectionMetadata(
182-
SsoProfileData(AmazonQLspConstants.AWS_BUILDER_ID_URL)
183-
)
184-
}
185-
}
171+
connection?.let { ConnectionMetadata.fromConnection(it) }
186172
}
187173

188174
override fun openTab(params: LSPAny): CompletableFuture<LSPAny> {

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,10 @@
44
package software.aws.toolkits.jetbrains.services.amazonq.lsp.auth
55

66
import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage
7+
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
78
import java.util.concurrent.CompletableFuture
89

910
interface AuthCredentialsService {
10-
fun updateTokenCredentials(accessToken: String, encrypted: Boolean): CompletableFuture<ResponseMessage>
11+
fun updateTokenCredentials(connection: ToolkitConnection, encrypted: Boolean): CompletableFuture<ResponseMessage>
1112
fun deleteTokenCredentials(): CompletableFuture<Unit>
1213
}

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

Lines changed: 24 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryp
2222
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.LspServerConfigurations
2323
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.UpdateConfigurationParams
2424
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.BearerCredentials
25+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.ConnectionMetadata
2526
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.UpdateCredentialsPayload
2627
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.UpdateCredentialsPayloadData
2728
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile
@@ -46,7 +47,7 @@ class DefaultAuthCredentialsService(
4647

4748
private val scheduler: ScheduledExecutorService = AppExecutorUtil.getAppScheduledExecutorService()
4849
private var tokenSyncTask: ScheduledFuture<*>? = null
49-
private val tokenSyncIntervalSeconds = 900L
50+
private val tokenSyncIntervalMinutes = 5L
5051

5152
init {
5253
project.messageBus.connect(serverInstance).apply {
@@ -98,14 +99,18 @@ class DefaultAuthCredentialsService(
9899
LOG.warn(e) { "Failed to sync bearer token to Flare" }
99100
}
100101
},
101-
tokenSyncIntervalSeconds,
102-
tokenSyncIntervalSeconds,
103-
TimeUnit.SECONDS
102+
tokenSyncIntervalMinutes,
103+
tokenSyncIntervalMinutes,
104+
TimeUnit.MINUTES
104105
)
105106
}
106107

107-
override fun updateTokenCredentials(accessToken: String, encrypted: Boolean): CompletableFuture<ResponseMessage> {
108-
val payload = createUpdateCredentialsPayload(accessToken, encrypted)
108+
override fun updateTokenCredentials(connection: ToolkitConnection, encrypted: Boolean): CompletableFuture<ResponseMessage> {
109+
val payload = try {
110+
createUpdateCredentialsPayload(connection, encrypted)
111+
} catch (e: Exception) {
112+
return CompletableFuture.failedFuture(e)
113+
}
109114

110115
return AmazonQLspService.executeIfRunning(project) { server ->
111116
server.updateTokenCredentials(payload)
@@ -142,35 +147,39 @@ class DefaultAuthCredentialsService(
142147
}
143148

144149
private fun updateTokenFromConnection(connection: ToolkitConnection): CompletableFuture<ResponseMessage> =
145-
(connection.getConnectionSettings() as? TokenConnectionSettings)
150+
updateTokenCredentials(connection, true)
151+
152+
override fun invalidate(providerId: String) {
153+
deleteTokenCredentials()
154+
}
155+
156+
private fun createUpdateCredentialsPayload(connection: ToolkitConnection, encrypted: Boolean): UpdateCredentialsPayload {
157+
val token = (connection.getConnectionSettings() as? TokenConnectionSettings)
146158
?.tokenProvider
147159
?.delegate
148160
?.let { it as? BearerTokenProvider }
149161
?.currentToken()
150162
?.accessToken
151-
?.let { token -> updateTokenCredentials(token, true) }
152-
?: CompletableFuture.failedFuture(IllegalStateException("Unable to get token from connection"))
163+
?: error("Unable to get token from connection")
153164

154-
override fun invalidate(providerId: String) {
155-
deleteTokenCredentials()
156-
}
157-
158-
private fun createUpdateCredentialsPayload(token: String, encrypted: Boolean): UpdateCredentialsPayload =
159-
if (encrypted) {
165+
return if (encrypted) {
160166
UpdateCredentialsPayload(
161167
data = encryptionManager.encrypt(
162168
UpdateCredentialsPayloadData(
163169
BearerCredentials(token)
164170
)
165171
),
172+
metadata = ConnectionMetadata.fromConnection(connection),
166173
encrypted = true
167174
)
168175
} else {
169176
UpdateCredentialsPayload(
170177
data = token,
178+
metadata = ConnectionMetadata.fromConnection(connection),
171179
encrypted = false
172180
)
173181
}
182+
}
174183

175184
override fun onProfileSelected(project: Project, profile: QRegionProfile?) {
176185
updateConfiguration()

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

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,29 @@
33

44
package software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials
55

6+
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
7+
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
8+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspConstants
9+
610
data class ConnectionMetadata(
711
val sso: SsoProfileData,
8-
)
12+
) {
13+
companion object {
14+
fun fromConnection(connection: ToolkitConnection) = when (connection) {
15+
is AwsBearerTokenConnection -> {
16+
ConnectionMetadata(
17+
SsoProfileData(connection.startUrl)
18+
)
19+
}
20+
else -> {
21+
// If no connection or not a bearer token connection return default builderID start url
22+
ConnectionMetadata(
23+
SsoProfileData(AmazonQLspConstants.AWS_BUILDER_ID_URL)
24+
)
25+
}
26+
}
27+
}
28+
}
929

1030
data class SsoProfileData(
1131
val startUrl: String,

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ package software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentia
55

66
data class UpdateCredentialsPayload(
77
val data: String,
8+
val metadata: ConnectionMetadata,
89
val encrypted: Boolean,
910
)
1011

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -331,7 +331,7 @@ class AmazonQLanguageClientImplTest {
331331
every { mockConnectionManager.activeConnectionForFeature(QConnection.getInstance()) } returns null
332332

333333
assertThat(sut.getConnectionMetadata().get())
334-
.isEqualTo(ConnectionMetadata(SsoProfileData(AmazonQLspConstants.AWS_BUILDER_ID_URL)))
334+
.isEqualTo(null)
335335
}
336336

337337
@Test

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

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,8 @@ import software.aws.toolkits.jetbrains.core.credentials.sso.bearer.InteractiveBe
3030
import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLanguageServer
3131
import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
3232
import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager
33+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.ConnectionMetadata
34+
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.SsoProfileData
3335
import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.UpdateCredentialsPayload
3436
import software.aws.toolkits.jetbrains.utils.isQConnected
3537
import software.aws.toolkits.jetbrains.utils.isQExpired
@@ -112,6 +114,7 @@ class DefaultAuthCredentialsServiceTest {
112114
connectionId: String = "test-connection-id",
113115
): AwsBearerTokenConnection = mockk {
114116
every { id } returns connectionId
117+
every { startUrl } returns "startUrl"
115118
every { getConnectionSettings() } returns createMockTokenSettings(accessToken)
116119
}
117120

@@ -192,17 +195,18 @@ class DefaultAuthCredentialsServiceTest {
192195

193196
@Test
194197
fun `test updateTokenCredentials unencrypted success`() {
195-
sut = DefaultAuthCredentialsService(project, mockEncryptionManager, mockk())
196-
197-
val token = "unencryptedToken"
198198
val isEncrypted = false
199+
sut = DefaultAuthCredentialsService(project, mockEncryptionManager, mockk())
199200

200-
sut.updateTokenCredentials(token, isEncrypted)
201+
sut.updateTokenCredentials(mockConnection, isEncrypted)
201202

202203
verify(exactly = 1) {
203204
mockLanguageServer.updateTokenCredentials(
204205
UpdateCredentialsPayload(
205-
token,
206+
"test-access-token",
207+
ConnectionMetadata(
208+
SsoProfileData("startUrl")
209+
),
206210
isEncrypted
207211
)
208212
)
@@ -214,17 +218,19 @@ class DefaultAuthCredentialsServiceTest {
214218
sut = DefaultAuthCredentialsService(project, mockEncryptionManager, mockk())
215219

216220
val encryptedToken = "encryptedToken"
217-
val decryptedToken = "decryptedToken"
218221
val isEncrypted = true
219222

220223
every { mockEncryptionManager.encrypt(any()) } returns encryptedToken
221224

222-
sut.updateTokenCredentials(decryptedToken, isEncrypted)
225+
sut.updateTokenCredentials(mockConnection, isEncrypted)
223226

224227
verify(atLeast = 1) {
225228
mockLanguageServer.updateTokenCredentials(
226229
UpdateCredentialsPayload(
227230
encryptedToken,
231+
ConnectionMetadata(
232+
SsoProfileData("startUrl")
233+
),
228234
isEncrypted
229235
)
230236
)

0 commit comments

Comments
 (0)