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

.changes/next-release/removal-758edd1c-4ebd-49c6-8ad8-0ec59a5d9f0b.json

@@ -0,0 +1,4 @@
+{
+  "type" : "removal",
+  "description" : "Amazon Q: Revert prefetch logic to enable more stable inline completion."
+}

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

@@ -77,10 +77,6 @@ interface CodeWhispererClientAdaptor : Disposable {
         firstRequest: GenerateCompletionsRequest,
     ): Sequence<GenerateCompletionsResponse>
 
-    fun generateCompletions(
-        firstRequest: GenerateCompletionsRequest,
-    ): GenerateCompletionsResponse
-
     fun createUploadUrl(
         request: CreateUploadUrlRequest,
     ): CreateUploadUrlResponse
@@ -327,8 +323,6 @@ 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,7 +7,6 @@ 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
@@ -28,8 +27,7 @@ class CodeWhispererPopupListener(private val states: InvocationContext) : JBPopu
             recommendationContext,
             CodeWhispererPopupManager.getInstance().sessionContext,
             event.isOk,
-            CodeWhispererInvocationStatus.getInstance().popupStartTimestamp?.let { Duration.between(it, Instant.now()) },
-            CodeWhispererService.getInstance().getNextInvocationContext()
+            CodeWhispererInvocationStatus.getInstance().popupStartTimestamp?.let { Duration.between(it, Instant.now()) }
         )
 
         CodeWhispererInvocationStatus.getInstance().setDisplaySessionActive(false)

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

@@ -68,7 +68,6 @@ 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
@@ -453,9 +452,6 @@ 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,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
@@ -25,7 +24,6 @@ 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
@@ -100,7 +98,6 @@ 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)
@@ -212,118 +209,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
         invokeCodeWhispererInBackground(requestContext)
     }
 
-    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
-        }
-
+    internal suspend fun invokeCodeWhispererInBackground(requestContext: RequestContext): Job {
         val popup = withContext(EDT) {
             CodeWhispererPopupManager.getInstance().initPopup().also {
                 Disposer.register(it) { CodeWhispererInvocationStatus.getInstance().finishInvocation() }
@@ -605,9 +491,6 @@ 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)
@@ -739,34 +622,6 @@ 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,
@@ -940,8 +795,6 @@ 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

@@ -457,7 +457,6 @@ class CodeWhispererTelemetryService {
         sessionContext: SessionContext,
         hasUserAccepted: Boolean,
         popupShownTime: Duration? = null,
-        nextInvocationContext: InvocationContext? = null,
     ) {
         val detailContexts = recommendationContext.details
         val decisions = mutableListOf<CodewhispererSuggestionState>()
@@ -505,19 +504,6 @@ 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

@@ -4,6 +4,7 @@
 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
@@ -19,12 +20,14 @@ 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
@@ -34,11 +37,14 @@ 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) {
@@ -80,7 +86,7 @@ class CodeWhispererSettingsTest : CodeWhispererTestBase() {
         whenever(stateManager.checkActiveCodeWhispererConnectionType(projectRule.project)).thenReturn(CodeWhispererLoginType.Logout)
         assertThat(isCodeWhispererEnabled(projectRule.project)).isFalse
         invokeCodeWhispererService()
-        verify(codewhispererService, never()).showRecommendationsInPopup(any(), any(), any())
+        verify(codewhispererServiceSpy, never()).showRecommendationsInPopup(any(), any(), any())
     }
 
     @Test
@@ -89,7 +95,7 @@ class CodeWhispererSettingsTest : CodeWhispererTestBase() {
         assertThat(stateManager.isAutoEnabled()).isFalse
         runInEdtAndWait {
             projectRule.fixture.type(':')
-            verify(codewhispererService, never()).showRecommendationsInPopup(any(), any(), any())
+            verify(codewhispererServiceSpy, never()).showRecommendationsInPopup(any(), any(), any())
         }
     }