config(amazonq): customization override should respect users' selection

.changes/next-release/bugfix-e0622832-a6ee-4113-99f2-0832e12705f5.json

@@ -0,0 +1,4 @@
+{
+  "type" : "bugfix",
+  "description" : "Amazon Q: Ensure server-pushed customization settings updates local value"
+}

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

@@ -16,7 +16,7 @@ interface CodeWhispererModelConfigurator {
 
     fun activeCustomization(project: Project): CodeWhispererCustomization?
 
-    fun switchCustomization(project: Project, newCustomization: CodeWhispererCustomization?)
+    fun switchCustomization(project: Project, newCustomization: CodeWhispererCustomization?, isOverride: Boolean = false)
 
     /**
      * This method is only used for invalidate a stale customization which was previously active but was removed, it will remove all usage of this customization

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

@@ -76,6 +76,8 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
 
     private val hasShownNewCustomizationNotification = AtomicBoolean(false)
 
+    private var serviceDefaultArn: String? = null
+
     override fun showConfigDialog(project: Project) {
         runInEdt {
             calculateIfIamIdentityCenterConnection(project) {
@@ -165,20 +167,19 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
     override fun activeCustomization(project: Project): CodeWhispererCustomization? {
         val selectedCustomization = calculateIfIamIdentityCenterConnection(project) { connectionIdToActiveCustomizationArn[it.id] }
 
-        if (selectedCustomization != null) {
-            return selectedCustomization
-        } else {
-            val customizationOverride = CodeWhispererFeatureConfigService.getInstance().getCustomizationFeature()
-            if (customizationOverride == null || customizationOverride.value.stringValue().isEmpty()) return null
-            return CodeWhispererCustomization(
-                arn = customizationOverride.value.stringValue(),
-                name = customizationOverride.variation,
-            )
-        }
+        return selectedCustomization
     }
 
-    override fun switchCustomization(project: Project, newCustomization: CodeWhispererCustomization?) {
+    /**
+     * Override happens when ALL following conditions are met
+     *  1. service returns non-empty override customization arn, refer to [CodeWhispererFeatureConfigService]
+     *  2. the override customization arn is different from the previous override customization if any. The purpose is to only do override once on users' behalf.
+     */
+    override fun switchCustomization(project: Project, newCustomization: CodeWhispererCustomization?, isOverride: Boolean) {
         calculateIfIamIdentityCenterConnection(project) {
+            if (isOverride && (newCustomization == null || newCustomization.arn.isEmpty() || serviceDefaultArn == newCustomization.arn)) {
+                return@calculateIfIamIdentityCenterConnection
+            }
             val oldCus = connectionIdToActiveCustomizationArn[it.id]
             if (oldCus != newCustomization) {
                 newCustomization?.let { newCus ->
@@ -191,6 +192,9 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
 
                 CodeWhispererCustomizationListener.notifyCustomUiUpdate()
             }
+            if (isOverride) {
+                serviceDefaultArn = newCustomization?.arn
+            }
         }
     }
 
@@ -240,6 +244,7 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
         val state = CodeWhispererCustomizationState()
         state.connectionIdToActiveCustomizationArn.putAll(this.connectionIdToActiveCustomizationArn)
         state.previousAvailableCustomizations.putAll(this.connectionToCustomizationsShownLastTime)
+        state.serviceDefaultArn = this.serviceDefaultArn
 
         return state
     }
@@ -250,6 +255,8 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
 
         connectionToCustomizationsShownLastTime.clear()
         connectionToCustomizationsShownLastTime.putAll(state.previousAvailableCustomizations)
+
+        this.serviceDefaultArn = state.serviceDefaultArn
     }
 
     override fun dispose() {}
