From 4874197744a4df92da315568f51aa3ee2106a11c Mon Sep 17 00:00:00 2001
From: Zoe Lin <60411978+zixlin7@users.noreply.github.com>
Date: Fri, 13 Oct 2023 12:42:31 -0700
Subject: [PATCH] codewhisperer: suppress token infilling for auto-trigger
 (#3914)

---
 .../editor/CodeWhispererEditorUtil.kt         | 13 +++++++
 .../editor/CodeWhispererEnterHandler.kt       |  5 +++
 .../editor/CodeWhispererTypedHandler.kt       |  6 +++
 .../service/CodeWhispererUserGroupSettings.kt |  8 +++-
 .../CodeWhispererUserActionsTest.kt           | 39 +++++++++++++++++++
 5 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEditorUtil.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEditorUtil.kt
index cf1849c7cf..454a40ba9d 100644
--- a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEditorUtil.kt
+++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEditorUtil.kt
@@ -16,6 +16,8 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.language.programmi
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.CaretContext
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.CaretPosition
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.FileContextInfo
+import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroup
+import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroupSettings
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.LEFT_CONTEXT_ON_CURRENT_LINE
 import java.awt.Point
@@ -90,6 +92,17 @@ object CodeWhispererEditorUtil {
         )
     }
 
+    fun shouldSkipInvokingBasedOnRightContext(editor: Editor): Boolean {
+        val caretContext = runReadAction { CodeWhispererEditorUtil.extractCaretContext(editor) }
+        val rightContextLines = caretContext.rightFileContext.split(Regex("\r?\n"))
+        val rightContextCurrentLine = if (rightContextLines.isEmpty()) "" else rightContextLines[0]
+
+        return CodeWhispererUserGroupSettings.getInstance().getUserGroup() == CodeWhispererUserGroup.RightContext &&
+            rightContextCurrentLine.isNotEmpty() &&
+            !rightContextCurrentLine.startsWith(" ") &&
+            rightContextCurrentLine.trim() != ("}")
+    }
+
     /**
      * Checks if the [otherRange] overlaps this TextRange. Note that the comparison is `<` because the endOffset of TextRange is exclusive.
      */
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEnterHandler.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEnterHandler.kt
index 8fac1cb2a7..db1fe23541 100644
--- a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEnterHandler.kt
+++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEnterHandler.kt
@@ -9,6 +9,7 @@ import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.editor.Caret
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.editor.actionSystem.EditorActionHandler
+import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorUtil.shouldSkipInvokingBasedOnRightContext
 import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererAutoTriggerService
 import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererAutomatedTriggerType
 
