feat: nep auto trigger

server/aws-lsp-codewhisperer/src/language-server/inline-completion/handler/editCompletionHandler.test.ts

@@ -6,6 +6,7 @@ import * as sinon from 'sinon'
 import { CodeWhispererSession, SessionData, SessionManager } from '../session/sessionManager'
 import { HELLO_WORLD_IN_CSHARP } from '../../../shared/testUtils'
 import { CodeWhispererServiceToken } from '../../../shared/codeWhispererService'
+import * as EditAutotrigger from '../auto-trigger/editPredictionAutoTrigger'
 
 describe('EditCompletionHandler', () => {
     let handler: EditCompletionHandler
@@ -54,7 +55,7 @@ describe('EditCompletionHandler', () => {
         }
         amazonQServiceManager = { getCodewhispererService: sinon.stub().returns(codeWhispererService) }
         cursorTracker = { trackPosition: sinon.stub() }
-        recentEditsTracker = {}
+        recentEditsTracker = { generateEditBasedContext: sinon.stub() }
         rejectedEditTracker = { isSimilarToRejected: sinon.stub().returns(false) }
         telemetry = { emitMetric: sinon.stub() }
         telemetryService = { emitUserTriggerDecision: sinon.stub() }
@@ -225,7 +226,7 @@ describe('EditCompletionHandler', () => {
         })
     })
 
-    describe('documentChanged', () => {
+    describe.skip('documentChanged', () => {
         it('should set hasDocumentChangedSinceInvocation when waiting', () => {
             handler['debounceTimeout'] = setTimeout(() => {}, 1000) as any
             handler['isWaiting'] = true
@@ -337,12 +338,23 @@ describe('EditCompletionHandler', () => {
             position: { line: 0, character: 0 },
             context: { triggerKind: InlineCompletionTriggerKind.Automatic },
         }
+
+        afterEach('teardown', function () {
+            sinon.restore()
+        })
+
+        function aTriggerStub(flag: boolean): EditAutotrigger.EditClassifier {
+            return {
+                shouldTriggerNep: sinon
+                    .stub()
+                    .returns({ score: 0, threshold: EditAutotrigger.EditClassifier.THRESHOLD, shouldTrigger: flag }),
+            } as any as EditAutotrigger.EditClassifier
+        }
+
         it('should return empty result when shouldTriggerEdits returns false', async () => {
             workspace.getWorkspaceFolder.returns(undefined)
 
-            const shouldTriggerEditsStub = sinon
-                .stub(require('../utils/triggerUtils'), 'shouldTriggerEdits')
-                .returns(false)
+            sinon.stub(EditAutotrigger, 'EditClassifier').returns(aTriggerStub(false))
 
             const result = await handler._invoke(
                 params as any,
@@ -354,15 +366,12 @@ describe('EditCompletionHandler', () => {
             )
 
             assert.deepEqual(result, EMPTY_RESULT)
-            shouldTriggerEditsStub.restore()
         })
 
         it('should create session and call generateSuggestions when trigger is valid', async () => {
             workspace.getWorkspaceFolder.returns(undefined)
 
-            const shouldTriggerEditsStub = sinon
-                .stub(require('../utils/triggerUtils'), 'shouldTriggerEdits')
-                .returns(true)
+            sinon.stub(EditAutotrigger, 'EditClassifier').returns(aTriggerStub(true))
             codeWhispererService.constructSupplementalContext.resolves(null)
             codeWhispererService.generateSuggestions.resolves({
                 suggestions: [{ itemId: 'item-1', content: 'test content' }],
@@ -380,7 +389,6 @@ describe('EditCompletionHandler', () => {
 
             assert.strictEqual(result.items.length, 1)
             sinon.assert.called(codeWhispererService.generateSuggestions)
-            shouldTriggerEditsStub.restore()
         })
 
         it('should handle active session and emit telemetry', async () => {
@@ -391,9 +399,7 @@ describe('EditCompletionHandler', () => {
             if (currentSession) {
                 sessionManager.activateSession(currentSession)
             }
-            const shouldTriggerEditsStub = sinon
-                .stub(require('../utils/triggerUtils'), 'shouldTriggerEdits')
-                .returns(true)
+            sinon.stub(EditAutotrigger, 'EditClassifier').returns(aTriggerStub(true))
             codeWhispererService.constructSupplementalContext.resolves(null)
             codeWhispererService.generateSuggestions.resolves({
                 suggestions: [{ itemId: 'item-1', content: 'test content' }],
@@ -410,15 +416,12 @@ describe('EditCompletionHandler', () => {
             )
 
             assert.strictEqual(currentSession?.state, 'DISCARD')
-            shouldTriggerEditsStub.restore()
         })
 
         it('should handle supplemental context when available', async () => {
             workspace.getWorkspaceFolder.returns(undefined)
 
-            const shouldTriggerEditsStub = sinon
-                .stub(require('../utils/triggerUtils'), 'shouldTriggerEdits')
-                .returns(true)
+            sinon.stub(EditAutotrigger, 'EditClassifier').returns(aTriggerStub(true))
             codeWhispererService.constructSupplementalContext.resolves({
                 items: [{ content: 'context', filePath: 'file.ts' }],
                 supContextData: { isUtg: false },
@@ -438,7 +441,6 @@ describe('EditCompletionHandler', () => {
             )
 
             sinon.assert.calledWith(codeWhispererService.generateSuggestions, sinon.match.has('supplementalContexts'))
-            shouldTriggerEditsStub.restore()
         })
     })
 

server/aws-lsp-codewhisperer/src/language-server/inline-completion/handler/editCompletionHandler.ts

@@ -22,7 +22,7 @@ import { CodeWhispererSession, SessionManager } from '../session/sessionManager'
 import { CursorTracker } from '../tracker/cursorTracker'
 import { CodewhispererLanguage, getSupportedLanguageId } from '../../../shared/languageDetection'
 import { WorkspaceFolderManager } from '../../workspaceContext/workspaceFolderManager'
-import { shouldTriggerEdits } from '../utils/triggerUtils'
+import { inferTriggerChar } from '../utils/triggerUtils'
 import {
     emitEmptyUserTriggerDecisionTelemetry,
     emitServiceInvocationFailure,
@@ -36,9 +36,10 @@ import { RejectedEditTracker } from '../tracker/rejectedEditTracker'
 import { getErrorMessage, hasConnectionExpired } from '../../../shared/utils'
 import { AmazonQError, AmazonQServiceConnectionExpiredError } from '../../../shared/amazonQServiceManager/errors'
 import { DocumentChangedListener } from '../documentChangedListener'
-import { EMPTY_RESULT, EDIT_DEBOUNCE_INTERVAL_MS } from '../contants/constants'
+import { EMPTY_RESULT } from '../contants/constants'
 import { StreakTracker } from '../tracker/streakTracker'
 import { processEditSuggestion } from '../utils/diffUtils'
+import { EditClassifier } from '../auto-trigger/editPredictionAutoTrigger'
 
 export class EditCompletionHandler {
     private readonly editsEnabled: boolean
@@ -79,14 +80,15 @@ export class EditCompletionHandler {
      * Also as a followup, ideally it should be a message/event publish/subscribe pattern instead of manual invocation like this
      */
     documentChanged() {
-        if (this.debounceTimeout) {
-            if (this.isWaiting) {
-                this.hasDocumentChangedSinceInvocation = true
-            } else {
-                this.logging.info(`refresh and debounce edits suggestion for another ${EDIT_DEBOUNCE_INTERVAL_MS}`)
-                this.debounceTimeout.refresh()
-            }
-        }
+        // TODO: Remove this entirely once we are sure we dont need debounce
+        // if (this.debounceTimeout) {
+        //     if (this.isWaiting) {
+        //         this.hasDocumentChangedSinceInvocation = true
+        //     } else {
+        //         this.logging.info(`refresh and debounce edits suggestion for another ${EDIT_DEBOUNCE_INTERVAL_MS}`)
+        //         this.debounceTimeout.refresh()
+        //     }
+        // }
     }
 
     async onEditCompletion(
@@ -171,40 +173,18 @@ export class EditCompletionHandler {
             }
         }
 
-        return new Promise<InlineCompletionListWithReferences>(async resolve => {
-            this.debounceTimeout = setTimeout(async () => {
-                try {
-                    this.isWaiting = true
-                    const result = await this._invoke(
-                        params,
-                        startPreprocessTimestamp,
-                        token,
-                        textDocument,
-                        inferredLanguageId,
-                        currentSession
-                    ).finally(() => {
-                        this.isWaiting = false
-                    })
-                    if (this.hasDocumentChangedSinceInvocation) {
-                        this.logging.info(
-                            'EditCompletionHandler - Document changed during execution, resolving empty result'
-                        )
-                        resolve({
-                            sessionId: SessionManager.getInstance('EDITS').getActiveSession()?.id ?? '',
-                            items: [],
-                        })
-                    } else {
-                        this.logging.info('EditCompletionHandler - No document changes, resolving result')
-                        resolve(result)
-                    }
-                } finally {
-                    this.debounceTimeout = undefined
-                    this.hasDocumentChangedSinceInvocation = false
-                }
-            }, EDIT_DEBOUNCE_INTERVAL_MS)
-        }).finally(() => {
+        try {
+            return await this._invoke(
+                params,
+                startPreprocessTimestamp,
+                token,
+                textDocument,
+                inferredLanguageId,
+                currentSession
+            )
+        } finally {
             this.isInProgress = false
-        })
+        }
     }
 
     async _invoke(
@@ -218,53 +198,57 @@ export class EditCompletionHandler {
         // Build request context
         const isAutomaticLspTriggerKind = params.context.triggerKind == InlineCompletionTriggerKind.Automatic
         const maxResults = isAutomaticLspTriggerKind ? 1 : 5
-        const fileContext = getFileContext({
+        const fileContextClss = getFileContext({
             textDocument,
             inferredLanguageId,
             position: params.position,
             workspaceFolder: this.workspace.getWorkspaceFolder(textDocument.uri),
         })
 
+        // TODO: Parametrize these to a util function, duplicate code as inineCompletionHandler
+        const triggerCharacters = inferTriggerChar(fileContextClss, params.documentChangeParams)
+
         const workspaceState = WorkspaceFolderManager.getInstance()?.getWorkspaceState()
         const workspaceId = workspaceState?.webSocketClient?.isConnected() ? workspaceState.workspaceId : undefined
 
-        const qEditsTrigger = shouldTriggerEdits(
-            this.codeWhispererService,
-            fileContext,
-            params,
-            this.cursorTracker,
-            this.recentEditsTracker,
-            this.sessionManager,
-            true
+        const recentEdits = await this.recentEditsTracker.generateEditBasedContext(textDocument)
+        const classifier = new EditClassifier(
+            {
+                fileContext: fileContextClss,
+                triggerChar: triggerCharacters,
+                recentEdits: recentEdits,
+                recentDecisions: this.sessionManager.userDecisionLog.map(it => it.decision),
+            },
+            this.logging
         )
 
-        if (!qEditsTrigger) {
+        const qEditsTrigger = classifier.shouldTriggerNep()
+
+        if (!qEditsTrigger.shouldTrigger) {
             return EMPTY_RESULT
         }
 
         const generateCompletionReq: GenerateSuggestionsRequest = {
-            fileContext: fileContext,
+            fileContext: fileContextClss.toServiceModel(),
             maxResults: maxResults,
             predictionTypes: ['EDITS'],
             workspaceId: workspaceId,
         }
 
-        if (qEditsTrigger) {
-            generateCompletionReq.editorState = {
-                document: {
-                    relativeFilePath: textDocument.uri,
-                    programmingLanguage: {
-                        languageName: generateCompletionReq.fileContext?.programmingLanguage?.languageName,
-                    },
-                    text: textDocument.getText(),
+        generateCompletionReq.editorState = {
+            document: {
+                relativeFilePath: textDocument.uri,
+                programmingLanguage: {
+                    languageName: generateCompletionReq.fileContext?.programmingLanguage?.languageName,
                 },
-                cursorState: {
-                    position: {
-                        line: params.position.line,
-                        character: params.position.character,
-                    },
+                text: textDocument.getText(),
+            },
+            cursorState: {
+                position: {
+                    line: params.position.line,
+                    character: params.position.character,
                 },
-            }
+            },
         }
 
         const supplementalContext = await this.codeWhispererService.constructSupplementalContext(
@@ -306,7 +290,7 @@ export class EditCompletionHandler {
             startPreprocessTimestamp: startPreprocessTimestamp,
             startPosition: params.position,
             triggerType: isAutomaticLspTriggerKind ? 'AutoTrigger' : 'OnDemand',
-            language: fileContext.programmingLanguage.languageName,
+            language: fileContextClss.programmingLanguage.languageName,
             requestContext: generateCompletionReq,
             autoTriggerType: undefined,
             triggerCharacter: '',

server/aws-lsp-codewhisperer/src/language-server/inline-completion/handler/inlineCompletionHandler.ts

@@ -191,7 +191,7 @@ export class InlineCompletionHandler {
                 inferredLanguageId,
                 position: params.position,
                 workspaceFolder: this.workspace.getWorkspaceFolder(textDocument.uri),
-            })
+            }).toServiceModel()
         }
 
         const workspaceState = WorkspaceFolderManager.getInstance()?.getWorkspaceState()

server/aws-lsp-codewhisperer/src/language-server/inline-completion/session/sessionManager.ts

@@ -309,6 +309,10 @@ export class SessionManager {
     private sessionsLog: CodeWhispererSession[] = []
     private maxHistorySize = 5
     // TODO, for user decision telemetry: accepted suggestions (not necessarily the full corresponding session) should be stored for 5 minutes
+    private _userDecisionLog: { sessionId: string; decision: UserTriggerDecision }[] = []
+    get userDecisionLog() {
+        return [...this._userDecisionLog]
+    }
 
     private constructor() {}
 
@@ -352,6 +356,19 @@ export class SessionManager {
 
     closeSession(session: CodeWhispererSession) {
         session.close()
+
+        // Note: it has to be called after session.close() as getAggregatedUserTriggerDecision() will return undefined if getAggregatedUserTriggerDecision() is called before session is closed
+        const decision = session.getAggregatedUserTriggerDecision()
+        // As we only care about AR here, pushing Accept/Reject only
+        if (decision === 'Accept' || decision === 'Reject') {
+            if (this._userDecisionLog.length === 5) {
+                this._userDecisionLog.shift()
+            }
+            this._userDecisionLog.push({
+                sessionId: session.codewhispererSessionId ?? 'undefined',
+                decision: decision,
+            })
+        }
     }
 
     discardSession(session: CodeWhispererSession) {