From a4619da36b142b3984b4434a6a4330c328ca9c0b Mon Sep 17 00:00:00 2001 From: Sam Fink Date: Tue, 7 Jan 2025 10:09:22 -0800 Subject: [PATCH] Change logic for customization override to not override if user has manually selected a customization --- ...-b6572d91-e5b4-4b2d-bb43-6a331aacc79c.json | 4 +++ .../CodeWhispererModelConfigurator.kt | 30 ++++++++++++------- .../CodeWhispererModelConfiguratorTest.kt | 23 +++++++++++--- 3 files changed, 43 insertions(+), 14 deletions(-) create mode 100644 .changes/next-release/bugfix-b6572d91-e5b4-4b2d-bb43-6a331aacc79c.json diff --git a/.changes/next-release/bugfix-b6572d91-e5b4-4b2d-bb43-6a331aacc79c.json b/.changes/next-release/bugfix-b6572d91-e5b4-4b2d-bb43-6a331aacc79c.json new file mode 100644 index 0000000000..82086fc5af --- /dev/null +++ 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" +} \ No newline at end of file diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt index 3cacc48325..2569da95b0 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/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/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererModelConfiguratorTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererModelConfiguratorTest.kt index 169473a97e..fc681b9e17 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererModelConfiguratorTest.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/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