Implement updated standard retry behavior

codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsDynamoDbRetryIntegration.java

@@ -0,0 +1,127 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package software.amazon.smithy.python.aws.codegen;
+
+import java.util.List;
+import java.util.Set;
+import software.amazon.smithy.aws.traits.ServiceTrait;
+import software.amazon.smithy.codegen.core.Symbol;
+import software.amazon.smithy.codegen.core.SymbolReference;
+import software.amazon.smithy.python.codegen.GenerationContext;
+import software.amazon.smithy.python.codegen.SmithyPythonDependency;
+import software.amazon.smithy.python.codegen.integrations.PythonIntegration;
+import software.amazon.smithy.python.codegen.integrations.RuntimeClientPlugin;
+
+/**
+ * Generates DynamoDB's retry defaults as a service-scoped client plugin on
+ * DynamoDB / DynamoDB Streams clients.
+ */
+public final class AwsDynamoDbRetryIntegration implements PythonIntegration {
+
+    private static final Set<String> DYNAMODB_SDK_IDS = Set.of("DynamoDB", "DynamoDB Streams");
+
+    public static final String DYNAMODB_RETRY_MODULE = """
+            _DYNAMODB_DEFAULT_MAX_ATTEMPTS = 4
+            _DYNAMODB_DEFAULT_BACKOFF_SCALE = 0.025
+
+
+            class _RetryConfig(Protocol):
+                retry_strategy: $1T | $2T | None
+
+
+            def dynamodb_retry_plugin(config: _RetryConfig) -> None:
+                \"\"\"Apply DynamoDB's standard-mode retry defaults for any option left unset.\"\"\"
+                retry_strategy = config.retry_strategy
+                if retry_strategy is not None and not isinstance(
+                    retry_strategy, $2T
+                ):
+                    return
+
+                if isinstance(retry_strategy, $2T):
+                    # Explicit options take precedence over separately resolved fields.
+                    retry_mode = retry_strategy.retry_mode
+                    max_attempts = retry_strategy.max_attempts
+                else:
+                    # Read independently resolved AsyncConfig fields when available. A legacy
+                    # Config has no scalar retry fields, so None represents an unset value.
+                    retry_mode = getattr(config, "retry_mode", None) or "standard"
+                    max_attempts = getattr(config, "max_attempts", None)
+                    source_of = getattr(config, "source_of", None)
+                    if (
+                        source_of is not None
+                        and source_of("max_attempts") == $4T.DEFAULT
+                    ):
+                        max_attempts = None
+
+                if retry_mode != "standard":
+                    return
+
+                config.retry_strategy = $3T(
+                    max_attempts=(
+                        max_attempts
+                        if max_attempts is not None
+                        else _DYNAMODB_DEFAULT_MAX_ATTEMPTS
+                    ),
+                    default_backoff_scale=_DYNAMODB_DEFAULT_BACKOFF_SCALE,
+                )
+            """;
+
+    @Override
+    public List<RuntimeClientPlugin> getClientPlugins(GenerationContext context) {
+        final String pluginFile = "retry";
+        final String moduleName = context.settings().moduleName();
+
+        final SymbolReference dynamodbRetryPlugin = SymbolReference.builder()
+                .symbol(Symbol.builder()
+                        .namespace(String.format("%s.%s", moduleName, pluginFile), ".")
+                        .definitionFile(String.format("./src/%s/%s.py", moduleName, pluginFile))
+                        .name("dynamodb_retry_plugin")
+                        .build())
+                .build();
+        final Symbol retryStrategy = Symbol.builder()
+                .namespace("smithy_core.aio.interfaces.retries", ".")
+                .name("RetryStrategy")
+                .build();
+        final Symbol retryStrategyOptions = Symbol.builder()
+                .namespace("smithy_core.retries", ".")
+                .name("RetryStrategyOptions")
+                .build();
+        final Symbol standardRetryStrategy = Symbol.builder()
+                .namespace("smithy_core.aio.retries", ".")
+                .name("StandardRetryStrategy")
+                .build();
+        final Symbol configSource = Symbol.builder()
+                .namespace("smithy_aws_core.config", ".")
+                .name("ConfigSource")
+                .build();
+
+        return List.of(
+                RuntimeClientPlugin.builder()
+                        .servicePredicate((model, service) -> service.getTrait(ServiceTrait.class)
+                                .map(trait -> DYNAMODB_SDK_IDS.contains(trait.getSdkId()))
+                                .orElse(false))
+                        .pythonPlugin(dynamodbRetryPlugin)
+                        .writeAdditionalFiles((c) -> {
+                            String filename = "src/%s/%s.py".formatted(moduleName, pluginFile);
+                            c.writerDelegator()
+                                    .useFileWriter(
+                                            filename,
+                                            moduleName + ".",
+                                            writer -> {
+                                                writer.addDependency(SmithyPythonDependency.SMITHY_CORE);
+                                                writer.addDependency(AwsPythonDependency.SMITHY_AWS_CORE);
+                                                writer.addStdlibImport("typing", "Protocol");
+                                                writer.write(
+                                                        DYNAMODB_RETRY_MODULE,
+                                                        retryStrategy,
+                                                        retryStrategyOptions,
+                                                        standardRetryStrategy,
+                                                        configSource);
+                                            });
+                            return List.of(filename);
+                        })
+                        .build());
+    }
+}

