From d95578464ef630d843368dbb5d0bec38dc624748 Mon Sep 17 00:00:00 2001 From: Steven Anton Date: Sat, 23 Aug 2025 21:52:18 -0700 Subject: [PATCH] Add support for AWS assume_role with a session token --- litellm/llms/bedrock/base_aws_llm.py | 3 +++ tests/llm_translation/test_aws_base_llm.py | 1 + .../test_litellm/llms/bedrock/test_base_aws_llm.py | 14 +++++++++++++- 3 files changed, 17 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 55aad673ced..ce196757f94 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -202,6 +202,7 @@ class BaseAWSLLM: credentials, _cache_ttl = self._auth_with_aws_role( aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, aws_role_name=aws_role_name, aws_session_name=aws_session_name, ) @@ -554,6 +555,7 @@ class BaseAWSLLM: self, aws_access_key_id: Optional[str], aws_secret_access_key: Optional[str], + aws_session_token: Optional[str], aws_role_name: str, aws_session_name: str, ) -> Tuple[Credentials, Optional[int]]: @@ -614,6 +616,7 @@ class BaseAWSLLM: "sts", aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, ) sts_response = sts_client.assume_role( diff --git a/tests/llm_translation/test_aws_base_llm.py b/tests/llm_translation/test_aws_base_llm.py index 5a299a76c07..7ce7f6ac0cb 100644 --- a/tests/llm_translation/test_aws_base_llm.py +++ b/tests/llm_translation/test_aws_base_llm.py @@ -88,6 +88,7 @@ def test_auth_with_aws_role(mock_boto3_client, base_aws_llm): credentials, ttl = base_aws_llm._auth_with_aws_role( aws_access_key_id="test_access", aws_secret_access_key="test_secret", + aws_session_token="test_token", aws_role_name="test_role", aws_session_name="test_session", ) diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index 01903737580..5effa6fa01a 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -576,6 +576,7 @@ def test_eks_irsa_ambient_credentials_used(): credentials, ttl = base_aws_llm._auth_with_aws_role( aws_access_key_id=None, aws_secret_access_key=None, + aws_session_token=None, aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole", aws_session_name="test-session" ) @@ -630,6 +631,7 @@ def test_explicit_credentials_used_when_provided(): credentials, ttl = base_aws_llm._auth_with_aws_role( aws_access_key_id="explicit-access-key", aws_secret_access_key="explicit-secret-key", + aws_session_token="assumed-session-token", aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole", aws_session_name="test-session" ) @@ -639,6 +641,7 @@ def test_explicit_credentials_used_when_provided(): "sts", aws_access_key_id="explicit-access-key", aws_secret_access_key="explicit-secret-key", + aws_session_token="assumed-session-token", ) # Should call assume_role @@ -687,6 +690,7 @@ def test_partial_credentials_still_use_ambient(): credentials, ttl = base_aws_llm._auth_with_aws_role( aws_access_key_id="AKIAEXAMPLE", aws_secret_access_key=None, + aws_session_token=None, aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole", aws_session_name="test-session" ) @@ -695,7 +699,8 @@ def test_partial_credentials_still_use_ambient(): mock_boto3_client.assert_called_once_with( "sts", aws_access_key_id="AKIAEXAMPLE", - aws_secret_access_key=None + aws_secret_access_key=None, + aws_session_token=None, ) # Should still call assume_role @@ -737,6 +742,7 @@ def test_cross_account_role_assumption(): credentials, ttl = base_aws_llm._auth_with_aws_role( aws_access_key_id=None, aws_secret_access_key=None, + aws_session_token=None, aws_role_name="arn:aws:iam::999999999999:role/CrossAccountRole", aws_session_name="cross-account-session" ) @@ -789,6 +795,7 @@ def test_role_assumption_with_custom_session_name(): credentials, ttl = base_aws_llm._auth_with_aws_role( aws_access_key_id=None, aws_secret_access_key=None, + aws_session_token=None, aws_role_name="arn:aws:iam::1111111111111:role/LitellmRole", aws_session_name="evals-bedrock-session" ) @@ -832,6 +839,7 @@ def test_role_assumption_ttl_calculation(): credentials, ttl = base_aws_llm._auth_with_aws_role( aws_access_key_id=None, aws_secret_access_key=None, + aws_session_token=None, aws_role_name="arn:aws:iam::1111111111111:role/LitellmRole", aws_session_name="ttl-test-session" ) @@ -858,6 +866,7 @@ def test_role_assumption_error_handling(): base_aws_llm._auth_with_aws_role( aws_access_key_id=None, aws_secret_access_key=None, + aws_session_token=None, aws_role_name="arn:aws:iam::1111111111111:role/UnauthorizedRole", aws_session_name="error-test-session" ) @@ -911,6 +920,7 @@ def test_multiple_role_assumptions_in_sequence(): credentials1, ttl1 = base_aws_llm._auth_with_aws_role( aws_access_key_id=None, aws_secret_access_key=None, + aws_session_token=None, aws_role_name="arn:aws:iam::1111111111111:role/LitellmRole", aws_session_name="session-1" ) @@ -919,6 +929,7 @@ def test_multiple_role_assumptions_in_sequence(): credentials2, ttl2 = base_aws_llm._auth_with_aws_role( aws_access_key_id=None, aws_secret_access_key=None, + aws_session_token=None, aws_role_name="arn:aws:iam::2222222222222:role/LitellmEvalBedrockRole", aws_session_name="session-2" ) @@ -980,6 +991,7 @@ def test_auth_with_aws_role_irsa_environment(): creds, ttl = base_llm._auth_with_aws_role( aws_access_key_id=None, aws_secret_access_key=None, + aws_session_token=None, aws_role_name='arn:aws:iam::222222222222:role/target-role', aws_session_name='test-session' )