Skip to content

Commit da381e9

Browse files
Will-ShaoHuaWill-ShaoHua
committed
Merge remote-tracking branch 'upstream/main' into throttle
2 parents ce9e49a + c146b0f commit da381e9

File tree

16 files changed

+206
-104
lines changed

16 files changed

+206
-104
lines changed

‎.changes/next-release/bugfix-7920ece4-99d6-4828-931e-29b66994dc10.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type" : "bugfix",
3+
"description" : "/dev: Fix missing Amazon Q feature dev auto build setting."
4+
}

‎.changes/next-release/feature-ae7d15f4-ae07-4faf-963a-6d83a85852e4.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type" : "feature",
3+
"description" : "Amazon Q: Support selecting customizations across all Q profiles with automatic profile switching for enterprise users"
4+
}

‎plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonqFeatureDev/controller/FeatureDevController.kt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -544,7 +544,8 @@ class FeatureDevController(
544544
}
545545

546546
private suspend fun handleDevCommandUserSetting(tabId: String, value: Boolean) {
547-
CodeWhispererSettings.getInstance().toggleAutoBuildFeature(context.project.basePath, value)
547+
val session = getSessionInfo(tabId)
548+
CodeWhispererSettings.getInstance().toggleAutoBuildFeature(session.context.workspaceRoot.path, value)
548549
messenger.sendAnswer(
549550
tabId = tabId,
550551
message = message("amazonqFeatureDev.chat_message.setting_updated"),

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ import software.aws.toolkits.core.utils.debug
3939
import software.aws.toolkits.core.utils.getLogger
4040
import software.aws.toolkits.jetbrains.core.credentials.sono.isInternalUser
4141
import software.aws.toolkits.jetbrains.services.amazonq.codeWhispererUserContext
42+
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile
4243
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
4344
import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization
4445
import software.aws.toolkits.jetbrains.services.codewhisperer.language.CodeWhispererProgrammingLanguage
@@ -83,7 +84,7 @@ interface CodeWhispererClientAdaptor {
8384

8485
fun getCodeFixJob(request: GetCodeFixJobRequest): GetCodeFixJobResponse
8586

86-
fun listAvailableCustomizations(): List<CodeWhispererCustomization>
87+
fun listAvailableCustomizations(profile: QRegionProfile): List<CodeWhispererCustomization>
8788

8889
fun startTestGeneration(uploadId: String, targetCode: List<TargetCode>, userInput: String): StartTestGenerationResponse
8990

@@ -287,9 +288,9 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW
287288
override fun getCodeFixJob(request: GetCodeFixJobRequest): GetCodeFixJobResponse = bearerClient().getCodeFixJob(request)
288289

289290
// DO NOT directly use this method to fetch customizations, use wrapper [CodeWhispererModelConfigurator.listCustomization()] instead
290-
override fun listAvailableCustomizations(): List<CodeWhispererCustomization> =
291-
bearerClient().listAvailableCustomizationsPaginator(
292-
ListAvailableCustomizationsRequest.builder().profileArn(QRegionProfileManager.getInstance().activeProfile(project)?.arn).build()
291+
override fun listAvailableCustomizations(profile: QRegionProfile): List<CodeWhispererCustomization> =
292+
QRegionProfileManager.getInstance().getQClient<CodeWhispererRuntimeClient>(project, profile).listAvailableCustomizationsPaginator(
293+
ListAvailableCustomizationsRequest.builder().profileArn(profile.arn).build()
293294
)
294295
.stream()
295296
.toList()
@@ -303,7 +304,8 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW
303304
CodeWhispererCustomization(
304305
arn = it.arn(),
305306
name = it.name(),
306-
description = it.description()
307+
description = it.description(),
308+
profile = profile
307309
)
308310
}
309311
}

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererCustomizationDialog.kt‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ import software.amazon.awssdk.arns.Arn
2626
import software.aws.toolkits.core.utils.debug
2727
import software.aws.toolkits.core.utils.getLogger
2828
import software.aws.toolkits.core.utils.tryOrNull
29+
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile
2930
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.Q_CUSTOM_LEARN_MORE_URI
3031
import software.aws.toolkits.jetbrains.ui.AsyncComboBox
3132
import software.aws.toolkits.jetbrains.utils.notifyInfo
@@ -34,7 +35,7 @@ import javax.swing.JComponent
3435
import javax.swing.JList
3536

3637
private val NoDataToDisplay = CustomizationUiItem(
37-
CodeWhispererCustomization("", message("codewhisperer.custom.dialog.option.no_data"), ""),
38+
CodeWhispererCustomization("", message("codewhisperer.custom.dialog.option.no_data"), "", QRegionProfile("", "")),
3839
false,
3940
false
4041
)
@@ -259,6 +260,10 @@ private object CustomizationRenderer : ColoredListCellRenderer<CustomizationUiIt
259260
}
260261
}
261262

