Skip to content

Commit 42aa641

Browse files
evanliu048evanliu048
committed
refactor
1 parent 593838f commit 42aa641

File tree

2 files changed

+39
-49
lines changed

2 files changed

+39
-49
lines changed

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

Lines changed: 7 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -190,19 +190,13 @@ class CodeWhispererCustomizationDialog(
190190
val activeCustomization = CodeWhispererModelConfigurator.getInstance().activeCustomization(project)
191191
val unsorted = myCustomizations ?: CodeWhispererModelConfigurator.getInstance().listCustomizations(project).orEmpty()
192192
val activeProfile = QRegionProfileManager.getInstance().activeProfile(project)
193-
// sort customization based on profile name first
194-
val baseSorted = unsorted.sortedWith(
195-
compareBy<CustomizationUiItem> {
196-
it.customization.profile?.profileName != activeProfile?.profileName
197-
}.thenBy { it.customization.profile?.profileName.orEmpty() }
198-
.thenBy { it.customization.name }
199-
)
200-
201-
val sorted = if (activeCustomization != null) {
202-
baseSorted.putPickedUpFront(setOf(activeCustomization))
203-
} else {
204-
baseSorted
205-
}
193+
// Group customizations by profile name (active profile first, then alphabetical), with the active customization on top
194+
val sorted = unsorted.sortedWith(
195+
compareBy<CustomizationUiItem> { it.customization.profile?.profileName != activeProfile?.profileName }
196+
.thenBy { it.customization.profile?.profileName.orEmpty() }
197+
.thenBy { it.customization.name }
198+
)
199+
.let { list -> activeCustomization?.let { list.putPickedUpFront(setOf(it)) } ?: list }
206200

207201
if (
208202
sorted.isNotEmpty() &&

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

Lines changed: 32 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,13 @@ private fun notifyNewCustomization(project: Project) {
6565
)
6666
}
6767

68+
/**
69+
* generate a stable composite key for profile & customization combination for diff/deduplicate/search
70+
* rule: `<profileArn>::<customizationArn>`
71+
*/
72+
private fun CodeWhispererCustomization.compositeKey(): String =
73+
"${profile?.arn.orEmpty()}::$arn"
74+
6875
@Service(Service.Level.APP)
6976
@State(name = "codewhispererCustomizationStates", storages = [Storage("aws.xml")])
7077
class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, PersistentStateComponent<CodeWhispererCustomizationState>, Disposable {
@@ -73,7 +80,7 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
7380
private val connectionIdToActiveCustomizationArn = Collections.synchronizedMap<String, CodeWhispererCustomization>(mutableMapOf())
7481

7582
// Map to store connectionId to its listAvailableCustomizations result last time
76-
// the customization has format profileArn::customizationArn
83+
// Format of the customization key: profileArn::customizationArn
7784
private val connectionToCustomizationsShownLastTime = mutableMapOf<String, MutableList<String>>()
7885

7986
private val connectionIdToIsAllowlisted = Collections.synchronizedMap<String, Boolean>(mutableMapOf())
@@ -111,35 +118,32 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
111118
@RequiresBackgroundThread
112119
override fun listCustomizations(project: Project, passive: Boolean): List<CustomizationUiItem>? =
113120
calculateIfIamIdentityCenterConnection(project) {
114-
// 1. invoke API and get result
121+
// 1. fetch all profiles, invoke fetch customizations API and get result for each profile and aggregate all the results
115122
val listAvailableProfilesResult = tryOrNull { QRegionProfileManager.getInstance().listRegionProfiles(project) }.orEmpty()
116123

117-
val aggregatedCustomizations = mutableListOf<CodeWhispererCustomization>()
118-
119-
for (profile in listAvailableProfilesResult) {
120-
try {
121-
val perProfileCustomizations = CodeWhispererClientAdaptor.getInstance(project).listAvailableCustomizations(profile)
122-
aggregatedCustomizations.addAll(perProfileCustomizations)
123-
} catch (e: Exception) {
124+
val aggregatedCustomizations = listAvailableProfilesResult.flatMap { profile ->
125+
runCatching {
126+
CodeWhispererClientAdaptor.getInstance(project).listAvailableCustomizations(profile)
127+
}.onFailure { e ->
124128
val requestId = (e as? CodeWhispererRuntimeException)?.requestId()
125-
val logMessage = if (CodeWhispererConstants.Customization.noAccessToCustomizationExceptionPredicate(e)) {
129+
val isNotAllowlisted = (e as? Exception)
130+
?.let { CodeWhispererConstants.Customization.noAccessToCustomizationExceptionPredicate(e) }
131+
?: false
132+
val logMessage = if (isNotAllowlisted) {
126133
// TODO: not required for non GP users
127134
"ListAvailableCustomizations: connection ${it.id} is not allowlisted, requestId: ${requestId.orEmpty()}"
128135
} else {
129136
"ListAvailableCustomizations: failed due to unknown error ${e.message}, requestId: ${requestId.orEmpty()}"
130137
}
131138
LOG.debug { logMessage }
132-
}
139+
}.getOrDefault(emptyList())
133140
}
134141

135142
// 2. get diff
136143
val previousCustomizationsShapshot = connectionToCustomizationsShownLastTime.getOrElse(it.id) { emptyList() }
137-
// calculate new profile+customization list using profile.arn :: customization.arn as the key
138-
val diff = aggregatedCustomizations.filterNot { customization ->
139-
previousCustomizationsShapshot.contains(
140-
"${customization.profile?.arn}::${customization.arn}"
141-
)
142-
}.toSet()
144+
// calculate new profile+customization list using compositeKey
145+
val diff = aggregatedCustomizations.filterNot { customization -> previousCustomizationsShapshot.contains(customization.compositeKey()) }.toSet()
146+
143147
// 3 if passive,
144148
// (1) update allowlisting
145149
// (2) prompt "You have New Customizations" toast notification (only show once)
@@ -152,39 +156,31 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
152156
notifyNewCustomization(project)
153157
}
154158
} else {
155-
aggregatedCustomizations.let { customizations ->
156-
connectionToCustomizationsShownLastTime[it.id] = customizations.map { customization,
157-
->
158-
"${customization.profile?.arn}::${customization.arn}"
159-
}.toMutableList()
160-
}
159+
connectionToCustomizationsShownLastTime[it.id] = aggregatedCustomizations.map { customization -> customization.compositeKey() }.toMutableList()
161160
}
162161

163162
// 4. invalidate selected customization if
164163
// (1) the API call failed
165164
// (2) the selected customization is not in the resultset of API call
165+
// (3) the q region profile associated with the selected customization does not match the currently active profile
166166
activeCustomization(project)?.let { activeCustom ->
167167
if (aggregatedCustomizations.isEmpty()) {
168168
invalidateSelectedAndNotify(project)
169-
} else if (!aggregatedCustomizations.any { latestCustom ->
170-
"${latestCustom.profile?.arn}::${latestCustom.arn}" == "${activeCustom.profile?.arn}::${activeCustom.arn}"
171-
} || activeCustom.profile != QRegionProfileManager.getInstance().activeProfile(project)
169+
} else if (!aggregatedCustomizations.any { latestCustom -> latestCustom.compositeKey() == activeCustom.compositeKey() } ||
170+
activeCustom.profile != QRegionProfileManager.getInstance().activeProfile(project)
172171
) {
173172
invalidateSelectedAndNotify(project)
174173
}
175174
}
176175

177176
// 5. transform result to UI items and return
178-
val customizationUiItems = aggregatedCustomizations.let { customizations ->
179-
val nameToCount = customizations.groupingBy { customization -> customization.name }.eachCount()
180-
181-
customizations.map { customization ->
182-
CustomizationUiItem(
183-
customization,
184-
isNew = diff.contains(customization),
185-
shouldPrefixAccountId = (nameToCount[customization.name] ?: 0) > 1
186-
)
187-
}
177+
val nameToCount = aggregatedCustomizations.groupingBy { customization -> customization.name }.eachCount()
178+
val customizationUiItems = aggregatedCustomizations.map { customization ->
179+
CustomizationUiItem(
180+
customization,
181+
isNew = diff.contains(customization),
182+
shouldPrefixAccountId = (nameToCount[customization.name] ?: 0) > 1
183+
)
188184
}
189185
connectionToCustomizationUiItems[it.id] = customizationUiItems
190186

0 commit comments

Comments
 (0)