feat(amazonq): Prefetch next inline recommendation (#5290)

evanliu048 · web-flow · commit 805cb399c5ba · 2025-02-04T13:01:09.000-08:00
This PR introduces the ability for users to see the next recommendation immediately after accepting the current one, provided there is a subsequent recommendation available. The key enhancements include:

    Automatic Prefetching: Introduced a function that calls the CodeWhisperer API to fetch the next recommendation while the current one is being reviewed.
    Session Promotion Added functionality to display the next recommendation instantly upon accepting the current recommendation (if there is one)
    Helper Utilities: Implemented helper functions to calculate the necessary file information for the next request and to send telemetry events related to the subsequent recommendation.
diff --git a/.changes/next-release/feature-af284b37-980e-4f98-a653-5fef6a71519b.json b/.changes/next-release/feature-af284b37-980e-4f98-a653-5fef6a71519b.json
@@ -0,0 +1,4 @@
+{
+  "type" : "feature",
+  "description" : "Inline suggestions: Pre-fetch recommendations to reduce suggestion latency."
+}
diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/credentials/CodeWhispererClientAdaptor.kt b/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)
diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererPopupListener.kt b/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().setDisplaySessionActive(false)
diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererPopupManager.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/popup/CodeWhispererPopupManager.kt
@@ -69,6 +69,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.popup.listeners.Co
 import software.aws.toolkits.jetbrains.services.codewhisperer.popup.listeners.CodeWhispererScrollListener
 import software.aws.toolkits.jetbrains.services.codewhisperer.popup.listeners.addIntelliSenseAcceptListener
 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
@@ -460,6 +461,9 @@ class CodeWhispererPopupManager {
                         CodeWhispererEditorManager.getInstance().updateEditorWithRecommendation(states, sessionContext)
                     }
                     closePopup(states.popup)
+                    if (sessionContext.selectedIndex == 0) {
+                        CodeWhispererService.getInstance().promoteNextInvocationIfAvailable()
+                    }
                 }
             }
         )
diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt b/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,118 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
         invokeCodeWhispererInBackground(requestContext)
     }
 
-    internal suspend fun invokeCodeWhispererInBackground(requestContext: RequestContext): Job {
+    internal suspend fun invokeCodeWhispererInBackground(
+        requestContext: RequestContext,
+        currStates: InvocationContext? = null,
+    ): Job {
+        // current states != null means that it's prefetch
+        if (currStates != null) {
+            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 = CaretPosition(
+                    line = requestContext.caretPosition.line + firstValidRecommendation.recommendation.content().count { it == '\n' },
+                    offset = requestContext.caretPosition.offset + firstValidRecommendation.recommendation.content().length
+                )
+
+                val nextFileContextInfo = requestContext.fileContextInfo.copy(
+                    caretContext = requestContext.fileContextInfo.caretContext.copy(
+                        leftFileContext = requestContext.fileContextInfo.caretContext.leftFileContext + firstValidRecommendation.recommendation.content()
+                    )
+                )
+
+                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,
+                            invocationSuccess = 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 +605,9 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
                 CodeWhispererPopupManager.getInstance().cancelPopup(popup)
                 return null
             }
+            cs.launch(getCoroutineBgContext()) {
+                invokeCodeWhispererInBackground(requestContext, nextStates)
+            }
         } else {
             // subsequent responses
             nextStates = updateStates(currStates, response)
@@ -616,6 +733,34 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
         CodeWhispererPopupManager.getInstance().changeStates(states, 0, recommendationAdded)
     }
 
+    fun promoteNextInvocationIfAvailable() {
+        val nextStates = nextInvocationContext ?: run {
+            LOG.debug { "No nextInvocationContext found, nothing to promote." }
+            return
+        }
+        nextInvocationContext?.popup?.let { Disposer.dispose(it) }
+        nextInvocationContext = null
+
+        cs.launch {
+            val newPopup = CodeWhispererPopupManager.getInstance().initPopup()
+            val updatedNextStates = nextStates.copy(popup = newPopup).also {
+                addPopupChildDisposables(it.requestContext.project, it.requestContext.editor, it.popup)
+                Disposer.register(newPopup, it)
+            }
+            CodeWhispererPopupManager.getInstance().initPopupListener(updatedNextStates)
+            withContext(EDT) {
+                CodeWhispererPopupManager.getInstance().changeStates(
+                    updatedNextStates,
+                    0,
+                    recommendationAdded = false
+                )
+            }
+            invokeCodeWhispererInBackground(updatedNextStates.requestContext, updatedNextStates)
+        }
+
+        LOG.debug { "Promoted nextInvocationContext to current session and displayed next recommendation." }
+    }
+
     private fun sendDiscardedUserDecisionEventForAll(
         requestContext: RequestContext,
         responseContext: ResponseContext,
@@ -782,6 +927,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
diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/telemetry/CodeWhispererTelemetryService.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/telemetry/CodeWhispererTelemetryService.kt
@@ -459,6 +459,7 @@ class CodeWhispererTelemetryService {
         sessionContext: SessionContext,
         hasUserAccepted: Boolean,
         popupShownTime: Duration? = null,
+        nextInvocationContext: InvocationContext? = null,
     ) {
         val detailContexts = recommendationContext.details
         val decisions = mutableListOf<CodewhispererSuggestionState>()
@@ -506,6 +507,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
+                    )
+                }
+            }
         }
     }
 
diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererSettingsTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererSettingsTest.kt
@@ -4,7 +4,6 @@
 package software.aws.toolkits.jetbrains.services.codewhisperer
 
 import com.intellij.analysis.problemsView.toolWindow.ProblemsView
-import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.components.service
 import com.intellij.openapi.wm.RegisterToolWindowTask
 import com.intellij.openapi.wm.ToolWindow
@@ -20,14 +19,12 @@ import org.junit.Ignore
 import org.junit.Test
 import org.mockito.kotlin.any
 import org.mockito.kotlin.never
-import org.mockito.kotlin.spy
 import org.mockito.kotlin.verify
 import org.mockito.kotlin.whenever
 import software.aws.toolkits.jetbrains.core.ToolWindowHeadlessManagerImpl
 import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererLoginType
 import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExploreActionState
 import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.isCodeWhispererEnabled
-import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService
 import software.aws.toolkits.jetbrains.services.codewhisperer.status.CodeWhispererStatusBarWidgetFactory
 import software.aws.toolkits.jetbrains.services.codewhisperer.toolwindow.CodeWhispererCodeReferenceToolWindowFactory
 import software.aws.toolkits.jetbrains.settings.CodeWhispererConfiguration
@@ -37,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) {
@@ -86,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
@@ -95,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())
         }
     }
 
diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTelemetryTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTelemetryTest.kt
diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTestBase.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTestBase.kt

-Original file line number
+Diff line change
@@ @@ -0,0 +1,4 @@ @@
 +{
 +  "type" : "feature",
 +  "description" : "Inline suggestions: Pre-fetch recommendations to reduce suggestion latency."
 +}
Original file line number	Diff line number	Diff line change
`@@ -69,6 +69,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.popup.listeners.Co`
`69`	`69`	`import software.aws.toolkits.jetbrains.services.codewhisperer.popup.listeners.CodeWhispererScrollListener`
`70`	`70`	`import software.aws.toolkits.jetbrains.services.codewhisperer.popup.listeners.addIntelliSenseAcceptListener`
`71`	`71`	`import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererInvocationStatus`
	`72`	`+import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService`
`72`	`73`	`import software.aws.toolkits.jetbrains.services.codewhisperer.telemetry.CodeWhispererTelemetryService`
`73`	`74`	`import software.aws.toolkits.jetbrains.services.codewhisperer.toolwindow.CodeWhispererCodeReferenceManager`
`74`	`75`	`import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererColorUtil.POPUP_DIM_HEX`
`@@ -460,6 +461,9 @@ class CodeWhispererPopupManager {`
`460`	`461`	`CodeWhispererEditorManager.getInstance().updateEditorWithRecommendation(states, sessionContext)`
`461`	`462`	`}`
`462`	`463`	`closePopup(states.popup)`
	`464`	`+ if (sessionContext.selectedIndex == 0) {`
	`465`	`+ CodeWhispererService.getInstance().promoteNextInvocationIfAvailable()`
	`466`	`+ }`
`463`	`467`	`}`
`464`	`468`	`}`
`465`	`469`	`)`