Prefetch next inline recommendation

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt

@@ -76,6 +76,10 @@ interface CodeWhispererClientAdaptor : Disposable {
         firstRequest: GenerateCompletionsRequest,
     ): Sequence<GenerateCompletionsResponse>
 
+    fun generateCompletions(
+        firstRequest: GenerateCompletionsRequest,
+    ): GenerateCompletionsResponse
+
     fun createUploadUrl(
         request: CreateUploadUrlRequest,
     ): CreateUploadUrlResponse
@@ -322,6 +326,8 @@ open class CodeWhispererClientAdaptorImpl(override val project: Project) : CodeW
             yield(response)
         } while (!nextToken.isNullOrEmpty())
     }
+    override fun generateCompletions(firstRequest: GenerateCompletionsRequest): GenerateCompletionsResponse =
+        bearerClient().generateCompletions(firstRequest)
 
     override fun createUploadUrl(request: CreateUploadUrlRequest): CreateUploadUrlResponse =
         bearerClient().createUploadUrl(request)

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

@@ -7,6 +7,7 @@ import com.intellij.openapi.ui.popup.JBPopupListener
 import com.intellij.openapi.ui.popup.LightweightWindowEvent
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext
 import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererInvocationStatus
+import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService
 import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererServiceNew
 import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeWhispererTelemetryService
 import java.time.Duration
@@ -27,7 +28,8 @@ class CodeWhispererPopupListener(private val states: InvocationContext) : JBPopu
             recommendationContext,
             CodeWhispererPopupManager.getInstance().sessionContext,
             event.isOk,
-            CodeWhispererInvocationStatus.getInstance().popupStartTimestamp?.let { Duration.between(it, Instant.now()) }
+            CodeWhispererInvocationStatus.getInstance().popupStartTimestamp?.let { Duration.between(it, Instant.now()) },
+            CodeWhispererService.getInstance().getNextInvocationContext()
         )
 
         CodeWhispererInvocationStatus.getInstance().setPopupActive(false)

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

@@ -68,6 +68,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.popup.listeners.Co
 import software.aws.toolkits.jetbrains.services.codewhisperer.popup.listeners.CodeWhispererPrevButtonActionListener
 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.CodeWhispererService
 import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeWhispererTelemetryService
 import software.aws.toolkits.jetbrains.services.codewhisperer.toolwindow.CodeWhispererCodeReferenceManager
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererColorUtil.POPUP_DIM_HEX
@@ -455,6 +456,9 @@ class CodeWhispererPopupManager {
                         CodeWhispererEditorManager.getInstance().updateEditorWithRecommendation(states, sessionContext)
                     }
                     closePopup(states.popup)
+                    if (sessionContext.selectedIndex == 0) {
+                        CodeWhispererService.getInstance().promoteNextInvocationIfAvailable()
+                    }
                 }
             }
         )

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

@@ -16,6 +16,7 @@ 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
@@ -24,6 +25,7 @@ import com.intellij.util.messages.Topic
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Deferred
 import kotlinx.coroutines.Job
+import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.async
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.isActive
@@ -98,6 +100,7 @@ import java.util.concurrent.TimeUnit
 class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
     private val codeInsightSettingsFacade = CodeInsightsSettingsFacade()
     private var refreshFailure: Int = 0
+    private var nextInvocationContext: InvocationContext? = null
 
     init {
         Disposer.register(this, codeInsightSettingsFacade)
@@ -209,7 +212,110 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
         invokeCodeWhispererInBackground(requestContext)
     }
 