@@ -16,6 +17,10 @@ class CodeWhispererEnterHandler(private val originalHandler: EditorActionHandler
     override fun executeWriteAction(editor: Editor, caret: Caret?, dataContext: DataContext?) {
         originalHandler.execute(editor, caret, dataContext)
 
+        if (shouldSkipInvokingBasedOnRightContext(editor)) {
+            return
+        }
+
         ApplicationManager.getApplication().executeOnPooledThread {
             CodeWhispererAutoTriggerService.getInstance().tryInvokeAutoTrigger(editor, CodeWhispererAutomatedTriggerType.Enter())
         }
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererTypedHandler.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererTypedHandler.kt
index 9069dbf0e1..1884569409 100644
--- a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererTypedHandler.kt
+++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererTypedHandler.kt
@@ -8,6 +8,7 @@ import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiFile
 import kotlinx.coroutines.Job
+import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorUtil.shouldSkipInvokingBasedOnRightContext
 import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererAutoTriggerService
 import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererAutomatedTriggerType
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants
@@ -17,6 +18,11 @@ class CodeWhispererTypedHandler : TypedHandlerDelegate() {
     override fun charTyped(c: Char, project: Project, editor: Editor, psiFiles: PsiFile): Result {
         triggerOnIdle?.cancel()
 
+        if (shouldSkipInvokingBasedOnRightContext(editor)
+        ) {
+            return Result.CONTINUE
+        }
+
         // Special Char
         if (CodeWhispererConstants.SPECIAL_CHARACTERS_LIST.contains(c.toString())) {
             CodeWhispererAutoTriggerService.getInstance().tryInvokeAutoTrigger(editor, CodeWhispererAutomatedTriggerType.SpecialChar(c))
diff --git a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererUserGroupSettings.kt b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererUserGroupSettings.kt
index a1d4828504..058dffaa2b 100644
--- a/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererUserGroupSettings.kt
+++ b/jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererUserGroupSettings.kt
@@ -78,7 +78,10 @@ class CodeWhispererUserGroupSettings : PersistentStateComponent<CodeWhispererUse
 
     @VisibleForTesting
     fun determineUserGroup(): CodeWhispererUserGroup {
-        val group = CodeWhispererUserGroup.Control
+        val randomNum = Math.random()
+        val group = if (randomNum < 1 / 2.0) {
+            CodeWhispererUserGroup.Control
+        } else CodeWhispererUserGroup.RightContext
 
         settings[USER_GROUP_KEY] = group.name
         version = AwsToolkit.PLUGIN_VERSION
@@ -125,7 +128,8 @@ interface CodeWhispererGroup
 enum class CodeWhispererUserGroup : CodeWhispererGroup {
     Control,
     CrossFile,
-    Classifier
+    Classifier,
+    RightContext,
 }
 
 enum class CodeWhispererExpThresholdGroup : CodeWhispererGroup {
diff --git a/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererUserActionsTest.kt b/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererUserActionsTest.kt
index 341c51cf46..734b5797df 100644
--- a/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererUserActionsTest.kt
+++ b/jetbrains-core/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererUserActionsTest.kt
@@ -13,9 +13,11 @@ import com.intellij.openapi.actionSystem.IdeActions.ACTION_EDITOR_MOVE_LINE_STAR
 import com.intellij.openapi.actionSystem.IdeActions.ACTION_EDITOR_SELECT_WORD_AT_CARET
 import com.intellij.openapi.actionSystem.IdeActions.ACTION_EDITOR_TEXT_END_WITH_SELECTION
 import com.intellij.openapi.actionSystem.IdeActions.ACTION_EDITOR_TEXT_START_WITH_SELECTION
+import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.command.WriteCommandAction
 import com.intellij.openapi.editor.event.VisibleAreaEvent
 import com.intellij.openapi.ui.popup.JBPopup
+import com.intellij.testFramework.replaceService
 import com.intellij.testFramework.runInEdtAndWait
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.Before
@@ -28,12 +30,15 @@ import org.mockito.kotlin.doReturn
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.timeout
 import org.mockito.kotlin.verify
+import org.mockito.kotlin.whenever
 import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.javaFileName
 import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonFileName
 import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonTestLeftContext
 import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
 import software.aws.toolkits.jetbrains.services.codewhisperer.popup.listeners.CodeWhispererScrollListener
 import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererInvocationStatus
+import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroup
+import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererUserGroupSettings
 import java.awt.Rectangle
 
 class CodeWhispererUserActionsTest : CodeWhispererTestBase() {
@@ -147,6 +152,40 @@ class CodeWhispererUserActionsTest : CodeWhispererTestBase() {
         }
     }
 
+    @Test
+    fun `test special characters should not trigger CodeWhisperer for user group when there is immediate right context`() {
+        testInputSpecialCharWithRightContext("add", false)
+    }
+
+    @Test
+    fun `test special characters should trigger CodeWhisperer for user group when it is not immediate or is single }`() {
+        testInputSpecialCharWithRightContext("}", true)
+        testInputSpecialCharWithRightContext(" add", true)
+        testInputSpecialCharWithRightContext("\nadd", true)
+    }
+
+    private fun testInputSpecialCharWithRightContext(rightContext: String, shouldtrigger: Boolean) {
+        val userGroupSetting = mock<CodeWhispererUserGroupSettings>()
+        ApplicationManager.getApplication().replaceService(CodeWhispererUserGroupSettings::class.java, userGroupSetting, disposableRule.disposable)
+
+        whenever(userGroupSetting.getUserGroup()).thenReturn(CodeWhispererUserGroup.RightContext)
+
+        CodeWhispererExplorerActionManager.getInstance().setAutoEnabled(true)
+        setFileContext(pythonFileName, "def", rightContext)
+        projectRule.fixture.type('{')
+        if (shouldtrigger) {
+            val popupCaptor = argumentCaptor<JBPopup>()
+            verify(popupManagerSpy, timeout(5000).atLeastOnce())
+                .showPopup(any(), any(), popupCaptor.capture(), any(), any())
+            runInEdtAndWait {
+                popupManagerSpy.closePopup(popupCaptor.lastValue)
+            }
+        } else {
+            verify(popupManagerSpy, times(0))
+                .showPopup(any(), any(), any(), any(), any())
+        }
+    }
+
     private fun testHittingEnterAfterWhitespaceCharsShouldTriggerCodeWhisperer(prompt: String, times: Int) {
         CodeWhispererExplorerActionManager.getInstance().setAutoEnabled(true)
         setFileContext(pythonFileName, prompt, "")