feat(shared): support AB testing aws#5443

ivikash · web-flow · commit e836e510440e · 2024-08-28T04:38:40.000-07:00
## Problem
Enable A/B testing for AWS Toolkits and Amazon Q by reusing
CodeWhisperer's A/B testing platform

## Solution
- Abstract out Feature Config from CodeWhisperer to Shared folder.
diff --git a/packages/amazonq/.changes/next-release/Feature-52b6b6ee-a123-4576-bce2-a19b07e93e99.json b/packages/amazonq/.changes/next-release/Feature-52b6b6ee-a123-4576-bce2-a19b07e93e99.json
@@ -0,0 +1,4 @@
+{
+	"type": "Feature",
+	"description": "Support AB testing"
+}
diff --git a/packages/core/src/codewhisperer/activation.ts b/packages/core/src/codewhisperer/activation.ts
@@ -39,7 +39,6 @@ import {
     connectWithCustomization,
     applySecurityFix,
     signoutCodeWhisperer,
-    fetchFeatureConfigsCmd,
     toggleCodeScans,
     registerToolkitApiCallback,
 } from './commands/basicCommands'
@@ -71,6 +70,7 @@ import { registerWebviewErrorHandler } from '../webviews/server'
 import { logAndShowError, logAndShowWebviewError } from '../shared/utilities/logAndShowUtils'
 import { openSettings } from '../shared/settings'
 import { telemetry } from '../shared/telemetry'
+import { FeatureConfigProvider } from '../shared/featureConfig'
 
 let localize: nls.LocalizeFunc
 
@@ -252,8 +252,6 @@ export async function activate(context: ExtContext): Promise<void> {
         selectCustomizationPrompt.register(),
         // notify new customizations
         notifyNewCustomizationsCmd.register(),
-        // fetch feature configs
-        fetchFeatureConfigsCmd.register(),
         /**
          * On recommendation acceptance
          */
@@ -594,6 +592,10 @@ export async function activate(context: ExtContext): Promise<void> {
         )
     }
 
+    void FeatureConfigProvider.instance.fetchFeatureConfigs().catch((error) => {
+        getLogger().error('Failed to fetch feature configs - %s', error)
+    })
+
     await Commands.tryExecute('aws.amazonq.refreshConnectionCallback')
     container.ready()
 }
diff --git a/packages/core/src/codewhisperer/client/codewhisperer.ts b/packages/core/src/codewhisperer/client/codewhisperer.ts
@@ -7,12 +7,7 @@ import { AWSError, Credentials, Service } from 'aws-sdk'
 import globals from '../../shared/extensionGlobals'
 import * as CodeWhispererClient from './codewhispererclient'
 import * as CodeWhispererUserClient from './codewhispereruserclient'
-import {
-    ListAvailableCustomizationsResponse,
-    ListFeatureEvaluationsRequest,
-    ListFeatureEvaluationsResponse,
-    SendTelemetryEventRequest,
-} from './codewhispereruserclient'
+import { ListAvailableCustomizationsResponse, SendTelemetryEventRequest } from './codewhispereruserclient'
 import { ServiceOptions } from '../../shared/awsClientBuilder'
 import { hasVendedIamCredentials } from '../../auth/auth'
 import { CodeWhispererSettings } from '../util/codewhispererSettings'
@@ -267,19 +262,6 @@ export class DefaultCodeWhispererClient {
         getLogger().debug(`codewhisperer: sendTelemetryEvent requestID: ${response.$response.requestId}`)
     }
 