-    internal suspend fun invokeCodeWhispererInBackground(requestContext: RequestContext): Job {
+    internal suspend fun invokeCodeWhispererInBackground(
+        requestContext: RequestContext,
+        isPrefetch: Boolean = false,
+        currStates: InvocationContext? = null,
+    ): Job {
+        if (isPrefetch && currStates?.recommendationContext?.details?.isNotEmpty() == true) {
+            val firstValidRecommendation = currStates.recommendationContext.details
+                .firstOrNull {
+                    !it.isDiscarded && it.recommendation.content().isNotEmpty()
+                } ?: return SupervisorJob().apply { complete() }
+            val job = cs.launch(getCoroutineBgContext()) {
+                val latencyContext = LatencyContext().apply {
+                    codewhispererPreprocessingStart = System.nanoTime()
+                    codewhispererEndToEndStart = System.nanoTime()
+                }
+
+                val nextCaretPosition = calculateNextCaretPosition(requestContext, firstValidRecommendation)
+                val nextFileContextInfo = createNextFileContextInfo(requestContext, firstValidRecommendation)
+
+                val nextRequestContext = requestContext.copy(
+                    caretPosition = nextCaretPosition,
+                    fileContextInfo = nextFileContextInfo,
+                    latencyContext = latencyContext
+                )
+                val newVisualPosition = withContext(EDT) {
+                    runReadAction {
+                        nextRequestContext.editor.offsetToVisualPosition(nextRequestContext.caretPosition.offset)
+                    }
+                }
+                try {
+                    val nextResponse = CodeWhispererClientAdaptor
+                        .getInstance(nextRequestContext.project)
+                        .generateCompletions(
+                            buildCodeWhispererRequest(
+                                nextRequestContext.fileContextInfo,
+                                nextRequestContext.awaitSupplementalContext(),
+                                nextRequestContext.customizationArn
+                            )
+                        )
+                    val startTime = System.nanoTime()
+                    nextRequestContext.latencyContext.codewhispererPreprocessingEnd = System.nanoTime()
+                    nextRequestContext.latencyContext.paginationAllCompletionsStart = System.nanoTime()
+                    CodeWhispererInvocationStatus.getInstance().setInvocationStart()
+                    nextResponse.let {
+                        val endTime = System.nanoTime()
+                        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]
+
+                        nextRequestContext.latencyContext.apply {
+                            codewhispererPostprocessingStart = System.nanoTime()
+                            paginationFirstCompletionTime = (endTime - codewhispererEndToEndStart).toDouble()
+                            firstRequestId = requestId
+                        }
+
+                        CodeWhispererInvocationStatus.getInstance().setInvocationSessionId(sessionId)
+
+                        val nextResponseContext = ResponseContext(sessionId)
+                        CodeWhispererTelemetryService.getInstance().sendServiceInvocationEvent(
+                            nextResponse.responseMetadata().requestId(),
+                            nextRequestContext,
+                            nextResponseContext,
+                            nextResponse.completions().size,
+                            true,
+                            latency,
+                            null
+                        )
+                        val validatedResponse = validateResponse(it)
+                        val detailContexts = withContext(EDT) {
+                            runReadAction {
+                                CodeWhispererRecommendationManager.getInstance().buildDetailContext(
+                                    nextRequestContext,
+                                    "",
+                                    validatedResponse.completions(),
+                                    validatedResponse.responseMetadata().requestId()
+                                )
+                            }
+                        }
+                        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) {
+                            CodeWhispererTelemetryService.getInstance().sendUserDecisionEventForAll(
+                                nextRequestContext,
+                                nextResponseContext,
+                                nextRecommendationContext,
+                                SessionContext(),
+                                false
+                            )
+                        } else {
+                            nextInvocationContext = InvocationContext(nextRequestContext, nextResponseContext, nextRecommendationContext, newPopup)
+                        }
+                        LOG.debug { "Prefetched next invocation stored in nextInvocationContext" }
+                    }
+                } catch (ex: Exception) {
+                    LOG.warn { "Failed to prefetch next codewhisperer invocation: ${ex.message}" }
+                }
+            }
+            return job
+        }
+
         val popup = withContext(EDT) {
             CodeWhispererPopupManager.getInstance().initPopup().also {
                 Disposer.register(it) { CodeWhispererInvocationStatus.getInstance().finishInvocation() }
@@ -491,6 +597,9 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
                 CodeWhispererPopupManager.getInstance().cancelPopup(popup)
                 return null
             }
+            cs.launch(getCoroutineBgContext()) {
+                invokeCodeWhispererInBackground(requestContext, true, nextStates)
+            }
         } else {
             // subsequent responses
             nextStates = updateStates(currStates, response)
@@ -621,6 +730,67 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
         CodeWhispererPopupManager.getInstance().changeStates(states, 0, "", true, recommendationAdded)
     }
 
+    fun promoteNextInvocationIfAvailable() {
+        val nextStates = nextInvocationContext ?: run {
+            LOG.debug { "No nextInvocationContext found, nothing to promote." }
+            return
+        }
+        nextInvocationContext?.popup?.let { Disposer.dispose(it) }
+        nextInvocationContext = null
+
+        runInEdt {
+            val newPopup = CodeWhispererPopupManager.getInstance().initPopup()
+            val updatedNextStates = nextStates.copy(popup = newPopup).also {
+                addPopupChildDisposables(it.popup)
+                Disposer.register(newPopup, it)
+            }
+            CodeWhispererPopupManager.getInstance().initPopupListener(updatedNextStates)
+            CodeWhispererPopupManager.getInstance().changeStates(
+                updatedNextStates,
+                0,
+                "",
+                typeaheadAdded = true,
+                recommendationAdded = false
+            )
+            cs.launch(getCoroutineBgContext()) {
+                invokeCodeWhispererInBackground(updatedNextStates.requestContext, true, updatedNextStates)
+            }
+        }
+
+        LOG.debug { "Promoted nextInvocationContext to current session and displayed next recommendation." }
+    }
+
+    private fun calculateNextCaretPosition(
+        currentRequestContext: RequestContext,
+        firstValidRecommendation: DetailContext,
+    ): CaretPosition {
+        val indent = currentRequestContext.fileContextInfo.caretContext.leftContextOnCurrentLine
+            .takeWhile { it.isWhitespace() }
+        val recommendedText = buildString {
+            append(indent)
+            append(firstValidRecommendation.recommendation.content())
+            if (!endsWith("\n")) {
+                append("\n")
+            }
+        }
+        val lineCount = recommendedText.count { it == '\n' }
+
+        return CaretPosition(
+            line = currentRequestContext.caretPosition.line + lineCount,
+            offset = currentRequestContext.caretPosition.offset + recommendedText.length
+        )
+    }
+
+    private fun createNextFileContextInfo(
+        requestContext: RequestContext,
+        firstValidRecommendation: DetailContext,
+    ): FileContextInfo {
+        val updatedCaretContext = requestContext.fileContextInfo.caretContext.copy(
+            leftFileContext = requestContext.fileContextInfo.caretContext.leftFileContext + firstValidRecommendation.recommendation.content()
+        )
+        return requestContext.fileContextInfo.copy(caretContext = updatedCaretContext)
+    }
+
     private fun sendDiscardedUserDecisionEventForAll(
         requestContext: RequestContext,
         responseContext: ResponseContext,
@@ -781,6 +951,8 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
 
     override fun dispose() {}
 
+    fun getNextInvocationContext(): InvocationContext? = nextInvocationContext
+
     companion object {
         private val LOG = getLogger<CodeWhispererService>()
         private const val MAX_REFRESH_ATTEMPT = 3

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/telemetry/CodeWhispererTelemetryService.kt

@@ -453,6 +453,7 @@ class CodeWhispererTelemetryService {
         sessionContext: SessionContext,
         hasUserAccepted: Boolean,
         popupShownTime: Duration? = null,
+        nextInvocationContext: InvocationContext? = null,
     ) {
         val detailContexts = recommendationContext.details
         val decisions = mutableListOf<CodewhispererSuggestionState>()
@@ -500,6 +501,19 @@ class CodeWhispererTelemetryService {
             previousUserTriggerDecisions.add(this)
             // we need this as well because AutotriggerService will reset the queue periodically
             CodeWhispererAutoTriggerService.getInstance().addPreviousDecision(this)
+            // send possible next session event if current action is reject and next popup haven't shown up
+            if (CodewhispererSuggestionState.from(this.toString()) == CodewhispererSuggestionState.Reject) {
+                nextInvocationContext?.let {
+                    sendUserDecisionEventForAll(
+                        it.requestContext,
+                        it.responseContext,
+                        it.recommendationContext,
+                        SessionContext(),
+                        false,
+                        nextInvocationContext = null
+                    )
+                }
+            }
         }
     }
 

plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererSettingsTest.kt

@@ -34,14 +34,11 @@ import kotlin.test.fail
 
 class CodeWhispererSettingsTest : CodeWhispererTestBase() {
 
-    private lateinit var codewhispererServiceSpy: CodeWhispererService
     private lateinit var toolWindowHeadlessManager: ToolWindowHeadlessManagerImpl
 
     @Before
     override fun setUp() {
         super.setUp()
-        codewhispererServiceSpy = spy(codewhispererService)
-        ApplicationManager.getApplication().replaceService(CodeWhispererService::class.java, codewhispererServiceSpy, disposableRule.disposable)
 
         // Create a mock ToolWindowManager with working implementation of setAvailable() and isAvailable()
         toolWindowHeadlessManager = object : ToolWindowHeadlessManagerImpl(projectRule.project) {
@@ -83,7 +80,7 @@ class CodeWhispererSettingsTest : CodeWhispererTestBase() {
         whenever(stateManager.checkActiveCodeWhispererConnectionType(projectRule.project)).thenReturn(CodeWhispererLoginType.Logout)
         assertThat(isCodeWhispererEnabled(projectRule.project)).isFalse
         invokeCodeWhispererService()
-        verify(codewhispererServiceSpy, never()).showRecommendationsInPopup(any(), any(), any())
+        verify(codewhispererService, never()).showRecommendationsInPopup(any(), any(), any())
     }
 
     @Test
@@ -92,7 +89,7 @@ class CodeWhispererSettingsTest : CodeWhispererTestBase() {
         assertThat(stateManager.isAutoEnabled()).isFalse
         runInEdtAndWait {
             projectRule.fixture.type(':')
-            verify(codewhispererServiceSpy, never()).showRecommendationsInPopup(any(), any(), any())
+            verify(codewhispererService, never()).showRecommendationsInPopup(any(), any(), any())
         }
     }