From 91052432b7d76d1292b20a1deff0263814b200e2 Mon Sep 17 00:00:00 2001 From: YIFAN LIU Date: Wed, 9 Apr 2025 10:16:56 -0700 Subject: [PATCH 1/4] call if --- .../services/amazonq/QLoginWebview.kt | 38 ++++++++++--------- .../amazonq/profile/QRegionProfileManager.kt | 2 +- 2 files changed, 22 insertions(+), 18 deletions(-) diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/QLoginWebview.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/QLoginWebview.kt index ceda1545ea1..e17b0ea5d50 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/QLoginWebview.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/QLoginWebview.kt @@ -264,23 +264,6 @@ class QWebviewBrowser(val project: Project, private val parentDisposable: Dispos // TODO: pass "REAUTH" if connection expires // Perform the potentially blocking AWS call outside the EDT to fetch available region profiles. ApplicationManager.getApplication().executeOnPooledThread { - var errorMessage: String? = null - val profiles: List = try { - QRegionProfileManager.getInstance().listRegionProfiles(project).orEmpty() - } catch (e: Exception) { - errorMessage = e.message - LOG.warn { "Failed to call listRegionProfiles API" } - val qConn = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance()) - Telemetry.amazonq.didSelectProfile.use { span -> - span.source(QProfileSwitchIntent.Auth.value) - .amazonQProfileRegion(QRegionProfileManager.getInstance().activeProfile(project)?.region ?: "not-set") - .ssoRegion((qConn as? AwsBearerTokenConnection)?.region) - .credentialStartUrl((qConn as? AwsBearerTokenConnection)?.startUrl) - .result(MetricResult.Failed) - .reason(e.message) - } - emptyList() - } val stage = if (isQExpired(project)) { "REAUTH" @@ -290,6 +273,27 @@ class QWebviewBrowser(val project: Project, private val parentDisposable: Dispos "START" } + var errorMessage: String? = null + var profiles: List = emptyList() + + if (stage == "PROFILE_SELECT") { + try { + profiles = QRegionProfileManager.getInstance().listRegionProfiles(project).orEmpty() + } catch (e: Exception) { + errorMessage = e.message + LOG.warn { "Failed to call listRegionProfiles API" } + val qConn = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance()) + Telemetry.amazonq.didSelectProfile.use { span -> + span.source(QProfileSwitchIntent.Auth.value) + .amazonQProfileRegion(QRegionProfileManager.getInstance().activeProfile(project)?.region ?: "not-set") + .ssoRegion((qConn as? AwsBearerTokenConnection)?.region) + .credentialStartUrl((qConn as? AwsBearerTokenConnection)?.startUrl) + .result(MetricResult.Failed) + .reason(e.message) + } + } + } + val jsonData = """ { stage: '$stage', diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileManager.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileManager.kt index a835a145a24..6b2659c198e 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileManager.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileManager.kt @@ -74,7 +74,7 @@ class QRegionProfileManager : PersistentStateComponent, Disposabl .awsClient() .listAvailableProfilesPaginator {} .profiles() - .map { p -> QRegionProfile(arn = p.arn(), profileName = p.profileName()) } + .map { p -> QRegionProfile(arn = p.arn(), profileName = p.profileName()?: "") } } if (mappedProfiles.size == 1) { switchProfile(project, mappedProfiles.first(), intent = QProfileSwitchIntent.Update) From ae016bd9ccde0bba99be2cbc1cb601c89f6a75f6 Mon Sep 17 00:00:00 2001 From: YIFAN LIU Date: Wed, 9 Apr 2025 10:41:33 -0700 Subject: [PATCH 2/4] add cache --- .../amazonq/profile/QProfileResources.kt | 39 +++++++++++++++++++ .../amazonq/profile/QRegionProfileManager.kt | 22 +++++------ 2 files changed, 49 insertions(+), 12 deletions(-) create mode 100644 plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QProfileResources.kt diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QProfileResources.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QProfileResources.kt new file mode 100644 index 00000000000..2421515c9bd --- /dev/null +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QProfileResources.kt @@ -0,0 +1,39 @@ +// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 + +package software.aws.toolkits.jetbrains.services.amazonq.profile + +import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient +import software.aws.toolkits.core.ClientConnectionSettings +import software.aws.toolkits.jetbrains.core.AwsClientManager +import software.aws.toolkits.jetbrains.core.region.AwsRegionProvider +import software.aws.toolkits.jetbrains.core.Resource +import java.time.Duration + +/** + * Save Amazon Q Profile Resource Cache + */ +object QProfileResources { + /** + * save available Q Profile list as cache with default duration 60 s。 + */ + val LIST_REGION_PROFILES = object : Resource.Cached>() { + override val id: String = "amazonq.allProfiles" + + override fun fetch(connectionSettings: ClientConnectionSettings<*>): List { + val mappedProfiles = QEndpoints.listRegionEndpoints().flatMap { (regionKey, _) -> + val awsRegion = AwsRegionProvider.getInstance()[regionKey] ?: return@flatMap emptyList() + val client = AwsClientManager + .getInstance() + .getClient(CodeWhispererRuntimeClient::class, connectionSettings.withRegion(awsRegion)) + + client.listAvailableProfilesPaginator {} + .profiles() + .map { p -> QRegionProfile(arn = p.arn(), profileName = p.profileName() ?: "") } + } + return mappedProfiles + } + + override fun expiry(): Duration = Duration.ofSeconds(60) + } +} diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileManager.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileManager.kt index 6b2659c198e..2d345cc6df0 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileManager.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QRegionProfileManager.kt @@ -14,14 +14,13 @@ import com.intellij.openapi.project.Project import com.intellij.util.xmlb.annotations.MapAnnotation import com.intellij.util.xmlb.annotations.Property import software.amazon.awssdk.core.SdkClient -import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient import software.aws.toolkits.core.TokenConnectionSettings import software.aws.toolkits.core.utils.debug import software.aws.toolkits.core.utils.getLogger import software.aws.toolkits.core.utils.tryOrNull import software.aws.toolkits.core.utils.warn import software.aws.toolkits.jetbrains.core.AwsClientManager -import software.aws.toolkits.jetbrains.core.awsClient +import software.aws.toolkits.jetbrains.core.AwsResourceCache import software.aws.toolkits.jetbrains.core.credentials.AwsBearerTokenConnection import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection @@ -31,6 +30,7 @@ import software.aws.toolkits.jetbrains.utils.notifyInfo import software.aws.toolkits.resources.AmazonQBundle.message import software.aws.toolkits.telemetry.MetricResult import software.aws.toolkits.telemetry.Telemetry +import java.time.Duration import java.util.Collections import kotlin.reflect.KClass @@ -66,16 +66,14 @@ class QRegionProfileManager : PersistentStateComponent, Disposabl fun listRegionProfiles(project: Project): List? { val connection = getIdcConnectionOrNull(project) ?: return null return try { - val mappedProfiles = QEndpoints.listRegionEndpoints() - .flatMap { (regionKey, _) -> - val awsRegion = AwsRegionProvider.getInstance()[regionKey] ?: return@flatMap emptyList() - connection.getConnectionSettings() - .withRegion(awsRegion) - .awsClient() - .listAvailableProfilesPaginator {} - .profiles() - .map { p -> QRegionProfile(arn = p.arn(), profileName = p.profileName()?: "") } - } + val connectionSettings = connection.getConnectionSettings() + val mappedProfiles = AwsResourceCache.getInstance().getResourceNow( + resource = QProfileResources.LIST_REGION_PROFILES, + connectionSettings = connectionSettings, + timeout = Duration.ofSeconds(30), + useStale = true, + forceFetch = false + ) if (mappedProfiles.size == 1) { switchProfile(project, mappedProfiles.first(), intent = QProfileSwitchIntent.Update) } From c1474fe789f65366e38f9128aab49de5f42a5331 Mon Sep 17 00:00:00 2001 From: YIFAN LIU Date: Wed, 9 Apr 2025 11:11:16 -0700 Subject: [PATCH 3/4] linter --- .../aws/toolkits/jetbrains/services/amazonq/QLoginWebview.kt | 1 - .../jetbrains/services/amazonq/profile/QProfileResources.kt | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/QLoginWebview.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/QLoginWebview.kt index e17b0ea5d50..bb56a5e06d1 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/QLoginWebview.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/QLoginWebview.kt @@ -264,7 +264,6 @@ class QWebviewBrowser(val project: Project, private val parentDisposable: Dispos // TODO: pass "REAUTH" if connection expires // Perform the potentially blocking AWS call outside the EDT to fetch available region profiles. ApplicationManager.getApplication().executeOnPooledThread { - val stage = if (isQExpired(project)) { "REAUTH" } else if (isQConnected(project) && QRegionProfileManager.getInstance().isPendingProfileSelection(project)) { diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QProfileResources.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QProfileResources.kt index 2421515c9bd..79cbdb38a27 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QProfileResources.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/profile/QProfileResources.kt @@ -6,8 +6,8 @@ package software.aws.toolkits.jetbrains.services.amazonq.profile import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient import software.aws.toolkits.core.ClientConnectionSettings import software.aws.toolkits.jetbrains.core.AwsClientManager -import software.aws.toolkits.jetbrains.core.region.AwsRegionProvider import software.aws.toolkits.jetbrains.core.Resource +import software.aws.toolkits.jetbrains.core.region.AwsRegionProvider import java.time.Duration /** From 07e63ca62bd127138ab006735980d806c3423cf6 Mon Sep 17 00:00:00 2001 From: YIFAN LIU Date: Wed, 9 Apr 2025 12:54:34 -0700 Subject: [PATCH 4/4] delete ut --- .../QRegionProfileManagerTest.kt | 31 +------------------ 1 file changed, 1 insertion(+), 30 deletions(-) diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/QRegionProfileManagerTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/QRegionProfileManagerTest.kt index 511e4e5272b..0453cf3a766 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/QRegionProfileManagerTest.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/QRegionProfileManagerTest.kt @@ -124,36 +124,7 @@ class QRegionProfileManagerTest { assertThat(cnt).isEqualTo(2) } - @Test - fun `listProfiles will call each client to get profiles`() { - val client = clientRule.create() - val mockResponse: SdkIterable = SdkIterable { - listOf( - Profile.builder().profileName("FOO").arn("foo").build(), - ).toMutableList().iterator() - } - - val mockResponse2: SdkIterable = SdkIterable { - listOf( - Profile.builder().profileName("BAR").arn("bar").build(), - ).toMutableList().iterator() - } - - val iterable: ListAvailableProfilesIterable = mock { - on { it.profiles() } doReturn mockResponse doReturn mockResponse2 - } - - // TODO: not sure if we can mock client with different region different response? - client.stub { - onGeneric { listAvailableProfilesPaginator(any>()) } doReturn iterable - } - - val r = sut.listRegionProfiles(project) - assertThat(r).hasSize(2) - - assertThat(r).contains(QRegionProfile("FOO", "foo")) - assertThat(r).contains(QRegionProfile("BAR", "bar")) - } + // TODO: Add two unit tests for listProfiles — one with cache hit, one without @Test fun `validateProfile should cross validate selected profile with latest API response for current project and remove it if its not longer accessible`() {