Repository: smithy-lang/smithy-rs
Pull request #4091: Incorporate account ID into service-specific endpoint params
Author: login: ysaito1001
Labels: name: needs-sdk-review
Created: 2025-04-12T02:37:23Z | Merged: 2025-04-17T15:14:40Z

## Motivation and Context
This is part 4 in the series of supporting account ID-based routing (continuation of https://github.com/smithy-lang/smithy-rs/pull/4047, https://github.com/smithy-lang/smithy-rs/pull/4057, and https://github.com/smithy-lang/smithy-rs/pull/4087).
 
## Description
In this PR, we update the orchestrator to pass an Identity, resolved using the resolve_identity function, to service-specific endpoint parameters in the code-generated client SDKs. To support this, we've made changes to the following APIs:
- `EndpointResolverParams::get_mut` provides a mutable reference to the underlying data
- `EndpointResolverParams` can set and get arbitrary properties keyed on a type ID
- `ConfigBag::get_mut` is refactored to extract the portion of the method that doesn't require `Clone` into a separate, reusable method
- The `ResolveEndpoint` trait introduces `finalize_params` method (defaulted to no-op) to allow the caller to incorporate a property stored in `EndpointResolverParams` into its underlying service-specific endpoint params.
 
The following describes the interaction and flow between the above APIs:
1. The orchestrator resolves an `Identity` with the `resolve_identity` function.
2. The orchestrator passes the `Identity` to `orchestrate_endpoint`.
3. The orchestrator stores the resolved Identity in the `EndpointResolverParams` by setting it as a property using `set_property`.
4. The orchestrator calls `SharedEndpointResolver::finalize_params` so that the concrete endpoint resolver struct (i.e. `DowncastParams`) incorporates the `Identity` into service-specific endpoint params.
   - If a service uses the `AccountID` built-in endpoint parameter in endpoint rules, the generated client SDK overrides the `DowncastParams::finalize_params` method to retrieve the account ID from the property of `Identity` and to set it to the service-specific endpoint parameters.
5. The orchestrator calls `SharedEndpointResolver::resolve_endpoint`, with the account ID included in endpoint parameters.
 
## Questions
I made `ResolveEndpoint::finalize_params` be async to allow for future extensibility. To be honest, I _currently_ cannot think of use cases where one wants to do something in an async manner, but I was concerned that if I made the method sync, that would prevent customers from doing so completely, unless we deprecate the method and introduce an async counterpart.
 
## Testing
If we generate the dynamodb SDK using the changes in this PR and run the following experimental test, 

<details>
<summary>Experimental test</summary>

```
#[tokio::test]
async fn experiment() {
    let (http_client, request) = aws_smithy_http_client::test_util::capture_request(None);
    let client = aws_sdk_dynamodb::Client::from_conf(
        aws_sdk_dynamodb::config::Config::builder()
            .region(aws_config::Region::new("us-east-1"))
            .credentials_provider(
                aws_credential_types::Credentials::builder()
                    .account_id("333333333333")
                    .access_key_id("ANOTREAL")
                    .secret_access_key("notrealrnrELgWzOk3IfjzDKtFBhDby")
                    .provider_name("test")
                    .build(),
            )
            .http_client(http_client)
            .build(),
    );
    let mut attr_v = std::collections::HashMap::new();
    attr_v.insert(
        ":s".to_string(),
        aws_sdk_dynamodb::types::AttributeValue::S("value".into()),
    );

    let mut kv = std::collections::HashMap::new();
    kv.insert(
        ":pk".to_string(),
        aws_sdk_dynamodb::types::AttributeValue::M(attr_v),
    );
    let _ = client
        .batch_get_item()
        .request_items(
            "arn:aws:dynamodb:us-east-1:333333333333:table/table_name",
            aws_sdk_dynamodb::types::KeysAndAttributes::builder()
                .keys(kv)
                .build()
                .unwrap(),
        )
        .send()
        .await;
    let _ = dbg!(request.expect_request());
}
``` 
</details>

`dbg!` above shows that the resulting request contains the endpoint URL with the account ID:
```
req = Request {
    ...
    uri: Uri {
        as_string: "https://333333333333.ddb.us-east-1.amazonaws.com/",
        parsed: H0(
            https://333333333333.ddb.us-east-1.amazonaws.com/,
        ),
    },
    ...
```
The DynamoDB endpoint tests include cases that validate behavior based on the account ID, like the one shown [here](https://github.com/awslabs/aws-sdk-rust/blob/d18b1585f06439909aede40fe206121379de0f20/sdk/dynamodb/tests/endpoint_tests.rs#L3027-L3058). Currently, we make these tests pass by explicitly setting the account ID on the client config builder. In the next PR, we plan to remove the `account_id` and `set_account_id` setters from the client config builder.

---
_By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice._


Diff:
diff --git a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/EndpointBuiltInsDecorator.kt b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/EndpointBuiltInsDecorator.kt
index c23e2a6b86c..a04c5effe6e 100644
--- a/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/EndpointBuiltInsDecorator.kt
+++ b/aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/EndpointBuiltInsDecorator.kt
@@ -20,6 +20,7 @@ import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext
 import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegenDecorator
 import software.amazon.smithy.rust.codegen.client.smithy.endpoint.EndpointCustomization
 import software.amazon.smithy.rust.codegen.client.smithy.endpoint.EndpointRulesetIndex
+import software.amazon.smithy.rust.codegen.client.smithy.endpoint.EndpointTypesGenerator
 import software.amazon.smithy.rust.codegen.client.smithy.endpoint.rustName
 import software.amazon.smithy.rust.codegen.client.smithy.endpoint.symbol
 import software.amazon.smithy.rust.codegen.client.smithy.generators.config.ConfigCustomization
@@ -36,6 +37,7 @@ import software.amazon.smithy.rust.codegen.core.rustlang.stripOuter
 import software.amazon.smithy.rust.codegen.core.rustlang.writable
 import software.amazon.smithy.rust.codegen.core.smithy.RuntimeConfig
 import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
+import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope
 import software.amazon.smithy.rust.codegen.core.smithy.customize.AdHocCustomization
 import software.amazon.smithy.rust.codegen.core.smithy.mapRustType
 import software.amazon.smithy.rust.codegen.core.util.PANIC
@@ -126,90 +128,158 @@ fun Model.sdkConfigSetter(
 }
 
 /**
- * Create a client codegen decorator that creates bindings for a builtIn parameter. Optionally, you can provide
- * [clientParam.Builder] which allows control over the config parameter that will be generated. You can also opt
- * to exclude including the extra sections that set the builtIn value on the SdkConfig. This is useful for builtIns
- * that are only minimally supported, like accountId and accountIdEndpointMode.
+ * A custom decorator that creates bindings for the `accountID` built-in parameter.
+ *
+ * The `accountID` parameter is special because:
+ * - It is not exposed in the `SdkConfig` setter.
+ * - It does not require customizations like `loadBuiltInFromServiceConfig`,
+ *  as it is not available when the `read_before_execution` method of the endpoint parameters interceptor is executed.
  */
-fun decoratorForBuiltIn(
-    builtIn: Parameter,
-    clientParamBuilder: ConfigParam.Builder? = null,
-    includeSdkConfigSetter: Boolean = true,
-): ClientCodegenDecorator {
-    val nameOverride = clientParamBuilder?.name
-    val name = nameOverride ?: builtIn.name.rustName()
-    return object : ClientCodegenDecorator {
-        override val name: String = "Auto${builtIn.builtIn.get()}"
-        override val order: Byte = 0
-
-        private fun rulesetContainsBuiltIn(codegenContext: ClientCodegenContext) =
-            codegenContext.getBuiltIn(builtIn) != null
-
-        override fun extraSections(codegenContext: ClientCodegenContext): List<AdHocCustomization> {
-            if (includeSdkConfigSetter) {
-                return listOfNotNull(
-                    codegenContext.model.sdkConfigSetter(
-                        codegenContext.serviceShape.id,
-                        builtIn,
-                        clientParamBuilder?.name,
-                    ),
-                )
-            } else {
-                return listOf()
-            }
-        }
+class DecoratorForAccountId(private val builtIn: Parameter) : DecoratorForBuiltIn(builtIn) {
+    // TODO(AccountIdBasedRouting): Override `configCustomizations` to avoid rendering `account_id` and `set_account_id`
+    //  on client config builder, and deprecate those that have already been exposed.
 
-        override fun configCustomizations(
-            codegenContext: ClientCodegenContext,
-            baseCustomizations: List<ConfigCustomization>,
-        ): List<ConfigCustomization> {
-            return baseCustomizations.extendIf(rulesetContainsBuiltIn(codegenContext)) {
-                standardConfigParam(
-                    clientParamBuilder?.toConfigParam(builtIn, codegenContext.runtimeConfig) ?: ConfigParam.Builder()
-                        .toConfigParam(builtIn, codegenContext.runtimeConfig),
-                )
-            }
-        }
+    override fun extraSections(codegenContext: ClientCodegenContext): List<AdHocCustomization> {
+        return emptyList()
+    }
 
-        override fun endpointCustomizations(codegenContext: ClientCodegenContext): List<EndpointCustomization> =
+    override fun endpointCustomizations(codegenContext: ClientCodegenContext): List<EndpointCustomization> =
+        // TODO(AccountIdBasedRouting): Remove `builtIn == AwsBuiltIns.ACCOUNT_ID` once accountID becomes the only
+        //  usage of this class
+        if (rulesetContainsBuiltIn(codegenContext) && builtIn == AwsBuiltIns.ACCOUNT_ID) {
             listOf(
                 object : EndpointCustomization {
-                    override fun loadBuiltInFromServiceConfig(
-                        parameter: Parameter,
-                        configRef: String,
-                    ): Writable? =
-                        when (parameter.builtIn) {
-                            builtIn.builtIn ->
-                                writable {
-                                    val newtype = configParamNewtype(parameter, name, codegenContext.runtimeConfig)
-                                    val symbol = parameter.symbol().mapRustType { t -> t.stripOuter<RustType.Option>() }
-                                    rustTemplate(
-                                        """$configRef.#{load_from_service_config_layer}""",
-                                        "load_from_service_config_layer" to loadFromConfigBag(symbol.name, newtype),
-                                    )
-                                }
-
-                            else -> null
-                        }
-
-                    override fun setBuiltInOnServiceConfig(
-                        name: String,
-                        value: Node,
-                        configBuilderRef: String,
+                    override fun serviceSpecificEndpointParamsFinalizer(
+                        codegenContext: ClientCodegenContext,
+                        params: String,
                     ): Writable? {
-                        if (name != builtIn.builtIn.get()) {
-                            return null
-                        }
+                        val runtimeConfig = codegenContext.runtimeConfig
                         return writable {
                             rustTemplate(
-                                "let $configBuilderRef = $configBuilderRef.${nameOverride ?: builtIn.name.rustName()}(#{value});",
-                                "value" to value.toWritable(),
+                                """
+                                // This is required to satisfy the borrow checker. By obtaining an `Option<Identity>`,
+                                // `params` is no longer mutably borrowed in the match expression below.
+                                // Furthermore, by using `std::mem::replace` with an empty `Identity`, we avoid
+                                // leaving the sensitive `Identity` inside `params` within `EndpointResolverParams`.
+                                let identity = $params
+                                    .get_property_mut::<#{Identity}>()
+                                    .map(|id| {
+                                        std::mem::replace(
+                                            id,
+                                            #{Identity}::new((), #{None}),
+                                        )
+                                    });
+                                match (
+                                    $params.get_mut::<#{Params}>(),
+                                    identity
+                                        .as_ref()
+                                        .and_then(|id| id.property::<#{AccountId}>()),
+                                ) {
+                                    (#{Some}(concrete_params), #{Some}(account_id)) => {
+                                        concrete_params.account_id = #{Some}(account_id.as_str().to_string());
+                                    }
+                                    (#{Some}(_), #{None}) => {
+                                        // No account ID; nothing to do.
+                                    }
+                                    (#{None}, _) => {
+                                        return #{Err}("service-specific endpoint params was not present".into());
+                                    }
+                                }
+                                """,
+                                *preludeScope,
+                                "AccountId" to
+                                    AwsRuntimeType.awsCredentialTypes(runtimeConfig)
+                                        .resolve("attributes::AccountId"),
+                                "Identity" to
+                                    RuntimeType.smithyRuntimeApiClient(runtimeConfig)
+                                        .resolve("client::identity::Identity"),
+                                "Params" to EndpointTypesGenerator.fromContext(codegenContext).paramsStruct(),
                             )
                         }
                     }
                 },
             )
+        } else {
+            emptyList()
+        }
+}
+
+/**
+ * A common client codegen decorator that creates bindings for a builtIn parameter. Optionally,
+ * you can provide [clientParam.Builder] which allows control over the config parameter that will be generated.
+ */
+open class DecoratorForBuiltIn(
+    private val builtIn: Parameter,
+    private val clientParamBuilder: ConfigParam.Builder? = null,
+) : ClientCodegenDecorator {
+    override val name: String = "Auto${builtIn.builtIn.get()}"
+    override val order: Byte = 0
+
+    val builtinParamName = clientParamBuilder?.name ?: builtIn.name.rustName()
+
+    protected fun rulesetContainsBuiltIn(codegenContext: ClientCodegenContext) =
+        codegenContext.getBuiltIn(builtIn) != null
+
+    override fun extraSections(codegenContext: ClientCodegenContext) =
+        listOfNotNull(
+            codegenContext.model.sdkConfigSetter(
+                codegenContext.serviceShape.id,
+                builtIn,
+                clientParamBuilder?.name,
+            ),
+        )
+
+    override fun configCustomizations(
+        codegenContext: ClientCodegenContext,
+        baseCustomizations: List<ConfigCustomization>,
+    ): List<ConfigCustomization> {
+        return baseCustomizations.extendIf(rulesetContainsBuiltIn(codegenContext)) {
+            standardConfigParam(
+                clientParamBuilder?.toConfigParam(builtIn, codegenContext.runtimeConfig) ?: ConfigParam.Builder()
+                    .toConfigParam(builtIn, codegenContext.runtimeConfig),
+            )
+        }
     }
+
+    override fun endpointCustomizations(codegenContext: ClientCodegenContext): List<EndpointCustomization> =
+        listOf(
+            object : EndpointCustomization {
+                override fun loadBuiltInFromServiceConfig(
+                    parameter: Parameter,
+                    configRef: String,
+                ): Writable? =
+                    when (parameter.builtIn) {
+                        builtIn.builtIn ->
+                            writable {
+                                val newtype =
+                                    configParamNewtype(parameter, builtinParamName, codegenContext.runtimeConfig)
+                                val symbol = parameter.symbol().mapRustType { t -> t.stripOuter<RustType.Option>() }
+                                rustTemplate(
+                                    """$configRef.#{load_from_service_config_layer}""",
+                                    "load_from_service_config_layer" to loadFromConfigBag(symbol.name, newtype),
+                                )
+                            }
+
+                        else -> null
+                    }
+
+                override fun setBuiltInOnServiceConfig(
+                    name: String,
+                    value: Node,
+                    configBuilderRef: String,
+                ): Writable? {
+                    if (name != builtIn.builtIn.get()) {
+                        return null
+                    }
+                    return writable {
+                        rustTemplate(
+                            "let $configBuilderRef = $configBuilderRef.$builtinParamName(#{value});",
+                            "value" to value.toWritable(),
+                        )
+                    }
+                }
+            },
+        )
 }
 
 private val endpointUrlDocs =
@@ -238,15 +308,17 @@ fun Node.toWritable(): Writable {
 
 val PromotedBuiltInsDecorators =
     listOf(
-        decoratorForBuiltIn(AwsBuiltIns.FIPS),
-        decoratorForBuiltIn(AwsBuiltIns.DUALSTACK),
-        decoratorForBuiltIn(
+        DecoratorForBuiltIn(AwsBuiltIns.FIPS),
+        DecoratorForBuiltIn(AwsBuiltIns.DUALSTACK),
+        DecoratorForBuiltIn(
             BuiltIns.SDK_ENDPOINT,
             ConfigParam.Builder()
                 .name("endpoint_url")
                 .type(RuntimeType.String.toSymbol())
                 .setterDocs(endpointUrlDocs),
         ),
-        decoratorForBuiltIn(AwsBuiltIns.ACCOUNT_ID_ENDPOINT_MODE, null, false),
-        decoratorForBuiltIn(AwsBuiltIns.ACCOUNT_ID, null, false),
+        // TODO(AccountIdBasedRouting): Switch to DecoratorForBuiltIn once account_id_endpoint_mode is exposed
+        //  in SdkConfig
+        DecoratorForAccountId(AwsBuiltIns.ACCOUNT_ID_ENDPOINT_MODE),
+        DecoratorForAccountId(AwsBuiltIns.ACCOUNT_ID),
     ).toTypedArray()
diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/EndpointsDecorator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/EndpointsDecorator.kt
index 222eaf41c3d..718543fca1b 100644
--- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/EndpointsDecorator.kt
+++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/EndpointsDecorator.kt
@@ -79,6 +79,34 @@ interface EndpointCustomization {
      * Provide a list of additional endpoints standard library functions that rules can use
      */
     fun customRuntimeFunctions(codegenContext: ClientCodegenContext): List<CustomRuntimeFunction> = listOf()
+
+    /**
+     * Allows overriding the `finalize_params` method  in the `ResolveEndpoint` trait within
+     * service-specific EndpointResolver, which is rendered in `serviceSpecificEndpointResolver`.
+     *
+     * `ResolveEndpoint::finalize_params` provides a default no-op implementation,
+     * and this customization enables the implementor to provide an alternative implementation.
+     *
+     * Example:
+     * ```kotlin
+     * override fun serviceSpecificEndpointParamsFinalizer(codegenContext: ClientCodegenContext, params: String): Writable? {
+     *     return writable {
+     *         rustTemplate("""
+     *             let identity = $params
+     *                 .get_property_mut::<#{Identity}>();
+     *             // do more things on `params`
+     *         """,
+     *         "Identity" to RuntimeType.smithyRuntimeApiClient(codegenContext.runtimeConfig)
+     *             .resolve("client::identity::Identity"),
+     *         )
+     *     }
+     * }
+     * ```
+     */
+    fun serviceSpecificEndpointParamsFinalizer(
+        codegenContext: ClientCodegenContext,
+        params: String,
+    ): Writable? = null
 }
 
 /**
diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/generators/EndpointResolverGenerator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/generators/EndpointResolverGenerator.kt
index 5ac72576a8e..5a16482bb95 100644
--- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/generators/EndpointResolverGenerator.kt
+++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/generators/EndpointResolverGenerator.kt
@@ -33,6 +33,7 @@ import software.amazon.smithy.rust.codegen.core.rustlang.comment
 import software.amazon.smithy.rust.codegen.core.rustlang.escape
 import software.amazon.smithy.rust.codegen.core.rustlang.join
 import software.amazon.smithy.rust.codegen.core.rustlang.rust
+import software.amazon.smithy.rust.codegen.core.rustlang.rustBlockTemplate
 import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.core.rustlang.writable
 import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
@@ -396,10 +397,22 @@ internal class EndpointResolverGenerator(
 }
 
 fun ClientCodegenContext.serviceSpecificEndpointResolver(): RuntimeType {
+    val ctx = this
     val generator = EndpointTypesGenerator.fromContext(this)
+    val endpointCustomizations = rootDecorator.endpointCustomizations(this)
     return RuntimeType.forInlineFun("ResolveEndpoint", ClientRustModule.Config.endpoint) {
-        val ctx =
-            arrayOf(*preludeScope, "Params" to generator.paramsStruct(), *Types(runtimeConfig).toArray(), "Debug" to RuntimeType.Debug)
+        val codegenScope =
+            arrayOf(
+                *preludeScope,
+                "BoxError" to RuntimeType.boxError(runtimeConfig),
+                "Params" to generator.paramsStruct(),
+                *Types(runtimeConfig).toArray(),
+                "Debug" to RuntimeType.Debug,
+            )
+        val paramsFinalizers =
+            endpointCustomizations.mapNotNull {
+                it.serviceSpecificEndpointParamsFinalizer(ctx, "params")
+            }
         rustTemplate(
             """
             /// Endpoint resolver trait specific to ${serviceShape.serviceNameOrDefault("this service")}
@@ -431,10 +444,31 @@ fun ClientCodegenContext.serviceSpecificEndpointResolver(): RuntimeType {
                     };
                     ep
                 }
+                #{finalize_endpoint_params:W}
             }
 
             """,
-            *ctx,
+            *codegenScope,
+            "finalize_endpoint_params" to finalizeEndpointParams(codegenScope, paramsFinalizers),
         )
     }
 }
+
+private fun finalizeEndpointParams(
+    codegenScope: Array<Pair<String, RuntimeType>>,
+    paramsFinalizers: List<Writable>,
+) = writable {
+    if (paramsFinalizers.isNotEmpty()) {
+        rustBlockTemplate(
+            """
+            fn finalize_params<'a>(&'a self, params: &'a mut #{EndpointResolverParams}) -> #{Result}<(), #{BoxError}>
+            """,
+            *codegenScope,
+        ) {
+            paramsFinalizers.forEach {
+                it(this)
+            }
+            rustTemplate("#{Ok}(())", *codegenScope)
+        }
+    }
+}
diff --git a/rust-runtime/Cargo.lock b/rust-runtime/Cargo.lock
index b7880b4e866..2b8981b1c43 100644
--- a/rust-runtime/Cargo.lock
+++ b/rust-runtime/Cargo.lock
@@ -462,7 +462,7 @@ dependencies = [
 
 [[package]]
 name = "aws-smithy-http-server"
-version = "0.65.0"
+version = "0.65.1"
 dependencies = [
  "aws-smithy-cbor",
  "aws-smithy-http",
@@ -645,7 +645,7 @@ dependencies = [
 
 [[package]]
 name = "aws-smithy-types"
-version = "1.3.0"
+version = "1.3.1"
 dependencies = [
  "base64 0.13.1",
  "base64-simd",
diff --git a/rust-runtime/aws-smithy-runtime-api/src/client/endpoint.rs b/rust-runtime/aws-smithy-runtime-api/src/client/endpoint.rs
index 64bc2dbd1fe..5d5870aa012 100644
--- a/rust-runtime/aws-smithy-runtime-api/src/client/endpoint.rs
+++ b/rust-runtime/aws-smithy-runtime-api/src/client/endpoint.rs
@@ -13,12 +13,14 @@ use aws_smithy_types::endpoint::Endpoint;
 use aws_smithy_types::type_erasure::TypeErasedBox;
 use error::InvalidEndpointError;
 use http_02x::uri::Authority;
+use std::any::TypeId;
+use std::collections::HashMap;
 use std::fmt;
 use std::str::FromStr;
 use std::sync::Arc;
 
 new_type_future! {
-    #[doc = "Future for [`EndpointResolver::resolve_endpoint`]."]
+    #[doc = "Future for [`ResolveEndpoint::resolve_endpoint`]."]
     pub struct EndpointFuture<'a, Endpoint, BoxError>;
 }
 
@@ -27,18 +29,53 @@ new_type_future! {
 /// The actual endpoint parameters are code generated from the Smithy model, and thus,
 /// are not known to the runtime crates. Hence, this struct is really a new-type around
 /// a [`TypeErasedBox`] that holds the actual concrete parameters in it.
+///
+/// This struct allows the caller to store and retrieve properties of arbitrary types.
+/// These arbitrary properties are intended to be incorporated into the concrete parameters
+/// by [`ResolveEndpoint::finalize_params`].
 #[derive(Debug)]
-pub struct EndpointResolverParams(TypeErasedBox);
+pub struct EndpointResolverParams {
+    inner: TypeErasedBox,
+    property: HashMap<TypeId, TypeErasedBox>,
+}
 
 impl EndpointResolverParams {
     /// Creates a new [`EndpointResolverParams`] from a concrete parameters instance.
     pub fn new<T: fmt::Debug + Send + Sync + 'static>(params: T) -> Self {
-        Self(TypeErasedBox::new(params))
+        Self {
+            inner: TypeErasedBox::new(params),
+            property: HashMap::new(),
+        }
     }
 
     /// Attempts to downcast the underlying concrete parameters to `T` and return it as a reference.
     pub fn get<T: fmt::Debug + Send + Sync + 'static>(&self) -> Option<&T> {
-        self.0.downcast_ref()
+        self.inner.downcast_ref()
+    }
+
+    /// Attempts to downcast the underlying concrete parameters to `T` and return it as a mutable reference.
+    pub fn get_mut<T: fmt::Debug + Send + Sync + 'static>(&mut self) -> Option<&mut T> {
+        self.inner.downcast_mut()
+    }
+
+    /// Sets property of an arbitrary type `T` for the endpoint resolver params.
+    pub fn set_property<T: fmt::Debug + Send + Sync + 'static>(&mut self, t: T) {
+        self.property
+            .insert(TypeId::of::<T>(), TypeErasedBox::new(t));
+    }
+
+    /// Attempts to retrieve a reference to property of a given type `T`.
+    pub fn get_property<T: fmt::Debug + Send + Sync + 'static>(&self) -> Option<&T> {
+        self.property
+            .get(&TypeId::of::<T>())
+            .and_then(|b| b.downcast_ref())
+    }
+
+    /// Attempts to retrieve a mutable reference to property of a given type `T`.
+    pub fn get_property_mut<T: fmt::Debug + Send + Sync + 'static>(&mut self) -> Option<&mut T> {
+        self.property
+            .get_mut(&TypeId::of::<T>())
+            .and_then(|b| b.downcast_mut())
     }
 }
 
@@ -50,6 +87,22 @@ impl Storable for EndpointResolverParams {
 pub trait ResolveEndpoint: Send + Sync + fmt::Debug {
     /// Asynchronously resolves an endpoint to use from the given endpoint parameters.
     fn resolve_endpoint<'a>(&'a self, params: &'a EndpointResolverParams) -> EndpointFuture<'a>;
+
+    /// Finalize the service-specific concrete parameters in `_params`.
+    ///
+    /// The `EndpointResolverParams` may need to include additional data at a later point,
+    /// after its creation in the `read_before_execution` method of an endpoint parameters interceptor.
+    /// Modifying it directly within the [`ResolveEndpoint::resolve_endpoint`] method is not feasible,
+    /// as `params` is passed by reference. This means that incorporating extra data would require
+    /// cloning `params` within the method. However, the return type `EndpointFuture` has a lifetime
+    /// tied to the input argument, making it impossible to return the cloned `params`, as its lifetime
+    /// is scoped to the method.
+    fn finalize_params<'a>(
+        &'a self,
+        _params: &'a mut EndpointResolverParams,
+    ) -> Result<(), BoxError> {
+        Ok(())
+    }
 }
 
 /// Shared endpoint resolver.
@@ -69,6 +122,13 @@ impl ResolveEndpoint for SharedEndpointResolver {
     fn resolve_endpoint<'a>(&'a self, params: &'a EndpointResolverParams) -> EndpointFuture<'a> {
         self.0.resolve_endpoint(params)
     }
+
+    fn finalize_params<'a>(
+        &'a self,
+        params: &'a mut EndpointResolverParams,
+    ) -> Result<(), BoxError> {
+        self.0.finalize_params(params)
+    }
 }
 
 impl ValidateConfig for SharedEndpointResolver {}
diff --git a/rust-runtime/aws-smithy-runtime/src/client/orchestrator.rs b/rust-runtime/aws-smithy-runtime/src/client/orchestrator.rs
index b87c60358ea..f0beea90d03 100644
--- a/rust-runtime/aws-smithy-runtime/src/client/orchestrator.rs
+++ b/rust-runtime/aws-smithy-runtime/src/client/orchestrator.rs
@@ -367,8 +367,7 @@ async fn try_attempt(
             cfg.interceptor_state().store_put(endpoint);
         }
         None => {
-            // TODO(AccountIdBasedRouting): Pass `identity` to `orchestrate_endpoint`.
-            halt_on_err!([ctx] => orchestrate_endpoint(ctx, runtime_components, cfg)
+            halt_on_err!([ctx] => orchestrate_endpoint(identity.clone(), ctx, runtime_components, cfg)
 				    .instrument(debug_span!("orchestrate_endpoint"))
 				    .await
 				    .map_err(OrchestratorError::other));
diff --git a/rust-runtime/aws-smithy-runtime/src/client/orchestrator/endpoints.rs b/rust-runtime/aws-smithy-runtime/src/client/orchestrator/endpoints.rs
index 77496a600b7..f8df38617e4 100644
--- a/rust-runtime/aws-smithy-runtime/src/client/orchestrator/endpoints.rs
+++ b/rust-runtime/aws-smithy-runtime/src/client/orchestrator/endpoints.rs
@@ -6,6 +6,7 @@
 use aws_smithy_runtime_api::client::endpoint::{
     error::ResolveEndpointError, EndpointFuture, EndpointResolverParams, ResolveEndpoint,
 };
+use aws_smithy_runtime_api::client::identity::Identity;
 use aws_smithy_runtime_api::client::interceptors::context::InterceptorContext;
 use aws_smithy_runtime_api::client::orchestrator::HttpRequest;
 use aws_smithy_runtime_api::client::runtime_components::RuntimeComponents;
@@ -68,6 +69,7 @@ impl From<StaticUriEndpointResolverParams> for EndpointResolverParams {
 }
 
 pub(super) async fn orchestrate_endpoint(
+    identity: Identity,
     ctx: &mut InterceptorContext,
     runtime_components: &RuntimeComponents,
     cfg: &mut ConfigBag,
@@ -75,13 +77,16 @@ pub(super) async fn orchestrate_endpoint(
     trace!("orchestrating endpoint resolution");
 
     let params = cfg
-        .load::<EndpointResolverParams>()
-        .expect("endpoint resolver params must be set");
+        .get_mut_from_interceptor_state::<EndpointResolverParams>()
+        .expect("endpoint resolver params must be set in the interceptor state");
+    params.set_property(identity);
+
+    let endpoint_resolver = runtime_components.endpoint_resolver();
+
+    endpoint_resolver.finalize_params(params)?;
+
     tracing::debug!(endpoint_params = ?params, "resolving endpoint");
-    let endpoint = runtime_components
-        .endpoint_resolver()
-        .resolve_endpoint(params)
-        .await?;
+    let endpoint = endpoint_resolver.resolve_endpoint(params).await?;
 
     apply_endpoint(&endpoint, ctx, cfg)?;
 
diff --git a/rust-runtime/aws-smithy-types/Cargo.toml b/rust-runtime/aws-smithy-types/Cargo.toml
index f252d5cff8d..177e1654a23 100644
--- a/rust-runtime/aws-smithy-types/Cargo.toml
+++ b/rust-runtime/aws-smithy-types/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "aws-smithy-types"
-version = "1.3.0"
+version = "1.3.1"
 authors = [
     "AWS Rust SDK Team <aws-sdk-rust@amazon.com>",
     "Russell Cohen <rcoh@amazon.com>",
diff --git a/rust-runtime/aws-smithy-types/src/config_bag.rs b/rust-runtime/aws-smithy-types/src/config_bag.rs
index 02bac2d2713..43427ea5567 100644
--- a/rust-runtime/aws-smithy-types/src/config_bag.rs
+++ b/rust-runtime/aws-smithy-types/src/config_bag.rs
@@ -572,7 +572,7 @@ impl ConfigBag {
         self
     }
 
-    /// Return a reference to the mutable interceptor state.
+    /// Return a mutable reference to the interceptor state.
     pub fn interceptor_state(&mut self) -> &mut Layer {
         &mut self.interceptor_state
     }
@@ -582,14 +582,34 @@ impl ConfigBag {
         self.sourced_get::<T::Storer>()
     }
 
-    /// Return a mutable reference to `T` if it is stored in the top layer of the bag
+    /// Return a mutable reference to `T` if found within the interceptor state.
+    ///
+    /// This method does not search the tail of the config bag (i.e., the collection of frozen layers).
+    ///
+    /// If the caller is unsure whether `T` exists in the interceptor state,
+    /// consider using [`Self::get_mut`] instead, which requires `T` to implement `Clone`, however.
+    pub fn get_mut_from_interceptor_state<T>(&mut self) -> Option<&mut T>
+    where
+        T: Storable<Storer = StoreReplace<T>> + Send + Sync + Debug + 'static,
+    {
+        match self.interceptor_state.get_mut::<StoreReplace<T>>() {
+            Some(Value::ExplicitlyUnset(_)) => None,
+            Some(Value::Set(t)) => Some(t),
+            None => None,
+        }
+    }
+
+    /// Return a mutable reference to `T` if found within the config bag
+    ///
+    /// This method requires `T` to implement `Clone` because, if the requested item is found in the tail of the config bag,
+    /// it will be cloned and moved to the head (i.e., the interceptor state).
+    ///
+    /// If the caller is certain that `T` already exists in the interceptor state,
+    /// consider using [`Self::get_mut_from_interceptor_state`] instead, which does not require `T` to be `Clone`.
     pub fn get_mut<T>(&mut self) -> Option<&mut T>
     where
         T: Storable<Storer = StoreReplace<T>> + Send + Sync + Debug + Clone + 'static,
     {
-        // this code looks weird to satisfy the borrow checker—we can't keep the result of `get_mut`
-        // alive (even in a returned branch) and then call `store_put`. So: drop the borrow immediately
-        // store, the value, then pull it right back
         if self
             .interceptor_state
             .get_mut::<StoreReplace<T>>()
@@ -600,17 +620,8 @@ impl ConfigBag {
                 None => return None,
             };
             self.interceptor_state.store_put(new_item);
-            self.get_mut()
-        } else if matches!(
-            self.interceptor_state.get::<StoreReplace<T>>(),
-            Some(Value::ExplicitlyUnset(_))
-        ) {
-            None
-        } else if let Some(Value::Set(t)) = self.interceptor_state.get_mut::<StoreReplace<T>>() {
-            Some(t)
-        } else {
-            unreachable!()
         }
+        self.get_mut_from_interceptor_state()
     }
 
     /// Returns a mutable reference to `T` if it is stored in the top layer of the bag
@@ -981,6 +992,39 @@ mod test {
         assert_eq!(bag.get_mut_or_default::<Foo>(), &Foo(0));
     }
 
+    #[test]
+    fn get_mut_from_interceptor_state() {
+        // excludes `Clone` to verify `get_mut_from_interceptor_state` works with types without the `Clone` trait
+        #[derive(Debug, PartialEq, Eq, Default)]
+        struct Foo(usize);
+        impl Storable for Foo {
+            type Storer = StoreReplace<Self>;
+        }
+
+        #[derive(Clone, Debug, PartialEq, Eq, Default)]
+        struct Bar(usize);
+        impl Storable for Bar {
+            type Storer = StoreReplace<Self>;
+        }
+
+        let mut bag = ConfigBag::base();
+        assert!(bag.get_mut_from_interceptor_state::<Foo>().is_none());
+
+        bag.interceptor_state().store_put(Foo(1));
+        assert_eq!(
+            bag.get_mut_from_interceptor_state::<Foo>(),
+            Some(&mut Foo(1))
+        );
+
+        // `Bar` is stored in the tail of the config bag.
+        let mut layer = Layer::new("test");
+        layer.store_put(Bar(1));
+        bag.push_layer(layer);
+
+        // the method under test should not find `Bar`, as it does not reside in the interceptor state.
+        assert!(bag.get_mut_from_interceptor_state::<Bar>().is_none());
+    }
+
     #[test]
     fn cloning_layers() {
         #[derive(Clone, Debug)]


Reviews:
user: login: aajtodd
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-15T13:40:47Z
user: login: ysaito1001
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-16T21:14:23Z
user: login: ysaito1001
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-16T21:15:48Z
user: login: aajtodd
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-17T12:53:36Z

Review comments:
diff_hunk: @@ -79,6 +79,14 @@ interface EndpointCustomization {
      * Provide a list of additional endpoints standard library functions that rules can use
      */
     fun customRuntimeFunctions(codegenContext: ClientCodegenContext): List<CustomRuntimeFunction> = listOf()
+
+    /**
+     * Allows overriding the `finalize_params` method  in the `ResolveEndpoint` trait.
+     *
+     * `ResolveEndpoint::finalize_params` provides a default no-op implementation,
+     * and this customization enables the implementor to provide an alternative implementation.
+     */
+    fun overrideFinalizeEndpointParams(codegenContext: ClientCodegenContext): Writable? = null
path: codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/EndpointsDecorator.kt
user: login: aajtodd
body: What happens if multiple customizations implement this?
created_at: 2025-04-15T13:38:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -13,32 +13,74 @@ use aws_smithy_types::endpoint::Endpoint;
 use aws_smithy_types::type_erasure::TypeErasedBox;
 use error::InvalidEndpointError;
 use http_02x::uri::Authority;
+use std::any::TypeId;
+use std::collections::HashMap;
 use std::fmt;
 use std::str::FromStr;
 use std::sync::Arc;
 
 new_type_future! {
-    #[doc = "Future for [`EndpointResolver::resolve_endpoint`]."]
+    #[doc = "Future for [`ResolveEndpoint::resolve_endpoint`]."]
     pub struct EndpointFuture<'a, Endpoint, BoxError>;
 }
 
+new_type_future! {
+    #[doc = "Future for [`ResolveEndpoint::finalize_params`]."]
+    pub struct FinalizeParamsFuture<'a, (), BoxError>;
path: rust-runtime/aws-smithy-runtime-api/src/client/endpoint.rs
user: login: aajtodd
body: I think we need to just make a design choice on whether we want to allow this to be async, I too can't think of a reason it needs to be.
created_at: 2025-04-15T13:39:37Z
author_association: CONTRIBUTOR
diff_hunk: @@ -79,6 +79,14 @@ interface EndpointCustomization {
      * Provide a list of additional endpoints standard library functions that rules can use
      */
     fun customRuntimeFunctions(codegenContext: ClientCodegenContext): List<CustomRuntimeFunction> = listOf()
+
+    /**
+     * Allows overriding the `finalize_params` method  in the `ResolveEndpoint` trait.
+     *
+     * `ResolveEndpoint::finalize_params` provides a default no-op implementation,
+     * and this customization enables the implementor to provide an alternative implementation.
+     */
+    fun overrideFinalizeEndpointParams(codegenContext: ClientCodegenContext): Writable? = null
path: codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/endpoint/EndpointsDecorator.kt
user: login: ysaito1001
body: Made customizations append-able and updated the doc in 280294c.
created_at: 2025-04-16T21:14:23Z
author_association: CONTRIBUTOR
diff_hunk: @@ -13,32 +13,74 @@ use aws_smithy_types::endpoint::Endpoint;
 use aws_smithy_types::type_erasure::TypeErasedBox;
 use error::InvalidEndpointError;
 use http_02x::uri::Authority;
+use std::any::TypeId;
+use std::collections::HashMap;
 use std::fmt;
 use std::str::FromStr;
 use std::sync::Arc;
 
 new_type_future! {
-    #[doc = "Future for [`EndpointResolver::resolve_endpoint`]."]
+    #[doc = "Future for [`ResolveEndpoint::resolve_endpoint`]."]
     pub struct EndpointFuture<'a, Endpoint, BoxError>;
 }
 
+new_type_future! {
+    #[doc = "Future for [`ResolveEndpoint::finalize_params`]."]
+    pub struct FinalizeParamsFuture<'a, (), BoxError>;
path: rust-runtime/aws-smithy-runtime-api/src/client/endpoint.rs
user: login: ysaito1001
body: Discussed offline and settled on a synchronous method. Values of other endpoint params are set in a synchronous interceptor method, so it would be consistent for this to be sync also. Updated in da38e79
created_at: 2025-04-16T21:15:48Z
author_association: CONTRIBUTOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-04-12T02:58:02Z
body: A new generated diff is ready to view.
- [AWS SDK](https://d2luzm2xt3nokh.cloudfront.net/codegen-diff/db06361c471660b1fd428db6b7ee2d81d59da2d2/6da840d6086c070f52de8361431acaf0ded2fa0a/aws-sdk/index.html) ([ignoring whitespace](https://d2luzm2xt3nokh.cloudfront.net/codegen-diff/db06361c471660b1fd428db6b7ee2d81d59da2d2/6da840d6086c070f52de8361431acaf0ded2fa0a/aws-sdk-ignore-whitespace/index.html))
- [Client Test](https://d2luzm2xt3nokh.cloudfront.net/codegen-diff/db06361c471660b1fd428db6b7ee2d81d59da2d2/6da840d6086c070f52de8361431acaf0ded2fa0a/client-test/index.html) ([ignoring whitespace](https://d2luzm2xt3nokh.cloudfront.net/codegen-diff/db06361c471660b1fd428db6b7ee2d81d59da2d2/6da840d6086c070f52de8361431acaf0ded2fa0a/client-test-ignore-whitespace/index.html))
- No codegen difference in the Server Test
- [Server Test Python](https://d2luzm2xt3nokh.cloudfront.net/codegen-diff/db06361c471660b1fd428db6b7ee2d81d59da2d2/6da840d6086c070f52de8361431acaf0ded2fa0a/server-test-python/index.html) ([ignoring whitespace](https://d2luzm2xt3nokh.cloudfront.net/codegen-diff/db06361c471660b1fd428db6b7ee2d81d59da2d2/6da840d6086c070f52de8361431acaf0ded2fa0a/server-test-python-ignore-whitespace/index.html))
- No codegen difference in the Server Test Typescript


A [new doc preview](https://d2luzm2xt3nokh.cloudfront.net/docs/6da840d6086c070f52de8361431acaf0ded2fa0a/index.html) is ready to view.


author_association: NONE
user: login: github-actions[bot]
created_at: 2025-04-14T19:35:11Z
body: A new generated diff is ready to view.
- [AWS SDK](https://d2luzm2xt3nokh.cloudfront.net/codegen-diff/db06361c471660b1fd428db6b7ee2d81d59da2d2/1a9fd8db6583f153f9d6818e263e60aa53c6f728/aws-sdk/index.html) ([ignoring whitespace](https://d2luzm2xt3nokh.cloudfront.net/codegen-diff/db06361c471660b1fd428db6b7ee2d81d59da2d2/1a9fd8db6583f153f9d6818e263e60aa53c6f728/aws-sdk-ignore-whitespace/index.html))
- [Client Test](https://d2luzm2xt3nokh.cloudfront.net/codegen-diff/db06361c471660b1fd428db6b7ee2d81d59da2d2/1a9fd8db6583f153f9d6818e263e60aa53c6f728/client-test/index.html) ([ignoring whitespace](https://d2luzm2xt3nokh.cloudfront.net/codegen-diff/db06361c471660b1fd428db6b7ee2d81d59da2d2/1a9fd8db6583f153f9d6818e263e60aa53c6f728/client-test-ignore-whitespace/index.html))
- No codegen difference in the Server Test
- [Server Test Python](https://d2luzm2xt3nokh.cloudfront.net/codegen-diff/db06361c471660b1fd428db6b7ee2d81d59da2d2/1a9fd8db6583f153f9d6818e263e60aa53c6f728/server-test-python/index.html) ([ignoring whitespace](https://d2luzm2xt3nokh.cloudfront.net/codegen-diff/db06361c471660b1fd428db6b7ee2d81d59da2d2/1a9fd8db6583f153f9d6818e263e60aa53c6f728/server-test-python-ignore-whitespace/index.html))
- No codegen difference in the Server Test Typescript


A [new doc preview](https://d2luzm2xt3nokh.cloudfront.net/docs/1a9fd8db6583f153f9d6818e263e60aa53c6f728/index.html) is ready to view.


author_association: NONE
user: login: github-actions[bot]
created_at: 2025-04-16T21:34:37Z
body: A new generated diff is ready to view.
- [AWS SDK](https://d2luzm2xt3nokh.cloudfront.net/codegen-diff/db06361c471660b1fd428db6b7ee2d81d59da2d2/280294c370ae545a89d55cd8994d893c77327965/aws-sdk/index.html) ([ignoring whitespace](https://d2luzm2xt3nokh.cloudfront.net/codegen-diff/db06361c471660b1fd428db6b7ee2d81d59da2d2/280294c370ae545a89d55cd8994d893c77327965/aws-sdk-ignore-whitespace/index.html))
- [Client Test](https://d2luzm2xt3nokh.cloudfront.net/codegen-diff/db06361c471660b1fd428db6b7ee2d81d59da2d2/280294c370ae545a89d55cd8994d893c77327965/client-test/index.html) ([ignoring whitespace](https://d2luzm2xt3nokh.cloudfront.net/codegen-diff/db06361c471660b1fd428db6b7ee2d81d59da2d2/280294c370ae545a89d55cd8994d893c77327965/client-test-ignore-whitespace/index.html))
- No codegen difference in the Server Test
- [Server Test Python](https://d2luzm2xt3nokh.cloudfront.net/codegen-diff/db06361c471660b1fd428db6b7ee2d81d59da2d2/280294c370ae545a89d55cd8994d893c77327965/server-test-python/index.html) ([ignoring whitespace](https://d2luzm2xt3nokh.cloudfront.net/codegen-diff/db06361c471660b1fd428db6b7ee2d81d59da2d2/280294c370ae545a89d55cd8994d893c77327965/server-test-python-ignore-whitespace/index.html))
- No codegen difference in the Server Test Typescript


A [new doc preview](https://d2luzm2xt3nokh.cloudfront.net/docs/280294c370ae545a89d55cd8994d893c77327965/index.html) is ready to view.


author_association: NONE