diff --git a/.changes/next-release/bugfix-21ad7a75-f4aa-4f92-a120-cff1715a6cd5.json b/.changes/next-release/bugfix-21ad7a75-f4aa-4f92-a120-cff1715a6cd5.json new file mode 100644 index 00000000000..9c013b203f6 --- /dev/null +++ b/.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)" +} \ No newline at end of file diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/actions/CodeWhispererRecommendationAction.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/actions/CodeWhispererRecommendationAction.kt index 62a6628f801..c5f4a546707 100644 --- a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/actions/CodeWhispererRecommendationAction.kt +++ b/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>("amazonq.codewhisperer.job") } } 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 index e8dd0e96562..13af5360d41 100644 --- 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 @@ -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 } 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 index bc5e09dd52d..55e1cb24f17 100644 --- 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 @@ -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() 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) + } } }