From 70a1bf92e24b186bae88e092feef113223db8694 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 10 Feb 2026 01:46:38 -0500 Subject: [PATCH] fix(bedrock): skip AssumeRole when ECS/EC2 already running as target IAM role MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When aws_role_name is configured but the environment (ECS task role, EC2 instance profile) is already running as that role, AssumeRole is unnecessary and can fail with AccessDenied. This adds same-role detection for ECS/EC2 (extending existing IRSA support) and a fallback to ambient credentials when AssumeRole returns AccessDenied. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude --- litellm/llms/bedrock/base_aws_llm.py | 98 +++++++-- litellm/utils.py | 6 + .../llms/bedrock/test_base_aws_llm.py | 196 ++++++++++++++++-- 3 files changed, 270 insertions(+), 30 deletions(-) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 1de1c40c43..506b8811d3 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -211,25 +211,13 @@ class BaseAWSLLM: aws_external_id=aws_external_id, ) elif aws_role_name is not None: - # Check if we're in IRSA and trying to assume the same role we already have - current_role_arn = os.getenv("AWS_ROLE_ARN") - web_identity_token_file = os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE") - - # In IRSA environments, we should skip role assumption if we're already running as the target role - # This is true when: - # 1. We have AWS_ROLE_ARN set (current role) - # 2. We have AWS_WEB_IDENTITY_TOKEN_FILE set (IRSA environment) - # 3. The current role matches the requested role - if ( - current_role_arn - and web_identity_token_file - and current_role_arn == aws_role_name - ): + # Check if we're already running as the target role and can skip assumption + # This handles IRSA (EKS), ECS task roles, and EC2 instance profiles + if self._is_already_running_as_role(aws_role_name): verbose_logger.debug( - "Using IRSA same-role optimization: calling _auth_with_env_vars" + "Already running as target role %s, using ambient credentials", + aws_role_name, ) - # We're already running as this role via IRSA, no need to assume it again - # Use the default boto3 credentials (which will use the IRSA credentials) credentials, _cache_ttl = self._auth_with_env_vars() else: verbose_logger.debug( @@ -553,6 +541,65 @@ class BaseAWSLLM: aws_region_name = "us-west-2" return aws_region_name + def _is_already_running_as_role(self, aws_role_name: str) -> bool: + """ + Check if the current environment is already running as the target IAM role. + + This handles multiple AWS environments: + - IRSA (EKS): AWS_ROLE_ARN + AWS_WEB_IDENTITY_TOKEN_FILE are set + - ECS task roles: Uses sts:GetCallerIdentity to check current role ARN + - EC2 instance profiles: Uses sts:GetCallerIdentity to check current role ARN + + Returns True if the current identity matches the target role, meaning + we can skip sts:AssumeRole and use ambient credentials directly. + """ + # Fast path: IRSA environment check (no API call needed) + current_role_arn = os.getenv("AWS_ROLE_ARN") + web_identity_token_file = os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE") + if current_role_arn and web_identity_token_file: + return current_role_arn == aws_role_name + + # For ECS/EC2: call sts:GetCallerIdentity to check if already running as the role + try: + import boto3 + + with tracer.trace("boto3.client(sts).get_caller_identity"): + sts_client = boto3.client("sts") + identity = sts_client.get_caller_identity() + caller_arn = identity.get("Arn", "") + + # The caller ARN for an ECS task role looks like: + # arn:aws:sts::123456789012:assumed-role/MyRole/session-name + # The target role ARN looks like: + # arn:aws:iam::123456789012:role/MyRole + # We need to compare the role name portion + if ":assumed-role/" in caller_arn: + # Extract role name from assumed-role ARN + # Format: arn:aws:sts::ACCOUNT:assumed-role/ROLE_NAME/SESSION + caller_role_name = caller_arn.split(":assumed-role/")[1].split("/")[0] + + # Extract role name from target role ARN + # Format: arn:aws:iam::ACCOUNT:role/ROLE_NAME or + # arn:aws:iam::ACCOUNT:role/path/ROLE_NAME + if ":role/" in aws_role_name: + target_role_name = aws_role_name.split(":role/")[-1].split("/")[-1] + else: + target_role_name = aws_role_name + + if caller_role_name == target_role_name: + verbose_logger.debug( + "Current identity already matches target role: %s", + aws_role_name, + ) + return True + + except Exception as e: + verbose_logger.debug( + "Could not determine current role identity: %s", str(e) + ) + + return False + @tracer.wrap() def _auth_with_web_identity_token( self, @@ -867,7 +914,22 @@ class BaseAWSLLM: if aws_external_id is not None: assume_role_params["ExternalId"] = aws_external_id - sts_response = sts_client.assume_role(**assume_role_params) + try: + sts_response = sts_client.assume_role(**assume_role_params) + except Exception as e: + error_str = str(e) + # If AssumeRole fails because the caller already IS the role + # (e.g., ECS task role, root account, or same-role scenario), + # fall back to using ambient credentials directly + if "AccessDenied" in error_str: + verbose_logger.warning( + "AssumeRole failed for %s (%s). " + "Falling back to ambient credentials (boto3 default chain).", + aws_role_name, + error_str, + ) + return self._auth_with_env_vars() + raise # Extract the credentials from the response and convert to Session Credentials sts_credentials = sts_response["Credentials"] diff --git a/litellm/utils.py b/litellm/utils.py index ed0d6ee930..733baeadd8 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6140,6 +6140,12 @@ def validate_environment( # noqa: PLR0915 if ( "AWS_ACCESS_KEY_ID" in os.environ and "AWS_SECRET_ACCESS_KEY" in os.environ + ) or ( + # IAM role, profile, or web identity auth don't require access keys + "AWS_ROLE_ARN" in os.environ + or "AWS_PROFILE" in os.environ + or "AWS_WEB_IDENTITY_TOKEN_FILE" in os.environ + or "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI" in os.environ # ECS task role ): keys_in_environment = True else: 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 77eda43251..b57c20fb7c 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -853,29 +853,66 @@ def test_role_assumption_ttl_calculation(): assert 3500 <= ttl <= 3600 # Allow some variance for test execution time -def test_role_assumption_error_handling(): +def test_role_assumption_access_denied_falls_back_to_env_vars(): """ - Test that role assumption errors are properly propagated. + Test that when AssumeRole fails with AccessDenied, we fall back to ambient credentials. + This handles ECS task roles, root accounts, and same-role scenarios where + AssumeRole is unnecessary because the caller already has the role's permissions. """ base_aws_llm = BaseAWSLLM() - - # Mock the boto3 STS client to raise an exception + + # Mock the boto3 STS client to raise AccessDenied mock_sts_client = MagicMock() - mock_sts_client.assume_role.side_effect = Exception("AccessDenied: User is not authorized to perform sts:AssumeRole") - + mock_sts_client.assume_role.side_effect = Exception( + "An error occurred (AccessDenied) when calling the AssumeRole operation: " + "Roles may not be assumed by root accounts." + ) + + # Mock _auth_with_env_vars to return fallback credentials + mock_creds = MagicMock() + mock_creds.access_key = "fallback-access-key" + mock_creds.secret_key = "fallback-secret-key" + + with patch("boto3.client", return_value=mock_sts_client): + with patch.object( + base_aws_llm, "_auth_with_env_vars", return_value=(mock_creds, None) + ) as mock_env_auth: + 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/UnauthorizedRole", + aws_session_name="error-test-session", + ) + + # Should have fallen back to env vars + mock_env_auth.assert_called_once() + assert credentials.access_key == "fallback-access-key" + + +def test_role_assumption_non_access_denied_error_propagated(): + """ + Test that non-AccessDenied errors from AssumeRole are still propagated. + """ + base_aws_llm = BaseAWSLLM() + + # Mock the boto3 STS client to raise a non-AccessDenied exception + mock_sts_client = MagicMock() + mock_sts_client.assume_role.side_effect = Exception( + "An error occurred (MalformedPolicyDocument) when calling the AssumeRole operation" + ) + with patch("boto3.client", return_value=mock_sts_client): - - # Should raise the exception with pytest.raises(Exception) as exc_info: 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" + aws_role_name="arn:aws:iam::1111111111111:role/BadPolicyRole", + aws_session_name="error-test-session", ) - - assert "AccessDenied" in str(exc_info.value) + + assert "MalformedPolicyDocument" in str(exc_info.value) def test_multiple_role_assumptions_in_sequence(): @@ -1195,3 +1232,138 @@ def test_converse_handler_external_id_extraction(): assert hasattr(mock_get_credentials, 'called_kwargs') assert "aws_external_id" in mock_get_credentials.called_kwargs assert mock_get_credentials.called_kwargs["aws_external_id"] == "TestExternalID123" + + +def test_is_already_running_as_role_irsa_same_role(): + """Test IRSA fast path: when AWS_ROLE_ARN matches target role.""" + base_aws_llm = BaseAWSLLM() + + with patch.dict(os.environ, { + "AWS_ROLE_ARN": "arn:aws:iam::123456789012:role/MyRole", + "AWS_WEB_IDENTITY_TOKEN_FILE": "/var/run/secrets/token", + }): + assert base_aws_llm._is_already_running_as_role( + "arn:aws:iam::123456789012:role/MyRole" + ) is True + + +def test_is_already_running_as_role_irsa_different_role(): + """Test IRSA fast path: when AWS_ROLE_ARN does NOT match target role.""" + base_aws_llm = BaseAWSLLM() + + with patch.dict(os.environ, { + "AWS_ROLE_ARN": "arn:aws:iam::123456789012:role/MyRole", + "AWS_WEB_IDENTITY_TOKEN_FILE": "/var/run/secrets/token", + }): + assert base_aws_llm._is_already_running_as_role( + "arn:aws:iam::999999999999:role/OtherRole" + ) is False + + +def test_is_already_running_as_role_ecs_task_role(): + """Test ECS/EC2 path: GetCallerIdentity shows assumed-role matching target.""" + base_aws_llm = BaseAWSLLM() + + mock_sts_client = MagicMock() + mock_sts_client.get_caller_identity.return_value = { + "Arn": "arn:aws:sts::123456789012:assumed-role/MyEcsTaskRole/ecs-task-id" + } + + with patch.dict(os.environ, {}, clear=False): + # Ensure no IRSA env vars + env = {k: v for k, v in os.environ.items() if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")} + with patch.dict(os.environ, env, clear=True): + with patch("boto3.client", return_value=mock_sts_client): + assert base_aws_llm._is_already_running_as_role( + "arn:aws:iam::123456789012:role/MyEcsTaskRole" + ) is True + + +def test_is_already_running_as_role_ecs_different_role(): + """Test ECS/EC2 path: GetCallerIdentity shows a different role.""" + base_aws_llm = BaseAWSLLM() + + mock_sts_client = MagicMock() + mock_sts_client.get_caller_identity.return_value = { + "Arn": "arn:aws:sts::123456789012:assumed-role/MyEcsTaskRole/ecs-task-id" + } + + with patch.dict(os.environ, {}, clear=False): + env = {k: v for k, v in os.environ.items() if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")} + with patch.dict(os.environ, env, clear=True): + with patch("boto3.client", return_value=mock_sts_client): + assert base_aws_llm._is_already_running_as_role( + "arn:aws:iam::999999999999:role/DifferentRole" + ) is False + + +def test_is_already_running_as_role_ecs_role_with_path(): + """Test ECS path with role that has a path prefix (e.g., /service-role/MyRole).""" + base_aws_llm = BaseAWSLLM() + + mock_sts_client = MagicMock() + mock_sts_client.get_caller_identity.return_value = { + "Arn": "arn:aws:sts::123456789012:assumed-role/MyEcsTaskRole/ecs-task-id" + } + + with patch.dict(os.environ, {}, clear=False): + env = {k: v for k, v in os.environ.items() if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")} + with patch.dict(os.environ, env, clear=True): + with patch("boto3.client", return_value=mock_sts_client): + # Role ARN with path + assert base_aws_llm._is_already_running_as_role( + "arn:aws:iam::123456789012:role/service-role/MyEcsTaskRole" + ) is True + + +def test_is_already_running_as_role_get_caller_identity_fails(): + """Test that when GetCallerIdentity fails, we return False (don't crash).""" + base_aws_llm = BaseAWSLLM() + + mock_sts_client = MagicMock() + mock_sts_client.get_caller_identity.side_effect = Exception("No credentials found") + + with patch.dict(os.environ, {}, clear=False): + env = {k: v for k, v in os.environ.items() if k not in ("AWS_ROLE_ARN", "AWS_WEB_IDENTITY_TOKEN_FILE")} + with patch.dict(os.environ, env, clear=True): + with patch("boto3.client", return_value=mock_sts_client): + assert base_aws_llm._is_already_running_as_role( + "arn:aws:iam::123456789012:role/SomeRole" + ) is False + + +def test_get_credentials_ecs_same_role_skips_assume_role(): + """ + End-to-end test: when running on ECS with the same role as aws_role_name, + get_credentials should use ambient credentials and NOT call AssumeRole. + """ + base_aws_llm = BaseAWSLLM() + + mock_creds = MagicMock() + mock_creds.access_key = "ecs-access-key" + mock_creds.secret_key = "ecs-secret-key" + mock_creds.token = "ecs-session-token" + + with patch.object( + base_aws_llm, + "_is_already_running_as_role", + return_value=True, + ): + with patch.object( + base_aws_llm, + "_auth_with_env_vars", + return_value=(mock_creds, None), + ) as mock_env_auth: + with patch.object( + base_aws_llm, + "_auth_with_aws_role", + ) as mock_role_auth: + credentials = base_aws_llm.get_credentials( + aws_role_name="arn:aws:iam::123456789012:role/MyEcsTaskRole", + aws_region_name="us-east-1", + ) + + # Should use env vars, NOT role assumption + mock_env_auth.assert_called_once() + mock_role_auth.assert_not_called() + assert credentials.access_key == "ecs-access-key"