fix: classifier last token sometimes fail to capture the right values

server/aws-lsp-codewhisperer/src/language-server/inline-completion/auto-trigger/autoTrigger.ts

@@ -3,6 +3,7 @@ import { Logging } from '@aws/language-server-runtimes/server-interface'
 import { FileContext } from '../../../shared/codeWhispererService'
 import typedCoefficients = require('./coefficients.json')
 import { TextDocumentContentChangeEvent } from 'vscode-languageserver-textdocument'
+import { lastTokenFromString } from '../utils/triggerUtils'
 
 type TypedCoefficients = typeof typedCoefficients
 type Coefficients = TypedCoefficients & {
@@ -223,8 +224,7 @@ export const autoTrigger = (
             classifierThreshold: TRIGGER_THRESHOLD,
         }
     }
-    const tokens = leftContextAtCurrentLine.trim().split(' ')
-    const lastToken = tokens[tokens.length - 1]
+    const lastToken = lastTokenFromString(fileContext.leftFileContent)
 
     const keyword = lastToken?.length > 1 ? lastToken : ''
 

server/aws-lsp-codewhisperer/src/language-server/inline-completion/auto-trigger/editPredictionAutoTrigger.ts

@@ -10,6 +10,7 @@ import { EditPredictionConfigManager } from './editPredictionConfig'
 import { CodeWhispererSupplementalContext } from '../../../shared/models/model'
 import { UserTriggerDecision } from '../session/sessionManager'
 import { Logging } from '@aws/language-server-runtimes/server-interface'
+import { lastTokenFromString } from '../utils/triggerUtils'
 
 // The sigmoid function to clamp the auto-trigger result to the (0, 1) range
 const sigmoid = (x: number) => {
@@ -360,8 +361,7 @@ ${params.recentEdits.supplementalContextItems.map(it => it.content)}`)
         const lang = params.fileContext.programmingLanguage
 
         // 7. Keywords
-        const tokens = params.fileContext.leftContextAtCurLine.trim().split(' ') // split(' ') Not strict enough?
-        const lastToken = tokens[tokens.length - 1]
+        const lastToken = lastTokenFromString(params.fileContext.leftFileContent)
 
         // 8. User AR for last 5
         // Cold start we assume 0.3 for AR

server/aws-lsp-codewhisperer/src/language-server/inline-completion/utils/triggerUtils.test.ts

@@ -1,6 +1,6 @@
 import * as assert from 'assert'
 import * as sinon from 'sinon'
-import { shouldTriggerEdits, NepTrigger, isDocumentChangedFromNewLine } from './triggerUtils'
+import { shouldTriggerEdits, NepTrigger, isDocumentChangedFromNewLine, lastTokenFromString } from './triggerUtils'
 import { SessionManager } from '../session/sessionManager'
 import { CursorTracker } from '../tracker/cursorTracker'
 import { RecentEditTracker } from '../tracker/codeEditTracker'
@@ -45,6 +45,50 @@ describe('triggerUtils', () => {
         sinon.restore()
     })
 
+    describe('lastTokenFromString', function () {
+        interface TestCase {
+            input: string
+            expected: string
+        }
+
+        const cases: TestCase[] = [
+            {
+                input: `line=str`,
+                expected: `str`,
+            },
+            {
+                input: `public class Main `,
+                expected: `Main`,
+            },
+            {
+                input: `const foo = 5`,
+                expected: `5`,
+            },
+            {
+                input: `const fooString = 'foo'`,
+                expected: `foo`,
+            },
+            {
+                input: `main(`,
+                expected: `main`,
+            },
+            {
+                input: `public class Main {
+                    // print foo
+                    `,
+                expected: `foo`,
+            },
+        ]
+
+        for (let i = 0; i < cases.length; i++) {
+            const c = cases[i]
+            it(`case ${i}`, function () {
+                const actual = lastTokenFromString(c.input)
+                assert.strictEqual(actual, c.expected)
+            })
+        }
+    })
+
     describe('isDocumentChangedFromNewLine', function () {
         interface TestCase {
             input: string

server/aws-lsp-codewhisperer/src/language-server/inline-completion/utils/triggerUtils.ts

@@ -115,3 +115,17 @@ export function inferTriggerChar(
 export function isDocumentChangedFromNewLine(s: string) {
     return /^\n[\s\t]+$/.test(s)
 }
+
+// TODO: Should refine the logic
+export function lastTokenFromString(str: string): string {
+    const tokens = str.trim().split(/\s+/)
+    if (tokens.length === 0) {
+        return ''
+    }
+
+    // This step can filter out most naive case however doesnt work for `line=str` such style without empty space
+    const candidate = tokens[tokens.length - 1]
+    // Only run regex against a small portion of string instead of the entire str
+    const finalCandate = candidate.match(/\w+/g) || []
+    return finalCandate.at(-1) ?? ''
+}