fix: addonly EDITS should be handled as COMPLETIONS

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

@@ -26,7 +26,7 @@ import {
     SuggestionType,
 } from '../../shared/codeWhispererService'
 import { CodewhispererLanguage, getSupportedLanguageId } from '../../shared/languageDetection'
-import { truncateOverlapWithRightContext } from './utils/mergeRightUtils'
+import { mergeSuggestionsWithRightContext, truncateOverlapWithRightContext } from './utils/mergeRightUtils'
 import { CodeWhispererSession, SessionManager } from './session/sessionManager'
 import { CodePercentageTracker } from './tracker/codePercentageTracker'
 import { getCompletionType, getEndPositionForAcceptedSuggestion, getErrorMessage, safeGet } from '../../shared/utils'
@@ -64,47 +64,6 @@ import { IdleWorkspaceManager } from '../workspaceContext/IdleWorkspaceManager'
 import { URI } from 'vscode-uri'
 import { isUsingIAMAuth } from '../../shared/utils'
 
-const mergeSuggestionsWithRightContext = (
-    rightFileContext: string,
-    suggestions: Suggestion[],
-    includeImportsWithSuggestions: boolean,
-    range?: Range
-): InlineCompletionItemWithReferences[] => {
-    return suggestions.map(suggestion => {
-        const insertText: string = truncateOverlapWithRightContext(rightFileContext, suggestion.content)
-        let references = suggestion.references
-            ?.filter(
-                ref =>
-                    !(
-                        ref.recommendationContentSpan?.start && insertText.length <= ref.recommendationContentSpan.start
-                    ) && insertText.length
-            )
-            .map(r => {
-                return {
-                    licenseName: r.licenseName,
-                    referenceUrl: r.url,
-                    referenceName: r.repository,
-                    position: r.recommendationContentSpan && {
-                        startCharacter: r.recommendationContentSpan.start,
-                        endCharacter: r.recommendationContentSpan.end
-                            ? Math.min(r.recommendationContentSpan.end, insertText.length - 1)
-                            : r.recommendationContentSpan.end,
-                    },
-                }
-            })
-
-        return {
-            itemId: suggestion.itemId,
-            insertText: insertText,
-            range,
-            references: references?.length ? references : undefined,
-            mostRelevantMissingImports: includeImportsWithSuggestions
-                ? suggestion.mostRelevantMissingImports
-                : undefined,
-        }
-    })
-}
-
 export const CodewhispererServerFactory =
     (serviceManager: (credentialsProvider?: any) => AmazonQBaseServiceManager): Server =>
     ({ credentialsProvider, lsp, workspace, telemetry, logging, runtime, sdkInitializator }) => {

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

@@ -38,6 +38,7 @@ import { AmazonQError, AmazonQServiceConnectionExpiredError } from '../../shared
 import { DocumentChangedListener } from './documentChangedListener'
 import { EMPTY_RESULT, EDIT_DEBOUNCE_INTERVAL_MS } from './contants/constants'
 import { StreakTracker } from './tracker/streakTracker'
+import { processEditSuggestion } from './utils/diffUtils'
 
 export class EditCompletionHandler {
     private readonly editsEnabled: boolean
@@ -396,6 +397,13 @@ export class EditCompletionHandler {
                         textDocument?.uri || ''
                     )
 
+                    const processedSuggestion = processEditSuggestion(
+                        suggestion.content,
+                        session.startPosition,
+                        session.document
+                    )
+                    const isInlineEdit = processedSuggestion.type === SuggestionType.EDIT
+
                     if (isSimilarToRejected) {
                         // Mark as rejected in the session
                         session.setSuggestionState(suggestion.itemId, 'Reject')
@@ -405,14 +413,14 @@ export class EditCompletionHandler {
                         // Return empty item that will be filtered out
                         return {
                             insertText: '',
-                            isInlineEdit: true,
+                            isInlineEdit: isInlineEdit,
                             itemId: suggestion.itemId,
                         }
                     }
 
                     return {
-                        insertText: suggestion.content,
-                        isInlineEdit: true,
+                        insertText: processedSuggestion.suggestionContent,
+                        isInlineEdit: isInlineEdit,
                         itemId: suggestion.itemId,
                     }
                 })

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

@@ -0,0 +1,258 @@
+import * as assert from 'assert'
+import { categorizeUnifieddiff, extractAdditions, removeOverlapCodeFromSuggestion } from './diffUtils'
+
+describe('extractAdditions', function () {
+    it('singleline', function () {
+        const udiff = `--- file:///Volumes/workplace/ide/sample_projects/Calculator/src/main/hello/MathUtil.java
++++ file:///Volumes/workplace/ide/sample_projects/Calculator/src/main/hello/MathUtil.java
+@@ -1,9 +1,10 @@
+ public class MathUtil {
+     // write a function to add 2 numbers
+     public static int add(int a, int b) {
+
++        return a + b;
+     }
+
+     // write a function to subtract 2 numbers
+     public static int subtract(int a, int b) {
+         return a - b;`
+
+        const r = extractAdditions(udiff)
+        assert.strictEqual(r, '        return a + b;')
+    })
+
+    it('multiline', function () {
+        const udiff = `--- file:///Volumes/workplace/ide/sample_projects/Calculator/src/main/hello/MathUtil.java
++++ file:///Volumes/workplace/ide/sample_projects/Calculator/src/main/hello/MathUtil.java
+@@ -1,9 +1,17 @@
+ public class MathUtil {
+     // write a function to add 2 numbers
+     public static int add(int a, int b) { 
+
++        if (a > Integer.MAX_VALUE - b){
++            throw new IllegalArgumentException("Overflow!");
++        }
++        else if (a < Integer.MIN_VALUE - b){
++            throw new IllegalArgumentException("Underflow");
++        }
++        else{
++            return a + b;
++        }
+     }
+
+     // write a function to subtract 2 numbers
+     public static int subtract(int a, int b) {
+         return a - b;`
+
+        const r = extractAdditions(udiff)
+        assert.strictEqual(
+            r,
+            `        if (a > Integer.MAX_VALUE - b){
+            throw new IllegalArgumentException("Overflow!");
+        }
+        else if (a < Integer.MIN_VALUE - b){
+            throw new IllegalArgumentException("Underflow");
+        }
+        else{
+            return a + b;
+        }`
+        )
+    })
+})
+
+describe('categorizeUnifieddiffv2 should return correct type (addOnly, edit, deleteOnly)', function () {
+    interface Case {
+        udiff: string
+    }
+
+    describe('addOnly', function () {
+        const addOnlyCases: Case[] = [
+            {
+                udiff: `--- file:///Volumes/workplace/ide/sample_projects/Calculator-2/src/main/hello/MathUtil.java
++++ file:///Volumes/workplace/ide/sample_projects/Calculator-2/src/main/hello/MathUtil.java
+@@ -6,7 +6,11 @@
+
+     // write a function to subtract 2 numbers
+      public static int subtract(int a, int b) {
+         return a - b;
+     }
+-    
++
++    // write a function to multiply 2 numbers
++    public static int multiply(int a, int b) {
++        return a * b;
++    }
+ }`,
+            },
+            {
+                udiff: `--- file:///Volumes/workplace/ide/sample_projects/Calculator-2/src/main/hello/MathUtil.java
++++ file:///Volumes/workplace/ide/sample_projects/Calculator-2/src/main/hello/MathUtil.java
+@@ -6,7 +6,11 @@
+
+     // write a function to subtract 2 numbers
+      public static int subtract(int a, int b) {
+         return a - b;
+     }
+-    
++
++
++    // write a function to multiply 2 numbers
++    public static int multiply(int a, int b) {
++        return a * b;
++    }
+ }`,
+            },
+            {
+                udiff: `--- file:///Volumes/workplace/ide/sample_projects/Calculator-2/src/main/hello/MathUtil.java
++++ file:///Volumes/workplace/ide/sample_projects/Calculator-2/src/main/hello/MathUtil.java
+@@ -6,7 +6,11 @@
+
+     // write a function to subtract 2 numbers
+      public static int subtract(int a, int b) {
+         return a - b;
+     }
+-    
++      
++    // write a function to multiply 2 numbers
++    public static int multiply(int a, int b) {
++        return a * b;
++    }
+ }`,
+            },
+            {
+                udiff: `--- file:///Volumes/workplace/ide/sample_projects/Calculator/src/main/hello/MathUtil.java
++++ file:///Volumes/workplace/ide/sample_projects/Calculator/src/main/hello/MathUtil.java
+@@ -1,9 +1,10 @@
+ public class MathUtil {
+     // write a function to add 2 numbers
+     public static int add(int a, int b) {
+
++        return a + b;
+     }
+
+     // write a function to subtract 2 numbers
+     public static int subtract(int a, int b) {
+         return a - b;`,
+            },
+            {
+                udiff: `--- file:///Volumes/workplace/ide/sample_projects/Calculator-2/src/main/hello/MathUtil.java
++++ file:///Volumes/workplace/ide/sample_projects/Calculator-2/src/main/hello/MathUtil.java
+@@ -3,7 +3,9 @@
+     public static int add(int a, int b) {
+         return a + b;
+     }
+
+     // write a function to subtract 2 numbers
+-    
++    public static int subtract(int a, int b) {
++        return a - b;
++    }
+ }`,
+            },
+            {
+                udiff: `--- file:///Volumes/workplace/ide/sample_projects/Calculator-2/src/main/hello/MathUtil.java
++++ file:///Volumes/workplace/ide/sample_projects/Calculator-2/src/main/hello/MathUtil.java
+@@ -4,8 +4,8 @@
+         return a + b;
+     }
+
+     // write a function to subtract 2 numbers
+     public static int subtract(int a, int b) {
+-        return 
++        return a - b;
+     }
+ }`,
+            },
+        ]
+
+        for (let i = 0; i < addOnlyCases.length; i++) {
+            it(`case ${i}`, function () {
+                const actual = categorizeUnifieddiff(addOnlyCases[i].udiff)
+                assert.strictEqual(actual, 'addOnly')
+            })
+        }
+    })
+
+    describe('edit', function () {
+        const cases: Case[] = [
+            {
+                udiff: `--- a/src/main/hello/MathUtil.java
++++ b/src/main/hello/MathUtil.java
+@@ -1,11 +1,11 @@
+ public class MathUtil {
+     // write a function to add 2 numbers
+-    public static int add(int a, int b) {
++    public static double add(double a, double b) {
+         return a + b;
+     }
+
+     // write a function to subtract 2 numbers
+     public static int subtract(int a, int b) {
+     public static double subtract(double a, double b) {
+         return a - b;
+     }   
+ }`,
+            },
+            {
+                udiff: `--- a/server/aws-lsp-codewhisperer/src/shared/codeWhispererService.ts
++++ b/server/aws-lsp-codewhisperer/src/shared/codeWhispererService.ts
+@@ -502,11 +502,7 @@ export class CodeWhispererServiceToken extends CodeWhispererServiceBase {
+             : undefined
+     }
+
+-    private withProfileArn<T extends object>(request: T): T {
+-        if (!this.profileArn) return request
+-
+-        return { ...request, profileArn: this.profileArn }
+-    }
++   // ddddddddddddddddd
+
+     async generateSuggestions(request: BaseGenerateSuggestionsRequest): Promise<GenerateSuggestionsResponse> {
+         // Cast is now safe because GenerateTokenSuggestionsRequest extends GenerateCompletionsRequest`,
+            },
+        ]
+
+        for (let i = 0; i < cases.length; i++) {
+            it(`case ${i}`, function () {
+                const actual = categorizeUnifieddiff(cases[i].udiff)
+                assert.strictEqual(actual, 'edit')
+            })
+        }
+    })
+})
+
+describe('removeOverlapCodeFromSuggestion', function () {
+    interface Case {
+        s1: string
+        s2: string
+        expected: string
+    }
+
+    const cases: Case[] = [
+        {
+            s1: `return`,
+            s2: `return a + b`,
+            expected: ` a + b`,
+        },
+        {
+            s1: `    `,
+            s2: `return a + b`,
+            expected: `return a + b`,
+        },
+        {
+            s1: `    System.out.print`,
+            s2: `println("a + b =", a + b)
+return a + b`,
+            expected: `ln("a + b =", a + b)
+return a + b`,
+        },
+    ]
+
+    for (let i = 0; i < cases.length; i++) {
+        it(`test case${i}`, function () {
+            const c = cases[i]
+            const actual = removeOverlapCodeFromSuggestion(c.s1, c.s2)
+            assert.strictEqual(actual, c.expected)
+        })
+    }
+})