Skip to content

Commit 736a831

Browse files
evanliu048evanliu048
authored
fix(amazonq): cache the listAvailableProfiles result (#5544)
Avoid calling ListAvailableProfiles at a higher frequency by: -Calling this API only when it's in the ProfileSelection stage -Using AwsResourceCache to cache the api result and retrieve it after the cache expired.
1 parent f5d161e commit 736a831

File tree

4 files changed

+71
-60
lines changed

4 files changed

+71
-60
lines changed

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/QLoginWebview.kt

Lines changed: 21 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -264,24 +264,6 @@ class QWebviewBrowser(val project: Project, private val parentDisposable: Dispos
264264
// TODO: pass "REAUTH" if connection expires
265265
// Perform the potentially blocking AWS call outside the EDT to fetch available region profiles.
266266
ApplicationManager.getApplication().executeOnPooledThread {
267-
var errorMessage: String? = null
268-
val profiles: List<QRegionProfile> = try {
269-
QRegionProfileManager.getInstance().listRegionProfiles(project).orEmpty()
270-
} catch (e: Exception) {
271-
errorMessage = e.message
272-
LOG.warn { "Failed to call listRegionProfiles API" }
273-
val qConn = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())
274-
Telemetry.amazonq.didSelectProfile.use { span ->
275-
span.source(QProfileSwitchIntent.Auth.value)
276-
.amazonQProfileRegion(QRegionProfileManager.getInstance().activeProfile(project)?.region ?: "not-set")
277-
.ssoRegion((qConn as? AwsBearerTokenConnection)?.region)
278-
.credentialStartUrl((qConn as? AwsBearerTokenConnection)?.startUrl)
279-
.result(MetricResult.Failed)
280-
.reason(e.message)
281-
}
282-
emptyList()
283-
}
284-
285267
val stage = if (isQExpired(project)) {
286268
"REAUTH"
287269
} else if (isQConnected(project) && QRegionProfileManager.getInstance().isPendingProfileSelection(project)) {
@@ -290,6 +272,27 @@ class QWebviewBrowser(val project: Project, private val parentDisposable: Dispos
290272
"START"
291273
}
292274

275+
var errorMessage: String? = null
276+
var profiles: List<QRegionProfile> = emptyList()
277+
278+
if (stage == "PROFILE_SELECT") {
279+
try {
280+
profiles = QRegionProfileManager.getInstance().listRegionProfiles(project).orEmpty()
281+
} catch (e: Exception) {
282+
errorMessage = e.message
283+
LOG.warn { "Failed to call listRegionProfiles API" }
284+
val qConn = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())
285+
Telemetry.amazonq.didSelectProfile.use { span ->
286+
span.source(QProfileSwitchIntent.Auth.value)
287+
.amazonQProfileRegion(QRegionProfileManager.getInstance().activeProfile(project)?.region ?: "not-set")
288+
.ssoRegion((qConn as? AwsBearerTokenConnection)?.region)
289+
.credentialStartUrl((qConn as? AwsBearerTokenConnection)?.startUrl)
290+
.result(MetricResult.Failed)
291+
.reason(e.message)
292+
}
293+
}
294+
}
295+
293296
val jsonData = """
294297
{
295298
stage: '$stage',

‎plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/QRegionProfileManagerTest.kt

Lines changed: 1 addition & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -124,36 +124,7 @@ class QRegionProfileManagerTest {
124124
assertThat(cnt).isEqualTo(2)
125125
}
126126

127-
@Test
128-
fun `listProfiles will call each client to get profiles`() {
129-
val client = clientRule.create<CodeWhispererRuntimeClient>()
130-
val mockResponse: SdkIterable<Profile> = SdkIterable<Profile> {
131-
listOf(
132-
Profile.builder().profileName("FOO").arn("foo").build(),
133-
).toMutableList().iterator()
134-
}
135-
136-
val mockResponse2: SdkIterable<Profile> = SdkIterable<Profile> {
137-
listOf(
138-
Profile.builder().profileName("BAR").arn("bar").build(),
139-
).toMutableList().iterator()
140-
}
141-
142-
val iterable: ListAvailableProfilesIterable = mock {
143-
on { it.profiles() } doReturn mockResponse doReturn mockResponse2
144-
}
145-
146-
// TODO: not sure if we can mock client with different region different response?
147-
client.stub {
148-
onGeneric { listAvailableProfilesPaginator(any<Consumer<ListAvailableProfilesRequest.Builder>>()) } doReturn iterable
149-
}
150-
151-
val r = sut.listRegionProfiles(project)
152-
assertThat(r).hasSize(2)
153-
154-
assertThat(r).contains(QRegionProfile("FOO", "foo"))
155-
assertThat(r).contains(QRegionProfile("BAR", "bar"))
156-
}
127+
// TODO: Add two unit tests for listProfiles — one with cache hit, one without
157128

158129
@Test
159130
fun `validateProfile should cross validate selected profile with latest API response for current project and remove it if its not longer accessible`() {

‎plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QProfileResources.kt

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package software.aws.toolkits.jetbrains.services.amazonq.profile
5+
6+
import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient
7+
import software.aws.toolkits.core.ClientConnectionSettings
8+
import software.aws.toolkits.jetbrains.core.AwsClientManager
9+
import software.aws.toolkits.jetbrains.core.Resource
10+
import software.aws.toolkits.jetbrains.core.region.AwsRegionProvider
11+
import java.time.Duration
12+
13+
/**
14+
* Save Amazon Q Profile Resource Cache
15+
*/
16+
object QProfileResources {
17+
/**
18+
* save available Q Profile list as cache with default duration 60 s。
19+
*/
20+
val LIST_REGION_PROFILES = object : Resource.Cached<List<QRegionProfile>>() {
21+
override val id: String = "amazonq.allProfiles"
22+
23+
override fun fetch(connectionSettings: ClientConnectionSettings<*>): List<QRegionProfile> {
24+
val mappedProfiles = QEndpoints.listRegionEndpoints().flatMap { (regionKey, _) ->
25+
val awsRegion = AwsRegionProvider.getInstance()[regionKey] ?: return@flatMap emptyList()
26+
val client = AwsClientManager
27+
.getInstance()
28+
.getClient(CodeWhispererRuntimeClient::class, connectionSettings.withRegion(awsRegion))
29+
30+
client.listAvailableProfilesPaginator {}
31+
.profiles()
32+
.map { p -> QRegionProfile(arn = p.arn(), profileName = p.profileName() ?: "<no name>") }
33+
}
34+
return mappedProfiles
35+
}
36+
37+
override fun expiry(): Duration = Duration.ofSeconds(60)
38+
}
39+
}

‎plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileManager.kt

Lines changed: 10 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -14,14 +14,13 @@ import com.intellij.openapi.project.Project
1414
import com.intellij.util.xmlb.annotations.MapAnnotation
1515
import com.intellij.util.xmlb.annotations.Property
1616
import software.amazon.awssdk.core.SdkClient
17-
import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient
1817
import software.aws.toolkits.core.TokenConnectionSettings
1918
import software.aws.toolkits.core.utils.debug
2019
import software.aws.toolkits.core.utils.getLogger
2120
import software.aws.toolkits.core.utils.tryOrNull
2221
import software.aws.toolkits.core.utils.warn
2322
import software.aws.toolkits.jetbrains.core.AwsClientManager
24-
import software.aws.toolkits.jetbrains.core.awsClient
23+
import software.aws.toolkits.jetbrains.core.AwsResourceCache
2524
import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection
2625
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
2726
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
@@ -31,6 +30,7 @@ import software.aws.toolkits.jetbrains.utils.notifyInfo
3130
import software.aws.toolkits.resources.AmazonQBundle.message
3231
import software.aws.toolkits.telemetry.MetricResult
3332
import software.aws.toolkits.telemetry.Telemetry
33+
import java.time.Duration
3434
import java.util.Collections
3535
import kotlin.reflect.KClass
3636

@@ -66,16 +66,14 @@ class QRegionProfileManager : PersistentStateComponent<QProfileState>, Disposabl
6666
fun listRegionProfiles(project: Project): List<QRegionProfile>? {
6767
val connection = getIdcConnectionOrNull(project) ?: return null
6868
return try {
69-
val mappedProfiles = QEndpoints.listRegionEndpoints()
70-
.flatMap { (regionKey, _) ->
71-
val awsRegion = AwsRegionProvider.getInstance()[regionKey] ?: return@flatMap emptyList()
72-
connection.getConnectionSettings()
73-
.withRegion(awsRegion)
74-
.awsClient<CodeWhispererRuntimeClient>()
75-
.listAvailableProfilesPaginator {}
76-
.profiles()
77-
.map { p -> QRegionProfile(arn = p.arn(), profileName = p.profileName() ?: "<no name>") }
78-
}
69+
val connectionSettings = connection.getConnectionSettings()
70+
val mappedProfiles = AwsResourceCache.getInstance().getResourceNow(
71+
resource = QProfileResources.LIST_REGION_PROFILES,
72+
connectionSettings = connectionSettings,
73+
timeout = Duration.ofSeconds(30),
74+
useStale = true,
75+
forceFetch = false
76+
)
7977
if (mappedProfiles.size == 1) {
8078
switchProfile(project, mappedProfiles.first(), intent = QProfileSwitchIntent.Update)
8179
}

0 commit comments

Comments
 (0)