Retry on IDPCommunicationError

services/sts/pom.xml

@@ -63,6 +63,12 @@
         </dependency>
 
         <!-- Test Dependencies -->
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>retries</artifactId>
+            <version>${awsjavasdk.version}</version>
+            <scope>test</scope>
+        </dependency>
         <dependency>
             <artifactId>iam</artifactId>
             <groupId>software.amazon.awssdk</groupId>

services/sts/src/main/java/software/amazon/awssdk/services/sts/internal/StsRetryStrategy.java

@@ -0,0 +1,68 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.services.sts.internal;
+
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.awscore.retry.AwsRetryStrategy;
+import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
+import software.amazon.awssdk.core.client.config.SdkClientOption;
+import software.amazon.awssdk.core.retry.NewRetries2026Resolver;
+import software.amazon.awssdk.core.retry.RetryMode;
+import software.amazon.awssdk.retries.api.RetryStrategy;
+import software.amazon.awssdk.services.sts.model.IdpCommunicationErrorException;
+
+/**
+ * Specialized retry strategy resolution for STS to enable retrying for {@link IdpCommunicationErrorException}.
+ */
+@SdkInternalApi
+public final class StsRetryStrategy {
+
+    private StsRetryStrategy() {
+    }
+
+    public static RetryStrategy resolveRetryStrategy(SdkClientConfiguration config) {
+        RetryStrategy configuredRetryStrategy = config.option(SdkClientOption.RETRY_STRATEGY);
+        if (configuredRetryStrategy != null) {
+            return configuredRetryStrategy;
+        }
+
+        // Just return null and let the normal retry strategy resolution occur
+        if (!isNewRetries2026Enabled(config)) {
+            return null;
+        }
+
+        RetryMode retryMode = resolveRetryMode(config);
+
+        return AwsRetryStrategy.forRetryMode(retryMode, true)
+                               .toBuilder()
+                               .retryOnException(IdpCommunicationErrorException.class)
+                               .build();
+    }
+
+    private static RetryMode resolveRetryMode(SdkClientConfiguration config) {
+        return RetryMode.resolver()
+                        .profileFile(config.option(SdkClientOption.PROFILE_FILE_SUPPLIER))
+                        .profileName(config.option(SdkClientOption.PROFILE_NAME))
+                        .defaultRetryMode(config.option(SdkClientOption.DEFAULT_RETRY_MODE))
+                        .defaultNewRetries2026(config.option(SdkClientOption.DEFAULT_NEW_RETRIES_2026))
+                        .resolve();
+    }
+
+    private static boolean isNewRetries2026Enabled(SdkClientConfiguration config) {
+        Boolean defaultNewRetries2026 = config.option(SdkClientOption.DEFAULT_NEW_RETRIES_2026);
+        return new NewRetries2026Resolver().defaultNewRetries2026(defaultNewRetries2026).resolve();
+    }
+}

services/sts/src/main/resources/codegen-resources/customization.config

@@ -21,6 +21,6 @@
         "UseGlobalEndpoint with legacy region `us-west-2`": "V2 does not support setting UseGlobalEndpoint. It's regional only/by default"
     },
 
-    "enableGenerateCompiledEndpointRules": true
-
+    "enableGenerateCompiledEndpointRules": true,
+    "customRetryStrategy" : "software.amazon.awssdk.services.sts.internal.StsRetryStrategy"
 }

services/sts/src/test/java/software/amazon/awssdk/services/sts/internal/StsRetryStrategyTest.java

