[CodeWhisperer] Update %code tracker acceptedToken logic (#3230)

Author: Will-ShaoHua

jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEditorListener.kt

@@ -23,7 +23,10 @@ class CodeWhispererEditorListener : EditorFactoryListener {
                     CodeWhispererInvocationStatus.getInstance().documentChanged()
                     editor.project?.let { project ->
                         PsiDocumentManager.getInstance(project).getPsiFile(editor.document)?.codeWhispererLanguage ?. let { language ->
-                            CodeWhispererCodeCoverageTracker.getInstance(language).documentChanged(event)
+                            CodeWhispererCodeCoverageTracker.getInstance(language).apply {
+                                activateTrackerIfNotActive()
+                                documentChanged(event)
+                            }
                         }
                     }
                 }

jetbrains-core/src/software/aws/toolkits/jetbrains/services/codewhisperer/telemetry/CodeWhispererCodeCoverageTracker.kt

@@ -5,10 +5,14 @@ package software.aws.toolkits.jetbrains.services.codewhisperer.telemetry
 
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.application.runReadAction
 import com.intellij.openapi.editor.RangeMarker
 import com.intellij.openapi.editor.event.DocumentEvent
+import com.intellij.openapi.util.Key
+import com.intellij.refactoring.suggested.range
 import com.intellij.util.Alarm
 import com.intellij.util.AlarmFactory
+import info.debatty.java.stringsimilarity.Levenshtein
 import org.jetbrains.annotations.TestOnly
 import software.aws.toolkits.core.utils.debug
 import software.aws.toolkits.core.utils.getLogger
@@ -41,33 +45,60 @@ abstract class CodeWhispererCodeCoverageTracker(
         get() = totalTokens.get()
     val acceptedRecommendationsCount: Int
         get() = rangeMarkers.size
+    private val isActive: AtomicBoolean = AtomicBoolean(false)
     private val alarm = AlarmFactory.getInstance().create(Alarm.ThreadToUse.POOLED_THREAD, this)
     private val isShuttingDown = AtomicBoolean(false)
     private var startTime: Instant = Instant.now()
 
-    init {
+    @Synchronized
+    fun activateTrackerIfNotActive() {
+        if (!isTelemetryEnabled() || isActive.get()) return
         val conn = ApplicationManager.getApplication().messageBus.connect()
         conn.subscribe(
             CodeWhispererPopupManager.CODEWHISPERER_USER_ACTION_PERFORMED,
             object : CodeWhispererUserActionListener {
                 override fun afterAccept(states: InvocationContext, sessionContext: SessionContext, rangeMarker: RangeMarker) {
                     if (states.requestContext.fileContextInfo.programmingLanguage.toCodeWhispererLanguage() != language) return
                     rangeMarkers.add(rangeMarker)
+                    val originalRecommendation = extractRangeMarkerString(rangeMarker)
+                    originalRecommendation?.let {
+                        rangeMarker.putUserData(KEY_REMAINING_RECOMMENDATION, it)
+                    }
                 }
             }
         )
+        startTime = Instant.now()
+        isActive.set(true)
         scheduleCodeWhispererCodeCoverageTracker()
     }
 
-    fun documentChanged(event: DocumentEvent) {
-        // Added this condition to filter out IDE reloading files
+    internal fun documentChanged(event: DocumentEvent) {
+        // When open a file for the first time, IDE will also emit DocumentEvent for loading with `isWholeTextReplaced = true`
+        // Added this condition to filter out those events
         if (event.isWholeTextReplaced) {
             LOG.debug { "event with isWholeTextReplaced flag: $event" }
             if (event.oldTimeStamp == 0L) return
         }
         addAndGetTotalTokens(event.newLength - event.oldLength)
     }
 
+    internal fun extractRangeMarkerString(rangeMarker: RangeMarker): String? = runReadAction {
+        rangeMarker.range?.let { myRange -> rangeMarker.document.getText(myRange) }
+    }
+
+    // With edit distance, complicate usermodification can be considered as simple edit(add, delete, replace),
+    // and thus the unmodified part of recommendation length can be deducted/approximated
+    // ex. (modified > original): originalRecom: foo -> modifiedRecom: fobarbarbaro, distance = 9, delta = 12 - 9 = 3
+    // ex. (modified == original): originalRecom: helloworld -> modifiedRecom: HelloWorld, distance = 2, delta = 10 - 2 = 8
+    // ex. (modified < original): originalRecom: CodeWhisperer -> modifiedRecom: CODE, distance = 12, delta = 13 - 12 = 1
+    internal fun getAcceptedTokensDelta(originalRecommendation: String, modifiedRecommendation: String): Int {
+        val editDistance = getEditDistance(modifiedRecommendation, originalRecommendation).toInt()
+        return maxOf(originalRecommendation.length, modifiedRecommendation.length) - editDistance
+    }
+
+    protected open fun getEditDistance(modifiedString: String, originalString: String): Double =
+        levenshteinChecker.distance(modifiedString, originalString)
+
     private fun flush() {
         try {
             if (isTelemetryEnabled()) emitCodeWhispererCodeContribution()
@@ -103,9 +134,21 @@ abstract class CodeWhispererCodeCoverageTracker(
     }
 
     private fun emitCodeWhispererCodeContribution() {
-        rangeMarkers.forEach {
-            if (!it.isValid) return@forEach
-            addAndGetAcceptedTokens(it.endOffset - it.startOffset)
+        rangeMarkers.forEach { rangeMarker ->
+            if (!rangeMarker.isValid) return@forEach
+            // if users add more code upon the recommendation generated from CodeWhisperer, we consider those added part as userToken but not CwsprTokens
+            val originalRecommendation = rangeMarker.getUserData(KEY_REMAINING_RECOMMENDATION)
+            val modifiedRecommendation = extractRangeMarkerString(rangeMarker)
+            if (originalRecommendation == null || modifiedRecommendation == null) {
+                LOG.debug {
+                    "failed to get accepted recommendation. " +
+                        "OriginalRecommendation is null: ${originalRecommendation == null}; " +
+                        "ModifiedRecommendation is null: ${modifiedRecommendation == null}"
+                }
+                return@forEach
+            }
+            val delta = getAcceptedTokensDelta(originalRecommendation, modifiedRecommendation)
+            addAndGetAcceptedTokens(delta)
         }
 
         // percentage == null means totalTokens == 0 and users are not editing the document, thus we shouldn't emit telemetry for this
@@ -137,6 +180,10 @@ abstract class CodeWhispererCodeCoverageTracker(
     }
 
     companion object {
+        @JvmStatic
+        protected val levenshteinChecker = Levenshtein()
+        private const val REMAINING_RECOMMENDATION = "remainingRecommendation"
+        private val KEY_REMAINING_RECOMMENDATION = Key<String>(REMAINING_RECOMMENDATION)
         private val LOG = getLogger<CodeWhispererCodeCoverageTracker>()
         private val instances: MutableMap<CodewhispererLanguage, CodeWhispererCodeCoverageTracker> = mutableMapOf()