Skip to content

Commit a9de9d6

Browse files
evanliu048evanliu048
committed
list all cus for all profiles and sort
1 parent da5315f commit a9de9d6

File tree

3 files changed

+86
-40
lines changed

3 files changed

+86
-40
lines changed

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

Lines changed: 21 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
1+
// Copyright 2023 Amazon.com, Inc. or it s affiliates. All Rights Reserved.
22
// SPDX-License-Identifier: Apache-2.0
33

44
package software.aws.toolkits.jetbrains.services.codewhisperer.customization
@@ -26,6 +26,8 @@ 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
30+
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
2931
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.Q_CUSTOM_LEARN_MORE_URI
3032
import software.aws.toolkits.jetbrains.ui.AsyncComboBox
3133
import software.aws.toolkits.jetbrains.utils.notifyInfo
@@ -34,7 +36,7 @@ import javax.swing.JComponent
3436
import javax.swing.JList
3537

3638
private val NoDataToDisplay = CustomizationUiItem(
37-
CodeWhispererCustomization("", message("codewhisperer.custom.dialog.option.no_data"), ""),
39+
CodeWhispererCustomization("", message("codewhisperer.custom.dialog.option.no_data"), "", QRegionProfile("","")),
3840
false,
3941
false
4042
)
@@ -179,11 +181,19 @@ class CodeWhispererCustomizationDialog(
179181
proposeModelUpdate { model ->
180182
val activeCustomization = CodeWhispererModelConfigurator.getInstance().activeCustomization(project)
181183
val unsorted = myCustomizations ?: CodeWhispererModelConfigurator.getInstance().listCustomizations(project).orEmpty()
182-
183-
val sorted = activeCustomization?.let {
184-
unsorted.putPickedUpFront(setOf(it))
185-
} ?: run {
186-
unsorted.sortedBy { it.customization.name }
184+
val activeProfile = QRegionProfileManager.getInstance().activeProfile(project)
185+
// sort customization based on profile name first
186+
val baseSorted = unsorted.sortedWith(
187+
compareBy<CustomizationUiItem> {
188+
it.customization.profile?.profileName != activeProfile?.profileName
189+
}.thenBy { it.customization.profile?.profileName ?: "" }
190+
.thenBy { it.customization.name }
191+
)
192+
193+
val sorted = if (activeCustomization != null) {
194+
baseSorted.putPickedUpFront(setOf(activeCustomization))
195+
} else {
196+
baseSorted
187197
}
188198

189199
if (
@@ -259,6 +269,10 @@ private object CustomizationRenderer : ColoredListCellRenderer<CustomizationUiIt
259269
}
260270
}
261271

272+
if (it.customization.profile?.profileName?.isNotEmpty() == true) {
273+
append(" [${it.customization.profile?.profileName}]", SimpleTextAttributes.REGULAR_ATTRIBUTES)
274+
}
275+
262276
if (it.isNew) {
263277
append(" New", SimpleTextAttributes.GRAYED_SMALL_ATTRIBUTES)
264278
}

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

Lines changed: 61 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -17,12 +17,20 @@ import com.intellij.openapi.project.Project
1717
import com.intellij.util.concurrency.annotations.RequiresBackgroundThread
1818
import com.intellij.util.xmlb.annotations.MapAnnotation
1919
import com.intellij.util.xmlb.annotations.Property
20+
import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient
2021
import software.amazon.awssdk.services.codewhispererruntime.model.CodeWhispererRuntimeException
22+
import software.amazon.awssdk.services.codewhispererruntime.model.ListAvailableCustomizationsRequest
2123
import software.aws.toolkits.core.utils.debug
2224
import software.aws.toolkits.core.utils.getLogger
25+
import software.aws.toolkits.core.utils.tryOrNull
26+
import software.aws.toolkits.jetbrains.core.AwsClientManager
27+
import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
28+
import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
29+
import software.aws.toolkits.jetbrains.core.region.AwsRegionProvider
2330
import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
2431
import software.aws.toolkits.jetbrains.services.amazonq.calculateIfIamIdentityCenterConnection
2532
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile
33+
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
2634
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileSelectedListener
2735
import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
2836
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
@@ -71,6 +79,7 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
7179
private val connectionIdToActiveCustomizationArn = Collections.synchronizedMap<String, CodeWhispererCustomization>(mutableMapOf())
7280

7381
// Map to store connectionId to its listAvailableCustomizations result last time
82+
// the customization has format profileArn::customizationArn
7483
private val connectionToCustomizationsShownLastTime = mutableMapOf<String, MutableList<String>>()
7584

7685
private val connectionIdToIsAllowlisted = Collections.synchronizedMap<String, Boolean>(mutableMapOf())
@@ -107,70 +116,89 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
107116

108117
@RequiresBackgroundThread
109118
override fun listCustomizations(project: Project, passive: Boolean): List<CustomizationUiItem>? =
110-
calculateIfIamIdentityCenterConnection(project) {
119+
calculateIfIamIdentityCenterConnection(project) { it ->
111120
// 1. invoke API and get result
112-
val listAvailableCustomizationsResult = try {
113-
CodeWhispererClientAdaptor.getInstance(project).listAvailableCustomizations()
114-
} catch (e: Exception) {
115-
val requestId = (e as? CodeWhispererRuntimeException)?.requestId()
116-
val logMessage = if (CodeWhispererConstants.Customization.noAccessToCustomizationExceptionPredicate(e)) {
117-
// TODO: not required for non GP users
118-
"ListAvailableCustomizations: connection ${it.id} is not allowlisted, requestId: ${requestId.orEmpty()}"
119-
} else {
120-
"ListAvailableCustomizations: failed due to unknown error ${e.message}, requestId: ${requestId.orEmpty()}"
121-
}
121+
val listAvailableProfilesResult = tryOrNull { QRegionProfileManager.getInstance().listRegionProfiles(project) } ?: emptyList()
122+
123+
val aggregatedCustomizations = mutableListOf<CodeWhispererCustomization>()
122124

123-
LOG.debug { logMessage }
124-
null
125+
for (profile in listAvailableProfilesResult) {
126+
try {
127+
val setting = AwsRegionProvider.getInstance()[profile.region]?.let { it1 ->
128+
ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())?.getConnectionSettings()
129+
?.withRegion(it1)
130+
}
131+
val client = setting?.let { it1 -> AwsClientManager.getInstance().getClient(CodeWhispererRuntimeClient::class, it1) }
132+
val perProfileCustomizations =
133+
client?.listAvailableCustomizationsPaginator(ListAvailableCustomizationsRequest.builder().profileArn(profile.arn).build())
134+
?.customizations()?.stream()?.toList()?.map {
135+
CodeWhispererCustomization(
136+
arn = it.arn(),
137+
name = it.name(),
138+
description = it.description(),
139+
profile = profile
140+
)
141+
}
142+
if (perProfileCustomizations != null) {
143+
aggregatedCustomizations.addAll(perProfileCustomizations)
144+
}
145+
} catch (e: Exception) {
146+
val requestId = (e as? CodeWhispererRuntimeException)?.requestId()
147+
val logMessage = if (CodeWhispererConstants.Customization.noAccessToCustomizationExceptionPredicate(e)) {
148+
// TODO: not required for non GP users
149+
"ListAvailableCustomizations: connection ${it.id} is not allowlisted, requestId: ${requestId.orEmpty()}"
150+
} else {
151+
"ListAvailableCustomizations: failed due to unknown error ${e.message}, requestId: ${requestId.orEmpty()}"
152+
}
153+
154+
LOG.debug { logMessage }
155+
}
125156
}
126157

127158
// 2. get diff
128159
val previousCustomizationsShapshot = connectionToCustomizationsShownLastTime.getOrElse(it.id) { emptyList() }
129-
val diff = listAvailableCustomizationsResult?.filterNot { customization -> previousCustomizationsShapshot.contains(customization.arn) }?.toSet()
130-
160+
// calculate new profile+customization list using profile.arn :: customization.arn as the key
161+
val diff = aggregatedCustomizations.filterNot { customization -> previousCustomizationsShapshot.contains("${customization.profile?.arn}::${customization.arn}") }.toSet()
131162
// 3 if passive,
132163
// (1) update allowlisting
133164
// (2) prompt "You have New Customizations" toast notification (only show once)
134165
//
135166
// if not passive,
136167
// (1) update the customization list snapshot (seen by users last time) if it will be displayed
137168
if (passive) {
138-
connectionIdToIsAllowlisted[it.id] = listAvailableCustomizationsResult != null
169+
connectionIdToIsAllowlisted[it.id] = aggregatedCustomizations.isNotEmpty()
139170
if (diff?.isNotEmpty() == true && !hasShownNewCustomizationNotification.getAndSet(true)) {
140171
notifyNewCustomization(project)
141172
}
142173
} else {
143-
listAvailableCustomizationsResult?.let { customizations ->
144-
connectionToCustomizationsShownLastTime[it.id] = customizations.map { customization -> customization.arn }.toMutableList()
174+
aggregatedCustomizations.let { customizations ->
175+
connectionToCustomizationsShownLastTime[it.id] = customizations.map { customization -> "${customization.profile?.arn}::${customization.arn}"}.toMutableList()
145176
}
146177
}
147178

148179
// 4. invalidate selected customization if
149180
// (1) the API call failed
150181
// (2) the selected customization is not in the resultset of API call
151182
activeCustomization(project)?.let { activeCustom ->
152-
if (listAvailableCustomizationsResult == null) {
183+
if (aggregatedCustomizations.isEmpty()) {
153184
invalidateSelectedAndNotify(project)
154-
} else if (!listAvailableCustomizationsResult.any { latestCustom -> latestCustom.arn == activeCustom.arn }) {
185+
} else if (!aggregatedCustomizations.any { latestCustom -> "${latestCustom.profile?.arn}::${latestCustom.arn}" == "${activeCustom.profile?.arn}::${activeCustom.arn}"}
186+
|| activeCustom.profile!= QRegionProfileManager.getInstance().activeProfile(project)){
155187
invalidateSelectedAndNotify(project)
156188
}
157189
}
158190

159191
// 5. transform result to UI items and return
160-
val customizationUiItems = if (diff != null) {
161-
listAvailableCustomizationsResult.let { customizations ->
162-
val nameToCount = customizations.groupingBy { customization -> customization.name }.eachCount()
163-
164-
customizations.map { customization ->
165-
CustomizationUiItem(
166-
customization,
167-
isNew = diff.contains(customization),
168-
shouldPrefixAccountId = (nameToCount[customization.name] ?: 0) > 1
169-
)
170-
}
192+
val customizationUiItems = aggregatedCustomizations.let { customizations ->
193+
val nameToCount = customizations.groupingBy { customization -> customization.name }.eachCount()
194+
195+
customizations.map { customization ->
196+
CustomizationUiItem(
197+
customization,
198+
isNew = diff.contains(customization),
199+
shouldPrefixAccountId = (nameToCount[customization.name] ?: 0) > 1
200+
)
171201
}
172-
} else {
173-
null
174202
}
175203
connectionToCustomizationUiItems[it.id] = customizationUiItems
176204

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
// SPDX-License-Identifier: Apache-2.0
33

44
package software.aws.toolkits.jetbrains.services.codewhisperer.customization
5+
import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfile
56

67
data class CodeWhispererCustomization(
78
@JvmField
@@ -12,4 +13,7 @@ data class CodeWhispererCustomization(
1213

1314
@JvmField
1415
var description: String? = null,
16+
17+
@JvmField
18+
var profile: QRegionProfile? = null
1519
)

0 commit comments

Comments
 (0)