fix(amazonq): prompt users to re-auth on invalid grant

Author: David Hasani

plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/panels/managers/CodeModernizerBottomWindowPanelManager.kt

@@ -8,6 +8,7 @@ import com.intellij.openapi.actionSystem.ActionGroup
 import com.intellij.openapi.actionSystem.ActionManager
 import com.intellij.openapi.actionSystem.ActionToolbar
 import com.intellij.openapi.application.invokeLater
+import com.intellij.openapi.application.runInEdt
 import com.intellij.openapi.components.service
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.projectRoots.JavaSdkVersion
@@ -62,14 +63,16 @@ class CodeModernizerBottomWindowPanelManager(private val project: Project) : JPa
     }
 
     private fun setUI(function: () -> Unit) {
-        lastShownProgressPanel = this.components.firstOrNull { it == fullSizeLoadingPanel || it == buildProgressSplitterPanelManager } ?: lastShownProgressPanel
-        removeAll()
-        add(BorderLayout.WEST, toolbar.component)
-        add(BorderLayout.NORTH, banner)
-        function.invoke()
-        updateRunTime()
-        revalidate()
-        repaint()
+        runInEdt {
+            lastShownProgressPanel = this.components.firstOrNull { it == fullSizeLoadingPanel || it == buildProgressSplitterPanelManager } ?: lastShownProgressPanel
+            removeAll()
+            add(BorderLayout.WEST, toolbar.component)
+            add(BorderLayout.NORTH, banner)
+            function.invoke()
+            updateRunTime()
+            revalidate()
+            repaint()
+        }
     }
 
     private fun updateRunTime(now: Instant? = null) {

plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/utils/CodeTransformApiUtils.kt

@@ -4,6 +4,7 @@
 package software.aws.toolkits.jetbrains.services.codemodernizer.utils
 
 import com.fasterxml.jackson.module.kotlin.readValue
+import com.intellij.notification.NotificationAction
 import com.intellij.openapi.project.Project
 import com.intellij.serviceContainer.AlreadyDisposedException
 import kotlinx.coroutines.delay
@@ -23,12 +24,14 @@ import software.aws.toolkits.core.utils.WaiterUnrecoverableException
 import software.aws.toolkits.core.utils.Waiters.waitUntil
 import software.aws.toolkits.jetbrains.services.codemodernizer.CodeTransformTelemetryManager
 import software.aws.toolkits.jetbrains.services.codemodernizer.client.GumbyClient
+import software.aws.toolkits.jetbrains.services.codemodernizer.commands.CodeTransformMessageListener
 import software.aws.toolkits.jetbrains.services.codemodernizer.constants.BILLING_RATE
 import software.aws.toolkits.jetbrains.services.codemodernizer.constants.JOB_STATISTICS_TABLE_KEY
 import software.aws.toolkits.jetbrains.services.codemodernizer.model.CodeModernizerArtifact.Companion.MAPPER
 import software.aws.toolkits.jetbrains.services.codemodernizer.model.CodeTransformType
 import software.aws.toolkits.jetbrains.services.codemodernizer.model.JobId
 import software.aws.toolkits.jetbrains.services.codemodernizer.model.PlanTable
+import software.aws.toolkits.jetbrains.utils.notifyStickyWarn
 import software.aws.toolkits.resources.message
 import java.time.Duration
 import java.util.Locale
@@ -64,8 +67,7 @@ suspend fun JobId.pollTransformationStatusAndPlan(
     val maxRefreshes = 10
     var numRefreshes = 0
 
-    // We refresh token at the start of polling, but for some long jobs that runs for 30 minutes plus, tokens may need to be
-    // refreshed again when AccessDeniedException or InvalidGrantException are caught.
+    // refresh token at start of polling since local build just prior can take a long time
     refreshToken(project)
 
     try {
@@ -113,8 +115,17 @@ suspend fun JobId.pollTransformationStatusAndPlan(
                 refreshToken(project)
                 return@waitUntil state
             } catch (e: InvalidGrantException) {
-                if (numRefreshes++ > maxRefreshes) throw e
-                refreshToken(project)
+                CodeTransformMessageListener.instance.onCheckAuth()
+                notifyStickyWarn(
+                    message("codemodernizer.notification.warn.expired_credentials.title"),
+                    message("codemodernizer.notification.warn.expired_credentials.content"),
+                    project,
+                    listOf(
+                        NotificationAction.createSimpleExpiring(message("codemodernizer.notification.warn.action.reauthenticate")) {
+                            CodeTransformMessageListener.instance.onReauthStarted()
+                        }
+                    )
+                )
                 return@waitUntil state
             } finally {
                 delay(sleepDurationMillis)

plugins/amazonq/codetransform/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codemodernizer/CodeWhispererCodeModernizerUtilsTest.kt

@@ -6,6 +6,7 @@ import io.mockk.every
 import io.mockk.just
 import io.mockk.mockkStatic
 import io.mockk.runs
+import io.mockk.verify
 import kotlinx.coroutines.runBlocking
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.Before
@@ -19,16 +20,19 @@ import software.amazon.awssdk.services.codewhispererruntime.model.AccessDeniedEx
 import software.amazon.awssdk.services.codewhispererruntime.model.TransformationProgressUpdate
 import software.amazon.awssdk.services.codewhispererruntime.model.TransformationStatus
 import software.amazon.awssdk.services.ssooidc.model.InvalidGrantException
+import software.aws.toolkits.jetbrains.services.codemodernizer.commands.CodeTransformMessageListener
 import software.aws.toolkits.jetbrains.services.codemodernizer.model.CodeTransformType
 import software.aws.toolkits.jetbrains.services.codemodernizer.utils.getBillingText
 import software.aws.toolkits.jetbrains.services.codemodernizer.utils.getTableMapping
 import software.aws.toolkits.jetbrains.services.codemodernizer.utils.parseBuildFile
 import software.aws.toolkits.jetbrains.services.codemodernizer.utils.pollTransformationStatusAndPlan
 import software.aws.toolkits.jetbrains.services.codemodernizer.utils.refreshToken
 import software.aws.toolkits.jetbrains.services.codemodernizer.utils.validateSctMetadata
+import software.aws.toolkits.jetbrains.utils.notifyStickyWarn
 import software.aws.toolkits.jetbrains.utils.rules.addFileToModule
 import java.util.concurrent.atomic.AtomicBoolean
 import kotlin.io.path.createTempFile
+import io.mockk.*
 
 class CodeWhispererCodeModernizerUtilsTest : CodeWhispererCodeModernizerTestBase() {
     @Before
@@ -114,24 +118,18 @@ class CodeWhispererCodeModernizerUtilsTest : CodeWhispererCodeModernizerTestBase
     }
 
     @Test
-    fun `refresh on invalid grant`() {
-        val mockInvalidGrantException = Mockito.mock(InvalidGrantException::class.java)
-
-        mockkStatic(::refreshToken)
-        every { refreshToken(any()) } just runs
+    fun `show re-auth notification on invalid grant exception`() {
+        mockkStatic(::notifyStickyWarn)
 
-        Mockito.doThrow(
-            mockInvalidGrantException
-        ).doReturn(
-            exampleGetCodeMigrationResponse,
-            exampleGetCodeMigrationResponse.replace(TransformationStatus.STARTED),
-            exampleGetCodeMigrationResponse.replace(TransformationStatus.COMPLETED), // Should stop before this point
-        ).whenever(clientAdaptorSpy).getCodeModernizationJob(any())
+        val mockMessageListener = mockk<CodeTransformMessageListener>()
+        mockkObject(CodeTransformMessageListener)
+        every { CodeTransformMessageListener.instance } returns mockMessageListener
+        every { mockMessageListener.onCheckAuth() } just runs
 
-        Mockito.doReturn(exampleGetCodeMigrationPlanResponse)
-            .whenever(clientAdaptorSpy).getCodeModernizationPlan(any())
+        val mockInvalidGrantException = Mockito.mock(InvalidGrantException::class.java)
+        Mockito.doThrow(mockInvalidGrantException)
+            .whenever(clientAdaptorSpy).getCodeModernizationJob(any())
 
-        val mutableList = mutableListOf<TransformationStatus>()
         runBlocking {
             jobId.pollTransformationStatusAndPlan(
                 CodeTransformType.LANGUAGE_UPGRADE,
@@ -142,17 +140,18 @@ class CodeWhispererCodeModernizerUtilsTest : CodeWhispererCodeModernizerTestBase
                 0,
                 AtomicBoolean(false),
                 project
-            ) { _, status, _ ->
-                mutableList.add(status)
-            }
+            ) { _, _, _ -> }
         }
-        val expected =
-            listOf<TransformationStatus>(
-                exampleGetCodeMigrationResponse.transformationJob().status(),
-                TransformationStatus.STARTED,
+
+        verify {
+            notifyStickyWarn(
+                "Your connection to Q has expired",
+                "Unable to check transformation status as your credentials expired.",
+                project,
+                any()
             )
-        assertThat(expected).isEqualTo(mutableList)
-        io.mockk.verify { refreshToken(any()) }
+        }
+        verify { mockMessageListener.onCheckAuth() }
     }
 
     @Test

plugins/core/resources/resources/software/aws/toolkits/resources/MessagesBundle.properties

@@ -778,6 +778,7 @@ codemodernizer.notification.warn.download_failed_invalid_artifact=Amazon Q was u
 codemodernizer.notification.warn.download_failed_other.content=Amazon Q ran into an issue while trying to download your {0}. Please try again. {1}
 codemodernizer.notification.warn.download_failed_ssl.content=Please make sure all your certificates for your proxy client have been set up correctly for your IDE.
 codemodernizer.notification.warn.download_failed_wildcard.content=Check your IDE proxy settings and remove any wildcard (*) references, and then try viewing the diff again.
+codemodernizer.notification.warn.expired_credentials.content=Unable to check transformation status as your credentials expired. Try signing out of Amazon Q and signing back in again if 'Reauthenticate' below does not work.
 codemodernizer.notification.warn.expired_credentials.title=Your connection to Q has expired
 codemodernizer.notification.warn.invalid_project.description.reason.missing_content_roots=None of your open modules are supported for code transformation with Amazon Q. Amazon Q can upgrade Java 8, Java 11, Java 17, and Java 21 projects built on Maven, with content roots configured.
 codemodernizer.notification.warn.invalid_project.description.reason.not_logged_in=Amazon Q cannot start the transformation as you are not logged in with Identity Center or Builder ID. Also ensure that you are not using IntelliJ version 232.8660.185 and that you are not developing on a remote host (uncommon).