@@ -0,0 +1,111 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.services.sts.internal;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.mock;
+
+import java.time.Duration;
+import java.util.stream.Stream;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+import software.amazon.awssdk.awscore.exception.AwsErrorDetails;
+import software.amazon.awssdk.core.SdkSystemSetting;
+import software.amazon.awssdk.core.client.config.SdkClientConfiguration;
+import software.amazon.awssdk.core.client.config.SdkClientOption;
+import software.amazon.awssdk.retries.AdaptiveRetryStrategy;
+import software.amazon.awssdk.retries.LegacyRetryStrategy;
+import software.amazon.awssdk.retries.StandardRetryStrategy;
+import software.amazon.awssdk.retries.api.AcquireInitialTokenRequest;
+import software.amazon.awssdk.retries.api.RefreshRetryTokenRequest;
+import software.amazon.awssdk.retries.api.RetryStrategy;
+import software.amazon.awssdk.retries.api.RetryToken;
+import software.amazon.awssdk.services.sts.model.IdpCommunicationErrorException;
+
+public class StsRetryStrategyTest {
+    @Test
+    void resolveRetryStrategy_preexistingStrategy_returnsPreexisting() {
+        RetryStrategy strategy = mock(RetryStrategy.class);
+        SdkClientConfiguration config = SdkClientConfiguration.builder()
+                                                              .option(SdkClientOption.RETRY_STRATEGY, strategy)
+                                                              .build();
+
+        assertThat(StsRetryStrategy.resolveRetryStrategy(config)).isSameAs(strategy);
+    }
+
+    @Test
+    void resolveRetryStrategy_defaultNewRetries2026False_returnsNull() {
+        SdkClientConfiguration config = SdkClientConfiguration.builder()
+                                                              .option(SdkClientOption.DEFAULT_NEW_RETRIES_2026, false)
+                                                              .build();
+
+        assertThat(StsRetryStrategy.resolveRetryStrategy(config)).isNull();
+    }
+
+    @Test
+    void resolveRetryStrategy_newRetries2026False_returnsNull() {
+        System.setProperty(SdkSystemSetting.AWS_NEW_RETRIES_2026.property(), "false");
+        try {
+            SdkClientConfiguration config = SdkClientConfiguration.builder().build();
+            assertThat(StsRetryStrategy.resolveRetryStrategy(config)).isNull();
+        } finally {
+            System.clearProperty(SdkSystemSetting.AWS_NEW_RETRIES_2026.property());
+        }
+    }
+
+    @ParameterizedTest
+    @MethodSource("retryModeResolutionCases")
+    void resolveRetryStrategy_returnsCorrectStrategyBasedOnMode(String mode, Class<?> expected) {
+        System.setProperty(SdkSystemSetting.AWS_RETRY_MODE.property(), mode);
+        System.setProperty(SdkSystemSetting.AWS_NEW_RETRIES_2026.property(), "true");
+        try {
+            SdkClientConfiguration config = SdkClientConfiguration.builder().build();
+            RetryStrategy resolved = StsRetryStrategy.resolveRetryStrategy(config);
+            assertThat(resolved).isInstanceOf(expected);
+            assertRetriesOnIdpCommunicationException(resolved);
+        } finally {
+            System.clearProperty(SdkSystemSetting.AWS_RETRY_MODE.property());
+            System.clearProperty(SdkSystemSetting.AWS_NEW_RETRIES_2026.property());
+        }
+    }
+
+    void assertRetriesOnIdpCommunicationException(RetryStrategy strategy) {
+        RetryToken token = strategy.acquireInitialToken(AcquireInitialTokenRequest.create("test")).token();
+
+        AwsErrorDetails errorDetails = AwsErrorDetails.builder()
+                                                      .errorCode("IDPCommunicationError")
+                                                      .build();
+
+        IdpCommunicationErrorException failure = IdpCommunicationErrorException.builder()
+                                                                               .awsErrorDetails(errorDetails)
+                                                                               .build();
+        RefreshRetryTokenRequest refresh = RefreshRetryTokenRequest.builder()
+                                                                   .token(token)
+                                                                   .failure(failure)
+                                                                   .build();
+        assertThat(strategy.refreshRetryToken(refresh).delay()).isGreaterThan(Duration.ZERO);
+    }
+
+    private static Stream<Arguments> retryModeResolutionCases() {
+        return Stream.of(
+            Arguments.of("standard", StandardRetryStrategy.class),
+            Arguments.of("legacy", LegacyRetryStrategy.class),
+            Arguments.of("adaptive", AdaptiveRetryStrategy.class)
+        );
+    }
+}