From 71e10dab493654935c7ee241dd663bd2e8f5d3f6 Mon Sep 17 00:00:00 2001 From: Andrew Yu Date: Tue, 25 Feb 2025 12:30:56 -0800 Subject: [PATCH 1/6] fix(amazonq): adjust feature config customization arn override logic 1. change poll interval from 30mins to 180 mins 2. Only call listAvailableCustomizations for new customization override for update 3. Use a new field customizationArnOverrideV2 for a clean override plan --- .../CodeWhispererModelConfigurator.kt | 5 ++ .../CodeWhispererModelConfigurator.kt | 14 ++- .../CodeWhispererProjectStartupActivity.kt | 8 +- .../util/CodeWhispererConstants.kt | 2 +- .../CodeWhispererFeatureConfigService.kt | 88 ++++++++++--------- 5 files changed, 73 insertions(+), 44 deletions(-) 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 53de2db3454..b95002fe19f 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 @@ -37,6 +37,11 @@ interface CodeWhispererModelConfigurator { */ fun getNewUpdate(connectionId: String): Collection? + /** + * Get any current persisted customization arn override config + */ + fun getPersistedCustomizationOverride(): String? + companion object { fun getInstance(): CodeWhispererModelConfigurator = service() } 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 077783d2c91..314d5d3c69d 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 @@ -78,6 +78,8 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe private var serviceDefaultArn: String? = null + private var customizationArnOverrideV2: String? = null + override fun showConfigDialog(project: Project) { runInEdt { calculateIfIamIdentityCenterConnection(project) { @@ -181,7 +183,7 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe */ override fun switchCustomization(project: Project, newCustomization: CodeWhispererCustomization?, isOverride: Boolean) { calculateIfIamIdentityCenterConnection(project) { - if (isOverride && (newCustomization == null || newCustomization.arn.isEmpty() || serviceDefaultArn == newCustomization.arn)) { + if (isOverride && (newCustomization == null || newCustomization.arn.isEmpty() || customizationArnOverrideV2 == newCustomization.arn)) { return@calculateIfIamIdentityCenterConnection } val oldCus = connectionIdToActiveCustomizationArn[it.id] @@ -197,7 +199,7 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe CodeWhispererCustomizationListener.notifyCustomUiUpdate() } if (isOverride) { - serviceDefaultArn = newCustomization?.arn + customizationArnOverrideV2 = newCustomization?.arn } } } @@ -249,6 +251,7 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe state.connectionIdToActiveCustomizationArn.putAll(this.connectionIdToActiveCustomizationArn) state.previousAvailableCustomizations.putAll(this.connectionToCustomizationsShownLastTime) state.serviceDefaultArn = this.serviceDefaultArn + state.customizationArnOverrideV2 = this.customizationArnOverrideV2 return state } @@ -261,8 +264,12 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe connectionToCustomizationsShownLastTime.putAll(state.previousAvailableCustomizations) this.serviceDefaultArn = state.serviceDefaultArn + this.customizationArnOverrideV2 = state.customizationArnOverrideV2 } + // current latest field is customizationArnOverrideV2 + override fun getPersistedCustomizationOverride() = customizationArnOverrideV2 + override fun dispose() {} private fun invalidateSelectedAndNotify(project: Project) { @@ -294,6 +301,9 @@ class CodeWhispererCustomizationState : BaseState() { @get:Property var serviceDefaultArn by string() + + @get:Property + var customizationArnOverrideV2 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 e2adcb8bdc6..1900239222a 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 @@ -74,12 +74,18 @@ class CodeWhispererProjectStartupActivity : StartupActivity.DumbAware { runOnce = true } - // Start a job that runs every 30 mins + // Start a job that runs every 180 minutes private fun initFeatureConfigPollingJob(project: Project) { projectCoroutineScope(project).launch { while (isActive) { CodeWhispererFeatureConfigService.getInstance().fetchFeatureConfigs(project) CodeWhispererFeatureConfigService.getInstance().getCustomizationFeature()?.let { customization -> + val persistedCustomizationOverride = CodeWhispererModelConfigurator.getInstance().getPersistedCustomizationOverride() + val latestCustomizationOverride = customization.value.stringValue() + if (persistedCustomizationOverride == latestCustomizationOverride) return@let + + // latest is different from the currently persisted, need update + CodeWhispererFeatureConfigService.getInstance().validateCustomizationOverride(project, customization) CodeWhispererModelConfigurator.getInstance().switchCustomization( project, CodeWhispererCustomization(arn = customization.value.stringValue(), name = customization.variation), diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt index 04952537aa4..68703549967 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt +++ b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/util/CodeWhispererConstants.kt @@ -88,7 +88,7 @@ object CodeWhispererConstants { const val TOTAL_MILLIS_IN_SECOND = 1000 const val TOTAL_SECONDS_IN_MINUTE: Long = 60L const val ACCOUNTLESS_START_URL = "accountless" - const val FEATURE_CONFIG_POLL_INTERVAL_IN_MS: Long = 30 * 60 * 1000L // 30 mins + const val FEATURE_CONFIG_POLL_INTERVAL_IN_MS: Long = 180 * 60 * 1000L // 180 mins const val USING: String = "using" const val GLOBAL_USING: String = "global using" const val STATIC: String = "static" diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/CodeWhispererFeatureConfigService.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/CodeWhispererFeatureConfigService.kt index 84a75038ad7..285f3b19989 100644 --- a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/CodeWhispererFeatureConfigService.kt +++ b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/CodeWhispererFeatureConfigService.kt @@ -42,47 +42,8 @@ class CodeWhispererFeatureConfigService { featureConfigs[it.feature()] = FeatureContext(it.feature(), it.variation(), it.value()) } - // Only apply new auto-trigger UX to BID users - val isNewAutoTriggerUX = getNewAutoTriggerUX() - if (isNewAutoTriggerUX) { - calculateIfIamIdentityCenterConnection(project) { - featureConfigs.remove(NEW_AUTO_TRIGGER_UX) - } - } - - val customizationArnOverride = featureConfigs[CUSTOMIZATION_ARN_OVERRIDE_NAME]?.value?.stringValue() - if (customizationArnOverride != null) { - // Double check if server-side wrongly returns a customizationArn to BID users - calculateIfBIDConnection(project) { - featureConfigs.remove(CUSTOMIZATION_ARN_OVERRIDE_NAME) - } - val availableCustomizations = - calculateIfIamIdentityCenterConnection(project) { - try { - connection.getConnectionSettings().awsClient().listAvailableCustomizationsPaginator( - ListAvailableCustomizationsRequest.builder().build() - ) - .stream() - .toList() - .flatMap { resp -> - resp.customizations().map { - it.arn() - } - } - } catch (e: Exception) { - LOG.debug(e) { "Failed to list available customizations" } - null - } - } + validateNewAutoTriggerUX(project) - // If customizationArn from A/B is not available in listAvailableCustomizations response, don't use this value - if (availableCustomizations?.contains(customizationArnOverride) == false) { - LOG.debug { - "Customization arn $customizationArnOverride not available in listAvailableCustomizations, not using" - } - featureConfigs.remove(CUSTOMIZATION_ARN_OVERRIDE_NAME) - } - } CodeWhispererFeatureConfigListener.notifyUiFeatureConfigsAvailable() } catch (e: Exception) { LOG.debug(e) { "Error when fetching feature configs" } @@ -132,6 +93,53 @@ class CodeWhispererFeatureConfigService { private fun connection(project: Project) = ToolkitConnectionManager.getInstance(project).activeConnectionForFeature(QConnection.getInstance()) + private fun validateNewAutoTriggerUX(project: Project) { + // Only apply new auto-trigger UX to BID users + val isNewAutoTriggerUX = getNewAutoTriggerUX() + if (isNewAutoTriggerUX) { + calculateIfIamIdentityCenterConnection(project) { + featureConfigs.remove(NEW_AUTO_TRIGGER_UX) + } + } + } + + fun validateCustomizationOverride(project: Project, customization: FeatureContext) { + val customizationArnOverride = customization.value.stringValue() + val connection = connection(project) ?: return + if (customizationArnOverride != null) { + // Double check if server-side wrongly returns a customizationArn to BID users + calculateIfBIDConnection(project) { + featureConfigs.remove(CUSTOMIZATION_ARN_OVERRIDE_NAME) + } + val availableCustomizations = + calculateIfIamIdentityCenterConnection(project) { + try { + connection.getConnectionSettings().awsClient().listAvailableCustomizationsPaginator( + ListAvailableCustomizationsRequest.builder().build() + ) + .stream() + .toList() + .flatMap { resp -> + resp.customizations().map { + it.arn() + } + } + } catch (e: Exception) { + LOG.debug(e) { "Failed to list available customizations" } + null + } + } + + // If customizationArn from A/B is not available in listAvailableCustomizations response, don't use this value + if (availableCustomizations?.contains(customizationArnOverride) == false) { + LOG.debug { + "Customization arn $customizationArnOverride not available in listAvailableCustomizations, not using" + } + featureConfigs.remove(CUSTOMIZATION_ARN_OVERRIDE_NAME) + } + } + } + companion object { fun getInstance(): CodeWhispererFeatureConfigService = service() private const val TEST_FEATURE_NAME = "testFeature" From 69c7f49603f6b176eae9490238fa94efd0dcba1a Mon Sep 17 00:00:00 2001 From: Andrew Yu Date: Tue, 25 Feb 2025 13:15:48 -0800 Subject: [PATCH 2/6] address comments and fix tests --- .../customization/CodeWhispererModelConfigurator.kt | 2 ++ .../codewhisperer/CodeWhispererModelConfiguratorTest.kt | 8 ++++---- 2 files changed, 6 insertions(+), 4 deletions(-) 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 314d5d3c69d..f19df1595bf 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,7 @@ class DefaultCodeWhispererModelConfigurator : CodeWhispererModelConfigurator, Pe private val hasShownNewCustomizationNotification = AtomicBoolean(false) + @Deprecated("Use customizationArnOverrideV2 for the latest arn override persistence") private var serviceDefaultArn: String? = null private var customizationArnOverrideV2: String? = null @@ -299,6 +300,7 @@ class CodeWhispererCustomizationState : BaseState() { @get:MapAnnotation val previousAvailableCustomizations by map>() + @Deprecated("Use customizationArnOverrideV2 for the latest arn override persistence") @get:Property var serviceDefaultArn by string() 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 baf42966096..a36a874fcd4 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 @@ -421,7 +421,7 @@ class CodeWhispererModelConfiguratorTest { ) ) - this.serviceDefaultArn = "arn:aws:codewhisperer:default" + this.customizationArnOverrideV2 = "arn:aws:codewhisperer:default" } XmlSerializer.serializeInto(state, element) @@ -470,7 +470,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() + assertThat(actual.customizationArnOverrideV2).isNull() } @Test @@ -509,7 +509,7 @@ class CodeWhispererModelConfiguratorTest { """ ) val actual = XmlSerializer.deserialize(element, CodeWhispererCustomizationState::class.java) - assertThat(actual.serviceDefaultArn).isEqualTo("arn:aws:codewhisperer:default") + assertThat(actual.customizationArnOverrideV2).isEqualTo("arn:aws:codewhisperer:default") assertThat(actual.connectionIdToActiveCustomizationArn).hasSize(1) assertThat(actual.connectionIdToActiveCustomizationArn["fake-sso-url"]).isEqualTo( CodeWhispererCustomization( @@ -546,7 +546,7 @@ class CodeWhispererModelConfiguratorTest { ) val actual = XmlSerializer.deserialize(element, CodeWhispererCustomizationState::class.java) assertThat(actual.connectionIdToActiveCustomizationArn).hasSize(0) - assertThat(actual.serviceDefaultArn).isNull() + assertThat(actual.customizationArnOverrideV2).isNull() assertThat(actual.previousAvailableCustomizations).hasSize(1) assertThat(actual.previousAvailableCustomizations["fake-sso-url"]).isEqualTo(listOf("arn_1", "arn_2", "arn_3")) } From d460b079a4e5c2c43b50a3094e5c5ac83839aefa Mon Sep 17 00:00:00 2001 From: Andrew Yu Date: Tue, 25 Feb 2025 13:51:39 -0800 Subject: [PATCH 3/6] fix tests --- .../codewhisperer/CodeWhispererModelConfiguratorTest.kt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 a36a874fcd4..0fef510b39b 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 @@ -453,7 +453,7 @@ class CodeWhispererModelConfiguratorTest { "" + "" + "" + - " - " + + "" + - "