-    public async listFeatureEvaluations(): Promise<ListFeatureEvaluationsResponse> {
-        const request: ListFeatureEvaluationsRequest = {
-            userContext: {
-                ideCategory: 'VSCODE',
-                operatingSystem: getOperatingSystem(),
-                product: 'CodeWhisperer', // TODO: update this?
-                clientId: getClientId(globals.globalState),
-                ideVersion: extensionVersion,
-            },
-        }
-        return (await this.createUserSdkClient()).listFeatureEvaluations(request).promise()
-    }
-
     /**
      * @description Use this function to start the transformation job.
      * @param request
diff --git a/packages/core/src/codewhisperer/commands/basicCommands.ts b/packages/core/src/codewhisperer/commands/basicCommands.ts
@@ -29,7 +29,6 @@ import { applyPatch } from 'diff'
 import { closeSecurityIssueWebview, showSecurityIssueWebview } from '../views/securityIssue/securityIssueWebview'
 import { Mutable } from '../../shared/utilities/tsUtils'
 import { CodeWhispererSource } from './types'
-import { FeatureConfigProvider } from '../service/featureConfigProvider'
 import { TelemetryHelper } from '../util/telemetryHelper'
 import { Auth, AwsConnection } from '../../auth'
 import { once } from '../../shared/utilities/functionUtils'
@@ -302,13 +301,6 @@ function focusQAfterDelay() {
     }, 1000)
 }
 
-export const fetchFeatureConfigsCmd = Commands.declare(
-    { id: 'aws.amazonq.fetchFeatureConfigs', logging: false },
-    () => async () => {
-        await FeatureConfigProvider.instance.fetchFeatureConfigs()
-    }
-)
-
 /**
  * Actually install Amazon Q.
  * Sometimes reload VS Code window after installation is necessary
diff --git a/packages/core/src/codewhisperer/index.ts b/packages/core/src/codewhisperer/index.ts
@@ -53,7 +53,7 @@ export { session } from './util/codeWhispererSession'
 export { onInlineAcceptance } from './commands/onInlineAcceptance'
 export { stopTransformByQ } from './commands/startTransformByQ'
 export { getCompletionItems, getCompletionItem, getLabel } from './service/completionProvider'
-export { featureDefinitions, FeatureConfigProvider } from './service/featureConfigProvider'
+export { featureDefinitions, FeatureConfigProvider } from '../shared/featureConfig'
 export { ReferenceInlineProvider } from './service/referenceInlineProvider'
 export { ReferenceHoverProvider } from './service/referenceHoverProvider'
 export { CWInlineCompletionItemProvider } from './service/inlineCompletionItemProvider'
diff --git a/packages/core/src/codewhisperer/service/transformByQ/transformApiHandler.ts b/packages/core/src/codewhisperer/service/transformByQ/transformApiHandler.ts
@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import * as vscode from 'vscode'
-import * as fs from 'fs-extra'
+import * as fs from 'fs'
 import * as path from 'path'
 import * as os from 'os'
 import * as codeWhisperer from '../../client/codewhisperer'
@@ -343,7 +343,7 @@ export async function zipCode({ dependenciesFolder, humanInTheLoopFlag, modulePa
             const sourceFiles = getFilesRecursively(modulePath, false)
             let sourceFilesSize = 0
             for (const file of sourceFiles) {
-                if ((await fs.stat(file)).isDirectory()) {
+                if (fs.statSync(file).isDirectory()) {
                     getLogger().info('CodeTransformation: Skipping directory, likely a symlink')
                     continue
                 }
diff --git a/packages/core/src/codewhisperer/service/transformByQ/transformationResultsViewProvider.ts b/packages/core/src/codewhisperer/service/transformByQ/transformationResultsViewProvider.ts
@@ -5,7 +5,7 @@
 
 import AdmZip from 'adm-zip'
 import os from 'os'
-import fs from 'fs-extra'
+import fs from 'fs'
 import { parsePatch, applyPatches, ParsedDiff } from 'diff'
 import path from 'path'
 import vscode from 'vscode'
@@ -129,7 +129,7 @@ export class DiffModel {
         changedFiles.forEach((file) => {
             const pathToTmpFile = path.join(pathToTmpSrcDir, file.oldFileName!.substring(2))
             // use mkdirsSync to create parent directories in pathToTmpFile too
-            fs.mkdirsSync(path.dirname(pathToTmpFile))
+            fs.mkdirSync(path.dirname(pathToTmpFile), { recursive: true })
             const pathToOldFile = path.join(pathToWorkspace, file.oldFileName!.substring(2))
             // pathToOldFile will not exist for new files such as summary.md
             if (fs.existsSync(pathToOldFile)) {
diff --git a/packages/core/src/codewhisperer/util/authUtil.ts b/packages/core/src/codewhisperer/util/authUtil.ts
@@ -139,7 +139,6 @@ export class AuthUtil {
             // To check valid connection
             if (this.isValidEnterpriseSsoInUse() || (this.isBuilderIdInUse() && !this.isConnectionExpired())) {
                 // start the feature config polling job
-                await vscode.commands.executeCommand('aws.amazonq.fetchFeatureConfigs')
                 await showAmazonQWalkthroughOnce()
             }
             await this.setVscodeContextProps()
diff --git a/packages/core/src/codewhisperer/util/customizationUtil.ts b/packages/core/src/codewhisperer/util/customizationUtil.ts
@@ -18,7 +18,7 @@ import { showMessageWithUrl } from '../../shared/utilities/messages'
 import { parse } from '@aws-sdk/util-arn-parser'
 import { Commands } from '../../shared/vscode/commands2'
 import { vsCodeState } from '../models/model'
-import { FeatureConfigProvider } from '../service/featureConfigProvider'
+import { FeatureConfigProvider } from '../../shared/featureConfig'
 
 /**
  *
diff --git a/packages/core/src/codewhisperer/util/telemetryHelper.ts b/packages/core/src/codewhisperer/util/telemetryHelper.ts
@@ -25,7 +25,7 @@ import { isAwsError } from '../../shared/errors'
 import { getLogger } from '../../shared/logger'
 import { session } from './codeWhispererSession'
 import { CodeWhispererSupplementalContext } from '../models/model'
-import { FeatureConfigProvider } from '../service/featureConfigProvider'
+import { FeatureConfigProvider } from '../../shared/featureConfig'
 import { CodeScanRemediationsEventType } from '../client/codewhispereruserclient'
 
 export class TelemetryHelper {
diff --git a/packages/core/src/codewhispererChat/controllers/chat/controller.ts b/packages/core/src/codewhispererChat/controllers/chat/controller.ts
@@ -47,7 +47,7 @@ import { randomUUID } from '../../../shared/crypto'
 import { LspController } from '../../../amazonq/lsp/lspController'
 import { CodeWhispererSettings } from '../../../codewhisperer/util/codewhispererSettings'
 import { getSelectedCustomization } from '../../../codewhisperer/util/customizationUtil'
-import { FeatureConfigProvider } from '../../../codewhisperer/service/featureConfigProvider'
+import { FeatureConfigProvider } from '../../../shared/featureConfig'
 import { getHttpStatusCode, getReasonFromSyntaxError } from '../../../shared/errors'
 
 export interface ChatControllerMessagePublishers {
diff --git a/packages/core/src/shared/featureConfig.ts b/packages/core/src/shared/featureConfig.ts
@@ -3,13 +3,21 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { Customization, FeatureValue } from '../client/codewhispereruserclient'
-import { codeWhispererClient as client } from '../client/codewhisperer'
-import { AuthUtil } from '../util/authUtil'
-import { getLogger } from '../../shared/logger'
-import { isBuilderIdConnection, isIdcSsoConnection } from '../../auth/connection'
-import { CodeWhispererSettings } from '../util/codewhispererSettings'
-import globals from '../../shared/extensionGlobals'
+import {
+    Customization,
+    FeatureValue,
+    ListFeatureEvaluationsRequest,
+    ListFeatureEvaluationsResponse,
+} from '../codewhisperer/client/codewhispereruserclient'
+import { codeWhispererClient as client } from '../codewhisperer/client/codewhisperer'
+import { AuthUtil } from '../codewhisperer/util/authUtil'
+import { getLogger } from './logger'
+import { isBuilderIdConnection, isIdcSsoConnection } from '../auth/connection'
+import { CodeWhispererSettings } from '../codewhisperer/util/codewhispererSettings'
+import globals from './extensionGlobals'
+import { getClientId, getOperatingSystem } from './telemetry/util'
+import { extensionVersion } from './vscode/env'
+import { telemetry } from './telemetry'
 
 export class FeatureContext {
     constructor(
@@ -56,22 +64,45 @@ export class FeatureConfigProvider {
         return this._isDataCollectionGroup
     }
 
+    public async listFeatureEvaluations(): Promise<ListFeatureEvaluationsResponse> {
+        const request: ListFeatureEvaluationsRequest = {
+            userContext: {
+                ideCategory: 'VSCODE',
+                operatingSystem: getOperatingSystem(),
+                product: 'CodeWhisperer', // TODO: update this?
+                clientId: getClientId(globals.globalState),
+                ideVersion: extensionVersion,
+            },
+        }
+        return (await client.createUserSdkClient()).listFeatureEvaluations(request).promise()
+    }
+
     async fetchFeatureConfigs(): Promise<void> {
         if (AuthUtil.instance.isConnectionExpired()) {
             return
         }
 
         getLogger().debug('amazonq: Fetching feature configs')
         try {
-            const response = await client.listFeatureEvaluations()
+            const response = await this.listFeatureEvaluations()
 
             // Overwrite feature configs from server response
             response.featureEvaluations.forEach((evaluation) => {
                 this.featureConfigs.set(
                     evaluation.feature,
                     new FeatureContext(evaluation.feature, evaluation.variation, evaluation.value)
                 )
+
+                telemetry.aws_featureConfig.run((span) => {
+                    span.record({
+                        id: evaluation.feature,
+                        featureVariation: evaluation.variation,
+                        featureValue: JSON.stringify(evaluation.value),
+                    })
+                })
             })
+            getLogger().info('AB Testing Cohort Assignments %s', JSON.stringify(response.featureEvaluations))
+
             const customizationArnOverride = this.featureConfigs.get(customizationArnOverrideName)?.value?.stringValue
             if (customizationArnOverride !== undefined) {
                 // Double check if server-side wrongly returns a customizationArn to BID users
@@ -153,4 +184,13 @@ export class FeatureConfigProvider {
     private getFeatureValueForKey(name: string): FeatureValue {
         return this.featureConfigs.get(name)?.value ?? featureDefinitions.get(name)?.value ?? { boolValue: true }
     }
+
+    /**
+     * Map of feature configurations.
+     *
+     * @returns {Map<string, FeatureContext>} A Map containing the feature configurations, where the keys are strings representing the feature names, and the values are FeatureContext objects.
+     */
+    public static getFeatureConfigs(): Map<string, FeatureContext> {
+        return FeatureConfigProvider.instance.featureConfigs
+    }
 }
