Skip to content

Commit 25f28ad

Browse files
dhasani23dhasani23
authored
fix(amazonq): add retries for /transform APIs (aws#5751)
Add back retries for /transform since they were being overridden to 0.
1 parent 8b3a635 commit 25f28ad

File tree

8 files changed

+121
-21
lines changed

8 files changed

+121
-21
lines changed

‎.changes/next-release/bugfix-ddfb0a9a-529f-4597-8b7a-519d4e838fa4.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type" : "bugfix",
3+
"description" : "/transform: handle InvalidGrantException properly when polling job status"
4+
}

‎plugins/amazonq/chat/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/clients/AmazonQStreamingClientTest.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,7 @@ class AmazonQStreamingClientTest : AmazonQTestBase() {
127127
amazonQStreamingClient.exportResultArchive("test-id", ExportIntent.TRANSFORMATION, null, {}, {})
128128
}
129129

130-
assertThat(attemptCount).isEqualTo(3)
130+
assertThat(attemptCount).isEqualTo(4)
131131
assertThat(thrown)
132132
.isInstanceOf(ValidationException::class.java)
133133
.hasMessage("Resource validation failed")

‎plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/CodeModernizerSession.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -547,7 +547,7 @@ class CodeModernizerSession(
547547
} catch (e: Exception) {
548548
return when (e) {
549549
is AlreadyDisposedException, is CancellationException -> {
550-
LOG.warn { "The session was disposed while polling for job details." }
550+
LOG.error(e) { "The session was disposed while polling for job details." }
551551
CodeModernizerJobCompletedResult.ManagerDisposed
552552
}
553553

‎plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/client/GumbyClient.kt‎

Lines changed: 31 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import com.intellij.openapi.components.Service
77
import com.intellij.openapi.components.service
88
import com.intellij.openapi.project.Project
99
import com.intellij.util.io.HttpRequests
10+
import software.amazon.awssdk.core.exception.SdkException
1011
import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient
1112
import software.amazon.awssdk.services.codewhispererruntime.model.CodeWhispererRuntimeResponse
1213
import software.amazon.awssdk.services.codewhispererruntime.model.ContentChecksumType
@@ -32,15 +33,18 @@ import software.amazon.awssdk.services.codewhispererruntime.model.UploadContext
3233
import software.amazon.awssdk.services.codewhispererruntime.model.UploadIntent
3334
import software.amazon.awssdk.services.codewhispererstreaming.model.ExportContext
3435
import software.amazon.awssdk.services.codewhispererstreaming.model.ExportIntent
36+
import software.amazon.awssdk.services.codewhispererstreaming.model.ThrottlingException
3537
import software.amazon.awssdk.services.codewhispererstreaming.model.TransformationDownloadArtifactType
3638
import software.amazon.awssdk.services.codewhispererstreaming.model.TransformationExportContext
39+
import software.amazon.awssdk.services.codewhispererstreaming.model.ValidationException
3740
import software.aws.toolkits.core.utils.error
3841
import software.aws.toolkits.core.utils.getLogger
3942
import software.aws.toolkits.core.utils.info
4043
import software.aws.toolkits.jetbrains.core.AwsClientManager
4144
import software.aws.toolkits.jetbrains.services.amazonq.APPLICATION_ZIP
4245
import software.aws.toolkits.jetbrains.services.amazonq.AWS_KMS
4346
import software.aws.toolkits.jetbrains.services.amazonq.CONTENT_SHA256
47+
import software.aws.toolkits.jetbrains.services.amazonq.RetryableOperation
4448
import software.aws.toolkits.jetbrains.services.amazonq.SERVER_SIDE_ENCRYPTION
4549
import software.aws.toolkits.jetbrains.services.amazonq.SERVER_SIDE_ENCRYPTION_AWS_KMS_KEY_ID
4650
import software.aws.toolkits.jetbrains.services.amazonq.clients.AmazonQStreamingClient
@@ -52,7 +56,9 @@ import software.aws.toolkits.jetbrains.services.codemodernizer.utils.calculateTo
5256
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.getTelemetryOptOutPreference
5357
import java.io.File
5458
import java.net.HttpURLConnection
59+
import java.net.SocketTimeoutException
5560
import java.time.Instant
61+
import java.util.concurrent.TimeoutException
5662

5763
@Service(Service.Level.PROJECT)
5864
class GumbyClient(private val project: Project) {
@@ -152,15 +158,34 @@ class GumbyClient(private val project: Project) {
152158
apiCall: () -> T,
153159
apiName: String,
154160
): T {
155-
var result: CodeWhispererRuntimeResponse? = null
161+
var result: T? = null
156162
try {
157-
result = apiCall()
158-
LOG.info { "$apiName request ID: ${result.responseMetadata()?.requestId()}" }
159-
return result
163+
RetryableOperation<Unit>().execute(
164+
operation = {
165+
result = apiCall()
166+
},
167+
isRetryable = { e ->
168+
when (e) {
169+
is ValidationException,
170+
is ThrottlingException,
171+
is SdkException,
172+
is TimeoutException,
173+
is SocketTimeoutException,
174+
-> true
175+
else -> false
176+
}
177+
},
178+
errorHandler = { e, attempts ->
179+
LOG.error(e) { "After $attempts attempts, $apiName failed: ${e.message}" }
180+
throw e
181+
}
182+
)
160183
} catch (e: Exception) {
161-
LOG.error(e) { "$apiName failed: ${e.message}" }
162-
throw e // pass along error to callee
184+
LOG.error(e) { "$apiName failed: ${e.message}; may have been retried up to 3 times" }
185+
throw e
163186
}
187+
LOG.info { "$apiName request ID: ${result?.responseMetadata()?.requestId()}" }
188+
return result ?: error("$apiName failed")
164189
}
165190

166191
suspend fun downloadExportResultArchive(

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

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@ package software.aws.toolkits.jetbrains.services.codemodernizer.utils
55

66
import com.fasterxml.jackson.module.kotlin.readValue
77
import com.intellij.grazie.utils.orFalse
8-
import com.intellij.notification.NotificationAction
98
import com.intellij.openapi.application.runInEdt
109
import com.intellij.openapi.application.runWriteAction
1110
import com.intellij.openapi.diff.impl.patch.PatchReader
@@ -144,23 +143,18 @@ suspend fun JobId.pollTransformationStatusAndPlan(
144143
refreshToken(project)
145144
return@waitUntil state
146145
} catch (e: InvalidGrantException) {
147-
CodeTransformMessageListener.instance.onCheckAuth()
146+
CodeTransformMessageListener.instance.onReauthStarted()
148147
notifyStickyWarn(
149148
message("codemodernizer.notification.warn.expired_credentials.title"),
150149
message("codemodernizer.notification.warn.expired_credentials.content"),
151-
project,
152-
listOf(
153-
NotificationAction.createSimpleExpiring(message("codemodernizer.notification.warn.action.reauthenticate")) {
154-
CodeTransformMessageListener.instance.onReauthStarted()
155-
}
156-
)
157150
)
158151
return@waitUntil state
159152
} finally {
160153
delay(sleepDurationMillis)
161154
}
162155
}
163156
} catch (e: Exception) {
157+
getLogger<CodeModernizerManager>().error(e) { "Error when polling for job status & plan" }
164158
// Still call onStateChange to update the UI
165159
onStateChange(state, TransformationStatus.FAILED, transformationPlan)
166160
when (e) {

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

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,10 +13,12 @@ import org.junit.Rule
1313
import org.junit.Test
1414
import org.mockito.kotlin.any
1515
import org.mockito.kotlin.argumentCaptor
16+
import org.mockito.kotlin.doAnswer
1617
import org.mockito.kotlin.doReturn
1718
import org.mockito.kotlin.eq
1819
import org.mockito.kotlin.mock
1920
import org.mockito.kotlin.stub
21+
import org.mockito.kotlin.times
2022
import org.mockito.kotlin.verify
2123
import org.mockito.kotlin.verifyNoInteractions
2224
import org.mockito.kotlin.verifyNoMoreInteractions
@@ -32,6 +34,7 @@ import software.amazon.awssdk.services.codewhispererruntime.model.StartTransform
3234
import software.amazon.awssdk.services.codewhispererruntime.model.StartTransformationResponse
3335
import software.amazon.awssdk.services.codewhispererruntime.model.StopTransformationRequest
3436
import software.amazon.awssdk.services.codewhispererruntime.model.StopTransformationResponse
37+
import software.amazon.awssdk.services.codewhispererruntime.model.ThrottlingException
3538
import software.amazon.awssdk.services.codewhispererruntime.model.TransformationLanguage
3639
import software.amazon.awssdk.services.codewhispererstreaming.CodeWhispererStreamingAsyncClient
3740
import software.amazon.awssdk.services.codewhispererstreaming.model.ExportIntent
@@ -135,6 +138,80 @@ class CodeWhispererCodeModernizerGumbyClientTest : CodeWhispererCodeModernizerTe
135138
}
136139
}
137140

141+
@Test
142+
fun `getCodeModernizationJob should retry on retryable exceptions`() {
143+
var callCount = 0
144+
bearerClient = mockClientManagerRule.create<CodeWhispererRuntimeClient>().stub {
145+
on { getTransformation(any<GetTransformationRequest>()) } doAnswer {
146+
callCount++
147+
when (callCount) {
148+
1 -> throw ThrottlingException.builder().message("Throttled 1").build()
149+
2 -> throw ThrottlingException.builder().message("Throttled 2").build()
150+
else -> exampleGetCodeMigrationResponse
151+
}
152+
}
153+
}
154+
val actual = gumbyClient.getCodeModernizationJob("jobId")
155+
argumentCaptor<GetTransformationRequest>().apply {
156+
// succeeds on 3rd attempt
157+
verify(bearerClient, times(3)).getTransformation(capture())
158+
verifyNoMoreInteractions(bearerClient)
159+
verifyNoInteractions(streamingBearerClient)
160+
assertThat(allValues).hasSize(3)
161+
allValues.forEach { request ->
162+
assertThat(request.transformationJobId()).isEqualTo("jobId")
163+
}
164+
assertThat(actual).isInstanceOf(GetTransformationResponse::class.java)
165+
assertThat(actual).usingRecursiveComparison()
166+
.comparingOnlyFields("jobId", "status", "transformationType", "source")
167+
.isEqualTo(exampleGetCodeMigrationResponse)
168+
}
169+
}
170+
171+
@Test
172+
fun `getCodeModernizationJob should fail immediately on non-retryable exception`() {
173+
val exception = IllegalArgumentException("Non-retryable error")
174+
bearerClient = mockClientManagerRule.create<CodeWhispererRuntimeClient>().stub {
175+
on { getTransformation(any<GetTransformationRequest>()) } doAnswer {
176+
throw exception
177+
}
178+
}
179+
val thrown = runCatching {
180+
gumbyClient.getCodeModernizationJob("jobId")
181+
}.exceptionOrNull()
182+
assertThat(thrown)
183+
.isInstanceOf(IllegalArgumentException::class.java)
184+
.hasMessage("Non-retryable error")
185+
// called just once since it fails immediately
186+
verify(bearerClient, times(1)).getTransformation(any<GetTransformationRequest>())
187+
verifyNoMoreInteractions(bearerClient)
188+
verifyNoInteractions(streamingBearerClient)
189+
}
190+
191+
@Test
192+
fun `getCodeModernizationJob should fail after max attempts`() {
193+
bearerClient = mockClientManagerRule.create<CodeWhispererRuntimeClient>().stub {
194+
on { getTransformation(any<GetTransformationRequest>()) } doAnswer {
195+
throw ThrottlingException.builder().message("Always throttled").build()
196+
}
197+
}
198+
val thrown = runCatching {
199+
gumbyClient.getCodeModernizationJob("jobId")
200+
}.exceptionOrNull()
201+
assertThat(thrown)
202+
.isInstanceOf(ThrottlingException::class.java)
203+
.hasMessage("Always throttled")
204+
argumentCaptor<GetTransformationRequest>().apply {
205+
// called 4 times since it always fails
206+
verify(bearerClient, times(4)).getTransformation(capture())
207+
allValues.forEach { request ->
208+
assertThat(request.transformationJobId()).isEqualTo("jobId")
209+
}
210+
}
211+
verifyNoMoreInteractions(bearerClient)
212+
verifyNoInteractions(streamingBearerClient)
213+
}
214+
138215
@Test
139216
fun `check startCodeModernization on JAVA_17 target`() {
140217
val actual = gumbyClient.startCodeModernization("jobId", TransformationLanguage.JAVA_8, TransformationLanguage.JAVA_17)

‎plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/RetryableOperation.kt‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,12 +30,12 @@ class RetryableOperation<T> {
3030
isRetryable: (Exception) -> Boolean = { it is RetryableException },
3131
errorHandler: (suspend (Exception, Int) -> Nothing),
3232
): T {
33-
while (attempts < MAX_RETRY_ATTEMPTS) {
33+
while (attempts < MAX_ATTEMPTS) {
3434
try {
3535
return operation()
3636
} catch (e: Exception) {
3737
attempts++
38-
if (attempts >= MAX_RETRY_ATTEMPTS || !isRetryable(e)) {
38+
if (attempts >= MAX_ATTEMPTS || !isRetryable(e)) {
3939
errorHandler.invoke(e, attempts)
4040
}
4141
delay(getJitteredDelay())
@@ -48,6 +48,6 @@ class RetryableOperation<T> {
4848
companion object {
4949
private const val INITIAL_DELAY = 100L
5050
private const val MAX_BACKOFF = 10000L
51-
private const val MAX_RETRY_ATTEMPTS = 3
51+
private const val MAX_ATTEMPTS = 4
5252
}
5353
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -626,7 +626,7 @@ codemodernizer.chat.form.user_selection.item.choose_sql_metadata_file=Okay, I ca
626626
codemodernizer.chat.form.user_selection.item.choose_target_version=Choose the target code version
627627
codemodernizer.chat.form.user_selection.title=Q - Code transformation
628628
codemodernizer.chat.message.absolute_path_detected=I detected {0} potential absolute file path(s) in your {1} file: **{2}**. Absolute file paths might cause issues when I build your code. Any errors will show up in the build log.
629-
codemodernizer.chat.message.auth_prompt=Follow instructions to re-authenticate ...
629+
codemodernizer.chat.message.auth_prompt=Follow instructions to re-authenticate. If you experience difficulty, sign out of the Amazon Q plugin and sign in again.
630630
codemodernizer.chat.message.button.cancel=Cancel
631631
codemodernizer.chat.message.button.confirm=Confirm
632632
codemodernizer.chat.message.button.hil_cancel=Cancel
@@ -780,7 +780,7 @@ codemodernizer.notification.warn.download_failed_invalid_artifact=Amazon Q was u
780780
codemodernizer.notification.warn.download_failed_other.content=Amazon Q ran into an issue while trying to download your {0}. Please try again. {1}
781781
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.
782782
codemodernizer.notification.warn.download_failed_wildcard.content=Check your IDE proxy settings and remove any wildcard (*) references, and then try viewing the diff again.
783-
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.
783+
codemodernizer.notification.warn.expired_credentials.content=Unable to check transformation status as your credentials expired. Try signing out of the Amazon Q plugin and signing in again.
784784
codemodernizer.notification.warn.expired_credentials.title=Your connection to Q has expired
785785
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.
786786
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).

0 commit comments

Comments
 (0)