Skip to content

Add initial support assume role credential resolver#721

Open
arandito wants to merge 1 commit into
smithy-lang:developfrom
arandito:assume-role-resolver
Open

Add initial support assume role credential resolver#721
arandito wants to merge 1 commit into
smithy-lang:developfrom
arandito:assume-role-resolver

Conversation

@arandito

@arandito arandito commented Jun 19, 2026

Copy link
Copy Markdown
Contributor

Important

This PR adds a standalone Assume Role credential resolver that can only be used programmatically. We do not currently support reading from a shared AWS config file and do not have a mechanism to create profile-based credential providers in the default chain. Once both land, users will be able to configure assume role credentials from their shared config file. However, this will immediately unblock anyone wishing to use assume role credentials.

Description

This PR adds an Assume Role credentials resolver the Smithy Python runtime that allows clients to source credentials from the Amazon Security Token Service (STS) using the AssumeRole API.

Changes

  • Introduces AssumeRoleCredentialsResolver in the identity module of the smithy-aws-core package
  • Executes an async STS Assume Role call using credentials resolved from a nested source_resolver that must be of the AWSCredentialsResolver type. It then converts the API output shape into an AWSCredentialsIdentity.
  • Reuses cached credentials when they are still valid and refreshes them when credentials have expired.
  • Adds a code-generated, slim STS client under _private/nested_clients/aws_sdk_sts/ that includes only the AssumeRole operation
  • Exposes a new [assume-role] optional dependency extra in smithy-aws-core that pulls in smithy-xml for the awsQuery protocol. Codegen now adds this extra to all AWS clients generated with Smithy Python to ensure they can use assume role credentials out of the box.
  • Supports MFA-protected roles via an optional mfa_serial and an async mfa_code_provider callback. The callback is invoked on every assume so a fresh, single-use token code is supplied on each credential refresh.

Nested Client Generation

The nested aws_sdk_sts client was generated with a trimmed build projection that removes unnecessary traits such as endpointTests and documentation to reduce the client size. The following is the smithy-build.json used:

{
    "version": "1.0",
    "sources": ["models"],
    "maven": {
        "dependencies": [
            "software.amazon.smithy:smithy-model:[1.71.0,2.0)",
            "software.amazon.smithy:smithy-aws-traits:[1.71.0,2.0)",
            "software.amazon.smithy:smithy-aws-endpoints:[1.71.0,2.0)",
            "software.amazon.smithy.python.codegen.aws:core:0.3.1"
        ]
    },
    "projections": {
        "client": {
            "transforms": [
                {
                    "name": "excludeShapesBySelector",
                    "args": {
                        "selector": "operation:not([id=com.amazonaws.sts#AssumeRole])"
                    }
                },
                {
                    "name": "excludeTraits",
                    "args": {
                        "traits": [
                            "smithy.rules#endpointTests",
                            "smithy.rules#endpointRuleSet",
                            "smithy.rules#endpointBdd",
                            "smithy.api#examples",
                            "smithy.api#documentation",
                            "smithy.api#externalDocumentation"
                        ]
                    }
                },
                {
                    "name": "removeUnusedShapes"
                }
            ],
            "plugins": {
                "python-client-codegen": {
                    "service": "com.amazonaws.sts#AWSSecurityTokenServiceV20110615",
                    "module": "aws_sdk_sts",
                    "moduleVersion": "0.1.0"
                }
            }
        }
    }
}

Testing

  • Added unit test coverage for resolver behavior
  • Tested E2E with aws-sdk-bedrock-runtime client.
    • Confirmed AssumeRoleCredentialResolver uses nested resolver to sign AssumeRole request
    • Confirmed multiple roles can be chained

Testing script:

import asyncio

from aws_sdk_bedrock_runtime.client import BedrockRuntimeClient
from aws_sdk_bedrock_runtime.config import Config
from aws_sdk_bedrock_runtime.models import ContentBlockText, ConverseInput, Message
from smithy_aws_core.identity import EnvironmentCredentialsResolver
from smithy_aws_core.identity import AssumeRoleCredentialsResolver


async def main():
    resolver = AssumeRoleCredentialsResolver(
        source_resolver=EnvironmentCredentialsResolver(),
        role_arn="arn:aws:iam::REDACTED:role/TestAssumeRole",
    )
    resolver2 = AssumeRoleCredentialsResolver(
        source_resolver=resolver,
        role_arn="arn:aws:iam::REDACTED:role/TestAssumeRole2",
    )
    client = BedrockRuntimeClient(
        config=Config(
            endpoint_uri="https://bedrock-runtime.us-west-2.amazonaws.com",
            region="us-west-2",
            aws_credentials_identity_resolver=resolver2,
        )
    )
    response = await client.converse(
        ConverseInput(
            model_id="us.anthropic.claude-opus-4-8",
            messages=[
                Message(
                    role="user",
                    content=[ContentBlockText(value="Tell me a joke.")],
                )
            ],
        )
    )
    print(response)


asyncio.run(main())

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

@arandito arandito requested a review from a team as a code owner June 19, 2026 18:02
)
self._external_id = external_id
self._duration_seconds = duration_seconds
self._region = region or DEFAULT_STS_REGION

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Open question: Should we expose region? We currently need this to configure the endpoint for the internal STS client and default to the global endpoint if its not specified. I think once we update our config resolution system, we can just use the top-level client's configured region instead of needing to pass it in. We can probably default to the global endpoint for now and remove this option. Curious what others think.

):
"""Resolves AWS credentials from an STS ``AssumeRole`` call."""

def __init__(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Just curious - it seems that the resolver exposes fewer AssumeRole inputs than the other SDKs. For example, policy / policy_arns aren't included here, while other SDKs such as Go and Kotlin already include them.

Not a blocker for initial support, but wondering if we have any plan to add more inputs in the future?

@Alan4506 Alan4506 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

The License Header Check is failing: the 8 generated nested STS client files are missing the Apache header (codegen only emits # Code generated...). Should we:

  • simply exclude _private/nested_clients/** in .licenserc.yaml, or
  • modify the codegen to add the license header to these files - maybe behind a flag so it only applies to vendored clients like this one?

Any preference on which way to go?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants