Q Code Transform: Refresh token if grant expires during transformation (#4638)

Author: tincheng

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

@@ -16,6 +16,7 @@ import software.amazon.awssdk.services.codewhispererruntime.model.Transformation
 import software.amazon.awssdk.services.codewhispererruntime.model.TransformationProgressUpdate
 import software.amazon.awssdk.services.codewhispererruntime.model.TransformationStatus
 import software.amazon.awssdk.services.codewhispererruntime.model.ValidationException
+import software.amazon.awssdk.services.ssooidc.model.InvalidGrantException
 import software.aws.toolkits.core.utils.WaiterUnrecoverableException
 import software.aws.toolkits.core.utils.Waiters.waitUntil
 import software.aws.toolkits.jetbrains.services.codemodernizer.CodeTransformTelemetryManager
@@ -53,6 +54,9 @@ suspend fun JobId.pollTransformationStatusAndPlan(
     var didSleepOnce = false
     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.
     refreshToken(project)
 
     try {
@@ -99,6 +103,10 @@ suspend fun JobId.pollTransformationStatusAndPlan(
                 if (numRefreshes++ > maxRefreshes) throw e
                 refreshToken(project)
                 return@waitUntil state
+            } catch (e: InvalidGrantException) {
+                if (numRefreshes++ > maxRefreshes) throw e
+                refreshToken(project)
+                return@waitUntil state
             } finally {
                 sleep(sleepDurationMillis)
             }
@@ -107,7 +115,7 @@ suspend fun JobId.pollTransformationStatusAndPlan(
         // Still call onStateChange to update the UI
         onStateChange(state, TransformationStatus.FAILED, transformationPlan)
         when (e) {
-            is WaiterUnrecoverableException, is AccessDeniedException -> {
+            is WaiterUnrecoverableException, is AccessDeniedException, is InvalidGrantException -> {
                 return PollingResult(false, transformationResponse?.transformationJob(), state, transformationPlan)
             }
             else -> throw e

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

@@ -2,6 +2,10 @@
 // SPDX-License-Identifier: Apache-2.0
 
 package software.aws.toolkits.jetbrains.services.codemodernizer
+import io.mockk.every
+import io.mockk.just
+import io.mockk.mockkStatic
+import io.mockk.runs
 import kotlinx.coroutines.runBlocking
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.Before
@@ -11,10 +15,13 @@ import org.mockito.kotlin.any
 import org.mockito.kotlin.times
 import org.mockito.kotlin.verify
 import org.mockito.kotlin.whenever
+import software.amazon.awssdk.services.codewhispererruntime.model.AccessDeniedException
 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.utils.getTableMapping
 import software.aws.toolkits.jetbrains.services.codemodernizer.utils.pollTransformationStatusAndPlan
+import software.aws.toolkits.jetbrains.services.codemodernizer.utils.refreshToken
 import java.util.concurrent.atomic.AtomicBoolean
 
 class CodeWhispererCodeModernizerUtilsTest : CodeWhispererCodeModernizerTestBase() {
@@ -57,6 +64,88 @@ class CodeWhispererCodeModernizerUtilsTest : CodeWhispererCodeModernizerTestBase
         assertThat(expected).isEqualTo(mutableList)
     }
 
+    @Test
+    fun `refresh on access denied`() {
+        val mockAccessDeniedException = Mockito.mock(AccessDeniedException::class.java)
+
+        mockkStatic(::refreshToken)
+        every { refreshToken(any()) } just runs
+
+        Mockito.doThrow(
+            mockAccessDeniedException
+        ).doReturn(
+            exampleGetCodeMigrationResponse,
+            exampleGetCodeMigrationResponse.replace(TransformationStatus.STARTED),
+            exampleGetCodeMigrationResponse.replace(TransformationStatus.COMPLETED), // Should stop before this point
+        ).whenever(clientAdaptorSpy).getCodeModernizationJob(any())
+
+        Mockito.doReturn(exampleGetCodeMigrationPlanResponse)
+            .whenever(clientAdaptorSpy).getCodeModernizationPlan(any())
+
+        val mutableList = mutableListOf<TransformationStatus>()
+        runBlocking {
+            jobId.pollTransformationStatusAndPlan(
+                setOf(TransformationStatus.STARTED),
+                setOf(TransformationStatus.FAILED),
+                clientAdaptorSpy,
+                0,
+                0,
+                AtomicBoolean(false),
+                project
+            ) { _, status, _ ->
+                mutableList.add(status)
+            }
+        }
+        val expected =
+            listOf<TransformationStatus>(
+                exampleGetCodeMigrationResponse.transformationJob().status(),
+                TransformationStatus.STARTED,
+            )
+        assertThat(expected).isEqualTo(mutableList)
+        io.mockk.verify { refreshToken(any()) }
+    }
+
+    @Test
+    fun `refresh on invalid grant`() {
+        val mockInvalidGrantException = Mockito.mock(InvalidGrantException::class.java)
+
+        mockkStatic(::refreshToken)
+        every { refreshToken(any()) } just runs
+
+        Mockito.doThrow(
+            mockInvalidGrantException
+        ).doReturn(
+            exampleGetCodeMigrationResponse,
+            exampleGetCodeMigrationResponse.replace(TransformationStatus.STARTED),
+            exampleGetCodeMigrationResponse.replace(TransformationStatus.COMPLETED), // Should stop before this point
+        ).whenever(clientAdaptorSpy).getCodeModernizationJob(any())
+
+        Mockito.doReturn(exampleGetCodeMigrationPlanResponse)
+            .whenever(clientAdaptorSpy).getCodeModernizationPlan(any())
+
+        val mutableList = mutableListOf<TransformationStatus>()
+        runBlocking {
+            jobId.pollTransformationStatusAndPlan(
+                setOf(TransformationStatus.STARTED),
+                setOf(TransformationStatus.FAILED),
+                clientAdaptorSpy,
+                0,
+                0,
+                AtomicBoolean(false),
+                project
+            ) { _, status, _ ->
+                mutableList.add(status)
+            }
+        }
+        val expected =
+            listOf<TransformationStatus>(
+                exampleGetCodeMigrationResponse.transformationJob().status(),
+                TransformationStatus.STARTED,
+            )
+        assertThat(expected).isEqualTo(mutableList)
+        io.mockk.verify { refreshToken(any()) }
+    }
+
     @Test
     fun `stops polling when status transitions to failOn`() {
         Mockito.doReturn(