fix(amazonq): adjust feature config customization arn override logic

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

@@ -37,6 +37,11 @@ interface CodeWhispererModelConfigurator {
      */
     fun getNewUpdate(connectionId: String): Collection<CustomizationUiItem>?
 
+    /**
+     * Get any current persisted customization arn override config
+     */
+    fun getPersistedCustomizationOverride(): String?
+
     companion object {
         fun getInstance(): CodeWhispererModelConfigurator = service()
     }

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

@@ -76,8 +76,11 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
 
     private val hasShownNewCustomizationNotification = AtomicBoolean(false)
 
+    @Deprecated("Use customizationArnOverrideV2 for the latest arn override persistence")
     private var serviceDefaultArn: String? = null
 
+    private var customizationArnOverrideV2: String? = null
+
     override fun showConfigDialog(project: Project) {
         runInEdt {
             calculateIfIamIdentityCenterConnection(project) {
@@ -181,7 +184,7 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
      */
     override fun switchCustomization(project: Project, newCustomization: CodeWhispererCustomization?, isOverride: Boolean) {
         calculateIfIamIdentityCenterConnection(project) {
-            if (isOverride && (newCustomization == null || newCustomization.arn.isEmpty() || serviceDefaultArn == newCustomization.arn)) {
+            if (isOverride && (newCustomization == null || newCustomization.arn.isEmpty() || customizationArnOverrideV2 == newCustomization.arn)) {
                 return@calculateIfIamIdentityCenterConnection
             }
             val oldCus = connectionIdToActiveCustomizationArn[it.id]
@@ -197,7 +200,7 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
                 CodeWhispererCustomizationListener.notifyCustomUiUpdate()
             }
             if (isOverride) {
-                serviceDefaultArn = newCustomization?.arn
+                customizationArnOverrideV2 = newCustomization?.arn
             }
         }
     }
@@ -249,6 +252,7 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
         state.connectionIdToActiveCustomizationArn.putAll(this.connectionIdToActiveCustomizationArn)
         state.previousAvailableCustomizations.putAll(this.connectionToCustomizationsShownLastTime)
         state.serviceDefaultArn = this.serviceDefaultArn
+        state.customizationArnOverrideV2 = this.customizationArnOverrideV2
 
         return state
     }
@@ -261,8 +265,12 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
         connectionToCustomizationsShownLastTime.putAll(state.previousAvailableCustomizations)
 
         this.serviceDefaultArn = state.serviceDefaultArn
+        this.customizationArnOverrideV2 = state.customizationArnOverrideV2
     }
 
+    // current latest field is customizationArnOverrideV2
+    override fun getPersistedCustomizationOverride() = customizationArnOverrideV2
+
     override fun dispose() {}
 
     private fun invalidateSelectedAndNotify(project: Project) {
@@ -292,8 +300,12 @@ class CodeWhispererCustomizationState : BaseState() {
     @get:MapAnnotation
     val previousAvailableCustomizations by map<String, MutableList<String>>()
 
+    @Deprecated("Use customizationArnOverrideV2 for the latest arn override persistence")
     @get:Property
     var serviceDefaultArn by string()
+
+    @get:Property
+    var customizationArnOverrideV2 by string()
 }
 
 data class CustomizationUiItem(

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

@@ -74,12 +74,18 @@ class CodeWhispererProjectStartupActivity : StartupActivity.DumbAware {
         runOnce = true
     }
 
-    // Start a job that runs every 30 mins
+    // Start a job that runs every 180 minutes
     private fun initFeatureConfigPollingJob(project: Project) {
         projectCoroutineScope(project).launch {
             while (isActive) {
                 CodeWhispererFeatureConfigService.getInstance().fetchFeatureConfigs(project)
                 CodeWhispererFeatureConfigService.getInstance().getCustomizationFeature()?.let { customization ->
+                    val persistedCustomizationOverride = CodeWhispererModelConfigurator.getInstance().getPersistedCustomizationOverride()
+                    val latestCustomizationOverride = customization.value.stringValue()
+                    if (persistedCustomizationOverride == latestCustomizationOverride) return@let
+
+                    // latest is different from the currently persisted, need update
+                    CodeWhispererFeatureConfigService.getInstance().validateCustomizationOverride(project, customization)
                     CodeWhispererModelConfigurator.getInstance().switchCustomization(
                         project,
                         CodeWhispererCustomization(arn = customization.value.stringValue(), name = customization.variation),

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

@@ -88,7 +88,7 @@ object CodeWhispererConstants {
     const val TOTAL_MILLIS_IN_SECOND = 1000
     const val TOTAL_SECONDS_IN_MINUTE: Long = 60L
     const val ACCOUNTLESS_START_URL = "accountless"
-    const val FEATURE_CONFIG_POLL_INTERVAL_IN_MS: Long = 30 * 60 * 1000L // 30 mins
+    const val FEATURE_CONFIG_POLL_INTERVAL_IN_MS: Long = 180 * 60 * 1000L // 180 mins
     const val USING: String = "using"
     const val GLOBAL_USING: String = "global using"
     const val STATIC: String = "static"

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

@@ -421,7 +421,7 @@ class CodeWhispererModelConfiguratorTest {
                 )
             )
 
-            this.serviceDefaultArn = "arn:aws:codewhisperer:default"
+            this.customizationArnOverrideV2 = "arn:aws:codewhisperer:default"
         }
 
         XmlSerializer.serializeInto(state, element)
@@ -441,6 +441,7 @@ class CodeWhispererModelConfiguratorTest {
             "</entry>" +
             "</map>" +
             "</option>" +
+            "<option name=\"customizationArnOverrideV2\" value=\"arn:aws:codewhisperer:default\" />" +
             "<option name=\"previousAvailableCustomizations\">" +
             "<map>" +
             "<entry key=\"fake-sso-url\">" +
@@ -453,7 +454,6 @@ class CodeWhispererModelConfiguratorTest {
             "</entry>" +
             "</map>" +
             "</option>" +
-            "<option name=\"serviceDefaultArn\" value=\"arn:aws:codewhisperer:default\" />" +
             "</component>"
 
         assertThat(actual).isEqualTo(expected)
@@ -470,7 +470,7 @@ class CodeWhispererModelConfiguratorTest {
         val actual = XmlSerializer.deserialize(element, CodeWhispererCustomizationState::class.java)
         assertThat(actual.connectionIdToActiveCustomizationArn).hasSize(0)
         assertThat(actual.previousAvailableCustomizations).hasSize(0)
-        assertThat(actual.serviceDefaultArn).isNull()
+        assertThat(actual.customizationArnOverrideV2).isNull()
     }
 
     @Test
@@ -504,12 +504,12 @@ class CodeWhispererModelConfiguratorTest {
                             </entry>
                         </map>
                     </option>
-                    <option name="serviceDefaultArn" value="arn:aws:codewhisperer:default"/>
+                    <option name="customizationArnOverrideV2" value="arn:aws:codewhisperer:default"/>
                 </component>
             """
         )
         val actual = XmlSerializer.deserialize(element, CodeWhispererCustomizationState::class.java)
-        assertThat(actual.serviceDefaultArn).isEqualTo("arn:aws:codewhisperer:default")
+        assertThat(actual.customizationArnOverrideV2).isEqualTo("arn:aws:codewhisperer:default")
         assertThat(actual.connectionIdToActiveCustomizationArn).hasSize(1)
         assertThat(actual.connectionIdToActiveCustomizationArn["fake-sso-url"]).isEqualTo(
             CodeWhispererCustomization(
@@ -546,7 +546,7 @@ class CodeWhispererModelConfiguratorTest {
         )
         val actual = XmlSerializer.deserialize(element, CodeWhispererCustomizationState::class.java)
         assertThat(actual.connectionIdToActiveCustomizationArn).hasSize(0)
-        assertThat(actual.serviceDefaultArn).isNull()
+        assertThat(actual.customizationArnOverrideV2).isNull()
         assertThat(actual.previousAvailableCustomizations).hasSize(1)
         assertThat(actual.previousAvailableCustomizations["fake-sso-url"]).isEqualTo(listOf("arn_1", "arn_2", "arn_3"))
     }

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

@@ -42,47 +42,8 @@ class CodeWhispererFeatureConfigService {
                 featureConfigs[it.feature()] = FeatureContext(it.feature(), it.variation(), it.value())
             }
 
-            // Only apply new auto-trigger UX to BID users
-            val isNewAutoTriggerUX = getNewAutoTriggerUX()
-            if (isNewAutoTriggerUX) {
-                calculateIfIamIdentityCenterConnection(project) {
-                    featureConfigs.remove(NEW_AUTO_TRIGGER_UX)
-                }
-            }
-
-            val customizationArnOverride = featureConfigs[CUSTOMIZATION_ARN_OVERRIDE_NAME]?.value?.stringValue()
-            if (customizationArnOverride != null) {
-                // Double check if server-side wrongly returns a customizationArn to BID users
-                calculateIfBIDConnection(project) {
-                    featureConfigs.remove(CUSTOMIZATION_ARN_OVERRIDE_NAME)
-                }
-                val availableCustomizations =
-                    calculateIfIamIdentityCenterConnection(project) {
-                        try {
-                            connection.getConnectionSettings().awsClient<CodeWhispererRuntimeClient>().listAvailableCustomizationsPaginator(
-                                ListAvailableCustomizationsRequest.builder().build()
-                            )
-                                .stream()
-                                .toList()
-                                .flatMap { resp ->
-                                    resp.customizations().map {
-                                        it.arn()
-                                    }
-                                }
-                        } catch (e: Exception) {
-                            LOG.debug(e) { "Failed to list available customizations" }
-                            null
-                        }
-                    }
+            validateNewAutoTriggerUX(project)
 
-                // If customizationArn from A/B is not available in listAvailableCustomizations response, don't use this value
-                if (availableCustomizations?.contains(customizationArnOverride) == false) {
-                    LOG.debug {
-                        "Customization arn $customizationArnOverride not available in listAvailableCustomizations, not using"
-                    }
-                    featureConfigs.remove(CUSTOMIZATION_ARN_OVERRIDE_NAME)
-                }
-            }
             CodeWhispererFeatureConfigListener.notifyUiFeatureConfigsAvailable()
         } catch (e: Exception) {
             LOG.debug(e) { "Error when fetching feature configs" }
@@ -132,6 +93,53 @@ class CodeWhispererFeatureConfigService {
     private fun connection(project: Project) =
         ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance())
 
+    private fun validateNewAutoTriggerUX(project: Project) {
+        // Only apply new auto-trigger UX to BID users
+        val isNewAutoTriggerUX = getNewAutoTriggerUX()
+        if (isNewAutoTriggerUX) {
+            calculateIfIamIdentityCenterConnection(project) {
+                featureConfigs.remove(NEW_AUTO_TRIGGER_UX)
+            }
+        }
+    }
+
+    fun validateCustomizationOverride(project: Project, customization: FeatureContext) {
+        val customizationArnOverride = customization.value.stringValue()
+        val connection = connection(project) ?: return
+        if (customizationArnOverride != null) {
+            // Double check if server-side wrongly returns a customizationArn to BID users
+            calculateIfBIDConnection(project) {
+                featureConfigs.remove(CUSTOMIZATION_ARN_OVERRIDE_NAME)
+            }
+            val availableCustomizations =
+                calculateIfIamIdentityCenterConnection(project) {
+                    try {
+                        connection.getConnectionSettings().awsClient<CodeWhispererRuntimeClient>().listAvailableCustomizationsPaginator(
+                            ListAvailableCustomizationsRequest.builder().build()
+                        )
+                            .stream()
+                            .toList()
+                            .flatMap { resp ->
+                                resp.customizations().map {
+                                    it.arn()
+                                }
+                            }
+                    } catch (e: Exception) {
+                        LOG.debug(e) { "Failed to list available customizations" }
+                        null
+                    }
+                }
+
+            // If customizationArn from A/B is not available in listAvailableCustomizations response, don't use this value
+            if (availableCustomizations?.contains(customizationArnOverride) == false) {
+                LOG.debug {
+                    "Customization arn $customizationArnOverride not available in listAvailableCustomizations, not using"
+                }
+                featureConfigs.remove(CUSTOMIZATION_ARN_OVERRIDE_NAME)
+            }
+        }
+    }
+
     companion object {
         fun getInstance(): CodeWhispererFeatureConfigService = service()
         private const val TEST_FEATURE_NAME = "testFeature"