feat(codewhisperer): improve file context fetching logic #3567

Improve codewhisperer file context fetching logic for generating recommendations.

Author: Will-ShaoHua

.changes/next-release/Feature-1f0018c9-e2b5-4f29-adcd-9eef2d4443fb.json

@@ -0,0 +1,4 @@
+{
+	"type": "Feature",
+	"description": "CodeWhisperer: Improve file context fetching logic"
+}

.changes/next-release/Feature-ce707655-611e-4a2f-997f-9d5b0472f7df.json

@@ -3422,7 +3422,7 @@
         "report": "nyc report --reporter=html --reporter=json"
     },
     "devDependencies": {
-        "@aws-toolkits/telemetry": "^1.0.129",
+        "@aws-toolkits/telemetry": "^1.0.130",
         "@cspotcode/source-map-support": "^0.8.1",
         "@sinonjs/fake-timers": "^10.0.2",
         "@types/adm-zip": "^0.4.34",

package-lock.json

@@ -62,9 +62,8 @@ export async function activate(context: ExtContext): Promise<void> {
 
     if (isCloud9()) {
         await enableDefaultConfigCloud9()
-    } else {
-        determineUserGroup()
     }
+
     /**
      * CodeWhisperer security panel
      */
@@ -215,23 +214,6 @@ export async function activate(context: ExtContext): Promise<void> {
         )
     }
 
-    function determineUserGroup() {
-        const userGroup = context.extensionContext.globalState.get<CodeWhispererConstants.UserGroup | undefined>(
-            CodeWhispererConstants.userGroupKey
-        )
-        if (userGroup === undefined) {
-            const randomNum = Math.random()
-            const result =
-                randomNum <= 1 / 3
-                    ? CodeWhispererConstants.UserGroup.Control
-                    : randomNum <= 2 / 3
-                    ? CodeWhispererConstants.UserGroup.CrossFile
-                    : CodeWhispererConstants.UserGroup.Classifier
-
-            context.extensionContext.globalState.update(CodeWhispererConstants.userGroupKey, result)
-        }
-    }
-
     function getAutoTriggerStatus(): boolean {
         return context.extensionContext.globalState.get<boolean>(CodeWhispererConstants.autoTriggerEnabledKey) || false
     }

package.json

@@ -47,6 +47,9 @@ export type GetCodeScanRequest = Readonly<
 export type ListCodeScanFindingsRequest = Readonly<
     CodeWhispererClient.ListCodeScanFindingsRequest | CodeWhispererUserClient.ListCodeAnalysisFindingsRequest
 >
+export type SupplementalContext = Readonly<
+    CodeWhispererClient.SupplementalContext | CodeWhispererUserClient.SupplementalContext
+>
 export type ArtifactType = Readonly<CodeWhispererClient.ArtifactType | CodeWhispererUserClient.ArtifactType>
 export type ArtifactMap = Readonly<CodeWhispererClient.ArtifactMap | CodeWhispererUserClient.ArtifactMap>
 export type ListCodeScanFindingsResponse =

src/codewhisperer/activation.ts

@@ -635,13 +635,48 @@
             "min": 0,
             "sensitive": true
         },