@@ -280,6 +287,9 @@ class CodeWhispererCustomizationState : BaseState() {
     @get:Property
     @get:MapAnnotation
     val previousAvailableCustomizations by map<String, MutableList<String>>()
+
+    @get:Property
+    var serviceDefaultArn by string()
 }
 
 data class CustomizationUiItem(

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

@@ -13,6 +13,7 @@ import software.aws.toolkits.jetbrains.core.coroutines.projectCoroutineScope
 import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
 import software.aws.toolkits.jetbrains.services.amazonq.calculateIfIamIdentityCenterConnection
 import software.aws.toolkits.jetbrains.services.codewhisperer.codescan.CodeWhispererCodeScanManager
+import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererCustomization
 import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator
 import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
 import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.isUserBuilderId
@@ -78,6 +79,14 @@ class CodeWhispererProjectStartupActivity : StartupActivity.DumbAware {
         projectCoroutineScope(project).launch {
             while (isActive) {
                 CodeWhispererFeatureConfigService.getInstance().fetchFeatureConfigs(project)
+                CodeWhispererFeatureConfigService.getInstance().getCustomizationFeature()?.let { customization ->
+                    CodeWhispererModelConfigurator.getInstance().switchCustomization(
+                        project,
+                        CodeWhispererCustomization(arn = customization.value.stringValue(), name = customization.variation),
+                        isOverride = true
+                    )
+                }
+
                 delay(FEATURE_CONFIG_POLL_INTERVAL_IN_MS)
             }
         }

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

@@ -143,9 +143,103 @@ class CodeWhispererModelConfiguratorTest {
                 FeatureContext("customizationArnOverride", "foo", FeatureValue.builder().stringValue("overrideArn").build())
             )
         }
+        abManager.getCustomizationFeature()?.let { customization ->
+            sut.switchCustomization(projectRule.project, CodeWhispererCustomization(arn = customization.value.stringValue(), name = customization.variation),
+                isOverride = true
+            )
+        }
         assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(CodeWhispererCustomization("overrideArn", "foo", null))
     }
 
+    @Test
+    fun `should update customization when user has never selected one`() {
+        val ssoConn = spy(LegacyManagedBearerSsoConnection(region = "us-east-1", startUrl = "url-1", scopes = Q_SCOPES))
+        ToolkitConnectionManager.getInstance(projectRule.project).switchConnection(ssoConn)
+
+        // Step 1: Server pushes first customization (arnOverride1)
+        abManager.stub {
+            on { getCustomizationFeature() }.thenReturn(
+                FeatureContext("customizationArnOverride", "foo", FeatureValue.builder().stringValue("arnOverride1").build())
+            )
+        }
+        abManager.getCustomizationFeature()?.let { customization ->
+            sut.switchCustomization(projectRule.project, CodeWhispererCustomization(arn = customization.value.stringValue(), name = customization.variation),
+                isOverride = true
+            )
+        }
+
+        // User should receive arnOverride1 from the server
+        assertThat(sut.activeCustomization(projectRule.project))
+            .isEqualTo(CodeWhispererCustomization("arnOverride1", "foo", null))
+
+        // Step 2: Server updates customization (arnOverride2)
+        abManager.stub {
+            on { getCustomizationFeature() }.thenReturn(
+                FeatureContext("customizationArnOverride", "foo", FeatureValue.builder().stringValue("arnOverride2").build())
+            )
+        }
+
+        abManager.getCustomizationFeature()?.let { customization ->
+            sut.switchCustomization(projectRule.project, CodeWhispererCustomization(arn = customization.value.stringValue(), name = customization.variation),
+                isOverride = true
+            )
+        }
+        // User should receive arnOverride2 from the server
+        assertThat(sut.activeCustomization(projectRule.project))
+            .isEqualTo(CodeWhispererCustomization("arnOverride2", "foo", null))
+    }
+
+    @Test
+    fun `should not override user selection when server updates customization`() {
+        val ssoConn = spy(LegacyManagedBearerSsoConnection(region = "us-east-1", startUrl = "url-1", scopes = Q_SCOPES))
+        ToolkitConnectionManager.getInstance(projectRule.project).switchConnection(ssoConn)
+
+        // Step 1: Server pushes first customization (arnOverride1)
+        abManager.stub {
+            on { getCustomizationFeature() }.thenReturn(
+                FeatureContext("customizationArnOverride", "foo", FeatureValue.builder().stringValue("arnOverride1").build())
+            )
+        }
+
+        abManager.getCustomizationFeature()?.let { customization ->
+            sut.switchCustomization(projectRule.project, CodeWhispererCustomization(arn = customization.value.stringValue(), name = customization.variation),
+                isOverride = true
+            )
+        }
+        // User should receive arnOverride1 from the server
+        assertThat(sut.activeCustomization(projectRule.project))
+            .isEqualTo(CodeWhispererCustomization("arnOverride1", "foo", null))
+
+        // Step 2: Server updates customization again (arnOverride2)
+        abManager.stub {
+            on { getCustomizationFeature() }.thenReturn(
+                FeatureContext("customizationArnOverride", "foo", FeatureValue.builder().stringValue("arnOverride2").build())
+            )
+        }
+
+        abManager.getCustomizationFeature()?.let { customization ->
+            sut.switchCustomization(projectRule.project, CodeWhispererCustomization(arn = customization.value.stringValue(), name = customization.variation),
+                isOverride = true
+            )
+        }
+        // Ensure server’s change is applied
+        assertThat(sut.activeCustomization(projectRule.project))
+            .isEqualTo(CodeWhispererCustomization("arnOverride2", "foo", null))
+
+        // Step 3: User manually selects a different customization (userSelectedArn)
+        val userCustomization = CodeWhispererCustomization("userSelectedArn", "userChoice", null)
+        sut.switchCustomization(projectRule.project, userCustomization)
+        abManager.getCustomizationFeature()?.let { customization ->
+            sut.switchCustomization(projectRule.project, CodeWhispererCustomization(arn = customization.value.stringValue(), name = customization.variation),
+                isOverride = true
+            )
+        }
+        // Ensure user selection is still respected (should not change to arnOverride2)
+        assertThat(sut.activeCustomization(projectRule.project))
+            .isEqualTo(userCustomization)
+    }
+
+
     @Test
     fun `loadState should load the correct values into memory`() {
         credManager.clear()
@@ -315,6 +409,8 @@ class CodeWhispererModelConfiguratorTest {
                     "fake-sso-url" to CodeWhispererCustomization(arn = "arn_2", name = "name_2", description = "description_2")
                 )
             )
