From ec9803f924523fa5fa0e19540bef37fe55b04a61 Mon Sep 17 00:00:00 2001 From: Sam Fink Date: Wed, 8 Jan 2025 11:43:13 -0800 Subject: [PATCH] Change logic for customization override to not override if user has manually selected a customization --- ...-d8c02c38-bdc4-492a-bf8c-6c35e35087be.json | 4 + packages/core/src/codewhisperer/index.ts | 3 +- .../codewhisperer/util/customizationUtil.ts | 35 ++++--- .../test/amazonq/customizationUtil.test.ts | 94 +++++++++++++++++++ 4 files changed, 121 insertions(+), 15 deletions(-) create mode 100644 packages/amazonq/.changes/next-release/Bug Fix-d8c02c38-bdc4-492a-bf8c-6c35e35087be.json create mode 100644 packages/core/src/test/amazonq/customizationUtil.test.ts diff --git a/packages/amazonq/.changes/next-release/Bug Fix-d8c02c38-bdc4-492a-bf8c-6c35e35087be.json b/packages/amazonq/.changes/next-release/Bug Fix-d8c02c38-bdc4-492a-bf8c-6c35e35087be.json new file mode 100644 index 00000000000..3f2f6330b2f --- /dev/null +++ b/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." +} diff --git a/packages/core/src/codewhisperer/index.ts b/packages/core/src/codewhisperer/index.ts index 54a1c508322..235da682ec9 100644 --- a/packages/core/src/codewhisperer/index.ts +++ b/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' diff --git a/packages/core/src/codewhisperer/util/customizationUtil.ts b/packages/core/src/codewhisperer/util/customizationUtil.ts index d2b14cb5247..e87d17cfdb3 100644 --- a/packages/core/src/codewhisperer/util/customizationUtil.ts +++ b/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, + } } } } diff --git a/packages/core/src/test/amazonq/customizationUtil.test.ts b/packages/core/src/test/amazonq/customizationUtil.test.ts new file mode 100644 index 00000000000..1c2e5de7e29 --- /dev/null +++ b/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 + 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) + }) +})