Change logic for customization override to not override if user has m…

…anually selected a customization
aws · Jan 7, 2025 · a4619da · a4619da
1 parent 370ed03
commit a4619da
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 14 deletions.
diff --git a/.changes/next-release/bugfix-b6572d91-e5b4-4b2d-bb43-6a331aacc79c.json b/.changes/next-release/bugfix-b6572d91-e5b4-4b2d-bb43-6a331aacc79c.json
@@ -0,0 +1,4 @@
+{
+  "type" : "bugfix",
+  "description" : "Change logic for customization override to not override if user has manually selected a customization"
+}
diff --git a/...toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt b/...toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt
@@ -18,9 +18,11 @@ import com.intellij.util.concurrency.annotations.RequiresBackgroundThread
 import com.intellij.util.xmlb.annotations.MapAnnotation
 import com.intellij.util.xmlb.annotations.Property
 import software.amazon.awssdk.services.codewhispererruntime.model.CodeWhispererRuntimeException
+import software.amazon.awssdk.services.codewhispererruntime.model.FeatureValue
 import software.aws.toolkits.core.utils.debug
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
+import software.aws.toolkits.jetbrains.services.amazonq.FeatureContext
 import software.aws.toolkits.jetbrains.services.amazonq.calculateIfIamIdentityCenterConnection
 import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
@@ -157,17 +159,25 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe
             return@calculateIfIamIdentityCenterConnection customizationUiItems
         }
 
+    /**
+     * Gets the active customization for a user. If a user has manually selected a customization,
+     * respect that choice. If a user has not selected a customization, check if they have a customization
+     * assigned to them via an AB feature. If so, use that customization.
+     */
     override fun activeCustomization(project: Project): CodeWhispererCustomization? {
-        val result = calculateIfIamIdentityCenterConnection(project) { connectionIdToActiveCustomizationArn[it.id] }
-
-        // A/B case
-        val customizationFeature = CodeWhispererFeatureConfigService.getInstance().getCustomizationFeature()
-        if (customizationFeature == null || customizationFeature.value.stringValue().isEmpty()) return result
-        return CodeWhispererCustomization(
-            arn = customizationFeature.value.stringValue(),
-            name = customizationFeature.variation,
-            description = result?.description
-        )
+        val selectedCustomization = calculateIfIamIdentityCenterConnection(project) { connectionIdToActiveCustomizationArn[it.id] }
+
+        if (selectedCustomization != null) {
+            return selectedCustomization
+        } else {
+            val customizationFeature = CodeWhispererFeatureConfigService.getInstance().getCustomizationFeature()
+            if (customizationFeature == null || customizationFeature.value.stringValue().isEmpty()) return selectedCustomization
+            return CodeWhispererCustomization(
+                arn = customizationFeature.value.stringValue(),
+                name = customizationFeature.variation,
+                description = selectedCustomization?.description
+            )
+        }
     }
 
     override fun switchCustomization(project: Project, newCustomization: CodeWhispererCustomization?) {

diff --git a/...tware/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererModelConfiguratorTest.kt b/...tware/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererModelConfiguratorTest.kt
@@ -113,17 +113,32 @@ class CodeWhispererModelConfiguratorTest {
     }
 
     @Test
-    fun `should override customization arn if there is one under AB test`() {
+    fun `should not override customization arn if there is one under AB test and manual selection has not been made`() {
         val ssoConn = spy(LegacyManagedBearerSsoConnection(region = "us-east-1", startUrl = "url 1", scopes = Q_SCOPES))
         ToolkitConnectionManager.getInstance(projectRule.project).switchConnection(ssoConn)
 
-        sut.switchCustomization(projectRule.project, CodeWhispererCustomization("foo", "customization_1", "description_1"))
-        assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(CodeWhispererCustomization("foo", "customization_1", "description_1"))
+        sut.switchCustomization(projectRule.project, CodeWhispererCustomization("selectedCustomizationArn", "customization_1", "description_1"))
+        assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(CodeWhispererCustomization("selectedCustomizationArn", "customization_1", "description_1"))
 
         abManager.stub {
             on { getCustomizationFeature() }.thenReturn(FeatureContext("customizationArnOverride", "foo", FeatureValue.builder().stringValue("bar").build()))
         }
-        assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(CodeWhispererCustomization("bar", "foo", "description_1"))
+        assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(CodeWhispererCustomization("selectedCustomizationArn", "customization_1", "description_1"))
+    }
+
+    @Test
+    fun `should override customization arn if there is one under AB test and manual selection has not been made`() {
+        val ssoConn = spy(LegacyManagedBearerSsoConnection(region = "us-east-1", startUrl = "url 1", scopes = Q_SCOPES))
+        ToolkitConnectionManager.getInstance(projectRule.project).switchConnection(ssoConn)
+
+        sut.switchCustomization(projectRule.project, CodeWhispererCustomization("selectedCustomizationArn", "customization_1", "description_1"))
+        assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(CodeWhispererCustomization("selectedCustomizationArn", "customization_1", "description_1"))
+        sut.invalidateCustomization("selectedCustomizationArn")
+
+        abManager.stub {
+            on { getCustomizationFeature() }.thenReturn(FeatureContext("customizationArnOverride", "foo", FeatureValue.builder().stringValue("overrideArn").build()))
+        }
+        assertThat(sut.activeCustomization(projectRule.project)).isEqualTo(CodeWhispererCustomization("overrideArn", "foo", null))
     }
 
     @Test