codewhisperer: suppress token filling from auto-trigger (#3892)

Author: zixlin7

src/codewhisperer/models/constants.ts

@@ -265,6 +265,7 @@ export enum UserGroup {
     Classifier = 'Classifier',
     CrossFile = 'CrossFile',
     Control = 'Control',
+    RightContext = 'RightContext',
 }
 
 export const isClassifierEnabledKey = 'CODEWHISPERER_CLASSIFIER_TRIGGER_ENABLED'

src/codewhisperer/service/keyStrokeHandler.ts

@@ -19,6 +19,8 @@ import { AuthUtil } from '../util/authUtil'
 import { ClassifierTrigger } from './classifierTrigger'
 import { isIamConnection } from '../../auth/connection'
 import { session } from '../util/codeWhispererSession'
+import { extractContextForCodeWhisperer } from '../util/editorContext'
+import { CodeWhispererUserGroupSettings } from '../util/userGroupUtil'
 
 const performance = globalThis.performance ?? require('perf_hooks').performance
 
@@ -111,6 +113,20 @@ export class KeyStrokeHandler {
                 }
             }
 
+            const { rightFileContent } = extractContextForCodeWhisperer(editor)
+            const rightContextLines = rightFileContent.split(/\r?\n/)
+            const rightContextAtCurrentLine = rightContextLines[0]
+            // we do not want to trigger when there is immediate right context on the same line
+            // with "}" being an exception because of IDE auto-complete
+            if (
+                CodeWhispererUserGroupSettings.getUserGroup() === CodeWhispererConstants.UserGroup.RightContext &&
+                rightContextAtCurrentLine.length &&
+                !rightContextAtCurrentLine.startsWith(' ') &&
+                rightContextAtCurrentLine.trim() !== '}'
+            ) {
+                return
+            }
+
             let triggerType: CodewhispererAutomatedTriggerType | undefined
             const changedSource = new DefaultDocumentChangedType(event.contentChanges).checkChangeSource()
 

src/codewhisperer/util/userGroupUtil.ts

@@ -56,7 +56,9 @@ export class CodeWhispererUserGroupSettings {
     }
 
     private guessUserGroup(): UserGroup {
-        return UserGroup.Control
+        const randomNum = Math.random()
+        const result = randomNum <= 1 / 2 ? UserGroup.Control : UserGroup.RightContext
+        return result
     }
 
     static #instance: CodeWhispererUserGroupSettings | undefined

src/test/codewhisperer/service/keyStrokeHandler.test.ts

@@ -20,6 +20,7 @@ import { RecommendationHandler } from '../../../codewhisperer/service/recommenda
 import { isInlineCompletionEnabled } from '../../../codewhisperer/util/commonUtil'
 import { ClassifierTrigger } from '../../../codewhisperer/service/classifierTrigger'
 import { CodeWhispererUserGroupSettings } from '../../../codewhisperer/util/userGroupUtil'
+import * as CodeWhispererConstants from '../../../codewhisperer/models/constants'
 
 const performance = globalThis.performance ?? require('perf_hooks').performance
 
@@ -213,6 +214,62 @@ describe('keyStrokeHandler', function () {
             await KeyStrokeHandler.instance.processKeyStroke(mockEvent, mockEditor, mockClient, config)
             assert.ok(invokeSpy.called)
         })
+
+        it('Should skip invoking if there is immediate right context on the same line and not a single } for the user group', async function () {
+            const casesForSuppressTokenFilling = [
+                {
+                    rightContext: 'add',
+                    shouldInvoke: false,
+                },
+                {
+                    rightContext: '}',
+                    shouldInvoke: true,
+                },
+                {
+                    rightContext: '}    ',
+                    shouldInvoke: true,
+                },
+                {
+                    rightContext: ' add',
+                    shouldInvoke: true,
+                },
+                {
+                    rightContext: '    ',
+                    shouldInvoke: true,
+                },
+                {
+                    rightContext: '\naddTwo',
+                    shouldInvoke: true,
+                },
+            ]
+            casesForSuppressTokenFilling.forEach(async ({ rightContext, shouldInvoke }) => {
+                await testIfRightContextShouldInvoke(
+                    rightContext,
+                    shouldInvoke,
+                    CodeWhispererConstants.UserGroup.RightContext
+                )
+            })
+        })
+
+        it('Should not skip invoking based on right context for control group', async function () {
+            await testIfRightContextShouldInvoke('add', true, CodeWhispererConstants.UserGroup.Control)
+        })
+
+        async function testIfRightContextShouldInvoke(
+            rightContext: string,
+            shouldTrigger: boolean,
+            userGroup: CodeWhispererConstants.UserGroup
+        ) {
+            const mockEditor = createMockTextEditor(rightContext, 'test.js', 'javascript', 1, 1)
+            const mockEvent: vscode.TextDocumentChangeEvent = createTextDocumentChangeEvent(
+                mockEditor.document,
+                new vscode.Range(new vscode.Position(0, 0), new vscode.Position(0, 1)),
+                '{'
+            )
+            CodeWhispererUserGroupSettings.instance.userGroup = userGroup
+            await KeyStrokeHandler.instance.processKeyStroke(mockEvent, mockEditor, mockClient, config)
+            assert.strictEqual(invokeSpy.called, shouldTrigger)
+        }
     })
 
     describe('invokeAutomatedTrigger', function () {