diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/migration/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/migration/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt index 5a155b5078d..53de2db3454 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/migration/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/migration/software/aws/toolkits/jetbrains/services/codewhisperer/customization/CodeWhispererModelConfigurator.kt @@ -18,6 +18,8 @@ interface CodeWhispererModelConfigurator { fun switchCustomization(project: Project, newCustomization: CodeWhispererCustomization?) + fun switchCustomization(project: Project, newCustomization: CodeWhispererCustomization?, isOverride: Boolean) + /** * 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 * but not limited to the specific connection. 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 f6535a4bba0..077783d2c91 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 @@ -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,23 @@ 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?) { + switchCustomization(project, newCustomization, false) + } + + /** + * 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 +196,9 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe CodeWhispererCustomizationListener.notifyCustomUiUpdate() } + if (isOverride) { + serviceDefaultArn = newCustomization?.arn + } } } @@ -240,6 +248,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 +259,8 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe connectionToCustomizationsShownLastTime.clear() connectionToCustomizationsShownLastTime.putAll(state.previousAvailableCustomizations) + + this.serviceDefaultArn = state.serviceDefaultArn } override fun dispose() {} @@ -280,6 +291,9 @@ class CodeWhispererCustomizationState : BaseState() { @get:Property @get:MapAnnotation val previousAvailableCustomizations by map>() + + @get:Property + var serviceDefaultArn by string() } data class CustomizationUiItem( diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/startup/CodeWhispererProjectStartupActivity.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/startup/CodeWhispererProjectStartupActivity.kt index cba7d4ab3e2..e2adcb8bdc6 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/startup/CodeWhispererProjectStartupActivity.kt +++ b/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) } } 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 d9608702402..baf42966096 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 @@ -143,9 +143,114 @@ 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 = 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 = 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 +420,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 +453,7 @@ class CodeWhispererModelConfiguratorTest { "" + "" + "" + + " +