[Codewhisperer] add serviceInvocation count field in the %code tracker (#3305)

Author: Will-ShaoHua

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

@@ -6,6 +6,7 @@ package software.aws.toolkits.jetbrains.services.codewhisperer.service
 import com.intellij.codeInsight.CodeInsightSettings
 import com.intellij.codeInsight.hint.HintManager
 import com.intellij.openapi.application.ApplicationInfo
+import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.application.runInEdt
 import com.intellij.openapi.components.service
 import com.intellij.openapi.editor.Editor
@@ -17,6 +18,7 @@ import com.intellij.openapi.util.Key
 import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.PsiFile
 import com.intellij.util.concurrency.annotations.RequiresEdt
+import com.intellij.util.messages.Topic
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.isActive
 import kotlinx.coroutines.launch
@@ -142,6 +144,8 @@ class CodeWhispererService {
                     val responseContext = ResponseContext(sessionId, completionType)
                     logServiceInvocation(requestId, requestContext, responseContext, response.recommendations(), latency, null)
                     lastRecommendationIndex += response.recommendations().size
+                    ApplicationManager.getApplication().messageBus.syncPublisher(CODEWHISPERER_CODE_COMPLETION_PERFORMED)
+                        .onSuccess(requestContext.fileContextInfo)
                     CodeWhispererTelemetryService.getInstance().sendServiceInvocationEvent(
                         requestId,
                         requestContext,
@@ -564,6 +568,10 @@ class CodeWhispererService {
 
     companion object {
         private val LOG = getLogger<CodeWhispererService>()
+        val CODEWHISPERER_CODE_COMPLETION_PERFORMED: Topic<CodeWhispererCodeCompletionServiceListener> = Topic.create(
+            "CodeWhisperer code completion service invoked",
+            CodeWhispererCodeCompletionServiceListener::class.java
+        )
         val KEY_CODEWHISPERER_METADATA: Key<CodeWhispererMetadata> = Key.create("codewhisperer.metadata")
         fun getInstance(): CodeWhispererService = service()
         const val KET_SESSION_ID = "x-amzn-SessionId"
@@ -601,3 +609,7 @@ data class ResponseContext(
     val sessionId: String,
     val completionType: CodewhispererCompletionType
 )
+
+interface CodeWhispererCodeCompletionServiceListener {
+    fun onSuccess(fileContextInfo: FileContextInfo) {}
+}

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

@@ -18,10 +18,13 @@ import org.jetbrains.annotations.TestOnly
 import software.aws.toolkits.core.utils.debug
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.jetbrains.services.codewhisperer.language.toCodeWhispererLanguage
+import software.aws.toolkits.jetbrains.services.codewhisperer.model.FileContextInfo
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.SessionContext
 import software.aws.toolkits.jetbrains.services.codewhisperer.popup.CodeWhispererPopupManager
 import software.aws.toolkits.jetbrains.services.codewhisperer.popup.CodeWhispererUserActionListener
+import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererCodeCompletionServiceListener
+import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.TOTAL_SECONDS_IN_MINUTE
 import software.aws.toolkits.telemetry.CodewhispererLanguage
 import software.aws.toolkits.telemetry.CodewhispererTelemetry
@@ -35,7 +38,8 @@ abstract class CodeWhispererCodeCoverageTracker(
     private val timeWindowInSec: Long,
     private val language: CodewhispererLanguage,
     private val rangeMarkers: MutableList<RangeMarker>,
-    private val fileToTokens: MutableMap<Document, CodeCoverageTokens>
+    private val fileToTokens: MutableMap<Document, CodeCoverageTokens>,
+    private val myServiceInvocationCount: AtomicInteger
 ) : Disposable {
     val percentage: Int?
         get() = if (totalTokensSize != 0) calculatePercentage(acceptedTokensSize, totalTokensSize) else null
@@ -53,6 +57,8 @@ abstract class CodeWhispererCodeCoverageTracker(
         }
     val acceptedRecommendationsCount: Int
         get() = rangeMarkers.size
+    val serviceInvocationCount: Int
+        get() = myServiceInvocationCount.get()
     private val isActive: AtomicBoolean = AtomicBoolean(false)
     private val alarm = AlarmFactory.getInstance().create(Alarm.ThreadToUse.POOLED_THREAD, this)
     private val isShuttingDown = AtomicBoolean(false)
@@ -77,6 +83,17 @@ abstract class CodeWhispererCodeCoverageTracker(
                 }
             }
         )
+
+        conn.subscribe(
+            CodeWhispererService.CODEWHISPERER_CODE_COMPLETION_PERFORMED,
+            object : CodeWhispererCodeCompletionServiceListener {
+                override fun onSuccess(fileContextInfo: FileContextInfo) {
+                    if (language == fileContextInfo.programmingLanguage.toCodeWhispererLanguage()) {
+                        myServiceInvocationCount.getAndIncrement()
+                    }
+                }
+            }
+        )
         startTime = Instant.now()
         isActive.set(true)
         scheduleCodeWhispererCodeCoverageTracker()
@@ -156,6 +173,7 @@ abstract class CodeWhispererCodeCoverageTracker(
         startTime = Instant.now()
         rangeMarkers.clear()
         fileToTokens.clear()
+        myServiceInvocationCount.set(0)
     }
 
     internal fun emitCodeWhispererCodeContribution() {
@@ -186,7 +204,7 @@ abstract class CodeWhispererCodeCoverageTracker(
                 language,
                 percentage,
                 totalTokensSize,
-                successCount = 0
+                successCount = myServiceInvocationCount.get()
             )
         }
     }
@@ -236,7 +254,8 @@ class DefaultCodeWhispererCodeCoverageTracker(language: CodewhispererLanguage) :
     5 * TOTAL_SECONDS_IN_MINUTE,
     language,
     mutableListOf(),
-    mutableMapOf()
+    mutableMapOf(),
+    AtomicInteger(0)
 )
 
 class CodeCoverageTokens(totalTokens: Int = 0, acceptedTokens: Int = 0) {

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

@@ -48,6 +48,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.model.ProgrammingL
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.RecommendationContext
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.SessionContext
 import software.aws.toolkits.jetbrains.services.codewhisperer.popup.CodeWhispererPopupManager.Companion.CODEWHISPERER_USER_ACTION_PERFORMED
+import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService
 import software.aws.toolkits.jetbrains.services.codewhisperer.service.RequestContext
 import software.aws.toolkits.jetbrains.services.codewhisperer.service.ResponseContext
 import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeCoverageTokens
@@ -62,14 +63,15 @@ import software.aws.toolkits.jetbrains.utils.rules.JavaCodeInsightTestFixtureRul
 import software.aws.toolkits.jetbrains.utils.rules.PythonCodeInsightTestFixtureRule
 import software.aws.toolkits.telemetry.CodewhispererCompletionType
 import software.aws.toolkits.telemetry.CodewhispererLanguage
+import java.util.concurrent.atomic.AtomicInteger
 
 internal abstract class CodeWhispererCodeCoverageTrackerTestBase(myProjectRule: CodeInsightTestFixtureRule) {
     protected class TestCodePercentageTracker(
         timeWindowInSec: Long,
         language: CodewhispererLanguage,
         rangeMarkers: MutableList<RangeMarker> = mutableListOf(),
         codeCoverageTokens: MutableMap<Document, CodeCoverageTokens> = mutableMapOf()
-    ) : CodeWhispererCodeCoverageTracker(timeWindowInSec, language, rangeMarkers, codeCoverageTokens)
+    ) : CodeWhispererCodeCoverageTracker(timeWindowInSec, language, rangeMarkers, codeCoverageTokens, AtomicInteger(0))
 
     protected class TestTelemetryService(
         publisher: TelemetryPublisher = NoOpPublisher(),
@@ -178,6 +180,22 @@ internal class CodeWhispererCodeCoverageTrackerTestPython : CodeWhispererCodeCov
         assertThat(CodeWhispererCodeCoverageTracker.getInstancesMap()).hasSize(3)
     }
 
+    @Test
+    fun `test tracker is listening to cwspr recommendation service invocation`() {
+        val pythonTracker = TestCodePercentageTracker(TOTAL_SECONDS_IN_MINUTE, CodewhispererLanguage.Python)
+        val jsxTracker = TestCodePercentageTracker(TOTAL_SECONDS_IN_MINUTE, CodewhispererLanguage.Jsx)
+        CodeWhispererCodeCoverageTracker.getInstancesMap()[CodewhispererLanguage.Python] = pythonTracker
+        CodeWhispererCodeCoverageTracker.getInstancesMap()[CodewhispererLanguage.Jsx] = jsxTracker
+        pythonTracker.activateTrackerIfNotActive()
+        assertThat(pythonTracker.serviceInvocationCount).isEqualTo(0)
+        assertThat(jsxTracker.serviceInvocationCount).isEqualTo(0)
+
+        val fileContextInfo = mock<FileContextInfo> { on { programmingLanguage } doReturn ProgrammingLanguage(CodewhispererLanguage.Python) }
+        ApplicationManager.getApplication().messageBus.syncPublisher(CodeWhispererService.CODEWHISPERER_CODE_COMPLETION_PERFORMED).onSuccess(fileContextInfo)
+        assertThat(pythonTracker.serviceInvocationCount).isEqualTo(1)
+        assertThat(jsxTracker.serviceInvocationCount).isEqualTo(0)
+    }
+
     @Test
     fun `test tracker is listening to document changes and increment totalTokens - add new code`() {
         val pythonTracker = spy(TestCodePercentageTracker(TOTAL_SECONDS_IN_MINUTE, CodewhispererLanguage.Python))