refactor(inline): Remove Dummy Popup Creation in Prefetch Flow and Ensure Session Validation

.changes/next-release/feature-5cdc6e55-3a3f-4a31-8499-543cf249c734.json

@@ -0,0 +1,4 @@
+{
+  "type" : "feature",
+  "description" : "Inline suggestions: Enable Pre-fetch recommendations to reduce suggestion latency for complex scenarios"
+}

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/inlay/CodeWhispererInlayManager.kt

@@ -22,7 +22,7 @@ class CodeWhispererInlayManager {
         clearInlays()
 
         chunks.forEach { chunk ->
-            createCodeWhispererInlays(editor, chunk.inlayOffset, chunk.text, states.popup)
+            states.popup?.let { createCodeWhispererInlays(editor, chunk.inlayOffset, chunk.text, it) }
         }
     }
 

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/model/CodeWhispererModel.kt

@@ -216,7 +216,7 @@ data class InvocationContext(
     val requestContext: RequestContext,
     val responseContext: ResponseContext,
     val recommendationContext: RecommendationContext,
-    val popup: JBPopup,
+    val popup: JBPopup? = null,
 ) : Disposable {
     override fun dispose() {}
 }

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererPopupManager.kt

@@ -140,15 +140,15 @@ class CodeWhispererPopupManager {
             val startOffset = states.requestContext.caretPosition.offset
             val currOffset = states.requestContext.editor.caretModel.offset
             if (startOffset > currOffset) {
-                cancelPopup(popup)
+                popup?.let { cancelPopup(popup) }
                 return
             }
 
             // userInput + typeahead
             val prefix = states.requestContext.editor.document.charsSequence
                 .substring(startOffset, currOffset)
             if (prefix.length < userInputOriginal.length) {
-                cancelPopup(popup)
+                popup?.let { cancelPopup(popup) }
                 return
             } else {
                 prefix.substring(userInputOriginal.length)
@@ -171,7 +171,7 @@ class CodeWhispererPopupManager {
         )
         if (selectedIndex == -1 || !isValidRecommendation(details[selectedIndex], userInput, typeaheadOriginal)) {
             LOG.debug { "None of the recommendation is valid at this point, cancelling the popup" }
-            cancelPopup(popup)
+            popup?.let { cancelPopup(popup) }
             return
         }
         val typeahead = resolveTypeahead(states, selectedIndex, typeaheadOriginal)
@@ -244,7 +244,7 @@ class CodeWhispererPopupManager {
                 states.requestContext.latencyContext.getPerceivedLatency(states.requestContext.triggerTypeInfo.triggerType)
         }
         if (!isRecommendationAdded) {
-            showPopup(states, sessionContext, states.popup, visible = sessionContext.isPopupShowing)
+            states.popup?.let { showPopup(states, sessionContext, it, visible = sessionContext.isPopupShowing) }
         }
     }
 
@@ -415,8 +415,8 @@ class CodeWhispererPopupManager {
 
     private fun addPopupListener(states: InvocationContext) {
         val listener = CodeWhispererPopupListener(states)
-        states.popup.addListener(listener)
-        Disposer.register(states) { states.popup.removeListener(listener) }
+        states.popup?.addListener(listener)
+        Disposer.register(states) { states.popup?.removeListener(listener) }
     }
 
     private fun addMessageSubscribers(states: InvocationContext) {
@@ -452,7 +452,7 @@ class CodeWhispererPopupManager {
                     dontClosePopupAndRun {
                         CodeWhispererEditorManager.getInstance().updateEditorWithRecommendation(states, sessionContext)
                     }
-                    closePopup(states.popup)
+                    states.popup?.let { closePopup(it) }
                     if (sessionContext.selectedIndex == 0) {
                         CodeWhispererService.getInstance().promoteNextInvocationIfAvailable()
                     }
@@ -504,7 +504,7 @@ class CodeWhispererPopupManager {
         val codewhispererSelectionListener: SelectionListener = object : SelectionListener {
             override fun selectionChanged(event: SelectionEvent) {
                 if (!allowTypingDuringSuggestionPreview && !allowIntelliSenseDuringSuggestionPreview) {
-                    cancelPopup(states.popup)
+                    states.popup?.let { cancelPopup(it) }
                 }
                 super.selectionChanged(event)
             }
@@ -520,7 +520,7 @@ class CodeWhispererPopupManager {
                 if (editor.caretModel.offset == event.offset) {
                     changeStates(states, 0)
                 } else if (!allowTypingDuringSuggestionPreview && !allowIntelliSenseDuringSuggestionPreview) {
-                    cancelPopup(states.popup)
+                    states.popup?.let { cancelPopup(it) }
                 }
             }
         }
@@ -529,7 +529,7 @@ class CodeWhispererPopupManager {
         val codewhispererCaretListener: CaretListener = object : CaretListener {
             override fun caretPositionChanged(event: CaretEvent) {
                 if (!allowTypingDuringSuggestionPreview && !allowIntelliSenseDuringSuggestionPreview) {
-                    cancelPopup(states.popup)
+                    states.popup?.let { cancelPopup(it) }
                 }
                 super.caretPositionChanged(event)
             }
@@ -542,11 +542,11 @@ class CodeWhispererPopupManager {
             val window = ComponentUtil.getWindow(editorComponent)
             val windowListener: ComponentListener = object : ComponentAdapter() {
                 override fun componentMoved(event: ComponentEvent) {
-                    cancelPopup(states.popup)
+                    states.popup?.let { cancelPopup(it) }
                 }
 
                 override fun componentShown(e: ComponentEvent?) {
-                    cancelPopup(states.popup)
+                    states.popup?.let { cancelPopup(it) }
                     super.componentShown(e)
                 }
             }
@@ -557,7 +557,7 @@ class CodeWhispererPopupManager {
         val suggestionHoverEnterListener: EditorMouseMotionListener = object : EditorMouseMotionListener {
             override fun mouseMoved(e: EditorMouseEvent) {
                 if (e.inlay != null) {
-                    showPopup(states, sessionContext, states.popup, visible = true)
+                    states.popup?.let { showPopup(states, sessionContext, it, visible = true) }
                 } else {
                     bringSuggestionInlayToFront(editor, states.popup, sessionContext, opposite = true)
                 }

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererUIChangeListener.kt

@@ -61,8 +61,10 @@ class CodeWhispererUIChangeListener : CodeWhispererPopupStateChangeListener {
                 },
                 HighlighterTargetArea.EXACT_RANGE
             )
-            Disposer.register(states.popup) {
-                editor.markupModel.removeHighlighter(rangeHighlighter)
+            states.popup?.let {
+                Disposer.register(it) {
+                    editor.markupModel.removeHighlighter(rangeHighlighter)
+                }
             }
             sessionContext.toBeRemovedHighlighter = rangeHighlighter
         }

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/handlers/CodeWhispererPopupEscHandler.kt

@@ -13,7 +13,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispe
 
 class CodeWhispererPopupEscHandler(states: InvocationContext) : CodeWhispererEditorActionHandler(states) {
     override fun doExecute(editor: Editor, caret: Caret?, dataContext: DataContext?) {
-        CodeWhispererPopupManager.getInstance().cancelPopup(states.popup)
+        states.popup?.let { CodeWhispererPopupManager.getInstance().cancelPopup(it) }
     }
 }
 

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt

@@ -16,7 +16,6 @@ import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.editor.VisualPosition
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.ui.popup.JBPopup
-import com.intellij.openapi.ui.popup.JBPopupFactory
 import com.intellij.openapi.util.Disposer
 import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.PsiFile
@@ -27,6 +26,7 @@ import kotlinx.coroutines.Deferred
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.async
+import kotlinx.coroutines.cancel
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.isActive
 import kotlinx.coroutines.launch
@@ -268,6 +268,11 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
                         val latency = TimeUnit.NANOSECONDS.toMillis(endTime - startTime).toDouble()
                         val requestId = nextResponse.responseMetadata().requestId()
                         val sessionId = nextResponse.sdkHttpResponse().headers().getOrDefault(KET_SESSION_ID, listOf(requestId))[0]
+                        if (sessionId != currStates.responseContext.sessionId) {
+                            LOG.debug { "Session id mismatch, expected: ${currStates.responseContext.sessionId}, actual: $sessionId" }
+                            this.cancel("Session mismatch")
+                            return@launch
+                        }
 
                         nextRequestContext.latencyContext.apply {
                             codewhispererPostprocessingStart = System.nanoTime()
@@ -299,12 +304,9 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
                             }
                         }
                         val nextRecommendationContext = RecommendationContext(detailContexts, "", "", newVisualPosition)
-                        val newPopup = withContext(EDT) {
-                            JBPopupFactory.getInstance().createMessage("Dummy popup")
-                        }
 
                         // send userDecision and trigger decision when next recommendation haven't been seen
-                        if (currStates.popup.isDisposed) {
+                        if (currStates.popup?.isDisposed == true) {
                             CodeWhispererTelemetryService.getInstance().sendUserDecisionEventForAll(
                                 nextRequestContext,
                                 nextResponseContext,
@@ -313,7 +315,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
                                 false
                             )
                         } else {
-                            nextInvocationContext = InvocationContext(nextRequestContext, nextResponseContext, nextRecommendationContext, newPopup)
+                            nextInvocationContext = InvocationContext(nextRequestContext, nextResponseContext, nextRecommendationContext, null)
                         }
                         LOG.debug { "Prefetched next invocation stored in nextInvocationContext" }
                     }
@@ -715,7 +717,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
         val updatedStates = states.copy(
             recommendationContext = recommendationContext.copy(details = details + newDetailContexts)
         )
-        Disposer.register(states.popup, updatedStates)
+        states.popup?.let { Disposer.register(it, updatedStates) }
         CodeWhispererPopupManager.getInstance().initPopupListener(updatedStates)
         return updatedStates
     }
@@ -750,7 +752,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
         cs.launch {
             val newPopup = CodeWhispererPopupManager.getInstance().initPopup()
             val updatedNextStates = nextStates.copy(popup = newPopup).also {
-                addPopupChildDisposables(it.requestContext.project, it.requestContext.editor, it.popup)
+                it.popup?.let { it1 -> addPopupChildDisposables(it.requestContext.project, it.requestContext.editor, it1) }
                 Disposer.register(newPopup, it)
             }
             CodeWhispererPopupManager.getInstance().initPopupListener(updatedNextStates)