fix unexpected right context merging causes recommednation gets trimmed (#4108)

Author: Will-ShaoHua

.changes/next-release/bugfix-71c6ba6c-d774-49b0-b7fb-81e14a1cb3f0.json

@@ -0,0 +1,4 @@
+{
+  "type" : "bugfix",
+  "description" : "Fix issue where CodeWhisperer suggestions are sometimes trimmed"
+}

jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererRecommendationManager.kt

@@ -4,6 +4,7 @@
 package software.aws.toolkits.jetbrains.services.codewhisperer.service
 
 import com.intellij.openapi.components.service
+import org.jetbrains.annotations.VisibleForTesting
 import software.amazon.awssdk.services.codewhispererruntime.model.Completion
 import software.amazon.awssdk.services.codewhispererruntime.model.Span
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.DetailContext
@@ -134,16 +135,21 @@ class CodeWhispererRecommendationManager {
         val caret = requestContext.editor.caretModel.primaryCaret
         val rightContext = document.charsSequence.subSequence(caret.offset, document.charsSequence.length).toString()
         val recommendationContent = recommendation.content()
+        return findRightContextOverlap(rightContext, recommendationContent)
+    }
+
+    @VisibleForTesting
+    fun findRightContextOverlap(rightContext: String, recommendationContent: String): String {
         val rightContextFirstLine = rightContext.substringBefore("\n")
         val overlap =
             if (rightContextFirstLine.isEmpty()) {
                 val tempOverlap = overlap(recommendationContent, rightContext)
-                if (tempOverlap.isEmpty()) overlap(recommendationContent.trimEnd(), rightContext.trimStart()) else tempOverlap
+                if (tempOverlap.isEmpty()) overlap(recommendationContent.trimEnd(), trimExtraPrefixNewLine(rightContext)) else tempOverlap
             } else {
                 // this is necessary to prevent display issue if first line of right context is not empty
                 var tempOverlap = overlap(recommendationContent, rightContext)
                 if (tempOverlap.isEmpty()) {
-                    tempOverlap = overlap(recommendationContent.trimEnd(), rightContext.trimStart())
+                    tempOverlap = overlap(recommendationContent.trimEnd(), trimExtraPrefixNewLine(rightContext))
                 }
                 if (recommendationContent.substring(0, recommendationContent.length - tempOverlap.length).none { it == '\n' }) {
                     tempOverlap
@@ -166,5 +172,33 @@ class CodeWhispererRecommendationManager {
 
     companion object {
         fun getInstance(): CodeWhispererRecommendationManager = service()
+
+        /**
+         * a function to trim extra prefixing new line character (only leave 1 new line character)
+         * example:
+         *  content = "\n\n\nfoo\n\nbar\nbaz"
+         *  return = "\nfoo\n\nbar\nbaz"
+         *
+         * example:
+         *  content = "\n\n\tfoobar\nbaz"
+         *  return = "\n\tfoobar\nbaz"
+         */
+        fun trimExtraPrefixNewLine(content: String): String {
+            if (content.isEmpty()) {
+                return ""
+            }
+
+            val firstChar = content.first()
+            if (firstChar != '\n') {
+                return content
+            }
+
+            var index = 1
+            while (index < content.length && content[index] == '\n') {
+                index++
+            }
+
+            return firstChar + content.substring(index)
+        }
     }
 }

jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererRecommendationManagerTest.kt

@@ -17,8 +17,10 @@ import org.mockito.kotlin.any
 import org.mockito.kotlin.doReturn
 import org.mockito.kotlin.spy
 import org.mockito.kotlin.stub
+import software.amazon.awssdk.services.codewhispererruntime.model.Completion
 import software.aws.toolkits.core.utils.test.aString
 import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererRecommendationManager
+import software.aws.toolkits.jetbrains.services.codewhisperer.service.RequestContext
 import software.aws.toolkits.jetbrains.utils.rules.PythonCodeInsightTestFixtureRule
 
 class CodeWhispererRecommendationManagerTest {
@@ -72,7 +74,7 @@ class CodeWhispererRecommendationManagerTest {
     fun `test duplicated recommendation after truncation will be discarded`() {
         val userInput = ""
         sut.stub {
-            onGeneric { findRightContextOverlap(any(), any()) } doReturn "}"
+            onGeneric { findRightContextOverlap(any<RequestContext>(), any<Completion>()) } doReturn "}"
             onGeneric { reformatReference(any(), any()) } doReturn aCompletion("def")
         }
         val detail = sut.buildDetailContext(
@@ -91,7 +93,7 @@ class CodeWhispererRecommendationManagerTest {
     fun `test blank recommendation after truncation will be discarded`() {
         val userInput = ""
         sut.stub {
-            onGeneric { findRightContextOverlap(any(), any()) } doReturn "}"
+            onGeneric { findRightContextOverlap(any<RequestContext>(), any<Completion>()) } doReturn "}"
         }
         val detail = sut.buildDetailContext(
             aRequestContext(project),
@@ -102,4 +104,104 @@ class CodeWhispererRecommendationManagerTest {
         assertThat(detail[0].isDiscarded).isTrue
         assertThat(detail[0].isTruncatedOnRight).isTrue
     }
+
+    @Test
+    fun `overlap calculation should trim new line character starting from second character (index 1 of a string)`() {
+        // recommendation is wrapped inside |recommendationContent|
+        /**
+         * public foo() {
+         *     re|turn foo
+         *}|
+         * public bar() {
+         *     return bar
+         * }
+         */
+        var overlap: String = sut.findRightContextOverlap(rightContext = " foo\n}\n\n\npublic bar () {\n\treturn bar\n}", recommendationContent = "turn foo\n}")
+        assertThat(overlap).isEqualTo(" foo\n}")
+
+        /**
+         * public foo() {
+         *     |return foo
+         * }|
+         *
+         * public bar() {
+         *     return bar
+         * }
+         */
+        overlap = sut.findRightContextOverlap(rightContext = "\n\n\n\npublic bar() {\n\treturn bar\n}", recommendationContent = "return foo\n}")
+        assertThat(overlap).isEqualTo("")
+
+        /**
+         * println(|world)|;
+         * String foo = "foo";
+         */
+        overlap = sut.findRightContextOverlap(rightContext = "ld);\nString foo = \"foo\";", recommendationContent = "world)")
+        assertThat(overlap).isEqualTo("ld)")
+
+        /**
+         * return |has_d_at_end|
+         *
+         * def foo:
+         *     pass
+         */
+        overlap = sut.findRightContextOverlap(rightContext = "\n\ndef foo():\n\tpass", recommendationContent = "has_d_at_end")
+        assertThat(overlap).isEqualTo("")
+
+        /**
+         * {
+         *    { foo: foo },
+         *    { bar: bar },
+         *    { |baz: baz }|
+         * }
+         *
+         */
+        overlap = sut.findRightContextOverlap(rightContext = "\n}", recommendationContent = "baz: baz }")
+        assertThat(overlap).isEqualTo("")
+
+        /**
+         * |
+         *
+         *     foo|
+         *
+         */
+        overlap = sut.findRightContextOverlap(rightContext = "\n\n\tfoo}", recommendationContent = "\n\tfoo")
+        assertThat(overlap).isEqualTo("\n\tfoo")
+
+        /** A case we can't cover
+         * def foo():
+         *   |print(foo)|
+         *
+         *
+         *   print(foo)
+         */
+        overlap = sut.findRightContextOverlap(rightContext = "\n\n\n\tprint(foo)", recommendationContent = "print(foo)")
+        assertThat(overlap).isEqualTo("")
+    }
+
+    @Test
+    fun `trim extra prefixing new line character`() {
+        var actual: String = CodeWhispererRecommendationManager.trimExtraPrefixNewLine("")
+        assertThat(actual).isEqualTo("")
+
+        actual = CodeWhispererRecommendationManager.trimExtraPrefixNewLine("f")
+        assertThat(actual).isEqualTo("f")
+
+        actual = CodeWhispererRecommendationManager.trimExtraPrefixNewLine("\n\n")
+        assertThat(actual).isEqualTo("\n")
+
+        actual = CodeWhispererRecommendationManager.trimExtraPrefixNewLine("foo")
+        assertThat(actual).isEqualTo("foo")
+
+        actual = CodeWhispererRecommendationManager.trimExtraPrefixNewLine("\nfoo")
+        assertThat(actual).isEqualTo("\nfoo")
+
+        actual = CodeWhispererRecommendationManager.trimExtraPrefixNewLine("\n\n\nfoo\nbar")
+        assertThat(actual).isEqualTo("\nfoo\nbar")
+
+        actual = CodeWhispererRecommendationManager.trimExtraPrefixNewLine("\n\n  foo\nbar")
+        assertThat(actual).isEqualTo("\n  foo\nbar")
+
+        actual = CodeWhispererRecommendationManager.trimExtraPrefixNewLine("\n\n\tfoo\nbar")
+        assertThat(actual).isEqualTo("\n\tfoo\nbar")
+    }
 }