@@ -7,6 +7,7 @@ import com.intellij.openapi.components.Service
77import com.intellij.openapi.components.service
88import com.intellij.openapi.project.Project
99import com.intellij.util.io.HttpRequests
10+ import software.amazon.awssdk.core.exception.SdkException
1011import software.amazon.awssdk.services.codewhispererruntime.CodeWhispererRuntimeClient
1112import software.amazon.awssdk.services.codewhispererruntime.model.CodeWhispererRuntimeResponse
1213import software.amazon.awssdk.services.codewhispererruntime.model.ContentChecksumType
@@ -32,15 +33,18 @@ import software.amazon.awssdk.services.codewhispererruntime.model.UploadContext
3233import software.amazon.awssdk.services.codewhispererruntime.model.UploadIntent
3334import software.amazon.awssdk.services.codewhispererstreaming.model.ExportContext
3435import software.amazon.awssdk.services.codewhispererstreaming.model.ExportIntent
36+ import software.amazon.awssdk.services.codewhispererstreaming.model.ThrottlingException
3537import software.amazon.awssdk.services.codewhispererstreaming.model.TransformationDownloadArtifactType
3638import software.amazon.awssdk.services.codewhispererstreaming.model.TransformationExportContext
39+ import software.amazon.awssdk.services.codewhispererstreaming.model.ValidationException
3740import software.aws.toolkits.core.utils.error
3841import software.aws.toolkits.core.utils.getLogger
3942import software.aws.toolkits.core.utils.info
4043import software.aws.toolkits.jetbrains.core.AwsClientManager
4144import software.aws.toolkits.jetbrains.services.amazonq.APPLICATION_ZIP
4245import software.aws.toolkits.jetbrains.services.amazonq.AWS_KMS
4346import software.aws.toolkits.jetbrains.services.amazonq.CONTENT_SHA256
47+ import software.aws.toolkits.jetbrains.services.amazonq.RetryableOperation
4448import software.aws.toolkits.jetbrains.services.amazonq.SERVER_SIDE_ENCRYPTION
4549import software.aws.toolkits.jetbrains.services.amazonq.SERVER_SIDE_ENCRYPTION_AWS_KMS_KEY_ID
4650import software.aws.toolkits.jetbrains.services.amazonq.clients.AmazonQStreamingClient
@@ -52,7 +56,9 @@ import software.aws.toolkits.jetbrains.services.codemodernizer.utils.calculateTo
5256import software.aws.toolkits.jetbrains.services.codewhisperer.util.CodeWhispererUtil.getTelemetryOptOutPreference
5357import java.io.File
5458import java.net.HttpURLConnection
59+ import java.net.SocketTimeoutException
5560import java.time.Instant
61+ import java.util.concurrent.TimeoutException
5662
5763@Service(Service .Level .PROJECT )
5864class 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 (
0 commit comments