Skip to content

fix(amazonq): add retries for /transform APIs #5751

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 12 commits into from
May 27, 2025
Merged

fix(amazonq): add retries for /transform APIs #5751

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
3805f06
fix(amazonq): add retries for /transform APIs
May 22, 2025
058d228
set attempts to 4
May 22, 2025
3679de3
handle re-auth
May 22, 2025
21d3966
text change
May 22, 2025
bd6cbd7
minor text change
May 22, 2025
db33b04
minor text change
May 22, 2025
5abc53b
remove unused import
May 22, 2025
bcfeb21
add changelog
May 22, 2025
d4f9edf
fix test
May 22, 2025
766a1c3
log error
May 23, 2025
9079590
add log
May 23, 2025
a6092d5
add tests for retry logic
May 23, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 31 additions & 6 deletions ...mmunity/src/software/aws/toolkits/jetbrains/services/codemodernizer/client/GumbyClient.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import com.intellij.openapi.components.Service
import com.intellij.openapi.components.service
import com.intellij.openapi.project.Project
import com.intellij.util.io.HttpRequests
import software.amazon.awssdk.core.exception.SdkException
import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient
import software.amazon.awssdk.services.codewhispererruntime.model.CodeWhispererRuntimeResponse
import software.amazon.awssdk.services.codewhispererruntime.model.ContentChecksumType
Expand All @@ -32,15 +33,18 @@ import software.amazon.awssdk.services.codewhispererruntime.model.UploadContext
import software.amazon.awssdk.services.codewhispererruntime.model.UploadIntent
import software.amazon.awssdk.services.codewhispererstreaming.model.ExportContext
import software.amazon.awssdk.services.codewhispererstreaming.model.ExportIntent
import software.amazon.awssdk.services.codewhispererstreaming.model.ThrottlingException
import software.amazon.awssdk.services.codewhispererstreaming.model.TransformationDownloadArtifactType
import software.amazon.awssdk.services.codewhispererstreaming.model.TransformationExportContext
import software.amazon.awssdk.services.codewhispererstreaming.model.ValidationException
import software.aws.toolkits.core.utils.error
import software.aws.toolkits.core.utils.getLogger
import software.aws.toolkits.core.utils.info
import software.aws.toolkits.jetbrains.core.AwsClientManager
import software.aws.toolkits.jetbrains.services.amazonq.APPLICATION_ZIP
import software.aws.toolkits.jetbrains.services.amazonq.AWS_KMS
import software.aws.toolkits.jetbrains.services.amazonq.CONTENT_SHA256
import software.aws.toolkits.jetbrains.services.amazonq.RetryableOperation
import software.aws.toolkits.jetbrains.services.amazonq.SERVER_SIDE_ENCRYPTION
import software.aws.toolkits.jetbrains.services.amazonq.SERVER_SIDE_ENCRYPTION_AWS_KMS_KEY_ID
import software.aws.toolkits.jetbrains.services.amazonq.clients.AmazonQStreamingClient
Expand All @@ -52,7 +56,9 @@ import software.aws.toolkits.jetbrains.services.codemodernizer.utils.calculateTo
import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.getTelemetryOptOutPreference
import java.io.File
import java.net.HttpURLConnection
import java.net.SocketTimeoutException
import java.time.Instant
import java.util.concurrent.TimeoutException

@Service(Service.Level.PROJECT)
class GumbyClient(private val project: Project) {
Expand Down Expand Up @@ -152,15 +158,34 @@ class GumbyClient(private val project: Project) {
apiCall: () -> T,
apiName: String,
): T {
var result: CodeWhispererRuntimeResponse? = null
var result: T? = null
try {
result = apiCall()
LOG.info { "$apiName request ID: ${result.responseMetadata()?.requestId()}" }
return result
RetryableOperation<Unit>().execute(
operation = {
result = apiCall()
},
isRetryable = { e ->
when (e) {
is ValidationException,
is ThrottlingException,
is SdkException,
is TimeoutException,
is SocketTimeoutException,
-> true
else -> false
Copy link
Contributor Author

@dhasani23 dhasani23 May 22, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Following the same retry strategy from:

aws-toolkit-jetbrains/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/clients/AmazonQStreamingClient.kt

Lines 82 to 89 in ebece9d

isRetryable = { e ->
when (e) {
is ValidationException,
is ThrottlingException,
is SdkException,
is TimeoutException,
-> true
else -> false

in the ExportResultArchive (download) API

and added SocketTimeoutException (which is not a type of TimeoutException, so it needs to be explicitly listed here)

}
},
errorHandler = { e, attempts ->
LOG.error(e) { "After $attempts attempts, $apiName failed: ${e.message}" }
throw e
}
)
} catch (e: Exception) {
LOG.error(e) { "$apiName failed: ${e.message}" }
throw e // pass along error to callee
LOG.error(e) { "$apiName failed: ${e.message}; may have been retried up to 3 times" }
throw e
}
LOG.info { "$apiName request ID: ${result?.responseMetadata()?.requestId()}" }
return result ?: error("$apiName failed")
}

suspend fun downloadExportResultArchive(
Expand Down
8 changes: 1 addition & 7 deletions ...rc/software/aws/toolkits/jetbrains/services/codemodernizer/utils/CodeTransformApiUtils.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -144,16 +144,10 @@ suspend fun JobId.pollTransformationStatusAndPlan(
refreshToken(project)
return@waitUntil state
} catch (e: InvalidGrantException) {
CodeTransformMessageListener.instance.onCheckAuth()
CodeTransformMessageListener.instance.onReauthStarted()
Copy link
Contributor Author

@dhasani23 dhasani23 May 22, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This onCheckAuth was disposing of our codeTransformationSessionContext which doesn't make sense to do. Better to just initiate the onReauthStarted flow which does:

aws-toolkit-jetbrains/plugins/amazonq/codetransform/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codemodernizer/controller/CodeTransformChatController.kt

Line 608 in ebece9d

* Invoking this is equivalent to user clicking "reauthenticate" in the Chat when credentials expired.

instead of waiting/hoping for the user to click on the Re-authenticate button in the notification.

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 {
Expand Down
6 changes: 3 additions & 3 deletions ...ains-community/src/software/aws/toolkits/jetbrains/services/amazonq/RetryableOperation.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,12 +30,12 @@ class RetryableOperation<T> {
isRetryable: (Exception) -> Boolean = { it is RetryableException },
errorHandler: (suspend (Exception, Int) -> Nothing),
): T {
while (attempts < MAX_RETRY_ATTEMPTS) {
while (attempts < MAX_ATTEMPTS) {
try {
return operation()
} catch (e: Exception) {
attempts++
if (attempts >= MAX_RETRY_ATTEMPTS || !isRetryable(e)) {
if (attempts >= MAX_ATTEMPTS || !isRetryable(e)) {
errorHandler.invoke(e, attempts)
}
delay(getJitteredDelay())
Expand All @@ -48,6 +48,6 @@ class RetryableOperation<T> {
companion object {
private const val INITIAL_DELAY = 100L
private const val MAX_BACKOFF = 10000L
private const val MAX_RETRY_ATTEMPTS = 3
private const val MAX_ATTEMPTS = 4
Copy link
Contributor Author

@dhasani23 dhasani23 May 22, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm guessing the intention was to do 3 retries (after the initial attempt), meaning this should technically be set to 4, and naming it MAX_ATTEMPTS is more accurate.

}
}
Loading