codegen/aws/core/src/main/java/software/amazon/smithy/python/aws/codegen/AwsLongPollingIntegration.java

@@ -0,0 +1,45 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package software.amazon.smithy.python.aws.codegen;
+
+import java.util.Map;
+import java.util.Set;
+import software.amazon.smithy.aws.traits.ServiceTrait;
+import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.shapes.OperationShape;
+import software.amazon.smithy.model.shapes.ServiceShape;
+import software.amazon.smithy.model.traits.LongPollTrait;
+import software.amazon.smithy.python.codegen.integrations.PythonIntegration;
+
+/**
+ * Marks long-polling operations so the generic client generator applies
+ * long-polling retry behavior to them.
+ *
+ * <p>An operation is long-polling if it carries the {@code smithy.api#longPoll}
+ * trait. Service models do not yet apply the trait, so the known operations are
+ * hard-coded as a fallback. Once the trait ships in the models, the fallback
+ * can be removed.
+ */
+public final class AwsLongPollingIntegration implements PythonIntegration {
+
+    private static final Map<String, Set<String>> LONG_POLLING_OPERATIONS = Map.of(
+            "SQS",
+            Set.of("ReceiveMessage"),
+            "SFN",
+            Set.of("GetActivityTask"),
+            "SWF",
+            Set.of("PollForActivityTask", "PollForDecisionTask"));
+
+    @Override
+    public boolean isLongPollingOperation(Model model, ServiceShape service, OperationShape operation) {
+        if (operation.hasTrait(LongPollTrait.class)) {
+            return true;
+        }
+        return service.getTrait(ServiceTrait.class)
+                .map(trait -> LONG_POLLING_OPERATIONS.get(trait.getSdkId()))
+                .map(operations -> operations.contains(operation.getId().getName()))
+                .orElse(false);
+    }
+}

codegen/aws/core/src/main/resources/META-INF/services/software.amazon.smithy.python.codegen.integrations.PythonIntegration

@@ -9,3 +9,5 @@ software.amazon.smithy.python.aws.codegen.AwsProtocolsIntegration
 software.amazon.smithy.python.aws.codegen.AwsServiceIdIntegration
 software.amazon.smithy.python.aws.codegen.AwsUserAgentIntegration
 software.amazon.smithy.python.aws.codegen.AwsStandardRegionalEndpointsIntegration
+software.amazon.smithy.python.aws.codegen.AwsDynamoDbRetryIntegration
+software.amazon.smithy.python.aws.codegen.AwsLongPollingIntegration

codegen/core/src/main/java/software/amazon/smithy/python/codegen/ClientGenerator.java

@@ -137,6 +137,15 @@ private void writeDefaultPlugins(PythonWriter writer, Collection<SymbolReference
         }
     }
 