diff --git a/packages/core/src/test/fake/mockFeatureConfigData.ts b/packages/core/src/test/fake/mockFeatureConfigData.ts
@@ -0,0 +1,14 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { FeatureEvaluation } from '../../codewhisperer'
+
+export const mockFeatureConfigsData: FeatureEvaluation[] = [
+    {
+        feature: 'testFeature',
+        variation: 'TREATMENT',
+        value: { stringValue: 'testValue' },
+    },
+]
diff --git a/packages/core/src/test/globalSetup.test.ts b/packages/core/src/test/globalSetup.test.ts
@@ -24,6 +24,8 @@ import { mapTestErrors, normalizeError, setRunnableTimeout } from './setupUtil'
 import { TelemetryDebounceInfo } from '../shared/vscode/commands2'
 import { disableAwsSdkWarning } from '../shared/awsClientBuilder'
 import { GlobalState } from '../shared/globalState'
+import { FeatureConfigProvider } from '../shared/featureConfig'
+import { mockFeatureConfigsData } from './fake/mockFeatureConfigData'
 
 disableAwsSdkWarning()
 
@@ -45,6 +47,10 @@ export async function mochaGlobalSetup(extensionId: string) {
         } catch (e) {}
         mkdirpSync(testReportDir)
 
+        sinon.stub(FeatureConfigProvider.prototype, 'listFeatureEvaluations').resolves({
+            featureEvaluations: mockFeatureConfigsData,
+        })
+
         // Shows the full error chain when tests fail
         mapTestErrors(this, normalizeError)
 
