diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 9591b962ab5..382a2700e6d 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -1289,7 +1289,7 @@ class BaseAWSLLM(SignsRequestsWithAWS): caller's key skips both. ``aws_session_name`` defaults inside this fetcher rather than in ``get_credentials`` so the cache key stays stable when the caller does not supply one. """ - if self._is_already_running_as_role(aws_role_name, ssl_verify=ssl_verify): + if not aws_session_tags and self._is_already_running_as_role(aws_role_name, ssl_verify=ssl_verify): verbose_logger.debug( "Already running as target role %s, using ambient credentials", aws_role_name, @@ -1409,7 +1409,15 @@ class BaseAWSLLM(SignsRequestsWithAWS): sts_response = sts_client.assume_role(**assume_role_params) except Exception as e: error_str: Final = str(e) - if "AccessDenied" in error_str: + if "AccessDenied" in error_str and aws_session_tags: + verbose_logger.error( + "AssumeRole AccessDenied for %s with aws_session_tags set. Ambient credentials cannot carry " + "session tags, so the role's trust policy must allow sts:AssumeRole and sts:TagSession for " + "the caller, even when the caller is that same role. Re-raising. Error: %s", + aws_role_name, + error_str, + ) + elif "AccessDenied" in error_str: # Only fall back to ambient credentials if we can positively # confirm the caller is already the target role (same account, # partition, and role name). This avoids silently using the diff --git a/tests/unit/llms/bedrock/test_base_aws_llm.py b/tests/unit/llms/bedrock/test_base_aws_llm.py index c220e39f4d1..0d4aa223dea 100644 --- a/tests/unit/llms/bedrock/test_base_aws_llm.py +++ b/tests/unit/llms/bedrock/test_base_aws_llm.py @@ -2363,13 +2363,19 @@ _TAGGED_ROLE_ARN = "arn:aws:iam::123456789012:role/TaggedRole" class _TagAwareSTSClient: """STS stand-in for a trust policy that only admits sessions carrying exactly the expected tags.""" - def __init__(self, expected_tags: tuple = (), access_key: str = "ASIATAGGEDSESSION") -> None: + def __init__( + self, + expected_tags: tuple = (), + access_key: str = "ASIATAGGEDSESSION", + caller_arn: str = "arn:aws:iam::111111111111:user/litellm-proxy-pod", + ) -> None: self.expected_tags = expected_tags self.access_key = access_key + self.caller_arn = caller_arn self.assume_role_calls: list = [] def get_caller_identity(self): - return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + return {"Arn": self.caller_arn} def assume_role_with_web_identity(self, **params): return { @@ -2511,6 +2517,57 @@ def test_irsa_same_account_assume_role_sends_session_tags(tmp_path): ] +def _running_as_tagged_role_env(tmp_path, runtime: str) -> dict: + match runtime: + case "ecs": + return _os_environ_without_aws_keys() + case "irsa": + return {**_os_environ_without_aws_keys(), **_irsa_env(tmp_path, _TAGGED_ROLE_ARN)} + case _: + raise AssertionError(runtime) + + +_TAGGED_ROLE_AMBIENT_SESSION_ARN = "arn:aws:sts::123456789012:assumed-role/TaggedRole/ambient-session" + + +@pytest.mark.parametrize("runtime", ["ecs", "irsa"]) +def test_get_credentials_assumes_its_own_role_to_apply_session_tags(tmp_path, runtime): + sts = _TagAwareSTSClient(expected_tags=_SORTED_SESSION_TAGS, caller_arn=_TAGGED_ROLE_AMBIENT_SESSION_ARN) + + with ( + patch.dict(os.environ, _running_as_tagged_role_env(tmp_path, runtime), clear=True), + patch("boto3.client", return_value=sts), + ): + credentials = BaseAWSLLM().get_credentials( + aws_role_name=_TAGGED_ROLE_ARN, + aws_session_name="team-session", + aws_session_tags=list(_SESSION_TAGS), + ) + + assert credentials.access_key == "ASIATAGGEDSESSION" + assert sts.assume_role_calls == [ + {"RoleArn": _TAGGED_ROLE_ARN, "RoleSessionName": "team-session", "Tags": _SORTED_SESSION_TAGS} + ] + + +@pytest.mark.parametrize("runtime", ["ecs", "irsa"]) +def test_get_credentials_surfaces_a_denied_tagged_self_assume_instead_of_dropping_the_tags(tmp_path, runtime): + sts = _TagAwareSTSClient(expected_tags=(), caller_arn=_TAGGED_ROLE_AMBIENT_SESSION_ARN) + + with ( + patch.dict(os.environ, _running_as_tagged_role_env(tmp_path, runtime), clear=True), + patch("boto3.client", return_value=sts), + pytest.raises(ClientError, match="sts:TagSession"), + ): + BaseAWSLLM().get_credentials( + aws_role_name=_TAGGED_ROLE_ARN, + aws_session_name="team-session", + aws_session_tags=list(_SESSION_TAGS), + ) + + assert [call["Tags"] for call in sts.assume_role_calls] == [_SORTED_SESSION_TAGS] + + def test_get_credentials_canonicalizes_session_tag_order_for_the_cache(): """Two deployments listing the same tags in a different order share one STS session.""" base_aws_llm = BaseAWSLLM()