+    private boolean isLongPollingOperation(OperationShape operation) {
+        for (PythonIntegration integration : context.integrations()) {
+            if (integration.isLongPollingOperation(model, service, operation)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     private void writeConstructorDocs(PythonWriter writer, String clientName) {
         writer.writeMultiLineDocs(() -> {
             writer.write("""
@@ -234,6 +243,7 @@ private void writeSharedOperationInit(
         }
 
         writer.putContext("operation", symbolProvider.toSymbol(operation));
+        writer.putContext("isLongPolling", isLongPollingOperation(operation));
         writer.addStdlibImport("copy", "deepcopy");
 
         writer.write("""
@@ -256,11 +266,14 @@ private void writeSharedOperationInit(
                     protocol=config.protocol,
                     transport=config.transport
                 )
-                call = $4T(
+                ${?isLongPolling}operation_context = $4T({"config": config})
+                operation_context[$5T] = True
+                ${/isLongPolling}call = $6T(
                     input=input,
                     operation=${operation:T},
-                    context=$5T({"config": config}),
-                    interceptor=$6T(config.interceptors),
+                    ${?isLongPolling}context=operation_context,
+                    ${/isLongPolling}${^isLongPolling}context=$4T({"config": config}),
+                    ${/isLongPolling}interceptor=$7T(config.interceptors),
                     auth_scheme_resolver=config.auth_scheme_resolver,
                     supported_auth_schemes=config.auth_schemes,
                     endpoint_resolver=config.endpoint_resolver,
@@ -270,8 +283,9 @@ private void writeSharedOperationInit(
                 writer.consumer(w -> writeDefaultPlugins(w, defaultPlugins)),
                 RuntimeTypes.EXPECTATION_NOT_MET_ERROR,
                 RuntimeTypes.REQUEST_PIPELINE,
-                RuntimeTypes.CLIENT_CALL,
                 RuntimeTypes.TYPED_PROPERTIES,
+                RuntimeTypes.LONG_POLLING,
+                RuntimeTypes.CLIENT_CALL,
                 RuntimeTypes.INTERCEPTOR_CHAIN);
 
     }

codegen/core/src/main/java/software/amazon/smithy/python/codegen/RuntimeTypes.java

@@ -68,6 +68,8 @@ public final class RuntimeTypes {
             createSymbol("aio.retries", "RetryStrategyResolver", SmithyPythonDependency.SMITHY_CORE);
     public static final Symbol RETRY_STRATEGY_OPTIONS =
             createSymbol("retries", "RetryStrategyOptions", SmithyPythonDependency.SMITHY_CORE);
+    public static final Symbol LONG_POLLING =
+            createSymbol("retries", "LONG_POLLING", SmithyPythonDependency.SMITHY_CORE);
     public static final Symbol SIMPLE_RETRY_STRATEGY =
             createSymbol("aio.retries", "SimpleRetryStrategy", SmithyPythonDependency.SMITHY_CORE);
 

codegen/core/src/main/java/software/amazon/smithy/python/codegen/integrations/PythonIntegration.java

@@ -8,6 +8,8 @@
 import java.util.List;
 import software.amazon.smithy.codegen.core.SmithyIntegration;
 import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.shapes.OperationShape;
+import software.amazon.smithy.model.shapes.ServiceShape;
 import software.amazon.smithy.python.codegen.GenerationContext;
 import software.amazon.smithy.python.codegen.PythonSettings;
 import software.amazon.smithy.python.codegen.generators.ProtocolGenerator;
@@ -44,6 +46,21 @@ default Model preprocessModel(Model model, PythonSettings settings) {
         return model;
     }
 
+    /**
+     * Determines whether the given operation is a long-polling operation, which
+     * must back off before returning even when the retry quota is exhausted. AWS
+     * integrations use this hook to identify these operations while the
+     * {@code smithy.api#longPoll} trait is not yet shipped in service models.
+     *
+     * @param model     Model the operation belongs to.
+     * @param service   Service the operation belongs to.
+     * @param operation Operation to test.
+     * @return Returns true if the operation is a long-polling operation.
+     */
+    default boolean isLongPollingOperation(Model model, ServiceShape service, OperationShape operation) {
+        return false;
+    }
+
     /**
      * Writes out all extra files required by runtime plugins.
      */

designs/retries.md

@@ -32,25 +32,33 @@ class RetryStrategy(Protocol):
     """Upper limit on total attempt count (initial attempt plus retries)."""
 
     async def acquire_initial_retry_token(
-        self, *, token_scope: str | None = None
+        self, *, token_scope: str | None = None, context: TypedProperties | None = None
     ) -> RetryToken:
         """Called before any retries (for the first attempt at the operation).
 
         :param token_scope: An arbitrary string accepted by the retry strategy to
             separate tokens into scopes.
+        :param context: The operation context, read for per-operation signals such as
+            :py:data:`smithy_core.retries.LONG_POLLING`.
         :returns: A retry token, to be used for determining the retry delay, refreshing
             the token after a failure, and recording success after success.
         :raises RetryError: If the retry strategy has no available tokens.
         """
         ...
 
     async def refresh_retry_token_for_retry(
-        self, *, token_to_renew: RetryToken, error: Exception
+        self,
+        *,
+        token_to_renew: RetryToken,
+        error: Exception,
+        context: TypedProperties | None = None,
     ) -> RetryToken:
         """Replace an existing retry token from a failed attempt with a new token.
 
         :param token_to_renew: The token used for the previous failed attempt.
         :param error: The error that triggered the need for a retry.
+        :param context: The operation context, read for per-operation signals such as
+            :py:data:`smithy_core.retries.LONG_POLLING`.
         :raises RetryError: If no further retry attempts are allowed.
         """
         ...
@@ -110,8 +118,9 @@ class HasFault(Protocol):
 
 `RetryStrategy` implementations MUST raise a `RetryError` if they receive an
 exception where `is_retry_safe` is `False` and SHOULD raise a `RetryError` if it
-is `None`. `RetryStrategy` implementations SHOULD use a delay that is at least
-as long as `retry_after` but MAY choose to wait longer.
+is `None`. `RetryStrategy` implementations SHOULD take `retry_after` into account
+when computing the delay, but MAY adjust it (for example, by clamping it to an
+upper bound).
 
 ### Backoff Strategy
 

packages/smithy-aws-core/.changes/next-release/smithy-aws-core-feature-881860b5bda049d1b000983dd2a3bd0e.json

@@ -0,0 +1,4 @@
+{
+  "type": "feature",
+  "description": "Added support for the `x-amz-retry-after` response header."
+}

packages/smithy-aws-core/src/smithy_aws_core/_private/query/errors.py

@@ -99,6 +99,7 @@ def create_aws_query_error(
     wrapper_elements: tuple[str, ...],
     status: int,
     context: TypedProperties,
+    retry_after: float | None = None,
 ) -> CallError:
     """Create a modeled or generic CallError from an awsQuery error response."""
     code = _parse_aws_query_error_code(body, wrapper_elements)
@@ -121,7 +122,10 @@ def create_aws_query_error(
             deserializer = XMLCodec().create_deserializer(
                 body, wrapper_elements=wrapper_elements
             )
-            return error_shape.deserialize(deserializer)
+            modeled_error = error_shape.deserialize(deserializer)
+            if retry_after is not None:
+                modeled_error.retry_after = retry_after
+            return modeled_error
 
     message = f"Unknown error for operation {operation.schema.id} - status: {status}"
     if code is not None:
@@ -137,4 +141,5 @@ def create_aws_query_error(
         is_throttling_error=is_throttle,
         is_timeout_error=is_timeout,
         is_retry_safe=is_throttle or is_timeout or None,
+        retry_after=retry_after,
     )

packages/smithy-aws-core/src/smithy_aws_core/aio/protocols.py

@@ -27,15 +27,18 @@
 from smithy_http import tuples_to_fields
 from smithy_http.aio import HTTPRequest as _HTTPRequest
 from smithy_http.aio.interfaces import HTTPErrorIdentifier, HTTPRequest, HTTPResponse
-from smithy_http.aio.protocols import HttpBindingClientProtocol, HttpClientProtocol
+from smithy_http.aio.protocols import (
+    HttpBindingClientProtocol,
+    HttpClientProtocol,
+)
 from smithy_http.deserializers import HTTPResponseDeserializer
 
 from .._private.query.errors import (
     create_aws_query_error,
 )
 from .._private.query.serializers import QueryShapeSerializer
 from ..traits import AwsQueryTrait, RestJson1Trait
-from ..utils import parse_document_discriminator, parse_error_code
+from ..utils import parse_document_discriminator, parse_error_code, parse_retry_after
 
 try:
     from smithy_json import JSONCodec, JSONDocument
@@ -166,6 +169,9 @@ def content_type(self) -> str:
     def error_identifier(self) -> HTTPErrorIdentifier:
         return self._error_identifier
 
+    def _retry_after(self, response: HTTPResponse) -> float | None:
+        return parse_retry_after(response)
+
     def _resolve_error_id(
         self,
         *,
@@ -364,6 +370,7 @@ async def _create_error(
             wrapper_elements=self._error_wrapper_elements(),
             status=response.status,
             context=context,
+            retry_after=parse_retry_after(response),
         )
 
     def _action_name(