Fix the issue when CodeWhisperer invocation is not set to complete when editor is disposed (#3762)

Author: andrewyuq

.changes/next-release/bugfix-6bbbb228-bbb0-4707-8117-5171de1add3e.json

@@ -0,0 +1,4 @@
+{
+  "type" : "bugfix",
+  "description" : "CodeWhisperer: Fix an issue where CodeWhisperer would stuck in the invocation state indefinitely"
+}

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

@@ -15,6 +15,7 @@ import java.util.concurrent.atomic.AtomicBoolean
 
 class CodeWhispererInvocationStatus {
     private val isInvokingCodeWhisperer: AtomicBoolean = AtomicBoolean(false)
+    private var invokingSessionId: String? = null
     private var timeAtLastInvocationComplete: Instant? = null
     var timeAtLastDocumentChanged: Instant = Instant.now()
         private set
@@ -25,7 +26,7 @@ class CodeWhispererInvocationStatus {
 
     fun checkExistingInvocationAndSet(): Boolean =
         if (isInvokingCodeWhisperer.getAndSet(true)) {
-            LOG.debug { "Have existing CodeWhisperer invocation" }
+            LOG.debug { "Have existing CodeWhisperer invocation, sessionId: $invokingSessionId" }
             true
         } else {
             ApplicationManager.getApplication().messageBus.syncPublisher(CODEWHISPERER_INVOCATION_STATE_CHANGED).invocationStateChanged(true)
@@ -39,6 +40,7 @@ class CodeWhispererInvocationStatus {
         if (isInvokingCodeWhisperer.compareAndSet(true, false)) {
             ApplicationManager.getApplication().messageBus.syncPublisher(CODEWHISPERER_INVOCATION_STATE_CHANGED).invocationStateChanged(false)
             LOG.debug { "Ending CodeWhisperer invocation" }
+            invokingSessionId = null
         }
     }
 
@@ -75,6 +77,11 @@ class CodeWhispererInvocationStatus {
         timeAtLastInvocationStart = Instant.now()
     }
 
+    fun setInvocationSessionId(sessionId: String?) {
+        LOG.debug { "Set current CodeWhisperer invocation sessionId: $sessionId" }
+        invokingSessionId = sessionId
+    }
+
     fun hasEnoughDelayToInvokeCodeWhisperer(): Boolean {
         val timeCanShowCodeWhisperer = timeAtLastInvocationStart?.plusMillis(CodeWhispererConstants.INVOCATION_INTERVAL) ?: return true
         return timeCanShowCodeWhisperer.isBefore(Instant.now())

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

@@ -170,15 +170,16 @@ class CodeWhispererService {
                     val latency = TimeUnit.NANOSECONDS.toMillis(endTime - startTime).toDouble()
                     startTime = endTime
                     val requestId = response.responseMetadata().requestId()
+                    val sessionId = response.sdkHttpResponse().headers().getOrDefault(KET_SESSION_ID, listOf(requestId))[0]
                     if (requestCount == 1) {
                         requestContext.latencyContext.codewhispererPostprocessingStart = System.nanoTime()
                         requestContext.latencyContext.paginationFirstCompletionTime = latency
                         requestContext.latencyContext.firstRequestId = requestId
+                        CodeWhispererInvocationStatus.getInstance().setInvocationSessionId(sessionId)
                     }
                     if (response.nextToken().isEmpty()) {
                         requestContext.latencyContext.paginationAllCompletionsEnd = System.nanoTime()
                     }
-                    val sessionId = response.sdkHttpResponse().headers().getOrDefault(KET_SESSION_ID, listOf(requestId))[0]
                     val emptyRecommendations = checkEmptyRecommendations(response.completions())
                     val completionType = checkCompletionType(response.completions(), emptyRecommendations)
                     val responseContext = ResponseContext(sessionId, completionType)
@@ -269,6 +270,7 @@ class CodeWhispererService {
                 }
                 val exceptionType = e::class.simpleName
                 val responseContext = ResponseContext(sessionId, CodewhispererCompletionType.Unknown)
+                CodeWhispererInvocationStatus.getInstance().setInvocationSessionId(sessionId)
                 logServiceInvocation(requestId, requestContext, responseContext, emptyList(), null, exceptionType)
                 CodeWhispererTelemetryService.getInstance().sendServiceInvocationEvent(
                     requestId,
@@ -332,6 +334,7 @@ class CodeWhispererService {
 
         if (requestContext.editor.isDisposed) {
             LOG.debug { "Stop showing CodeWhisperer recommendations since editor is disposed. RequestId: $requestId" }
+            CodeWhispererPopupManager.getInstance().cancelPopup(popup)
             return null
         }