+        "SupplementalContext": {
+            "type": "structure",
+            "required": ["filePath", "content"],
+            "members": {
+                "filePath": {
+                    "shape": "SupplementalContextFilePathString"
+                },
+                "content": {
+                    "shape": "SupplementalContextContentString"
+                }
+            }
+        },
+        "SupplementalContextFilePathString": {
+            "type": "string",
+            "max": 1024,
+            "min": 1,
+            "sensitive": true
+        },
+        "SupplementalContextContentString": {
+            "type": "string",
+            "max": 5120,
+            "min": 0,
+            "sensitive": true
+        },
+        "SupplementalContextList": {
+            "type": "list",
+            "member": {
+                "shape": "SupplementalContext"
+            },
+            "max": 10,
+            "min": 0
+        },
         "GenerateRecommendationsRequest": {
             "type": "structure",
             "required": ["fileContext"],
             "members": {
                 "fileContext": {
                     "shape": "FileContext"
                 },
+                "supplementalContexts": {
+                    "shape": "SupplementalContextList"
+                },
                 "maxResults": {
                     "shape": "GenerateRecommendationsRequestMaxResultsInteger"
                 },
@@ -855,6 +890,9 @@
                 "maxResults": {
                     "shape": "ListRecommendationsRequestMaxResultsInteger"
                 },
+                "supplementalContexts": {
+                    "shape": "SupplementalContextList"
+                },
                 "nextToken": {
                     "shape": "ListRecommendationsRequestNextTokenString"
                 },

src/codewhisperer/client/codewhisperer.ts

@@ -330,6 +330,38 @@
             "min": 0,
             "sensitive": true
         },
+        "SupplementalContext": {
+            "type": "structure",
+            "required": ["filePath", "content"],
+            "members": {
+                "filePath": {
+                    "shape": "SupplementalContextFilePathString"
+                },
+                "content": {
+                    "shape": "SupplementalContextContentString"
+                }
+            }
+        },
+        "SupplementalContextFilePathString": {
+            "type": "string",
+            "max": 1024,
+            "min": 1,
+            "sensitive": true
+        },
+        "SupplementalContextContentString": {
+            "type": "string",
+            "max": 5120,
+            "min": 0,
+            "sensitive": true
+        },
+        "SupplementalContextList": {
+            "type": "list",
+            "member": {
+                "shape": "SupplementalContext"
+            },
+            "max": 10,
+            "min": 0
+        },
         "GenerateCompletionsRequest": {
             "type": "structure",
             "required": ["fileContext"],
@@ -345,6 +377,9 @@
                 },
                 "referenceTrackerConfiguration": {
                     "shape": "ReferenceTrackerConfiguration"
+                },
+                "supplementalContexts": {
+                    "shape": "SupplementalContextList"
                 }
             }
         },

src/codewhisperer/client/service-2.json

@@ -51,6 +51,8 @@ export const lineBreak = '\n'
 
 export const lineBreakWin = '\r\n'
 
+export const supplementalContextTimeoutInMs = 50
+
 /**
  * Ux of recommendations
  */
@@ -261,3 +263,7 @@ export enum UserGroup {
     CrossFile = 'CrossFile',
     Control = 'Control',
 }
+
+export const isClassifierEnabledKey = 'CODEWHISPERER_CLASSIFIER_TRIGGER_ENABLED'
+
+export const supplemetalContextFetchingTimeoutMsg = 'codewhisperer supplemental context fetching timeout'

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

@@ -5,13 +5,13 @@
 
 import * as os from 'os'
 import * as vscode from 'vscode'
-import globals from '../../shared/extensionGlobals'
 import { CodewhispererLanguage, CodewhispererAutomatedTriggerType } from '../../shared/telemetry/telemetry'
 import { extractContextForCodeWhisperer } from '../util/editorContext'
 import { TelemetryHelper } from '../util/telemetryHelper'
 import * as CodeWhispererConstants from '../models/constants'
 import { runtimeLanguageContext } from '../util/runtimeLanguageContext'
 import { ProgrammingLanguage } from '../client/codewhispereruserclient'
+import { CodeWhispererUserGroupSettings } from '../util/userGroupUtil'
 
 interface normalizedCoefficients {
     readonly cursor: number
@@ -318,12 +318,8 @@ export class ClassifierTrigger {
         this.lastInvocationLineNumber = lineNumber
     }
 
-    public isClassifierEnabled() {
-        const userGroup = globals.context.globalState.get<CodeWhispererConstants.UserGroup | undefined>(
-            CodeWhispererConstants.userGroupKey
-        )
-
-        return userGroup === CodeWhispererConstants.UserGroup.Classifier
+    public isClassifierEnabled(): boolean {
+        return CodeWhispererUserGroupSettings.getUserGroup() === CodeWhispererConstants.UserGroup.Classifier
     }
 
     public getThreshold() {