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 updates"
+}

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

@@ -35,6 +35,11 @@ interface CodeWhispererModelConfigurator {
      */
     fun getNewUpdate(connectionId: String): Collection<CustomizationUiItem>?
 
+    /**
+     * Refreshes the default customization ARN by retrieving the latest value from the feature configuration service.
+     */
+    fun refreshDefaultCustomizationArn(project: Project)
+
     companion object {
         fun getInstance(): CodeWhispererModelConfigurator = service()
     }

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,16 +167,7 @@ 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?) {
@@ -234,12 +227,22 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
         } ?: false
     }
 
+    override fun refreshDefaultCustomizationArn(project: Project) {
+        CodeWhispererFeatureConfigService.getInstance().getCustomizationFeature()?.let { customization ->
+            if (customization.value.stringValue() != serviceDefaultArn) {
+                serviceDefaultArn = customization.value.stringValue()
+                switchCustomization(project, CodeWhispererCustomization(arn = customization.value.stringValue(), name = customization.variation))
+            }
+        }
+    }
+
     override fun getNewUpdate(connectionId: String) = connectionToCustomizationUiItems[connectionId]
 
     override fun getState(): CodeWhispererCustomizationState {
         val state = CodeWhispererCustomizationState()
         state.connectionIdToActiveCustomizationArn.putAll(this.connectionIdToActiveCustomizationArn)
         state.previousAvailableCustomizations.putAll(this.connectionToCustomizationsShownLastTime)
+        state.serviceDefaultArn = this.serviceDefaultArn
 
         return state
     }
@@ -250,6 +253,8 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
 
         connectionToCustomizationsShownLastTime.clear()
         connectionToCustomizationsShownLastTime.putAll(state.previousAvailableCustomizations)
+
+        this.serviceDefaultArn = state.serviceDefaultArn
     }
 
     override fun dispose() {}
@@ -280,6 +285,10 @@ class CodeWhispererCustomizationState : BaseState() {
     @get:Property
     @get:MapAnnotation
     val previousAvailableCustomizations by map<String, MutableList<String>>()
+
+    @get:Property
+    @get:MapAnnotation
+    var serviceDefaultArn by string()
 }
 
 data class CustomizationUiItem(

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

@@ -63,6 +63,7 @@ class CodeWhispererProjectStartupActivity : StartupActivity.DumbAware {
         calculateIfIamIdentityCenterConnection(project) {
             pluginAwareExecuteOnPooledThread {
                 CodeWhispererModelConfigurator.getInstance().listCustomizations(project, passive = true)
+                CodeWhispererModelConfigurator.getInstance().refreshDefaultCustomizationArn(project)
             }
         }
 

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

@@ -143,9 +143,79 @@ class CodeWhispererModelConfiguratorTest {
                 FeatureContext("customizationArnOverride", "foo", FeatureValue.builder().stringValue("overrideArn").build())
             )
         }
+        sut.refreshDefaultCustomizationArn(projectRule.project)
         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())
+            )
+        }
+        sut.refreshDefaultCustomizationArn(projectRule.project)
+
+        // 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())
+            )
+        }
+        sut.refreshDefaultCustomizationArn(projectRule.project)
+
+        // 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())
+            )
+        }
+        sut.refreshDefaultCustomizationArn(projectRule.project)
+
+        // 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())
+            )
+        }
+        sut.refreshDefaultCustomizationArn(projectRule.project)
+
+        // 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)
+
+        // 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()