Merge pull request #13919 from stevenmanton/fix-aws-assume-role-with-token

Add support for AWS assume_role with a session token

Author: krrishdholakia

litellm/llms/bedrock/base_aws_llm.py

@@ -202,6 +202,7 @@ def get_credentials(
                 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 @@ def _auth_with_aws_role(
         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 @@ def _auth_with_aws_role(
                     "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(

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",
     )

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'
                 )