feat(codewhisperer): ignore whitespace in right context #3255

* do not trim if there is no overlap
* do not trim suggestion
* update right context range

Author: zixlin7

.changes/next-release/Feature-c7b6b67a-45c3-452e-b791-5392cef57962.json

@@ -0,0 +1,4 @@
+{
+	"type": "Feature",
+	"description": "CodeWhisperer improves right context handling"
+}

src/codewhisperer/service/inlineCompletionService.ts

@@ -88,16 +88,18 @@ export class CWInlineCompletionItemProvider implements vscode.InlineCompletionIt
     }
 
     truncateOverlapWithRightContext(document: vscode.TextDocument, suggestion: string): string {
-        let rightContextRange: vscode.Range | undefined = undefined
+        const trimmedSuggestion = suggestion.trim()
         const pos = vscode.window.activeTextEditor?.selection.active || RecommendationHandler.instance.startPos
-        if (suggestion.split(/\r?\n/).length > 1) {
-            rightContextRange = new vscode.Range(pos, document.positionAt(document.offsetAt(pos) + suggestion.length))
+        // limit of 5000 for right context matching
+        const rightContext = document.getText(new vscode.Range(pos, document.positionAt(document.offsetAt(pos) + 5000)))
+        const overlap = getPrefixSuffixOverlap(trimmedSuggestion, rightContext.trim())
+        const overlapIndex = suggestion.lastIndexOf(overlap)
+        if (overlapIndex >= 0) {
+            const truncated = suggestion.slice(0, overlapIndex)
+            return truncated.trim().length ? truncated : ''
         } else {
-            rightContextRange = new vscode.Range(pos, document.lineAt(pos).range.end)
+            return suggestion
         }
-        const rightContext = document.getText(rightContextRange)
-        const overlap = getPrefixSuffixOverlap(suggestion, rightContext)
-        return suggestion.slice(0, suggestion.length - overlap.length)
     }
 
     getInlineCompletionItem(

src/test/codewhisperer/service/inlineCompletionService.test.ts

@@ -101,6 +101,62 @@ describe('inlineCompletionService', function () {
         })
     })
 
+    describe('truncateOverlapWithRightContext', function () {
+        const fileName = 'test.py'
+        const language = 'python'
+        let rightContext = 'return target\n'
+        const doc = `import math\ndef two_sum(nums, target):\n`
+        const provider = new CWInlineCompletionItemProvider(0, 0)
+
+        it('removes overlap with right context from suggestion', async function () {
+            const mockSuggestion = 'return target\n'
+            const mockEditor = createMockTextEditor(`${doc}${rightContext}`, fileName, language)
+            const result = provider.truncateOverlapWithRightContext(mockEditor.document, mockSuggestion)
+            assert.strictEqual(result, '')
+        })
+
+        it('only removes the overlap part from suggestion', async function () {
+            const mockSuggestion = 'print(nums)\nreturn target\n'
+            const mockEditor = createMockTextEditor(`${doc}${rightContext}`, fileName, language)
+            const result = provider.truncateOverlapWithRightContext(mockEditor.document, mockSuggestion)
+            assert.strictEqual(result, 'print(nums)\n')
+        })
+
+        it('only removes the last overlap pattern from suggestion', async function () {
+            const mockSuggestion = 'return target\nprint(nums)\nreturn target\n'
+            const mockEditor = createMockTextEditor(`${doc}${rightContext}`, fileName, language)
+            const result = provider.truncateOverlapWithRightContext(mockEditor.document, mockSuggestion)
+            assert.strictEqual(result, 'return target\nprint(nums)\n')
+        })
+
+        it('returns empty string if the remaining suggestion only contains white space', async function () {
+            const mockSuggestion = 'return target\n     '
+            const mockEditor = createMockTextEditor(`${doc}${rightContext}`, fileName, language)
+            const result = provider.truncateOverlapWithRightContext(mockEditor.document, mockSuggestion)
+            assert.strictEqual(result, '')
+        })
+
+        it('returns the original suggestion if no match found', async function () {
+            const mockSuggestion = 'import numpy\n'
+            const mockEditor = createMockTextEditor(`${doc}${rightContext}`, fileName, language)
+            const result = provider.truncateOverlapWithRightContext(mockEditor.document, mockSuggestion)
+            assert.strictEqual(result, 'import numpy\n')
+        })
+        it('ignores the space at the end of recommendation', async function () {
+            const mockSuggestion = 'return target\n\n\n\n\n'
+            const mockEditor = createMockTextEditor(`${doc}${rightContext}`, fileName, language)
+            const result = provider.truncateOverlapWithRightContext(mockEditor.document, mockSuggestion)
+            assert.strictEqual(result, '')
+        })
+        it('ignores the space at the start of right context', async function () {
+            rightContext = '\n\n\n\nreturn target'
+            const mockEditor = createMockTextEditor(`${doc}${rightContext}`, fileName, language)
+            const mockSuggestion = 'return target\n'
+            const result = provider.truncateOverlapWithRightContext(mockEditor.document, mockSuggestion)
+            assert.strictEqual(result, '')
+        })
+    })
+
     describe('on event change', async function () {
         beforeEach(function () {
             const fakeReferences = [