263+
if (it.customization.profile?.profileName?.isNotEmpty() == true) {
264+
append(" [${it.customization.profile?.profileName}]", SimpleTextAttributes.REGULAR_ATTRIBUTES)
265+
}
266+
262267
if (it.isNew) {
263268
append(" New", SimpleTextAttributes.GRAYED_SMALL_ATTRIBUTES)
264269
}

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt‎

Lines changed: 52 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -22,10 +22,11 @@ import software.aws.toolkits.core.utils.debug
2222
import software.aws.toolkits.core.utils.getLogger
2323
import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
2424
import software.aws.toolkits.jetbrains.services.amazonq.calculateIfIamIdentityCenterConnection
25+
import software.aws.toolkits.jetbrains.services.amazonq.profile.QProfileSwitchIntent
2526
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile
27+
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
2628
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener
2729
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
28-
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
2930
import software.aws.toolkits.jetbrains.utils.notifyInfo
3031
import software.aws.toolkits.jetbrains.utils.notifyWarn
3132
import software.aws.toolkits.jetbrains.utils.pluginAwareExecuteOnPooledThread
@@ -106,25 +107,24 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
106107
@RequiresBackgroundThread
107108
override fun listCustomizations(project: Project, passive: Boolean): List<CustomizationUiItem>? =
108109
calculateIfIamIdentityCenterConnection(project) {
109-
// 1. invoke API and get result
110-
val listAvailableCustomizationsResult = try {
111-
CodeWhispererClientAdaptor.getInstance(project).listAvailableCustomizations()
112-
} catch (e: Exception) {
113-
val requestId = (e as? CodeWhispererRuntimeException)?.requestId()
114-
val logMessage = if (CodeWhispererConstants.Customization.noAccessToCustomizationExceptionPredicate(e)) {
115-
// TODO: not required for non GP users
116-
"ListAvailableCustomizations: connection ${it.id} is not allowlisted, requestId: ${requestId.orEmpty()}"
117-
} else {
118-
"ListAvailableCustomizations: failed due to unknown error ${e.message}, requestId: ${requestId.orEmpty()}"
119-
}
120-
121-
LOG.debug { logMessage }
122-
null
110+
// 1. fetch all profiles, invoke fetch customizations API and get result for each profile and aggregate all the results
111+
val profiles = QRegionProfileManager.getInstance().listRegionProfiles(project)
112+
?: error("Attempted to fetch profiles while there does not exist")
113+
114+
val customizations = profiles.flatMap { profile ->
115+
runCatching {
116+
CodeWhispererClientAdaptor.getInstance(project).listAvailableCustomizations(profile)
117+
}.onFailure { e ->
118+
val requestId = (e as? CodeWhispererRuntimeException)?.requestId()
119+
val logMessage = "ListAvailableCustomizations: failed due to unknown error ${e.message}, " +
120+
"requestId: ${requestId.orEmpty()}, profileName: ${profile.profileName}"
121+
LOG.debug { logMessage }
122+
}.getOrDefault(emptyList())
123123
}
124124

125125
// 2. get diff
126126
val previousCustomizationsShapshot = connectionToCustomizationsShownLastTime.getOrElse(it.id) { emptyList() }
127-
val diff = listAvailableCustomizationsResult?.filterNot { customization -> previousCustomizationsShapshot.contains(customization.arn) }?.toSet()
127+
val diff = customizations.filterNot { customization -> previousCustomizationsShapshot.contains(customization.arn) }.toSet()
128128

129129
// 3 if passive,
130130
// (1) update allowlisting
@@ -133,42 +133,45 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
133133
// if not passive,
134134
// (1) update the customization list snapshot (seen by users last time) if it will be displayed
135135
if (passive) {
136-
connectionIdToIsAllowlisted[it.id] = listAvailableCustomizationsResult != null
137-
if (diff?.isNotEmpty() == true && !hasShownNewCustomizationNotification.getAndSet(true)) {
136+
connectionIdToIsAllowlisted[it.id] = customizations.isNotEmpty()
137+
if (diff.isNotEmpty() && !hasShownNewCustomizationNotification.getAndSet(true)) {
138138
notifyNewCustomization(project)
139139
}
140140
} else {
141-
listAvailableCustomizationsResult?.let { customizations ->
142-
connectionToCustomizationsShownLastTime[it.id] = customizations.map { customization -> customization.arn }.toMutableList()
143-
}
141+
connectionToCustomizationsShownLastTime[it.id] = customizations.map { customization -> customization.arn }.toMutableList()
144142
}
145143

146144
// 4. invalidate selected customization if
147145
// (1) the API call failed
148146
// (2) the selected customization is not in the resultset of API call
147+
// (3) the existing q region profile associated with the selected customization does not match the currently active profile
149148
activeCustomization(project)?.let { activeCustom ->
150-
if (listAvailableCustomizationsResult == null) {
149+
if (customizations.isEmpty()) {
151150
invalidateSelectedAndNotify(project)
152-
} else if (!listAvailableCustomizationsResult.any { latestCustom -> latestCustom.arn == activeCustom.arn }) {
151+
} else if (customizations.none { latestCustom -> latestCustom.arn == activeCustom.arn }) {
153152
invalidateSelectedAndNotify(project)
153+
} else {
154+
// for backward compatibility, previous schema didn't have profile arn, so backfill profile here if it's null
155+
if (activeCustom.profile == null) {
156+
customizations.find { c -> c.arn == activeCustom.arn }?.profile?.let { p ->
157+
activeCustom.profile = p
158+
}
159+
}
160+
161+
if (activeCustom.profile != null && activeCustom.profile != QRegionProfileManager.getInstance().activeProfile(project)) {
162+
invalidateSelectedAndNotify(project)
163+
}
154164
}
155165
}
156166

157167
// 5. transform result to UI items and return
158-
val customizationUiItems = if (diff != null) {
159-
listAvailableCustomizationsResult.let { customizations ->
160-
val nameToCount = customizations.groupingBy { customization -> customization.name }.eachCount()
161-
162-
customizations.map { customization ->
163-
CustomizationUiItem(
164-
customization,
165-
isNew = diff.contains(customization),
166-
shouldPrefixAccountId = (nameToCount[customization.name] ?: 0) > 1
167-
)
168-
}
169-
}
170-
} else {
171-
null
168+
val nameToCount = customizations.groupingBy { customization -> customization.name }.eachCount()
169+
val customizationUiItems = customizations.map { customization ->
170+
CustomizationUiItem(
171+
customization,
172+
isNew = diff.contains(customization),
173+
shouldPrefixAccountId = (nameToCount[customization.name] ?: 0) > 1
174+
)
172175
}
173176
connectionToCustomizationUiItems[it.id] = customizationUiItems
174177

@@ -210,6 +213,18 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
210213

211214
LOG.debug { "Switch from customization $oldCus to $newCustomization" }
212215

216+
// Switch profile if it doesn't match the customization's profile.
217+
// Customizations are profile-scoped and must be used under the correct context.
218+
newCustomization?.profile?.let { p ->
219+
if (p.arn != QRegionProfileManager.getInstance().activeProfile(project)?.arn) {
220+
QRegionProfileManager.getInstance().switchProfile(
221+
project,
222+
p,
223+
QProfileSwitchIntent.Customization
224+
)
225+
}
226+
}
227+
213228
CodeWhispererCustomizationListener.notifyCustomUiUpdate()
214229
}
215230
if (isOverride) {

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/startup/CodeWhispererProjectStartupActivity.kt‎

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@ import software.aws.toolkits.jetbrains.core.coroutines.projectCoroutineScope
1313
import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
1414
import software.aws.toolkits.jetbrains.services.amazonq.calculateIfIamIdentityCenterConnection
1515
import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanManager
16-
import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization
1716
import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator
1817
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
1918
import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.isUserBuilderId
@@ -79,18 +78,16 @@ class CodeWhispererProjectStartupActivity : StartupActivity.DumbAware {
7978
projectCoroutineScope(project).launch {
8079
while (isActive) {
8180
CodeWhispererFeatureConfigService.getInstance().fetchFeatureConfigs(project)
82-
CodeWhispererFeatureConfigService.getInstance().getCustomizationFeature()?.let { customization ->
81+
CodeWhispererFeatureConfigService.getInstance().getCustomizationFeature()?.let { overrideContext ->
8382
val persistedCustomizationOverride = CodeWhispererModelConfigurator.getInstance().getPersistedCustomizationOverride()
84-
val latestCustomizationOverride = customization.value.stringValue()
83+
val latestCustomizationOverride = overrideContext.value.stringValue()
8584
if (persistedCustomizationOverride == latestCustomizationOverride) return@let
8685

8786
// latest is different from the currently persisted, need update
88-
CodeWhispererFeatureConfigService.getInstance().validateCustomizationOverride(project, customization)
89-
CodeWhispererModelConfigurator.getInstance().switchCustomization(
90-
project,
91-
CodeWhispererCustomization(arn = customization.value.stringValue(), name = customization.variation),
92-
isOverride = true
93-
)
87+
val customization = CodeWhispererFeatureConfigService.getInstance().validateCustomizationOverride(project, overrideContext)
88+
if (customization != null) {
89+
CodeWhispererModelConfigurator.getInstance().switchCustomization(project, customization, isOverride = true)
90+
}
9491
}
9592

9693
delay(FEATURE_CONFIG_POLL_INTERVAL_IN_MS)

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

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@ import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
6565
import software.aws.toolkits.jetbrains.core.credentials.sono.Q_SCOPES
6666
import software.aws.toolkits.jetbrains.core.credentials.sono.SONO_REGION
6767
import software.aws.toolkits.jetbrains.services.amazonq.FEATURE_EVALUATION_PRODUCT_NAME
68+
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile
6869
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.metadata
6970
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonRequest
7071
import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonResponseWithToken
@@ -189,13 +190,13 @@ class CodeWhispererClientAdaptorTest {
189190
on { client.listAvailableCustomizationsPaginator(any<ListAvailableCustomizationsRequest>()) } doReturn sdkIterable
190191
}
191192

192-
val actual = sut.listAvailableCustomizations()
193+
val actual = sut.listAvailableCustomizations(QRegionProfile("fake_profile", "fake arn"))
193194
assertThat(actual).hasSize(3)
194195
assertThat(actual).isEqualTo(
195196
listOf(
196-
CodeWhispererCustomization(name = "custom-1", arn = "arn-1"),
197-
CodeWhispererCustomization(name = "custom-2", arn = "arn-2"),
198-
CodeWhispererCustomization(name = "custom-3", arn = "arn-3")
197+
CodeWhispererCustomization(name = "custom-1", arn = "arn-1", profile = QRegionProfile("fake_profile", "fake arn")),
198+
CodeWhispererCustomization(name = "custom-2", arn = "arn-2", profile = QRegionProfile("fake_profile", "fake arn")),
199+
CodeWhispererCustomization(name = "custom-3", arn = "arn-3", profile = QRegionProfile("fake_profile", "fake arn"))
199200
)
200201
)
201202
}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
3535
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
3636
import software.aws.toolkits.jetbrains.core.credentials.sono.SONO_URL
3737
import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
38+
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile
3839
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
3940
import kotlin.reflect.full.memberFunctions
4041
import kotlin.test.Test
@@ -78,7 +79,7 @@ class CodeWhispererFeatureConfigServiceTest {
7879

7980
projectRule.project.replaceService(
8081
QRegionProfileManager::class.java,
81-
mock<QRegionProfileManager> { on { getQClient(any<Project>(), eq(CodeWhispererRuntimeClient::class)) } doReturn mockClient },
82+
mock<QRegionProfileManager> { on { getQClient(any<Project>(), eq(QRegionProfile()), eq(CodeWhispererRuntimeClient::class)) } doReturn mockClient },
8283
disposableRule.disposable
8384
)
8485

0 commit comments

Comments
 (0)