fix(amazonq): fix inline pagination logic and corountine scope logic (#5717)

1. Show the configurable tab action only on new auto trigger UX(A/B)
2. remove disposableCoroutineScope usage and directly check the dispose state of the popup(session validity)
3. use partialResultToken instead of workdoneToken for pagination logic with flare

Author: andrewyuq

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/editor/CodeWhispererEditorManager.kt

@@ -12,6 +12,7 @@ import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.options.ShowSettingsUtil
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiDocumentManager
+import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.CaretPosition
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.InvocationContext
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.SessionContext
@@ -66,7 +67,9 @@ class CodeWhispererEditorManager {
         }
 
         // Display tab accept priority once when the first accept is made
-        if (!CodeWhispererSettings.getInstance().isTabAcceptPriorityNotificationShownOnce()) {
+        if (!CodeWhispererSettings.getInstance().isTabAcceptPriorityNotificationShownOnce() &&
+            CodeWhispererFeatureConfigService.getInstance().getNewAutoTriggerUX()
+        ) {
             notifyInfo(
                 "Amazon Q",
                 message("codewhisperer.inline.settings.tab_priority.notification.text"),

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

@@ -22,7 +22,6 @@ import com.intellij.util.concurrency.annotations.RequiresEdt
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.future.await
-import kotlinx.coroutines.isActive
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
 import org.eclipse.lsp4j.Position
@@ -36,7 +35,6 @@ import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.info
 import software.aws.toolkits.core.utils.warn
 import software.aws.toolkits.jetbrains.core.coroutines.EDT
-import software.aws.toolkits.jetbrains.core.coroutines.disposableCoroutineScope
 import software.aws.toolkits.jetbrains.core.coroutines.getCoroutineBgContext
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
@@ -181,15 +179,9 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
             }
         }
 
-        // When popup is disposed we will cancel this coroutine. The only places popup can get disposed should be
-        // from CodeWhispererPopupManager.cancelPopup() and CodeWhispererPopupManager.closePopup().
-        // It's possible and ok that coroutine will keep running until the next time we check it's state.
-        // As long as we don't show to the user extra info we are good.
-        val coroutineScope = disposableCoroutineScope(popup)
-
         var states: InvocationContext? = null
 
-        val job = coroutineScope.launch {
+        val job = cs.launch {
             try {
                 var startTime = System.nanoTime()
                 CodeWhispererInvocationStatus.getInstance().setInvocationStart()
@@ -211,13 +203,11 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
                     runInEdt {
                         states = processCodeWhispererUI(workerContext, states)
                     }
-                    if (!isActive) {
-                        // If job is cancelled before we do another request, don't bother making
-                        // another API call to save resources
+                    if (popup.isDisposed) {
                         LOG.debug { "Skipping sending remaining requests on CodeWhisperer session exit" }
                         return@launch
                     }
-                } while (nextToken != null)
+                } while (nextToken != null && nextToken.left.isNotEmpty())
             } catch (e: Exception) {
                 // TODO flare: flare doesn't return exceptions
                 val sessionId = ""
@@ -302,7 +292,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
             return null
         }
 
-        if (completions.partialResultToken == null) {
+        if (completions.partialResultToken?.left.isNullOrEmpty()) {
             CodeWhispererInvocationStatus.getInstance().finishInvocation()
         }
 
@@ -336,7 +326,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
         // If there are no recommendations at all in this session, we need to manually send the user decision event here
         // since it won't be sent automatically later
         if (!hasAtLeastOneValid) {
-            if (completions.partialResultToken == null) {
+            if (completions.partialResultToken?.left.isNullOrEmpty()) {
                 LOG.debug { "None of the recommendations are valid, exiting CodeWhisperer session" }
                 CodeWhispererPopupManager.getInstance().cancelPopup(popup)
                 return null
@@ -531,7 +521,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
                     editor.caretModel.primaryCaret.logicalPosition.column
                 )
                 if (nextToken != null) {
-                    workDoneToken = nextToken
+                    partialResultToken = nextToken
                 }
             }
         }

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

@@ -16,6 +16,7 @@ import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocume
 import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.javaFileName
 import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.javaResponse
 import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.javaTestContext
+import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.testNextToken
 import software.aws.toolkits.jetbrains.services.codewhisperer.actions.CodeWhispererActionPromoter
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.QInlineActionId.qInlineAcceptActionId
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererConstants.QInlineActionId.qInlineForceAcceptActionId
@@ -89,7 +90,7 @@ class CodeWhispererAcceptTest : CodeWhispererTestBase() {
                     )
                 ),
                 sessionId = "sessionId",
-                partialResultToken = null
+                partialResultToken = testNextToken
             )
         )
 

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

@@ -50,7 +50,7 @@ object CodeWhispererTestUtil {
     const val codeWhispererRecommendationActionId = "CodeWhispererRecommendationAction"
     const val codeWhispererCodeScanActionId = "codewhisperer.toolbar.security.scan"
     const val testValidAccessToken = "test_valid_access_token"
-    val testNextToken: Either<String, Int> = Either.forLeft("test_next_token")
+    val testNextToken: Either<String, Int> = Either.forLeft("")
     val metadata: DefaultAwsResponseMetadata = DefaultAwsResponseMetadata.create(
         mapOf(AwsHeader.AWS_REQUEST_ID to testRequestId)
     )
@@ -67,6 +67,7 @@ object CodeWhispererTestUtil {
             InlineCompletionItem("item4", "test recommendation 5"),
         ),
         sessionId = "sessionId",
+        partialResultToken = testNextToken
     )
     val javaResponse: InlineCompletionListWithReferences = InlineCompletionListWithReferences(
         items = listOf(
@@ -77,6 +78,7 @@ object CodeWhispererTestUtil {
             InlineCompletionItem("item5", "test recommendation 5"),
         ),
         sessionId = "sessionId",
+        partialResultToken = testNextToken
     )
     const val pythonFileName = "test.py"
     const val javaFileName = "test.java"

plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/textDocument/InlineCompletionWithReferencesParams.kt

@@ -3,8 +3,8 @@
 
 package software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.textDocument
 
-import org.eclipse.lsp4j.TextDocumentPositionAndWorkDoneProgressParams
+import org.eclipse.lsp4j.TextDocumentPositionAndWorkDoneProgressAndPartialResultParams
 
 data class InlineCompletionWithReferencesParams(
     var context: InlineCompletionContext,
-) : TextDocumentPositionAndWorkDoneProgressParams()
+) : TextDocumentPositionAndWorkDoneProgressAndPartialResultParams()