Merge main into feature/dev-execution

aws-toolkit-automation · web-flow · commit 15c8a17bf9ad · 2025-02-04T22:01:16.000+01:00
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`	`)`