CodeWhisperer: ab testing feature (#4109)

* CodeWhisperer: Add OptOut field to GenerateCompletions API (#3989)

Add OptOut field to GenerateCompletions API

* CodeWhisperer: Add feature fetching component (#4047)

* CodeWhisperer: Add feature fetching component

1. Add a component to fetch feature assignments every 30 mins,
calling ListFeatureEvaluations API and cache values in memory.

2. Add a field in UserTriggerDecision codewhispererFeatureEvaluations
to record all the feature configs in a string, which can be queried
on Kibana

* remove unused function

* implment a fetchFeatureConfigs command to resolve cicular dependency

* add a test for getFeatureConfigsTelemetry

* Add UserContext, suggestionCount and clientId into STE (#4108)

* use const instead of let (#4115)

---------

Co-authored-by: aws-toolkit-automation <[email protected]>
Co-authored-by: Justin M. Keyes <[email protected]>

Author: 3 people

src/codewhisperer/activation.ts

@@ -39,6 +39,7 @@ import {
     applySecurityFix,
     signoutCodeWhisperer,
     showManageCwConnections,
+    fetchFeatureConfigsCmd,
 } from './commands/basicCommands'
 import { sleep } from '../shared/utilities/timeoutUtils'
 import { ReferenceLogViewProvider } from './service/referenceLogViewProvider'
@@ -206,6 +207,8 @@ export async function activate(context: ExtContext): Promise<void> {
         selectCustomizationPrompt.register(),
         // notify new customizations
         notifyNewCustomizationsCmd.register(),
+        // fetch feature configs
+        fetchFeatureConfigsCmd.register(),
         /**
          * On recommendation acceptance
          */

src/codewhisperer/client/codewhisperer.ts

@@ -7,7 +7,12 @@ import { AWSError, Credentials, Service } from 'aws-sdk'
 import globals from '../../shared/extensionGlobals'
 import * as CodeWhispererClient from './codewhispererclient'
 import * as CodeWhispererUserClient from './codewhispereruserclient'
-import { ListAvailableCustomizationsResponse, SendTelemetryEventRequest } from './codewhispereruserclient'
+import {
+    ListAvailableCustomizationsResponse,
+    ListFeatureEvaluationsRequest,
+    ListFeatureEvaluationsResponse,
+    SendTelemetryEventRequest,
+} from './codewhispereruserclient'
 import * as CodeWhispererConstants from '../models/constants'
 import { ServiceOptions } from '../../shared/awsClientBuilder'
 import { hasVendedIamCredentials } from '../../auth/auth'
@@ -22,6 +27,9 @@ import { session } from '../util/codeWhispererSession'
 import { getLogger } from '../../shared/logger'
 import { indent } from '../../shared/utilities/textUtilities'
 import { keepAliveHeader } from './agent'
+import { getOptOutPreference } from '../util/commonUtil'
+import * as os from 'os'
+import { getClientId } from '../../shared/telemetry/util'
 
 export type ProgrammingLanguage = Readonly<
     CodeWhispererClient.ProgrammingLanguage | CodeWhispererUserClient.ProgrammingLanguage
@@ -219,17 +227,46 @@ export class DefaultCodeWhispererClient {
     }
 
     public async sendTelemetryEvent(request: SendTelemetryEventRequest) {
-        const requestWithOptOut: SendTelemetryEventRequest = {
+        const requestWithCommonFields: SendTelemetryEventRequest = {
             ...request,
-            optOutPreference: globals.telemetry.telemetryEnabled ? 'OPTIN' : 'OPTOUT',
+            optOutPreference: getOptOutPreference(),
+            userContext: {
+                ideCategory: 'VSCODE',
+                operatingSystem: this.getOperatingSystem(),
+                product: 'CodeWhisperer',
+                clientId: await getClientId(globals.context.globalState),
+            },
         }
         if (!AuthUtil.instance.isValidEnterpriseSsoInUse() && !globals.telemetry.telemetryEnabled) {
             return
         }
-        const response = await (await this.createUserSdkClient()).sendTelemetryEvent(requestWithOptOut).promise()
+        const response = await (await this.createUserSdkClient()).sendTelemetryEvent(requestWithCommonFields).promise()
         getLogger().debug(`codewhisperer: sendTelemetryEvent requestID: ${response.$response.requestId}`)
     }
 
+    public async listFeatureEvaluations(): Promise<ListFeatureEvaluationsResponse> {
+        const request: ListFeatureEvaluationsRequest = {
+            userContext: {
+                ideCategory: 'VSCODE',
+                operatingSystem: this.getOperatingSystem(),
+                product: 'CodeWhisperer',
+                clientId: await getClientId(globals.context.globalState),
+            },
+        }
+        return (await this.createUserSdkClient()).listFeatureEvaluations(request).promise()
+    }
+
+    private getOperatingSystem(): string {
+        const osId = os.platform() // 'darwin', 'win32', 'linux', etc.
+        if (osId === 'darwin') {
+            return 'MAC'
+        } else if (osId === 'win32') {
+            return 'WINDOWS'
+        } else {
+            return 'LINUX'
+        }
+    }
+
     /**
      * @description Use this function to start the transformation job.
      * @param request

src/codewhisperer/client/user-service-2.json

@@ -53,6 +53,8 @@
             "output": { "shape": "CreateUploadUrlResponse" },
             "errors": [
                 { "shape": "ThrottlingException" },
+                { "shape": "ConflictException" },
+                { "shape": "ResourceNotFoundException" },
                 { "shape": "InternalServerException" },
                 { "shape": "ValidationException" },
                 { "shape": "AccessDeniedException" }
@@ -116,6 +118,8 @@
             "output": { "shape": "GetTaskAssistCodeGenerationResponse" },
             "errors": [
                 { "shape": "ThrottlingException" },
+                { "shape": "ConflictException" },
+                { "shape": "ResourceNotFoundException" },
                 { "shape": "InternalServerException" },
                 { "shape": "ValidationException" },
                 { "shape": "AccessDeniedException" }
@@ -241,6 +245,8 @@
             "output": { "shape": "StartTaskAssistCodeGenerationResponse" },
             "errors": [
                 { "shape": "ThrottlingException" },
+                { "shape": "ConflictException" },
+                { "shape": "ResourceNotFoundException" },
                 { "shape": "InternalServerException" },
                 { "shape": "ValidationException" },
                 { "shape": "AccessDeniedException" }
@@ -474,6 +480,7 @@
                 "contentMd5": { "shape": "CreateUploadUrlRequestContentMd5String" },
                 "contentChecksum": { "shape": "CreateUploadUrlRequestContentChecksumString" },
                 "contentChecksumType": { "shape": "ContentChecksumType" },
+                "contentLength": { "shape": "CreateUploadUrlRequestContentLengthLong" },
                 "artifactType": { "shape": "ArtifactType" },
                 "uploadIntent": { "shape": "UploadIntent" },
                 "uploadContext": { "shape": "UploadContext" }
@@ -485,6 +492,11 @@
             "min": 1,
             "sensitive": true
         },
+        "CreateUploadUrlRequestContentLengthLong": {
+            "type": "long",
+            "box": true,
+            "min": 1
+        },
         "CreateUploadUrlRequestContentMd5String": {
             "type": "string",
             "max": 128,
@@ -521,7 +533,7 @@
             "type": "string",
             "max": 950,
             "min": 0,
-            "pattern": "$|^arn:[-.a-z0-9]{1,63}:codewhisperer:([-.a-z0-9]{0,63}:){2}([a-zA-Z0-9-_:/]){1,1023}"
+            "pattern": "arn:[-.a-z0-9]{1,63}:codewhisperer:([-.a-z0-9]{0,63}:){2}([a-zA-Z0-9-_:/]){1,1023}"
         },
         "CustomizationName": {
             "type": "string",
@@ -955,7 +967,8 @@
         "PreSignedUrl": {
             "type": "string",
             "max": 2048,
-            "min": 1
+            "min": 1,
+            "sensitive": true
         },
         "PrimitiveInteger": { "type": "integer" },
         "ProgrammingLanguage": {
@@ -1442,7 +1455,8 @@
             "members": {
                 "ideCategory": { "shape": "IdeCategory" },
                 "operatingSystem": { "shape": "OperatingSystem" },
-                "product": { "shape": "UserContextProductString" }
+                "product": { "shape": "UserContextProductString" },
+                "clientId": { "shape": "UUID" }
             }
         },
         "UserContextProductString": {
@@ -1518,7 +1532,8 @@
                 "recommendationLatencyMilliseconds": { "shape": "Double" },
                 "timestamp": { "shape": "Timestamp" },
                 "suggestionReferenceCount": { "shape": "PrimitiveInteger" },
-                "generatedLine": { "shape": "PrimitiveInteger" }
+                "generatedLine": { "shape": "PrimitiveInteger" },
+                "numberOfRecommendations": { "shape": "PrimitiveInteger" }
             }
         },
         "ValidationException": {

src/codewhisperer/commands/basicCommands.ts

@@ -32,6 +32,7 @@ import { FileSystemCommon } from '../../srcShared/fs'
 import { Mutable } from '../../shared/utilities/tsUtils'
 import { CodeWhispererSource } from './types'
 import { showManageConnections } from '../../auth/ui/vue/show'
+import { FeatureConfigProvider } from '../service/featureConfigProvider'
 
 export const toggleCodeSuggestions = Commands.declare(
     { id: 'aws.codeWhisperer.toggleCodeSuggestion', compositeKey: { 1: 'source' } },
@@ -246,6 +247,13 @@ export const notifyNewCustomizationsCmd = Commands.declare(
     }
 )
 
+export const fetchFeatureConfigsCmd = Commands.declare(
+    { id: 'aws.codeWhisperer.fetchFeatureConfigs', logging: false },
+    () => () => {
+        FeatureConfigProvider.instance.fetchFeatureConfigs()
+    }
+)
+
 export const applySecurityFix = Commands.declare(
     'aws.codeWhisperer.applySecurityFix',
     () => async (issue: CodeScanIssue, filePath: string, source: Component) => {

src/codewhisperer/service/completionProvider.ts

@@ -8,7 +8,6 @@ import * as CodeWhispererConstants from '../models/constants'
 import { runtimeLanguageContext } from '../util/runtimeLanguageContext'
 import { Recommendation } from '../client/codewhisperer'
 import { LicenseUtil } from '../util/licenseUtil'
-import { TelemetryHelper } from '../util/telemetryHelper'
 import { RecommendationHandler } from './recommendationHandler'
 import { session } from '../util/codeWhispererSession'
 /**
@@ -60,7 +59,7 @@ export function getCompletionItem(
             recommendation,
             RecommendationHandler.instance.requestId,
             session.sessionId,
-            TelemetryHelper.instance.triggerType,
+            session.triggerType,
             session.getCompletionType(recommendationIndex),
             languageContext.language,
             references,

src/codewhisperer/service/featureConfigProvider.ts

@@ -0,0 +1,88 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { FeatureValue } from '../client/codewhispereruserclient'
+import { codeWhispererClient as client } from '../client/codewhisperer'
+import { AuthUtil } from '../util/authUtil'
+import { getLogger } from '../../shared/logger'
+
+export class FeatureContext {
+    constructor(public name: string, public variation: string, public value: FeatureValue) {}
+}
+
+const testFeatureName = 'testFeature'
+const featureConfigPollIntervalInMs = 30 * 60 * 1000 // 30 mins
+
+// TODO: add real feature later
+export const featureDefinitions = new Map([
+    [testFeatureName, new FeatureContext(testFeatureName, 'CONTROL', { stringValue: 'testValue' })],
+])
+
+export class FeatureConfigProvider {
+    private featureConfigs = new Map<string, FeatureContext>()
+
+    static #instance: FeatureConfigProvider
+
+    constructor() {
+        this.fetchFeatureConfigs()
+
+        setInterval(this.fetchFeatureConfigs.bind(this), featureConfigPollIntervalInMs)
+    }
+
+    public static get instance() {
+        return (this.#instance ??= new this())
+    }
+
+    async fetchFeatureConfigs(): Promise<void> {
+        if (AuthUtil.instance.isConnectionExpired()) {
+            return
+        }
+
+        getLogger().debug('CodeWhisperer: Fetching feature configs')
+        try {
+            const response = await client.listFeatureEvaluations()
+
+            // Overwrite feature configs from server response
+            response.featureEvaluations.forEach(evaluation => {
+                this.featureConfigs.set(
+                    evaluation.feature,
+                    new FeatureContext(evaluation.feature, evaluation.variation, evaluation.value)
+                )
+            })
+        } catch (e) {
+            getLogger().debug('CodeWhisperer: Error when fetching feature configs', e)
+        }
+        getLogger().debug(`CodeWhisperer: Current feature configs: ${this.getFeatureConfigsTelemetry()}`)
+    }
+
+    // Sample format: "{testFeature: CONTROL}""
+    getFeatureConfigsTelemetry(): string {
+        return `{${Array.from(this.featureConfigs.entries())
+            .map(([name, context]) => `${name}: ${context.variation}`)
+            .join(', ')}}`
+    }
+
+    // TODO: for all feature variations, define a contract that can be enforced upon the implementation of
+    // the business logic.
+    // When we align on a new feature config, client-side will implement specific business logic to utilize
+    // these values by:
+    // 1) Add an entry in featureDefinitions, which is <feature_name> to <feature_context>.
+    // 2) Add a function with name `getXXX`, where XXX refers to the feature name.
+    // 3) Specify the return type: One of the return type string/boolean/Long/Double should be used here.
+    // 4) Specify the key for the `getFeatureValueForKey` helper function which is the feature name.
+    // 5) Specify the corresponding type value getter for the `FeatureValue` class. For example,
+    // if the return type is Long, then the corresponding type value getter is `longValue()`.
+    // 6) Add a test case for this feature.
+    // 7) In case `getXXX()` returns undefined, it should be treated as a default/control group.
+    getTestFeature(): string | undefined {
+        return this.getFeatureValueForKey(testFeatureName).stringValue
+    }
+
+    // Get the feature value for the given key.
+    // In case of a misconfiguration, it will return a default feature value of Boolean true.
+    private getFeatureValueForKey(name: string): FeatureValue {
+        return this.featureConfigs.get(name)?.value ?? featureDefinitions.get(name)?.value ?? { boolValue: true }
+    }
+}

src/codewhisperer/service/inlineCompletionItemProvider.ts

@@ -114,7 +114,7 @@ export class CWInlineCompletionItemProvider implements vscode.InlineCompletionIt
                     truncatedSuggestion,
                     this.requestId,
                     session.sessionId,
-                    TelemetryHelper.instance.triggerType,
+                    session.triggerType,
                     session.getCompletionType(index),
                     runtimeLanguageContext.getLanguageContext(document.languageId).language,
                     r.references,

src/codewhisperer/service/recommendationHandler.ts

@@ -67,6 +67,7 @@ const lock = new AsyncLock({ maxPending: 1 })
 
 export class RecommendationHandler {
     public lastInvocationTime: number
+    // TODO: remove this requestId
     public requestId: string
     private nextToken: string
     private cancellationToken: vscode.CancellationTokenSource
@@ -192,9 +193,9 @@ export class RecommendationHandler {
             } else {
                 session.requestContext = {
                     request: {
-                        fileContext: session.requestContext.request.fileContext,
+                        ...session.requestContext.request,
+                        // Putting nextToken assignment in the end so it overwrites the existing nextToken
                         nextToken: this.nextToken,
-                        supplementalContexts: session.requestContext.request.supplementalContexts,
                     },
                     supplementalMetadata: session.requestContext.supplementalMetadata,
                 }
@@ -207,7 +208,10 @@ export class RecommendationHandler {
         // set start pos for non pagination call or first pagination call
         if (!pagination || (pagination && page === 0)) {
             session.startPos = editor.selection.active
+            session.startCursorOffset = editor.document.offsetAt(session.startPos)
             session.leftContextOfCurrentLine = EditorContext.getLeftContext(editor, session.startPos.line)
+            session.triggerType = triggerType
+            session.autoTriggerType = autoTriggerType
 
             /**
              * Validate request
@@ -248,7 +252,6 @@ export class RecommendationHandler {
                 recommendations = (resp && resp.completions) || []
             }
             invocationResult = 'Succeeded'
-            TelemetryHelper.instance.triggerType = triggerType
             requestId = resp?.$response && resp?.$response?.requestId
             nextToken = resp?.nextToken ? resp?.nextToken : ''
             sessionId = resp?.$response?.httpResponse?.headers['x-amzn-sessionid']
@@ -343,11 +346,8 @@ export class RecommendationHandler {
                     requestId,
                     sessionId,
                     session.recommendations.length + recommendations.length - 1,
-                    triggerType,
-                    autoTriggerType,
                     invocationResult,
                     latency,
-                    session.startPos.line,
                     session.language,
                     session.taskType,
                     reason,
@@ -690,7 +690,7 @@ export class RecommendationHandler {
             telemetry.codewhisperer_perceivedLatency.emit({
                 codewhispererRequestId: this.requestId,
                 codewhispererSessionId: session.sessionId,
-                codewhispererTriggerType: TelemetryHelper.instance.triggerType,
+                codewhispererTriggerType: session.triggerType,
                 codewhispererCompletionType: session.getCompletionType(0),
                 codewhispererLanguage: languageContext.language,
                 duration: performance.now() - this.lastInvocationTime,