Fix UI freeze when trying to refresh tokens offline (#4880)

`CodeWhispererService#showRecommendationsInPopup` was previously running on EDT, which caused issues when we need to wait for the fairly expensive call chain for `promptReauth` if the connection is expired. This is very visible in the offline case.

Fix by moving all the operations onto a background worker and limiting invocations to one at a time. This caused some issues with the test suite since they assumed that `showRecommendationsInPopup` would block until background operations have initialized, which was worked around by passing along job tracking

Author: rli

.changes/next-release/bugfix-21ad7a75-f4aa-4f92-a120-cff1715a6cd5.json

@@ -0,0 +1,4 @@
+{
+  "type" : "bugfix",
+  "description" : "Fix UI slowdown when Amazon Q Inline Suggestions are enabled, but token cannot be refreshed (#4868)"
+}

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/actions/CodeWhispererRecommendationAction.kt

@@ -8,12 +8,15 @@ import com.intellij.openapi.actionSystem.AnAction
 import com.intellij.openapi.actionSystem.AnActionEvent
 import com.intellij.openapi.actionSystem.CommonDataKeys
 import com.intellij.openapi.project.DumbAware
+import com.intellij.openapi.util.Key
+import kotlinx.coroutines.Job
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.LatencyContext
 import software.aws.toolkits.jetbrains.services.codewhisperer.model.TriggerTypeInfo
 import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererAutomatedTriggerType
 import software.aws.toolkits.jetbrains.services.codewhisperer.service.CodeWhispererService
 import software.aws.toolkits.resources.message
 import software.aws.toolkits.telemetry.CodewhispererTriggerType
+import java.util.concurrent.atomic.AtomicReference
 
 class CodeWhispererRecommendationAction : AnAction(message("codewhisperer.trigger.service")), DumbAware {
     override fun getActionUpdateThread() = ActionUpdateThread.BGT
@@ -32,6 +35,12 @@ class CodeWhispererRecommendationAction : AnAction(message("codewhisperer.trigge
         }
 
         val triggerType = TriggerTypeInfo(CodewhispererTriggerType.OnDemand, CodeWhispererAutomatedTriggerType.Unknown())
-        CodeWhispererService.getInstance().showRecommendationsInPopup(editor, triggerType, latencyContext)
+        val job = CodeWhispererService.getInstance().showRecommendationsInPopup(editor, triggerType, latencyContext)
+
+        e.getData(CommonDataKeys.EDITOR)?.getUserData(ACTION_JOB_KEY)?.set(job)
+    }
+
+    companion object {
+        val ACTION_JOB_KEY = Key.create<AtomicReference<Job?>>("amazonq.codewhisperer.job")
     }
 }

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

@@ -28,6 +28,7 @@ import kotlinx.coroutines.async
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.isActive
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
 import software.amazon.awssdk.core.exception.SdkServiceException
 import software.amazon.awssdk.core.util.DefaultSdkAutoConstructList
 import software.amazon.awssdk.services.codewhisperer.model.CodeWhispererException
@@ -46,6 +47,7 @@ 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.disposableCoroutineScope
+import software.aws.toolkits.jetbrains.core.coroutines.getCoroutineBgContext
 import software.aws.toolkits.jetbrains.core.coroutines.projectCoroutineScope
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
@@ -82,7 +84,6 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.util.FileContextPr
 import software.aws.toolkits.jetbrains.utils.isInjectedText
 import software.aws.toolkits.jetbrains.utils.isQExpired
 import software.aws.toolkits.jetbrains.utils.notifyWarn
-import software.aws.toolkits.jetbrains.utils.pluginAwareExecuteOnPooledThread
 import software.aws.toolkits.resources.message
 import software.aws.toolkits.telemetry.CodewhispererCompletionType
 import software.aws.toolkits.telemetry.CodewhispererSuggestionState
@@ -98,26 +99,46 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
         Disposer.register(this, codeInsightSettingsFacade)
     }
 
+    private var job: Job? = null
     fun showRecommendationsInPopup(
         editor: Editor,
         triggerTypeInfo: TriggerTypeInfo,
         latencyContext: LatencyContext
+    ): Job? {
+        if (job == null || job?.isCompleted == true) {
+            job = cs.launch(getCoroutineBgContext()) {
+                doShowRecommendationsInPopup(editor, triggerTypeInfo, latencyContext)
+            }
+        }
+
+        // did some wrangling, but compiler didn't believe this can't be null
+        return job
+    }
+
+    private suspend fun doShowRecommendationsInPopup(
+        editor: Editor,
+        triggerTypeInfo: TriggerTypeInfo,
+        latencyContext: LatencyContext
     ) {
         val project = editor.project ?: return
         if (!isCodeWhispererEnabled(project)) return
 
         latencyContext.credentialFetchingStart = System.nanoTime()
 
+        // try to refresh automatically if possible, otherwise ask user to login again
         if (isQExpired(project)) {
+            // consider changing to only running once a ~minute since this is relatively expensive
             // say the connection is un-refreshable if refresh fails for 3 times
             val shouldReauth = if (refreshFailure < MAX_REFRESH_ATTEMPT) {
-                pluginAwareExecuteOnPooledThread {
+                val attempt = withContext(getCoroutineBgContext()) {
                     promptReAuth(project)
-                }.get().also { success ->
-                    if (!success) {
-                        refreshFailure++
-                    }
                 }
+
+                if (!attempt) {
+                    refreshFailure++
+                }
+
+                attempt
             } else {
                 true
             }

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

@@ -8,6 +8,10 @@ import com.intellij.testFramework.DisposableRule
 import com.intellij.testFramework.RuleChain
 import com.intellij.testFramework.replaceService
 import com.intellij.testFramework.runInEdtAndWait
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.test.runTest
+import kotlinx.coroutines.yield
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.After
 import org.junit.Before
@@ -32,6 +36,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestU
 import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonResponse
 import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.pythonTestLeftContext
 import software.aws.toolkits.jetbrains.services.codewhisperer.CodeWhispererTestUtil.testValidAccessToken
+import software.aws.toolkits.jetbrains.services.codewhisperer.actions.CodeWhispererRecommendationAction
 import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
 import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererLoginType
 import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorManager
@@ -51,6 +56,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.settings.CodeWhisp
 import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererColorUtil.POPUP_DIM_HEX
 import software.aws.toolkits.jetbrains.utils.rules.PythonCodeInsightTestFixtureRule
 import software.aws.toolkits.resources.message
+import java.util.concurrent.atomic.AtomicReference
 
 // TODO: restructure testbase, too bulky and hard to debug
 open class CodeWhispererTestBase {
@@ -180,12 +186,26 @@ open class CodeWhispererTestBase {
         runInEdtAndWait {}
     }
 
+    /**
+     * Block until manual action has either failed or completed
+     */
     fun invokeCodeWhispererService() {
+        val jobRef = AtomicReference<Job?>()
         runInEdtAndWait {
+            projectRule.fixture.editor.putUserData(CodeWhispererRecommendationAction.ACTION_JOB_KEY, jobRef)
+            // does not block, so we need to extract something to track the async task
             projectRule.fixture.performEditorAction(codeWhispererRecommendationActionId)
         }
-        while (CodeWhispererInvocationStatus.getInstance().hasExistingInvocation()) {
-            Thread.sleep(10)
+
+        runTest {
+            // wait for CodeWhispererService#showRecommendationsInPopup to complete, if started
+            jobRef.get()?.join()
+
+            // wait for subsequent background operations to be complete
+            while (CodeWhispererInvocationStatus.getInstance().hasExistingInvocation()) {
+                yield()
+                delay(10)
+            }
         }
     }