diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 92b32d32dc0..edd2e88f95c 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -17,6 +17,7 @@ AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset( "aws_web_identity_token", "aws_sts_endpoint", "aws_external_id", + "aws_session_tags", "aws_bedrock_runtime_endpoint", "aws_bedrock_project_id", } diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index a197702921f..ab6c4ab85d6 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -4,13 +4,14 @@ import json import os import re import urllib.parse -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence from datetime import datetime from threading import Lock from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast, get_args, overload import httpx -from pydantic import BaseModel, ValidationError +from pydantic import BaseModel, TypeAdapter, ValidationError +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.caching.caching import DualCache @@ -26,6 +27,7 @@ from litellm.constants import ( from litellm.litellm_core_utils.aws_partition import contains_bedrock_arn, get_aws_dns_suffix from litellm.litellm_core_utils.dd_tracing import tracer from litellm.secret_managers.main import get_secret, get_secret_str +from litellm.types.llms.bedrock import AwsSessionTag if TYPE_CHECKING: from botocore.awsrequest import AWSPreparedRequest @@ -47,6 +49,47 @@ _STS_REGION_FROM_ENDPOINT_PATTERN: Final = re.compile( SIGV4_COMPUTED_HEADERS: Final = frozenset({"authorization", "x-amz-date", "x-amz-security-token", "date"}) +_AWS_SESSION_TAGS_ADAPTER: Final[TypeAdapter[tuple[AwsSessionTag, ...]]] = TypeAdapter(tuple[AwsSessionTag, ...]) + + +def _canonical_aws_session_tags(raw_tags: object) -> tuple[AwsSessionTag, ...] | None: + if raw_tags is None: + return None + try: + validated: Final = _AWS_SESSION_TAGS_ADAPTER.validate_python(raw_tags) + except ValidationError as e: + raise ValueError( + "Invalid 'aws_session_tags' value. Expected a list of {'Key': , 'Value': } dicts, " + f"e.g. [{{'Key': 'team', 'Value': 'genai'}}]. Got: {raw_tags!r}" + ) from e + return tuple(sorted(validated, key=lambda tag: tag["Key"])) + + +class _AssumeRoleParams(TypedDict): + RoleArn: ReadOnly[str] + RoleSessionName: ReadOnly[str] + ExternalId: ReadOnly[NotRequired[str]] + Tags: ReadOnly[NotRequired[tuple[AwsSessionTag, ...]]] + + +def _assume_role_params( + aws_role_name: str, + aws_session_name: str, + aws_external_id: str | None, + aws_session_tags: Sequence[AwsSessionTag] | None, +) -> _AssumeRoleParams: + match (aws_external_id, tuple(aws_session_tags or ())): + case (None, ()): + return _AssumeRoleParams(RoleArn=aws_role_name, RoleSessionName=aws_session_name) + case (None, tags): + return _AssumeRoleParams(RoleArn=aws_role_name, RoleSessionName=aws_session_name, Tags=tags) + case (external_id, ()): + return _AssumeRoleParams(RoleArn=aws_role_name, RoleSessionName=aws_session_name, ExternalId=external_id) + case (external_id, tags): + return _AssumeRoleParams( + RoleArn=aws_role_name, RoleSessionName=aws_session_name, ExternalId=external_id, Tags=tags + ) + class BedrockRequestTarget(BaseModel): aws_region_name: str @@ -120,6 +163,7 @@ class BaseAWSLLM: "aws_sts_endpoint", "aws_bedrock_runtime_endpoint", "aws_external_id", + "aws_session_tags", ] def _get_ssl_verify(self, ssl_verify: bool | str | None = None): @@ -137,7 +181,7 @@ class BaseAWSLLM: return get_ssl_verify(ssl_verify=ssl_verify) - def get_cache_key(self, credential_args: Mapping[str, str | bool | None]) -> str: + def get_cache_key(self, credential_args: Mapping[str, str | bool | tuple[AwsSessionTag, ...] | None]) -> str: """ Generate a unique cache key based on the credential arguments. """ @@ -147,7 +191,7 @@ class BaseAWSLLM: def _get_or_set_cached_credentials( self, - credential_args: Mapping[str, str | bool | None], + credential_args: Mapping[str, str | bool | tuple[AwsSessionTag, ...] | None], credential_fetcher: Callable[[], tuple[Credentials, int | None]], ) -> Any: """ @@ -222,6 +266,7 @@ class BaseAWSLLM: aws_web_identity_token: str | None = None, aws_sts_endpoint: str | None = None, aws_external_id: str | None = None, + aws_session_tags: Sequence[AwsSessionTag] | None = None, ssl_verify: bool | str | None = None, ): """ @@ -258,6 +303,7 @@ class BaseAWSLLM: (aws_external_id, "AWS_EXTERNAL_ID"), ) ) + session_tags: Final = _canonical_aws_session_tags(aws_session_tags) verbose_logger.debug( "in get credentials\n" @@ -270,7 +316,8 @@ class BaseAWSLLM: "aws_role_name=%s\n" "aws_web_identity_token=[set=%s]\n" "aws_sts_endpoint=%s\n" - "aws_external_id=%s", + "aws_external_id=%s\n" + "aws_session_tags=%s", aws_access_key_id is not None, aws_secret_access_key is not None, aws_session_token is not None, @@ -281,6 +328,7 @@ class BaseAWSLLM: aws_web_identity_token is not None, aws_sts_endpoint, aws_external_id, + session_tags, ) args: Final = { @@ -294,6 +342,7 @@ class BaseAWSLLM: "aws_web_identity_token": aws_web_identity_token, "aws_sts_endpoint": aws_sts_endpoint, "aws_external_id": aws_external_id, + "aws_session_tags": session_tags, "ssl_verify": ssl_verify, } @@ -336,6 +385,7 @@ class BaseAWSLLM: aws_region_name=aws_region_name, aws_sts_endpoint=aws_sts_endpoint, aws_external_id=aws_external_id, + aws_session_tags=session_tags, ssl_verify=ssl_verify, ), ) @@ -980,6 +1030,7 @@ class BaseAWSLLM: aws_sts_endpoint: str | None = None, ssl_verify: bool | str | None = None, aws_region_name: str | None = None, + aws_session_tags: Sequence[AwsSessionTag] | None = None, ) -> dict: """Handle cross-account role assumption for IRSA.""" import boto3 @@ -1032,16 +1083,9 @@ class BaseAWSLLM: # Now assume the target role verbose_logger.debug("Attempting to assume target role: %s with session: %s", aws_role_name, aws_session_name) - assume_role_params: Final = { - "RoleArn": aws_role_name, - "RoleSessionName": aws_session_name, - } - - # Add ExternalId parameter if provided - if aws_external_id is not None: - assume_role_params["ExternalId"] = aws_external_id - - return sts_client_with_creds.assume_role(**assume_role_params) + return sts_client_with_creds.assume_role( + **_assume_role_params(aws_role_name, aws_session_name, aws_external_id, aws_session_tags) + ) def _handle_irsa_same_account( self, @@ -1051,6 +1095,7 @@ class BaseAWSLLM: aws_sts_endpoint: str | None = None, ssl_verify: bool | str | None = None, aws_region_name: str | None = None, + aws_session_tags: Sequence[AwsSessionTag] | None = None, ) -> dict: """Handle same-account role assumption for IRSA.""" import boto3 @@ -1074,16 +1119,9 @@ class BaseAWSLLM: # Assume the role verbose_logger.debug("Attempting to assume role: %s with session: %s", aws_role_name, aws_session_name) - assume_role_params: Final = { - "RoleArn": aws_role_name, - "RoleSessionName": aws_session_name, - } - - # Add ExternalId parameter if provided - if aws_external_id is not None: - assume_role_params["ExternalId"] = aws_external_id - - return sts_client.assume_role(**assume_role_params) + return sts_client.assume_role( + **_assume_role_params(aws_role_name, aws_session_name, aws_external_id, aws_session_tags) + ) def _extract_credentials_and_ttl(self, sts_response: dict) -> tuple[Credentials, int | None]: """Extract credentials and TTL from STS response. @@ -1118,6 +1156,7 @@ class BaseAWSLLM: aws_region_name: str | None, aws_sts_endpoint: str | None, aws_external_id: str | None, + aws_session_tags: tuple[AwsSessionTag, ...] | None, ssl_verify: bool | str | None, ) -> tuple[Credentials, int | None]: """ @@ -1144,6 +1183,7 @@ class BaseAWSLLM: aws_region_name=aws_region_name, aws_sts_endpoint=aws_sts_endpoint, aws_external_id=aws_external_id, + aws_session_tags=aws_session_tags, ssl_verify=ssl_verify, ) @@ -1159,6 +1199,7 @@ class BaseAWSLLM: aws_sts_endpoint: str | None = None, aws_external_id: str | None = None, ssl_verify: bool | str | None = None, + aws_session_tags: Sequence[AwsSessionTag] | None = None, ) -> tuple[Credentials, int | None]: """ Authenticate with AWS Role @@ -1189,6 +1230,7 @@ class BaseAWSLLM: aws_sts_endpoint=aws_sts_endpoint, ssl_verify=ssl_verify, aws_region_name=aws_region_name, + aws_session_tags=aws_session_tags, ) else: sts_response = self._handle_irsa_same_account( @@ -1198,6 +1240,7 @@ class BaseAWSLLM: aws_sts_endpoint=aws_sts_endpoint, ssl_verify=ssl_verify, aws_region_name=aws_region_name, + aws_session_tags=aws_session_tags, ) return self._extract_credentials_and_ttl(sts_response) @@ -1234,14 +1277,9 @@ class BaseAWSLLM: **sts_client_kwargs, ) - assume_role_params: Final = { - "RoleArn": aws_role_name, - "RoleSessionName": aws_session_name, - } - - # Add ExternalId parameter if provided - if aws_external_id is not None: - assume_role_params["ExternalId"] = aws_external_id + assume_role_params: Final = _assume_role_params( + aws_role_name, aws_session_name, aws_external_id, aws_session_tags + ) try: sts_response = sts_client.assume_role(**assume_role_params) @@ -1460,6 +1498,7 @@ class BaseAWSLLM: "aws_bedrock_runtime_endpoint", None ) # https://bedrock-runtime.{region_name}.amazonaws.com aws_external_id: Final = optional_params.pop("aws_external_id", None) + aws_session_tags: Final = optional_params.pop("aws_session_tags", None) if bearer_token is not None: return BearerRequestTarget( @@ -1478,6 +1517,7 @@ class BaseAWSLLM: aws_web_identity_token=aws_web_identity_token, aws_sts_endpoint=aws_sts_endpoint, aws_external_id=aws_external_id, + aws_session_tags=aws_session_tags, ) return Boto3CredentialsInfo( credentials=credentials, @@ -1623,6 +1663,7 @@ class BaseAWSLLM: aws_web_identity_token: Final = optional_params.get("aws_web_identity_token", None) aws_sts_endpoint: Final = optional_params.get("aws_sts_endpoint", None) aws_external_id: Final = optional_params.get("aws_external_id", None) + aws_session_tags: Final = optional_params.get("aws_session_tags", None) aws_region_name: Final = self._get_aws_region_name(optional_params=optional_params, model=model) credentials: Final[Credentials] = self.get_credentials( @@ -1636,6 +1677,7 @@ class BaseAWSLLM: aws_web_identity_token=aws_web_identity_token, aws_sts_endpoint=aws_sts_endpoint, aws_external_id=aws_external_id, + aws_session_tags=aws_session_tags, ) sigv4: Final = SigV4Auth(credentials, service_name, aws_region_name) diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index 4b500897642..f73d5f5d3cb 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -283,7 +283,7 @@ class BedrockBatchesHandler: ``aws_session_token``, ``aws_profile_name``, ``aws_role_name``, ``aws_session_name``, ``aws_web_identity_token``, ``aws_sts_endpoint``, - ``aws_external_id``). Unknown keys are ignored. + ``aws_external_id``, ``aws_session_tags``). Unknown keys are ignored. Returns: ``LiteLLMBatch`` shaped like an OpenAI Batch resource. @@ -317,6 +317,7 @@ class BedrockBatchesHandler: aws_web_identity_token=kwargs.get("aws_web_identity_token"), aws_sts_endpoint=kwargs.get("aws_sts_endpoint"), aws_external_id=kwargs.get("aws_external_id"), + aws_session_tags=kwargs.get("aws_session_tags"), ) client: Final = boto3.client( diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 984ba371898..573fc629383 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -355,6 +355,7 @@ class BedrockConverseLLM(BaseAWSLLM): aws_web_identity_token: Final = optional_params.pop("aws_web_identity_token", None) aws_sts_endpoint: Final = optional_params.pop("aws_sts_endpoint", None) aws_external_id: Final = optional_params.pop("aws_external_id", None) + aws_session_tags: Final = optional_params.pop("aws_session_tags", None) optional_params.pop("aws_region_name", None) litellm_params["aws_region_name"] = aws_region_name # [DO NOT DELETE] important for async calls @@ -373,6 +374,7 @@ class BedrockConverseLLM(BaseAWSLLM): aws_web_identity_token=aws_web_identity_token, aws_sts_endpoint=aws_sts_endpoint, aws_external_id=aws_external_id, + aws_session_tags=aws_session_tags, ) ) diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index be4f0f32689..cb2c70e74c8 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -93,6 +93,7 @@ _BEDROCK_AWS_AUTH_PARAMETER_KEYS: Final[tuple[str, ...]] = ( "aws_web_identity_token", "aws_sts_endpoint", "aws_external_id", + "aws_session_tags", ) @@ -1663,6 +1664,7 @@ class CommonBatchFilesUtils: aws_web_identity_token=optional_params.get("aws_web_identity_token"), aws_sts_endpoint=optional_params.get("aws_sts_endpoint"), aws_external_id=optional_params.get("aws_external_id"), + aws_session_tags=optional_params.get("aws_session_tags"), ) # Prepare the request data diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index e249feb9ff3..8ddc67c7f0e 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -73,6 +73,7 @@ class BedrockEmbedding(BaseAWSLLM): aws_web_identity_token: Final = optional_params.pop("aws_web_identity_token", None) aws_sts_endpoint: Final = optional_params.pop("aws_sts_endpoint", None) aws_external_id: Final = optional_params.pop("aws_external_id", None) + aws_session_tags: Final = optional_params.pop("aws_session_tags", None) ### SET REGION NAME ### if aws_region_name is None: @@ -103,6 +104,7 @@ class BedrockEmbedding(BaseAWSLLM): aws_web_identity_token=aws_web_identity_token, aws_sts_endpoint=aws_sts_endpoint, aws_external_id=aws_external_id, + aws_session_tags=aws_session_tags, ) ) return credentials, aws_region_name diff --git a/litellm/llms/sagemaker/chat/handler.py b/litellm/llms/sagemaker/chat/handler.py index 3f62b7276df..38978300c52 100644 --- a/litellm/llms/sagemaker/chat/handler.py +++ b/litellm/llms/sagemaker/chat/handler.py @@ -36,6 +36,7 @@ class SagemakerChatHandler(BaseAWSLLM): aws_web_identity_token: Final = optional_params.pop("aws_web_identity_token", None) aws_sts_endpoint: Final = optional_params.pop("aws_sts_endpoint", None) aws_external_id: Final = optional_params.pop("aws_external_id", None) + aws_session_tags: Final = optional_params.pop("aws_session_tags", None) ### SET REGION NAME ### if aws_region_name is None: @@ -63,6 +64,7 @@ class SagemakerChatHandler(BaseAWSLLM): aws_web_identity_token=aws_web_identity_token, aws_sts_endpoint=aws_sts_endpoint, aws_external_id=aws_external_id, + aws_session_tags=aws_session_tags, ) return credentials, aws_region_name diff --git a/litellm/llms/sagemaker/completion/handler.py b/litellm/llms/sagemaker/completion/handler.py index fb8074d3682..fad0a460647 100644 --- a/litellm/llms/sagemaker/completion/handler.py +++ b/litellm/llms/sagemaker/completion/handler.py @@ -59,6 +59,7 @@ class SagemakerLLM(BaseAWSLLM): aws_web_identity_token: Final = optional_params.pop("aws_web_identity_token", None) aws_sts_endpoint: Final = optional_params.pop("aws_sts_endpoint", None) aws_external_id: Final = optional_params.pop("aws_external_id", None) + aws_session_tags: Final = optional_params.pop("aws_session_tags", None) ### SET REGION NAME ### if aws_region_name is None: @@ -86,6 +87,7 @@ class SagemakerLLM(BaseAWSLLM): aws_web_identity_token=aws_web_identity_token, aws_sts_endpoint=aws_sts_endpoint, aws_external_id=aws_external_id, + aws_session_tags=aws_session_tags, ) return credentials, aws_region_name diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index f78c4221f5a..f498f2fe8f5 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -316,6 +316,7 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = ( "aws_profile_name", "aws_session_name", "aws_external_id", + "aws_session_tags", "vertex_credentials", # Azure managed-identity / federated-auth token. The Azure provider # transformer reads ``azure_ad_token`` (top-level or via diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 9f93886a9c6..76756ac35bb 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -1107,6 +1107,11 @@ class BedrockTag(TypedDict): value: str +class AwsSessionTag(TypedDict): + Key: str # writable-ok: boto3's STS stubs type assume_role Tags as writable TagTypeDef, which rejects ReadOnly + Value: str # writable-ok: boto3's STS stubs type assume_role Tags as writable TagTypeDef, which rejects ReadOnly + + class BedrockCreateBatchRequest(TypedDict, total=False): """ Request structure for creating a Bedrock batch inference job. diff --git a/litellm/types/router.py b/litellm/types/router.py index 5c9eab30f3d..30ee68f5545 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -4,7 +4,7 @@ litellm.Router Types - includes RouterConfig, UpdateRouterConfig, ModelInfo etc import datetime import enum -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, TypeVar, get_type_hints @@ -21,6 +21,7 @@ if TYPE_CHECKING: from .completion import CompletionRequest from .embedding import EmbeddingRequest +from .llms.bedrock import AwsSessionTag from .llms.openai import OpenAIFileObject from .search import SearchProvider from .utils import ( @@ -288,6 +289,7 @@ class CredentialLiteLLMParams(BaseModel): aws_web_identity_token: str | None = None aws_sts_endpoint: str | None = None aws_external_id: str | None = None + aws_session_tags: Sequence[AwsSessionTag] | None = None aws_bedrock_runtime_endpoint: str | None = None aws_bedrock_project_id: str | None = None s3_bucket_name: str | None = None diff --git a/tests/llm_translation/test_bedrock_dynamic_auth_params_unit_tests.py b/tests/llm_translation/test_bedrock_dynamic_auth_params_unit_tests.py index dad2fdbf065..6c059423f74 100644 --- a/tests/llm_translation/test_bedrock_dynamic_auth_params_unit_tests.py +++ b/tests/llm_translation/test_bedrock_dynamic_auth_params_unit_tests.py @@ -194,6 +194,7 @@ class DummyCredentials: ("aws_web_identity_token", "dummy_web_identity_token"), ("aws_sts_endpoint", "dummy_sts_endpoint"), ("aws_external_id", "dummy_external_id"), + ("aws_session_tags", [{"Key": "team", "Value": "genai"}]), ], ) def test_dynamic_aws_params_propagation(model, param_name, param_value): diff --git a/tests/test_litellm/llms/bedrock/batches/test_handler.py b/tests/test_litellm/llms/bedrock/batches/test_handler.py index 2a9b7a6d138..4ccddc33e77 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_handler.py +++ b/tests/test_litellm/llms/bedrock/batches/test_handler.py @@ -8,7 +8,7 @@ the tests don't hit AWS. from __future__ import annotations -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from unittest.mock import MagicMock, patch import pytest @@ -480,3 +480,53 @@ def test_litellm_cancel_batch_dispatches_to_bedrock(patched_boto3): fake_client.stop_model_invocation_job.assert_called_once_with(jobIdentifier=JOB_ARN) assert batch.status == "cancelled" + + +def test_handle_model_invocation_job_status_builds_the_client_from_the_tagged_session(monkeypatch): + """Status polling must assume the role with the deployment's session tags, like every other call.""" + from botocore.exceptions import ClientError + + monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False) + monkeypatch.delenv("AWS_ROLE_ARN", raising=False) + tags = [{"Key": "team", "Value": "genai"}] + bedrock_client_kwargs: list[dict] = [] + fake_bedrock = MagicMock() + fake_bedrock.get_model_invocation_job.return_value = _fake_boto3_response() + + class FakeSTSClient: + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + if list(params.get("Tags", ())) != tags: + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:TagSession"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": "ASIABATCHSTATUSTAGGED", + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.now(timezone.utc) + timedelta(minutes=30), + } + } + + def boto3_client(service_name, **kwargs): + if service_name == "sts": + return FakeSTSClient() + bedrock_client_kwargs.append(kwargs) + return fake_bedrock + + with patch("boto3.client", side_effect=boto3_client): + batch = BedrockBatchesHandler._handle_model_invocation_job_status( + batch_id=JOB_ARN, + aws_access_key_id="AKIABATCHSTATUSCALLER", + aws_secret_access_key="pod-caller-secret", + aws_role_name="arn:aws:iam::999999999999:role/litellm-batch-role", + aws_session_name="litellm-batch-session", + aws_session_tags=tags, + ) + + assert batch.status == "completed" + assert [kwargs["aws_access_key_id"] for kwargs in bedrock_client_kwargs] == ["ASIABATCHSTATUSTAGGED"] diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py index 531f334e460..bcd1df26020 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py @@ -15,6 +15,7 @@ AWS_AUTH_PARAMS = { "aws_sts_endpoint": "https://sts.us-west-2.amazonaws.com", "aws_bedrock_runtime_endpoint": "https://bedrock-runtime.us-west-2.amazonaws.com", "aws_external_id": "external", + "aws_session_tags": [{"Key": "team", "Value": "genai"}], } diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index cbf160c451f..bcba4bf7711 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -57,6 +57,7 @@ def test_aws_params_filtered_from_request_body(): "aws_sts_endpoint": "https://sts.amazonaws.com", "aws_bedrock_runtime_endpoint": "https://bedrock-runtime.us-west-2.amazonaws.com", "aws_external_id": "external-id-123", + "aws_session_tags": [{"Key": "team", "Value": "genai"}], } # Transform the request @@ -105,6 +106,9 @@ def test_aws_params_filtered_from_request_body(): assert ( "aws_external_id" not in result_json ), "AWS external ID should not be in request body" + assert ( + "aws_session_tags" not in result_json + ), "AWS session tags should not be in request body" # Also check that the sensitive values themselves are not in the response assert ( diff --git a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py index f34b8eb1fb9..92bd6d48aae 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py +++ b/tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py @@ -6,12 +6,15 @@ extension, and AWS credential resolution is stubbed so nothing reaches STS. from __future__ import annotations +from datetime import datetime, timedelta, timezone from unittest.mock import MagicMock, patch +import boto3 import httpx import pytest from botocore.credentials import Credentials +from botocore.exceptions import ClientError from litellm.llms.bedrock.chat.converse_handler import BedrockConverseLLM from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.rust_bridge import chat_completions as bridge @@ -541,3 +544,54 @@ def test_bearer_token_auth_never_runs_the_sigv4_credential_chain(monkeypatch, co assert response.choices[0].message.content == "hi" assert client.post.call_args.kwargs["headers"]["Authorization"] == "Bearer bedrock-bearer-token" + + +def test_session_tags_sign_the_request_and_stay_out_of_the_body(monkeypatch): + """The tagged STS session signs the Converse call and the tags never reach the request body (#34069).""" + monkeypatch.setenv("LITELLM_RUST", "0") + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False) + monkeypatch.delenv("AWS_ROLE_ARN", raising=False) + tags = [{"Key": "team", "Value": "genai"}] + + class FakeSTSClient: + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + if list(params.get("Tags", ())) != tags: + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:TagSession"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": "ASIACONVERSETAGGED", + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.now(timezone.utc) + timedelta(minutes=30), + } + } + + client = _sync_client_returning_converse_response() + with patch.object(boto3, "client", return_value=FakeSTSClient()): + response = BedrockConverseLLM().completion( + **_completion_kwargs( + optional_params={ + "maxTokens": 16, + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIACONVERSECALLER", + "aws_secret_access_key": "pod-caller-secret", + "aws_role_name": "arn:aws:iam::999999999999:role/litellm-converse-role", + "aws_session_name": "litellm-converse-session", + "aws_session_tags": tags, + }, + litellm_params={}, + client=client, + ) + ) + + assert response.choices[0].message.content == "hi" + sent = client.post.call_args.kwargs + assert "Credential=ASIACONVERSETAGGED/" in sent["headers"]["Authorization"] + assert "aws_session_tags" not in sent["data"] diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py index bcd1a29d0e8..1f8c9e77165 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py @@ -1036,6 +1036,63 @@ def test_load_credentials_assumes_role_with_external_id(monkeypatch): assert "aws_external_id" not in optional_params +def test_embedding_session_tags_sign_the_request_and_stay_out_of_the_body(monkeypatch): + """The tagged STS session signs the InvokeModel call and the tags never reach the body (#34069).""" + import datetime + + import boto3 + from botocore.exceptions import ClientError + + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) + monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False) + monkeypatch.delenv("AWS_ROLE_ARN", raising=False) + tags = [{"Key": "team", "Value": "genai"}] + + class FakeSTSClient: + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + if list(params.get("Tags", ())) != tags: + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:TagSession"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": "ASIAEMBEDTAGGED", + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30), + } + } + + client = HTTPHandler() + with patch.object(boto3, "client", return_value=FakeSTSClient()), patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(titan_embedding_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_post.return_value = mock_response + + response = litellm.embedding( + model="bedrock/amazon.titan-embed-text-v1", + input=test_input, + client=client, + aws_region_name="us-east-1", + aws_access_key_id="AKIAEMBEDCALLERKEY", + aws_secret_access_key="pod-caller-secret", + aws_role_name="arn:aws:iam::999999999999:role/litellm-embed-role", + aws_session_name="litellm-embed-session", + aws_session_tags=tags, + ) + + assert response.data[0]["embedding"] == titan_embedding_response["embedding"] + sent = mock_post.call_args.kwargs + assert "Credential=ASIAEMBEDTAGGED/" in sent["headers"]["Authorization"] + assert "aws_session_tags" not in sent["data"] + + def test_bedrock_embedding_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch): """The deployment's AWS profile does not exist, so resolving SigV4 credentials raises; a bearer-token deployment must still serve the request, since the 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 f854d806bdc..1aa9cb3a9db 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -15,7 +15,7 @@ from unittest.mock import MagicMock, patch from botocore.awsrequest import AWSPreparedRequest, AWSRequest from botocore.auth import SigV4Auth from botocore.credentials import Credentials -from botocore.exceptions import NoCredentialsError +from botocore.exceptions import ClientError, NoCredentialsError import litellm from litellm.llms.bedrock.base_aws_llm import ( @@ -2395,6 +2395,283 @@ def test_assume_role_without_external_id(): ) +_SESSION_TAGS = ({"Key": "team", "Value": "genai"}, {"Key": "env", "Value": "prod"}) +_SORTED_SESSION_TAGS = ({"Key": "env", "Value": "prod"}, {"Key": "team", "Value": "genai"}) +_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: + self.expected_tags = expected_tags + self.access_key = access_key + self.assume_role_calls: list = [] + + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role_with_web_identity(self, **params): + return { + "Credentials": { + "AccessKeyId": "ASIAIRSATEMP", + "SecretAccessKey": "irsa-temp-secret-key", + "SessionToken": "irsa-temp-session-token", + "Expiration": datetime.now(timezone.utc) + timedelta(hours=1), + } + } + + def assume_role(self, **params): + self.assume_role_calls.append(params) + if tuple(params.get("Tags", ())) != self.expected_tags: + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:TagSession"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": self.access_key, + "SecretAccessKey": "assumed-secret-key", + "SessionToken": "assumed-session-token", + "Expiration": datetime.now(timezone.utc) + timedelta(hours=1), + } + } + + +def _irsa_env(tmp_path, irsa_role_arn: str) -> dict: + token_file = tmp_path / "web-identity-token" + token_file.write_text("test-web-identity-token") + return { + "AWS_WEB_IDENTITY_TOKEN_FILE": str(token_file), + "AWS_ROLE_ARN": irsa_role_arn, + "AWS_REGION": "us-east-1", + } + + +def test_assume_role_sends_session_tags(): + """The STS session carries the configured tags, so a trust policy gated on sts:TagSession admits it.""" + sts = _TagAwareSTSClient(expected_tags=_SESSION_TAGS) + + with patch("boto3.client", return_value=sts): + credentials, _ttl = BaseAWSLLM()._auth_with_aws_role( + aws_access_key_id=None, + aws_secret_access_key=None, + aws_session_token=None, + aws_role_name=_TAGGED_ROLE_ARN, + aws_session_name="test-session", + aws_session_tags=list(_SESSION_TAGS), + ) + + assert credentials.access_key == "ASIATAGGEDSESSION" + assert sts.assume_role_calls == [ + {"RoleArn": _TAGGED_ROLE_ARN, "RoleSessionName": "test-session", "Tags": _SESSION_TAGS} + ] + + +def test_assume_role_sends_session_tags_alongside_external_id(): + sts = _TagAwareSTSClient(expected_tags=_SESSION_TAGS) + + with patch("boto3.client", return_value=sts): + credentials, _ttl = BaseAWSLLM()._auth_with_aws_role( + aws_access_key_id=None, + aws_secret_access_key=None, + aws_session_token=None, + aws_role_name=_TAGGED_ROLE_ARN, + aws_session_name="test-session", + aws_external_id="UniqueExternalID123", + aws_session_tags=_SESSION_TAGS, + ) + + assert credentials.access_key == "ASIATAGGEDSESSION" + assert sts.assume_role_calls == [ + { + "RoleArn": _TAGGED_ROLE_ARN, + "RoleSessionName": "test-session", + "ExternalId": "UniqueExternalID123", + "Tags": _SESSION_TAGS, + } + ] + + +@pytest.mark.parametrize("aws_session_tags", [None, [], ()], ids=["none", "empty-list", "empty-tuple"]) +def test_assume_role_omits_the_tags_key_without_session_tags(aws_session_tags): + """Nothing configured means the AssumeRole request looks exactly as it did before tags existed.""" + sts = _TagAwareSTSClient(expected_tags=()) + + with patch("boto3.client", return_value=sts): + credentials, _ttl = BaseAWSLLM()._auth_with_aws_role( + aws_access_key_id=None, + aws_secret_access_key=None, + aws_session_token=None, + aws_role_name=_TAGGED_ROLE_ARN, + aws_session_name="test-session", + aws_session_tags=aws_session_tags, + ) + + assert credentials.access_key == "ASIATAGGEDSESSION" + assert sts.assume_role_calls == [{"RoleArn": _TAGGED_ROLE_ARN, "RoleSessionName": "test-session"}] + + +def test_irsa_cross_account_assume_role_sends_session_tags(tmp_path): + irsa_role_arn = "arn:aws:iam::111111111111:role/eks-service-account-role" + sts = _TagAwareSTSClient(expected_tags=_SESSION_TAGS) + + with patch.dict(os.environ, _irsa_env(tmp_path, irsa_role_arn)), patch("boto3.client", return_value=sts): + credentials, _ttl = BaseAWSLLM()._auth_with_aws_role( + aws_access_key_id=None, + aws_secret_access_key=None, + aws_session_token=None, + aws_role_name=_TAGGED_ROLE_ARN, + aws_session_name="test-session", + aws_session_tags=_SESSION_TAGS, + ) + + assert credentials.access_key == "ASIATAGGEDSESSION" + assert sts.assume_role_calls == [ + {"RoleArn": _TAGGED_ROLE_ARN, "RoleSessionName": "test-session", "Tags": _SESSION_TAGS} + ] + + +def test_irsa_same_account_assume_role_sends_session_tags(tmp_path): + sts = _TagAwareSTSClient(expected_tags=_SESSION_TAGS) + + with patch.dict(os.environ, _irsa_env(tmp_path, _TAGGED_ROLE_ARN)), patch("boto3.client", return_value=sts): + credentials, _ttl = BaseAWSLLM()._auth_with_aws_role( + aws_access_key_id=None, + aws_secret_access_key=None, + aws_session_token=None, + aws_role_name=_TAGGED_ROLE_ARN, + aws_session_name="test-session", + aws_session_tags=_SESSION_TAGS, + ) + + assert credentials.access_key == "ASIATAGGEDSESSION" + assert sts.assume_role_calls == [ + {"RoleArn": _TAGGED_ROLE_ARN, "RoleSessionName": "test-session", "Tags": _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() + sts = _TagAwareSTSClient(expected_tags=_SORTED_SESSION_TAGS) + + with patch.dict(os.environ, _os_environ_without_aws_keys(), clear=True), patch("boto3.client", return_value=sts): + first = base_aws_llm.get_credentials( + aws_role_name=_TAGGED_ROLE_ARN, + aws_session_name="team-session", + aws_session_tags=list(_SESSION_TAGS), + ) + second = base_aws_llm.get_credentials( + aws_role_name=_TAGGED_ROLE_ARN, + aws_session_name="team-session", + aws_session_tags=list(reversed(_SESSION_TAGS)), + ) + + assert first.access_key == second.access_key == "ASIATAGGEDSESSION" + assert sts.assume_role_calls == [ + {"RoleArn": _TAGGED_ROLE_ARN, "RoleSessionName": "team-session", "Tags": _SORTED_SESSION_TAGS} + ] + + +def test_get_credentials_scopes_the_cache_per_session_tag_set(): + """Different tag sets are different principals to AWS, so each gets its own STS session.""" + base_aws_llm = BaseAWSLLM() + mock_sts_client = _assume_role_sts_mock() + mock_sts_client.assume_role.side_effect = [ + { + "Credentials": { + "AccessKeyId": f"assumed-access-key-{team}", + "SecretAccessKey": "assumed-secret-key", + "SessionToken": f"assumed-session-token-{team}", + "Expiration": datetime.now(timezone.utc) + timedelta(hours=1), + } + } + for team in ("genai", "platform") + ] + + with ( + patch.dict(os.environ, _os_environ_without_aws_keys(), clear=True), + patch("boto3.client", return_value=mock_sts_client), + ): + genai = base_aws_llm.get_credentials( + aws_role_name=_TAGGED_ROLE_ARN, + aws_session_name="team-session", + aws_session_tags=[{"Key": "team", "Value": "genai"}], + ) + platform = base_aws_llm.get_credentials( + aws_role_name=_TAGGED_ROLE_ARN, + aws_session_name="team-session", + aws_session_tags=[{"Key": "team", "Value": "platform"}], + ) + + assert genai.access_key == "assumed-access-key-genai" + assert platform.access_key == "assumed-access-key-platform" + assert [call.kwargs["Tags"] for call in mock_sts_client.assume_role.call_args_list] == [ + ({"Key": "team", "Value": "genai"},), + ({"Key": "team", "Value": "platform"},), + ] + + +@pytest.mark.parametrize( + "aws_session_tags", + [ + "team=genai", + {"team": "genai"}, + [["team", "genai"]], + [{"key": "team", "value": "genai"}], + [{"Key": "team"}], + [{"Key": 1, "Value": "genai"}], + ], + ids=["string", "flat-dict", "pair-list", "lowercase-keys", "missing-value", "non-string-key"], +) +def test_get_credentials_rejects_malformed_session_tags(aws_session_tags): + with pytest.raises(ValueError, match="Invalid 'aws_session_tags' value"): + BaseAWSLLM().get_credentials( + aws_role_name=_TAGGED_ROLE_ARN, + aws_session_name="team-session", + aws_session_tags=aws_session_tags, + ) + + +def test_get_boto_credentials_from_optional_params_consumes_session_tags(): + """Tags feed the STS call and must not linger in optional_params to be serialized into the body.""" + sts = _TagAwareSTSClient(expected_tags=_SORTED_SESSION_TAGS) + optional_params = { + "aws_region_name": "us-east-1", + "aws_role_name": _TAGGED_ROLE_ARN, + "aws_session_name": "team-session", + "aws_session_tags": list(_SESSION_TAGS), + } + + with patch.dict(os.environ, _os_environ_without_aws_keys(), clear=True), patch("boto3.client", return_value=sts): + target = BaseAWSLLM()._get_boto_credentials_from_optional_params(optional_params) + + assert target.credentials.access_key == "ASIATAGGEDSESSION" + assert "aws_session_tags" not in optional_params + + +def test_sign_request_signs_with_the_tagged_sts_session(): + sts = _TagAwareSTSClient(expected_tags=_SORTED_SESSION_TAGS) + optional_params = { + "aws_region_name": "us-east-1", + "aws_role_name": _TAGGED_ROLE_ARN, + "aws_session_name": "team-session", + "aws_session_tags": list(_SESSION_TAGS), + } + + with patch.dict(os.environ, _os_environ_without_aws_keys(), clear=True), patch("boto3.client", return_value=sts): + headers, _body = BaseAWSLLM()._sign_request( + service_name="bedrock", + headers={}, + optional_params=optional_params, + request_data={"prompt": "hi"}, + api_base="https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-opus-5/invoke", + ) + + assert "Credential=ASIATAGGEDSESSION/" in headers["Authorization"] + + def test_converse_handler_external_id_extraction(): """Test that BedrockConverseLLM properly extracts and passes aws_external_id parameter""" from litellm.llms.bedrock.chat.converse_handler import BedrockConverseLLM diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py index 3f03305423a..a8a21e2cd37 100644 --- a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py @@ -486,6 +486,7 @@ def test_merge_bedrock_aws_request_params_strips_caller_identity_when_deployment "aws_role_name": "arn:aws:iam::123456789012:role/caller", "aws_session_token": "caller-token", "aws_web_identity_token": "caller-web-identity", + "aws_session_tags": [{"Key": "team", "Value": "caller-chosen"}], "timeout": 600, }, ) @@ -500,6 +501,7 @@ def test_merge_bedrock_aws_request_params_strips_caller_identity_when_deployment "aws_role_name", "aws_session_token", "aws_web_identity_token", + "aws_session_tags", ): assert stripped not in merged @@ -616,6 +618,60 @@ def test_sign_aws_request_assumes_role_with_external_id(monkeypatch): assert signed_data == b'{"jobName": "litellm-batch-job"}' +def test_sign_aws_request_assumes_role_with_session_tags(monkeypatch): + """Batch and file signing must carry the deployment's session tags into the AssumeRole call too.""" + import datetime + from unittest.mock import patch + + import boto3 + from botocore.exceptions import ClientError + + from litellm.llms.bedrock.common_utils import CommonBatchFilesUtils + + monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False) + monkeypatch.delenv("AWS_ROLE_ARN", raising=False) + tags = [{"Key": "team", "Value": "genai"}] + + class FakeSTSClient: + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + if list(params.get("Tags", ())) != tags: + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:TagSession"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": "ASIABATCHSIGNTAGGED", + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30), + } + } + + optional_params = { + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIABATCHSIGNCALLER", + "aws_secret_access_key": "pod-caller-secret", + "aws_role_name": "arn:aws:iam::999999999999:role/litellm-batch-sign-role", + "aws_session_name": "litellm-batch-sign-session", + "aws_session_tags": tags, + } + + with patch.object(boto3, "client", return_value=FakeSTSClient()): + signed_headers, _signed_data = CommonBatchFilesUtils().sign_aws_request( + service_name="bedrock", + data={"jobName": "litellm-batch-job"}, + endpoint_url="https://bedrock.us-east-1.amazonaws.com/model-invocation-job", + optional_params=optional_params, + ) + + authorization = {key.lower(): value for key, value in signed_headers.items()}["authorization"] + assert "Credential=ASIABATCHSIGNTAGGED/" in authorization + + # --------------------------------------------------------------------------- # # Provider error headers (LIT-5428) # # --------------------------------------------------------------------------- # diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_handler.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_handler.py index aa1a59d0e5c..bb891c06fa2 100644 --- a/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_handler.py +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_handler.py @@ -46,3 +46,45 @@ def test_load_credentials_assumes_role_with_external_id(monkeypatch): assert credentials.token == "assumed-session-token" assert aws_region_name == "us-east-1" assert "aws_external_id" not in optional_params + + +def test_load_credentials_assumes_role_with_session_tags(monkeypatch): + """A trust policy gated on sts:TagSession only admits the session when the deployment's tags are sent.""" + monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False) + monkeypatch.delenv("AWS_ROLE_ARN", raising=False) + tags = [{"Key": "team", "Value": "genai"}] + + class FakeSTSClient: + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + if list(params.get("Tags", ())) != tags: + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:TagSession"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": "ASIASMCHATTAGGED", + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30), + } + } + + optional_params = { + "aws_access_key_id": "AKIASMCHATCALLERKEY", + "aws_secret_access_key": "pod-caller-secret", + "aws_region_name": "us-east-1", + "aws_role_name": "arn:aws:iam::999999999999:role/litellm-sm-chat-role", + "aws_session_name": "litellm-sm-chat-session", + "aws_session_tags": tags, + } + + with patch.object(boto3, "client", return_value=FakeSTSClient()): + credentials, aws_region_name = SagemakerChatHandler()._load_credentials(optional_params) + + assert credentials.access_key == "ASIASMCHATTAGGED" + assert aws_region_name == "us-east-1" + assert "aws_session_tags" not in optional_params diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py index 881bac096b1..9a5f36e2081 100644 --- a/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py @@ -219,3 +219,51 @@ def test_load_credentials_assumes_role_with_external_id(monkeypatch): assert credentials.token == "assumed-session-token" assert aws_region_name == "us-east-1" assert "aws_external_id" not in optional_params + + +def test_load_credentials_assumes_role_with_session_tags(monkeypatch): + """A trust policy gated on sts:TagSession only admits the session when the deployment's tags are sent.""" + import datetime + + import boto3 + from botocore.exceptions import ClientError + from unittest.mock import patch + + monkeypatch.delenv("AWS_WEB_IDENTITY_TOKEN_FILE", raising=False) + monkeypatch.delenv("AWS_ROLE_ARN", raising=False) + tags = [{"Key": "team", "Value": "genai"}] + + class FakeSTSClient: + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + if list(params.get("Tags", ())) != tags: + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:TagSession"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": "ASIASMCOMPTAGGED", + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30), + } + } + + optional_params = { + "aws_access_key_id": "AKIASMCOMPCALLERKEY", + "aws_secret_access_key": "pod-caller-secret", + "aws_region_name": "us-east-1", + "aws_role_name": "arn:aws:iam::999999999999:role/litellm-sm-completion-role", + "aws_session_name": "litellm-sm-completion-session", + "aws_session_tags": tags, + } + + with patch.object(boto3, "client", return_value=FakeSTSClient()): + credentials, aws_region_name = SagemakerLLM()._load_credentials(optional_params) + + assert credentials.access_key == "ASIASMCOMPTAGGED" + assert aws_region_name == "us-east-1" + assert "aws_session_tags" not in optional_params diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index aaf630ad29b..97cda730e04 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -3497,7 +3497,7 @@ class TestIsRequestBodySafeBlocksAwsIdentitySelectors: @pytest.mark.parametrize( "selector", - ["aws_profile_name", "aws_session_name", "aws_external_id"], + ["aws_profile_name", "aws_session_name", "aws_external_id", "aws_session_tags"], ) def test_aws_identity_selector_in_batch_body_is_rejected(self, selector): with pytest.raises(ValueError, match=selector): @@ -3516,7 +3516,7 @@ class TestIsRequestBodySafeBlocksAwsIdentitySelectors: @pytest.mark.parametrize( "selector", - ["aws_profile_name", "aws_session_name", "aws_external_id"], + ["aws_profile_name", "aws_session_name", "aws_external_id", "aws_session_tags"], ) def test_aws_identity_selector_under_extra_body_is_rejected(self, selector): with pytest.raises(ValueError, match=selector): diff --git a/tests/test_litellm/types/test_router.py b/tests/test_litellm/types/test_router.py index fd933a9d993..7fcb76b638f 100644 --- a/tests/test_litellm/types/test_router.py +++ b/tests/test_litellm/types/test_router.py @@ -1,6 +1,7 @@ import logging import pytest +from pydantic import ValidationError from litellm.types.router import ( SPECIAL_MODEL_INFO_PARAMS, @@ -122,3 +123,26 @@ def test_drop_params_flags_and_strings_log_nothing(value, caplog): with caplog.at_level(logging.WARNING, logger="LiteLLM"): GenericLiteLLMParams(drop_params=value) assert caplog.text == "" + + +def test_aws_session_tags_round_trip_as_sts_shaped_pairs(): + """The deployment field keeps the exact Key/Value shape STS AssumeRole expects.""" + params = LiteLLM_Params( + model="bedrock/anthropic.claude-opus-5", + aws_session_tags=[{"Key": "team", "Value": "genai"}, {"Key": "env", "Value": "prod"}], + ) + + assert params.model_dump(exclude_none=True)["aws_session_tags"] == [ + {"Key": "team", "Value": "genai"}, + {"Key": "env", "Value": "prod"}, + ] + + +@pytest.mark.parametrize( + "aws_session_tags", + ["team=genai", {"team": "genai"}, [{"key": "team", "value": "genai"}], [{"Key": "team"}]], + ids=["string", "flat-dict", "lowercase-keys", "missing-value"], +) +def test_aws_session_tags_reject_shapes_sts_would_refuse(aws_session_tags): + with pytest.raises(ValidationError, match="aws_session_tags"): + LiteLLM_Params(model="bedrock/anthropic.claude-opus-5", aws_session_tags=aws_session_tags) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 83b0d58f2b2..1009655bedc 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -23708,6 +23708,15 @@ export interface components { /** @description The decision record this request would have written to its log row */ routing_decision: components["schemas"]["StandardLoggingRoutingDecision"]; }; + /** AwsSessionTag */ + AwsSessionTag: { + /** Key */ + Key: string; + /** Value */ + Value: string; + } & { + [key: string]: unknown; + }; /** BaseLitellmParams */ BaseLitellmParams: { /** @@ -29426,6 +29435,8 @@ export interface components { aws_secret_access_key?: string | null; /** Aws Session Name */ aws_session_name?: string | null; + /** Aws Session Tags */ + aws_session_tags?: components["schemas"]["AwsSessionTag"][] | null; /** Aws Session Token */ aws_session_token?: string | null; /** Aws Sts Endpoint */ @@ -39617,6 +39628,8 @@ export interface components { aws_secret_access_key?: string | null; /** Aws Session Name */ aws_session_name?: string | null; + /** Aws Session Tags */ + aws_session_tags?: components["schemas"]["AwsSessionTag"][] | null; /** Aws Session Token */ aws_session_token?: string | null; /** Aws Sts Endpoint */