Skip to content

Commit 1f9a366

Browse files
andrewyuqrli
andrewyuq
and
rli
authored
fix(amazonq): adjust feature config customization arn override logic (#5406)
* fix(amazonq): adjust feature config customization arn override logic 1. change poll interval from 30mins to 180 mins 2. Only call listAvailableCustomizations for new customization override for update 3. Use a new field customizationArnOverrideV2 for a clean override plan * address comments and fix tests * fix tests * fix tests * Update plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/CodeWhispererFeatureConfigService.kt Co-authored-by: Richard Li <742829+rli@users.noreply.github.com> * Update CodeWhispererFeatureConfigService.kt --------- Co-authored-by: Richard Li <742829+rli@users.noreply.github.com>
1 parent ca7295f commit 1f9a366

File tree

6 files changed

+77
-51
lines changed

6 files changed

+77
-51
lines changed

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,11 @@ interface CodeWhispererModelConfigurator {
3737
*/
3838
fun getNewUpdate(connectionId: String): Collection<CustomizationUiItem>?
3939

40+
/**
41+
* Get any current persisted customization arn override config
42+
*/
43+
fun getPersistedCustomizationOverride(): String?
44+
4045
companion object {
4146
fun getInstance(): CodeWhispererModelConfigurator = service()
4247
}

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

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -76,8 +76,11 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
7676

7777
private val hasShownNewCustomizationNotification = AtomicBoolean(false)
7878

79+
@Deprecated("Use customizationArnOverrideV2 for the latest arn override persistence")
7980
private var serviceDefaultArn: String? = null
8081

82+
private var customizationArnOverrideV2: String? = null
83+
8184
override fun showConfigDialog(project: Project) {
8285
runInEdt {
8386
calculateIfIamIdentityCenterConnection(project) {
@@ -181,7 +184,7 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
181184
*/
182185
override fun switchCustomization(project: Project, newCustomization: CodeWhispererCustomization?, isOverride: Boolean) {
183186
calculateIfIamIdentityCenterConnection(project) {
184-
if (isOverride && (newCustomization == null || newCustomization.arn.isEmpty() || serviceDefaultArn == newCustomization.arn)) {
187+
if (isOverride && (newCustomization == null || newCustomization.arn.isEmpty() || customizationArnOverrideV2 == newCustomization.arn)) {
185188
return@calculateIfIamIdentityCenterConnection
186189
}
187190
val oldCus = connectionIdToActiveCustomizationArn[it.id]
@@ -197,7 +200,7 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
197200
CodeWhispererCustomizationListener.notifyCustomUiUpdate()
198201
}
199202
if (isOverride) {
200-
serviceDefaultArn = newCustomization?.arn
203+
customizationArnOverrideV2 = newCustomization?.arn
201204
}
202205
}
203206
}
@@ -249,6 +252,7 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
249252
state.connectionIdToActiveCustomizationArn.putAll(this.connectionIdToActiveCustomizationArn)
250253
state.previousAvailableCustomizations.putAll(this.connectionToCustomizationsShownLastTime)
251254
state.serviceDefaultArn = this.serviceDefaultArn
255+
state.customizationArnOverrideV2 = this.customizationArnOverrideV2
252256

253257
return state
254258
}
@@ -261,8 +265,12 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
261265
connectionToCustomizationsShownLastTime.putAll(state.previousAvailableCustomizations)
262266

263267
this.serviceDefaultArn = state.serviceDefaultArn
268+
this.customizationArnOverrideV2 = state.customizationArnOverrideV2
264269
}
265270

271+
// current latest field is customizationArnOverrideV2
272+
override fun getPersistedCustomizationOverride() = customizationArnOverrideV2
273+
266274
override fun dispose() {}
267275

268276
private fun invalidateSelectedAndNotify(project: Project) {
@@ -292,8 +300,12 @@ class CodeWhispererCustomizationState : BaseState() {
292300
@get:MapAnnotation
293301
val previousAvailableCustomizations by map<String, MutableList<String>>()
294302

303+
@Deprecated("Use customizationArnOverrideV2 for the latest arn override persistence")
295304
@get:Property
296305
var serviceDefaultArn by string()
306+
307+
@get:Property
308+
var customizationArnOverrideV2 by string()
297309
}
298310

299311
data class CustomizationUiItem(

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

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -74,12 +74,18 @@ class CodeWhispererProjectStartupActivity : StartupActivity.DumbAware {
7474
runOnce = true
7575
}
7676

77-
// Start a job that runs every 30 mins
77+
// Start a job that runs every 180 minutes
7878
private fun initFeatureConfigPollingJob(project: Project) {
7979
projectCoroutineScope(project).launch {
8080
while (isActive) {
8181
CodeWhispererFeatureConfigService.getInstance().fetchFeatureConfigs(project)
8282
CodeWhispererFeatureConfigService.getInstance().getCustomizationFeature()?.let { customization ->
83+
val persistedCustomizationOverride = CodeWhispererModelConfigurator.getInstance().getPersistedCustomizationOverride()
84+
val latestCustomizationOverride = customization.value.stringValue()
85+
if (persistedCustomizationOverride == latestCustomizationOverride) return@let
86+
87+
// latest is different from the currently persisted, need update
88+
CodeWhispererFeatureConfigService.getInstance().validateCustomizationOverride(project, customization)
8389
CodeWhispererModelConfigurator.getInstance().switchCustomization(
8490
project,
8591
CodeWhispererCustomization(arn = customization.value.stringValue(), name = customization.variation),

‎plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,7 @@ object CodeWhispererConstants {
8888
const val TOTAL_MILLIS_IN_SECOND = 1000
8989
const val TOTAL_SECONDS_IN_MINUTE: Long = 60L
9090
const val ACCOUNTLESS_START_URL = "accountless"
91-
const val FEATURE_CONFIG_POLL_INTERVAL_IN_MS: Long = 30 * 60 * 1000L // 30 mins
91+
const val FEATURE_CONFIG_POLL_INTERVAL_IN_MS: Long = 180 * 60 * 1000L // 180 mins
9292
const val USING: String = "using"
9393
const val GLOBAL_USING: String = "global using"
9494
const val STATIC: String = "static"

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

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -421,7 +421,7 @@ class CodeWhispererModelConfiguratorTest {
421421
)
422422
)
423423

424-
this.serviceDefaultArn = "arn:aws:codewhisperer:default"
424+
this.customizationArnOverrideV2 = "arn:aws:codewhisperer:default"
425425
}
426426

427427
XmlSerializer.serializeInto(state, element)
@@ -441,6 +441,7 @@ class CodeWhispererModelConfiguratorTest {
441441
"</entry>" +
442442
"</map>" +
443443
"</option>" +
444+
"<option name=\"customizationArnOverrideV2\" value=\"arn:aws:codewhisperer:default\" />" +
444445
"<option name=\"previousAvailableCustomizations\">" +
445446
"<map>" +
446447
"<entry key=\"fake-sso-url\">" +
@@ -453,7 +454,6 @@ class CodeWhispererModelConfiguratorTest {
453454
"</entry>" +
454455
"</map>" +
455456
"</option>" +
456-
"<option name=\"serviceDefaultArn\" value=\"arn:aws:codewhisperer:default\" />" +
457457
"</component>"
458458

459459
assertThat(actual).isEqualTo(expected)
@@ -470,7 +470,7 @@ class CodeWhispererModelConfiguratorTest {
470470
val actual = XmlSerializer.deserialize(element, CodeWhispererCustomizationState::class.java)
471471
assertThat(actual.connectionIdToActiveCustomizationArn).hasSize(0)
472472
assertThat(actual.previousAvailableCustomizations).hasSize(0)
473-
assertThat(actual.serviceDefaultArn).isNull()
473+
assertThat(actual.customizationArnOverrideV2).isNull()
474474
}
475475

476476
@Test
@@ -504,12 +504,12 @@ class CodeWhispererModelConfiguratorTest {
504504
</entry>
505505
</map>
506506
</option>
507-
<option name="serviceDefaultArn" value="arn:aws:codewhisperer:default"/>
507+
<option name="customizationArnOverrideV2" value="arn:aws:codewhisperer:default"/>
508508
</component>
509509
"""
510510
)
511511
val actual = XmlSerializer.deserialize(element, CodeWhispererCustomizationState::class.java)
512-
assertThat(actual.serviceDefaultArn).isEqualTo("arn:aws:codewhisperer:default")
512+
assertThat(actual.customizationArnOverrideV2).isEqualTo("arn:aws:codewhisperer:default")
513513
assertThat(actual.connectionIdToActiveCustomizationArn).hasSize(1)
514514
assertThat(actual.connectionIdToActiveCustomizationArn["fake-sso-url"]).isEqualTo(
515515
CodeWhispererCustomization(
@@ -546,7 +546,7 @@ class CodeWhispererModelConfiguratorTest {
546546
)
547547
val actual = XmlSerializer.deserialize(element, CodeWhispererCustomizationState::class.java)
548548
assertThat(actual.connectionIdToActiveCustomizationArn).hasSize(0)
549-
assertThat(actual.serviceDefaultArn).isNull()
549+
assertThat(actual.customizationArnOverrideV2).isNull()
550550
assertThat(actual.previousAvailableCustomizations).hasSize(1)
551551
assertThat(actual.previousAvailableCustomizations["fake-sso-url"]).isEqualTo(listOf("arn_1", "arn_2", "arn_3"))
552552
}

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

Lines changed: 44 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@ import com.intellij.openapi.project.Project
99
import com.intellij.util.concurrency.annotations.RequiresBackgroundThread
1010
import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient
1111
import software.amazon.awssdk.services.codewhispererruntime.model.FeatureValue
12-
import software.amazon.awssdk.services.codewhispererruntime.model.ListAvailableCustomizationsRequest
1312
import software.aws.toolkits.core.utils.debug
1413
import software.aws.toolkits.core.utils.error
1514
import software.aws.toolkits.core.utils.getLogger
@@ -42,47 +41,8 @@ class CodeWhispererFeatureConfigService {
4241
featureConfigs[it.feature()] = FeatureContext(it.feature(), it.variation(), it.value())
4342
}
4443

45-
// Only apply new auto-trigger UX to BID users
46-
val isNewAutoTriggerUX = getNewAutoTriggerUX()
47-
if (isNewAutoTriggerUX) {
48-
calculateIfIamIdentityCenterConnection(project) {
49-
featureConfigs.remove(NEW_AUTO_TRIGGER_UX)
50-
}
51-
}
52-
53-
val customizationArnOverride = featureConfigs[CUSTOMIZATION_ARN_OVERRIDE_NAME]?.value?.stringValue()
54-
if (customizationArnOverride != null) {
55-
// Double check if server-side wrongly returns a customizationArn to BID users
56-
calculateIfBIDConnection(project) {
57-
featureConfigs.remove(CUSTOMIZATION_ARN_OVERRIDE_NAME)
58-
}
59-
val availableCustomizations =
60-
calculateIfIamIdentityCenterConnection(project) {
61-
try {
62-
connection.getConnectionSettings().awsClient<CodeWhispererRuntimeClient>().listAvailableCustomizationsPaginator(
63-
ListAvailableCustomizationsRequest.builder().build()
64-
)
65-
.stream()
66-
.toList()
67-
.flatMap { resp ->
68-
resp.customizations().map {
69-
it.arn()
70-
}
71-
}
72-
} catch (e: Exception) {
73-
LOG.debug(e) { "Failed to list available customizations" }
74-
null
75-
}
76-
}
44+
validateNewAutoTriggerUX(project)
7745

78-
// If customizationArn from A/B is not available in listAvailableCustomizations response, don't use this value
79-
if (availableCustomizations?.contains(customizationArnOverride) == false) {
80-
LOG.debug {
81-
"Customization arn $customizationArnOverride not available in listAvailableCustomizations, not using"
82-
}
83-
featureConfigs.remove(CUSTOMIZATION_ARN_OVERRIDE_NAME)
84-
}
85-
}
8646
CodeWhispererFeatureConfigListener.notifyUiFeatureConfigsAvailable()
8747
} catch (e: Exception) {
8848
LOG.debug(e) { "Error when fetching feature configs" }
@@ -132,6 +92,49 @@ class CodeWhispererFeatureConfigService {
13292
private fun connection(project: Project) =
13393
ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())
13494

95+
private fun validateNewAutoTriggerUX(project: Project) {
96+
// Only apply new auto-trigger UX to BID users
97+
val isNewAutoTriggerUX = getNewAutoTriggerUX()
98+
if (isNewAutoTriggerUX) {
99+
calculateIfIamIdentityCenterConnection(project) {
100+
featureConfigs.remove(NEW_AUTO_TRIGGER_UX)
101+
}
102+
}
103+
}
104+
105+
fun validateCustomizationOverride(project: Project, customization: FeatureContext) {
106+
val customizationArnOverride = customization.value.stringValue()
107+
val connection = connection(project) ?: return
108+
if (customizationArnOverride != null) {
109+
// Double check if server-side wrongly returns a customizationArn to BID users
110+
calculateIfBIDConnection(project) {
111+
featureConfigs.remove(CUSTOMIZATION_ARN_OVERRIDE_NAME)
112+
}
113+
val availableCustomizations =
114+
calculateIfIamIdentityCenterConnection(project) {
115+
try {
116+
connection.getConnectionSettings().awsClient<CodeWhispererRuntimeClient>().listAvailableCustomizationsPaginator {}
117+
.flatMap { resp ->
118+
resp.customizations().map {
119+
it.arn()
120+
}
121+
}
122+
} catch (e: Exception) {
123+
LOG.debug(e) { "Failed to list available customizations" }
124+
null
125+
}
126+
}
127+
128+
// If customizationArn from A/B is not available in listAvailableCustomizations response, don't use this value
129+
if (availableCustomizations?.contains(customizationArnOverride) == false) {
130+
LOG.debug {
131+
"Customization arn $customizationArnOverride not available in listAvailableCustomizations, not using"
132+
}
133+
featureConfigs.remove(CUSTOMIZATION_ARN_OVERRIDE_NAME)
134+
}
135+
}
136+
}
137+
135138
companion object {
136139
fun getInstance(): CodeWhispererFeatureConfigService = service()
137140
private const val TEST_FEATURE_NAME = "testFeature"

0 commit comments

Comments
 (0)