+
+            this.serviceDefaultArn = "arn:aws:codewhisperer:default"
         }
 
         XmlSerializer.serializeInto(state, element)
@@ -346,6 +442,7 @@ class CodeWhispererModelConfiguratorTest {
             "</entry>" +
             "</map>" +
             "</option>" +
+            "<option name=\"serviceDefaultArn\" value=\"arn:aws:codewhisperer:default\" />" +
             "</component>"
 
         assertThat(actual).isEqualTo(expected)
@@ -362,6 +459,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()
     }
 
     @Test
@@ -395,10 +493,12 @@ class CodeWhispererModelConfiguratorTest {
                             </entry>
                         </map>
                     </option>
+                    <option name="serviceDefaultArn" value="arn:aws:codewhisperer:default"/>
                 </component>
             """
         )
         val actual = XmlSerializer.deserialize(element, CodeWhispererCustomizationState::class.java)
+        assertThat(actual.serviceDefaultArn).isEqualTo("arn:aws:codewhisperer:default")
         assertThat(actual.connectionIdToActiveCustomizationArn).hasSize(1)
         assertThat(actual.connectionIdToActiveCustomizationArn["fake-sso-url"]).isEqualTo(
             CodeWhispererCustomization(
@@ -435,6 +535,7 @@ class CodeWhispererModelConfiguratorTest {
         )
         val actual = XmlSerializer.deserialize(element, CodeWhispererCustomizationState::class.java)
         assertThat(actual.connectionIdToActiveCustomizationArn).hasSize(0)
+        assertThat(actual.serviceDefaultArn).isNull()
         assertThat(actual.previousAvailableCustomizations).hasSize(1)
         assertThat(actual.previousAvailableCustomizations["fake-sso-url"]).isEqualTo(listOf("arn_1", "arn_2", "arn_3"))
     }

plugins/amazonq/src/test/kotlin/PluginAmazonQJvmBinaryCompatabilityTest.kt

@@ -71,6 +71,6 @@ class PluginAmazonQJvmBinaryCompatabilityTest {
         assertThat(customizationUiItem.getMethod("getCustomization").returnType).isEqualTo(customizationClazz)
 
         // CodeWhispererModelConfigurator.switchCustomization(...)
-        assertThat(modelConfiguratorClazz.getMethod("switchCustomization", Class.forName("com.intellij.openapi.project.Project"), customizationClazz).returnType).isEqualTo(Void.TYPE)
+        assertThat(modelConfiguratorClazz.getMethod("switchCustomization", Class.forName("com.intellij.openapi.project.Project"), customizationClazz, java.lang.Boolean.TYPE).returnType).isEqualTo(Void.TYPE)
     }
 }