diff --git a/packages/core/src/test/index.ts b/packages/core/src/test/index.ts
@@ -23,6 +23,7 @@ export * from './codewhisperer/testUtil'
 export * from './credentials/testUtil'
 export * from './testUtil'
 export * from './amazonqFeatureDev/utils'
+export * from './fake/mockFeatureConfigData'
 
 import request from '../shared/request'
 import { stub } from 'sinon'
diff --git a/packages/core/src/test/shared/featureConfig.test.ts b/packages/core/src/test/shared/featureConfig.test.ts
@@ -4,17 +4,14 @@
  */
 
 import assert from 'assert'
-import {
-    FeatureConfigProvider,
-    featureDefinitions,
-    FeatureEvaluation,
-    ListFeatureEvaluationsResponse,
-} from 'aws-core-vscode/codewhisperer'
-import { createSpyClient } from 'aws-core-vscode/test'
 import sinon from 'sinon'
 import { AWSError, Request } from 'aws-sdk'
+import { FeatureConfigProvider, featureDefinitions } from '../../shared/featureConfig'
+import { ListFeatureEvaluationsResponse } from '../../codewhisperer'
+import { createSpyClient } from '../codewhisperer/testUtil'
+import { mockFeatureConfigsData } from '../fake/mockFeatureConfigData'
 
