fix(amazonq): Prevent customization override if user has manually selected a customization

packages/amazonq/.changes/next-release/Bug Fix-d8c02c38-bdc4-492a-bf8c-6c35e35087be.json

@@ -0,0 +1,4 @@
+{
+	"type": "Bug Fix",
+	"description": "User-selected customizations are sometimes not being persisted."
+}

packages/core/src/codewhisperer/index.ts

@@ -26,6 +26,7 @@ export type {
     SendTelemetryEventResponse,
     TelemetryEvent,
     InlineChatEvent,
+    Customization,
 } from './client/codewhispereruserclient.d.ts'
 export type { default as CodeWhispererUserClient } from './client/codewhispereruserclient.d.ts'
 export { SecurityPanelViewProvider } from './views/securityPanelViewProvider'
@@ -98,6 +99,6 @@ export * as diagnosticsProvider from './service/diagnosticsProvider'
 export * from './ui/codeWhispererNodes'
 export { SecurityScanError } from '../codewhisperer/models/errors'
 export * as CodeWhispererConstants from '../codewhisperer/models/constants'
-export { getSelectedCustomization } from './util/customizationUtil'
+export { getSelectedCustomization, setSelectedCustomization, baseCustomization } from './util/customizationUtil'
 export { Container } from './service/serviceContainer'
 export * from './util/gitUtil'

packages/core/src/codewhisperer/util/customizationUtil.ts

@@ -91,6 +91,12 @@ export const baseCustomization = {
     ),
 }
 
+/**
+ * Gets the customization that should be used for user requests. If a user has manually selected
+ * a customization, always respect that choice. If not, check if the user is part of an AB
+ * group assigned a specific customization. If so, use that customization. If not, use the
+ * base customization.
+ */
 export const getSelectedCustomization = (): Customization => {
     if (
         !AuthUtil.instance.isCustomizationFeatureEnabled ||
@@ -105,21 +111,22 @@ export const getSelectedCustomization = (): Customization => {
         Object,
         {}
     )
-    const result = selectedCustomizationArr[AuthUtil.instance.conn.label] || baseCustomization
-
-    // A/B case
-    const customizationFeature = FeatureConfigProvider.getFeature(Features.customizationArnOverride)
-    const arnOverride = customizationFeature?.value.stringValue
-    const customizationOverrideName = customizationFeature?.variation
-    if (arnOverride === undefined || arnOverride === '') {
-        return result
+    const selectedCustomization = selectedCustomizationArr[AuthUtil.instance.conn.label]
+
+    if (selectedCustomization && selectedCustomization.name !== baseCustomization.name) {
+        return selectedCustomization
     } else {
-        // A trick to prioritize arn from A/B over user's currently selected(for request and telemetry)
-        // but still shows customization info of user's currently selected.
-        return {
-            arn: arnOverride,
-            name: customizationOverrideName,
-            description: result.description,
+        const customizationFeature = FeatureConfigProvider.getFeature(Features.customizationArnOverride)
+        const arnOverride = customizationFeature?.value.stringValue
+        const customizationOverrideName = customizationFeature?.variation
+        if (arnOverride === undefined) {
+            return baseCustomization
+        } else {
+            return {
+                arn: arnOverride,
+                name: customizationOverrideName,
+                description: baseCustomization.description,
+            }
         }
     }
 }

packages/core/src/test/amazonq/customizationUtil.test.ts

@@ -0,0 +1,94 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as sinon from 'sinon'
+import assert from 'assert'
+import { tryRegister } from '../testUtil'
+import {
+    amazonQScopes,
+    AuthUtil,
+    baseCustomization,
+    Customization,
+    FeatureConfigProvider,
+    getSelectedCustomization,
+    refreshStatusBar,
+    setSelectedCustomization,
+} from '../../codewhisperer'
+import { FeatureContext, globals } from '../../shared'
+import { resetCodeWhispererGlobalVariables } from '../codewhisperer/testUtil'
+import { createSsoProfile, createTestAuth } from '../credentials/testUtil'
+import { SsoConnection } from '../../auth'
+
+const enterpriseSsoStartUrl = 'https://enterprise.awsapps.com/start'
+
+describe('CodeWhisperer-customizationUtils', function () {
+    let auth: ReturnType<typeof createTestAuth>
+    let ssoConn: SsoConnection
+    let featureCustomization: FeatureContext
+
+    before(async function () {
+        createTestAuth(globals.globalState)
+        tryRegister(refreshStatusBar)
+    })
+
+    beforeEach(async function () {
+        auth = createTestAuth(globals.globalState)
+        ssoConn = await auth.createInvalidSsoConnection(
+            createSsoProfile({ startUrl: enterpriseSsoStartUrl, scopes: amazonQScopes })
+        )
+        featureCustomization = {
+            name: 'featureCustomizationName',
+            value: {
+                stringValue: 'featureCustomizationArn',
+            },
+            variation: 'featureCustomizationName',
+        }
+        sinon.stub(FeatureConfigProvider, 'getFeature').returns(featureCustomization)
+
+        sinon.stub(AuthUtil.instance, 'isConnectionExpired').returns(false)
+        sinon.stub(AuthUtil.instance, 'isConnected').returns(true)
+        sinon.stub(AuthUtil.instance, 'isCustomizationFeatureEnabled').value(true)
+        sinon.stub(AuthUtil.instance, 'conn').value(ssoConn)
+
+        await resetCodeWhispererGlobalVariables()
+    })
+
+    afterEach(function () {
+        sinon.restore()
+    })
+
+    it('Returns baseCustomization when not SSO', async function () {
+        sinon.stub(AuthUtil.instance, 'isValidEnterpriseSsoInUse').returns(false)
+        const customization = getSelectedCustomization()
+
+        assert.strictEqual(customization.name, baseCustomization.name)
+    })
+
+    it('Returns selectedCustomization when customization manually selected', async function () {
+        sinon.stub(AuthUtil.instance, 'isValidEnterpriseSsoInUse').returns(true)
+
+        const selectedCustomization: Customization = {
+            arn: 'selectedCustomizationArn',
+            name: 'selectedCustomizationName',
+            description: 'selectedCustomizationDescription',
+        }
+
+        await setSelectedCustomization(selectedCustomization)
+
+        const actualCustomization = getSelectedCustomization()
+
+        assert.strictEqual(actualCustomization.name, selectedCustomization.name)
+    })
+
+    it('Returns AB customization', async function () {
+        sinon.stub(AuthUtil.instance, 'isValidEnterpriseSsoInUse').returns(true)
+
+        await setSelectedCustomization(baseCustomization)
+
+        const returnedCustomization = getSelectedCustomization()
+
+        assert.strictEqual(returnedCustomization.name, featureCustomization.name)
+    })
+})