-describe('CodeWhispererFeatureConfigServiceTest', () => {
+describe('FeatureConfigProvider', () => {
     afterEach(function () {
         sinon.restore()
     })
@@ -35,27 +32,18 @@ describe('CodeWhispererFeatureConfigServiceTest', () => {
     })
 
     it('test getFeatureConfigsTelemetry will return expected string', async () => {
-        const testFeatureContext = {
-            feature: 'testFeature',
-            variation: 'TREATMENT',
-            value: 'testValue',
-        } as FeatureEvaluation
-
         const clientSpy = await createSpyClient()
         sinon.stub(clientSpy, 'listFeatureEvaluations').returns({
             promise: () =>
                 Promise.resolve({
                     $response: {
                         requestId: '',
                     },
-                    featureEvaluations: [testFeatureContext],
+                    featureEvaluations: mockFeatureConfigsData,
                 }),
         } as Request<ListFeatureEvaluationsResponse, AWSError>)
 
         await FeatureConfigProvider.instance.fetchFeatureConfigs()
-        assert.strictEqual(
-            FeatureConfigProvider.instance.getFeatureConfigsTelemetry(),
-            `{${testFeatureContext.feature}: ${testFeatureContext.variation}}`
-        )
+        assert.strictEqual(FeatureConfigProvider.instance.getFeatureConfigsTelemetry(), `{testFeature: TREATMENT}`)
     })
 })

-Original file line number
+Diff line change
@@ @@ -0,0 +1,4 @@ @@
 +{
 +	"type": "Feature",
 +	"description": "Support AB testing"
 +}
Original file line number	Diff line number	Diff line change
`@@ -139,7 +139,6 @@ export class AuthUtil {`
`139`	`139`	`// To check valid connection`
`140`	`140`	`if (this.isValidEnterpriseSsoInUse() \|\| (this.isBuilderIdInUse() && !this.isConnectionExpired())) {`
`141`	`141`	`// start the feature config polling job`
`142`		`- await vscode.commands.executeCommand('aws.amazonq.fetchFeatureConfigs')`
`143`	`142`	`await showAmazonQWalkthroughOnce()`
`144`	`143`	`}`
`145`	`144`	`await this.setVscodeContextProps()`
Original file line number	Diff line number	Diff line change
`@@ -18,7 +18,7 @@ import { showMessageWithUrl } from '../../shared/utilities/messages'`
`18`	`18`	`import { parse } from '@aws-sdk/util-arn-parser'`
`19`	`19`	`import { Commands } from '../../shared/vscode/commands2'`
`20`	`20`	`import { vsCodeState } from '../models/model'`
`21`		`-import { FeatureConfigProvider } from '../service/featureConfigProvider'`
	`21`	`+import { FeatureConfigProvider } from '../../shared/featureConfig'`
`22`	`22`
`23`	`23`	`/**`
`24`	`24`	`*`