From 7b43977460249021a89ddfb285c83ce361ff90b1 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 9 Sep 2026 12:37:36 -0700 Subject: [PATCH] feat(bedrock): thread aws_session_tags into STS AssumeRole Operators can now set aws_session_tags on a Bedrock or SageMaker deployment and every AssumeRole call carries them as STS session tags, so trust policies gated on sts:TagSession admit the session and CloudTrail and cost reports see the tags. Tags are validated up front, sorted into the credential cache key so tag order does not fork sessions, stripped from invoke and embedding request bodies, and blocked from client request bodies like the other AWS identity params. Based on #34073. Co-authored-by: Satya Yedida <222552552+satya-rubrik@users.noreply.github.com> --- .../litellm_core_utils/get_litellm_params.py | 1 + litellm/llms/bedrock/base_aws_llm.py | 108 ++++--- litellm/llms/bedrock/batches/handler.py | 3 +- litellm/llms/bedrock/chat/converse_handler.py | 2 + litellm/llms/bedrock/common_utils.py | 2 + litellm/llms/bedrock/embed/embedding.py | 2 + litellm/llms/sagemaker/chat/handler.py | 2 + litellm/llms/sagemaker/completion/handler.py | 2 + litellm/proxy/auth/auth_utils.py | 1 + litellm/types/llms/bedrock.py | 5 + litellm/types/router.py | 4 +- ..._bedrock_dynamic_auth_params_unit_tests.py | 1 + .../llms/bedrock/batches/test_handler.py | 52 +++- .../test_amazon_moonshot_transformation.py | 1 + ...ations_anthropic_claude3_transformation.py | 4 + .../chat/test_bedrock_converse_handler.py | 54 ++++ .../bedrock/embed/test_bedrock_embedding.py | 57 ++++ .../llms/bedrock/test_base_aws_llm.py | 279 +++++++++++++++++- .../llms/bedrock/test_bedrock_common_utils.py | 56 ++++ .../sagemaker/test_sagemaker_chat_handler.py | 42 +++ .../test_sagemaker_completion_handler.py | 48 +++ .../proxy/auth/test_auth_utils.py | 4 +- tests/test_litellm/types/test_router.py | 24 ++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 13 + 24 files changed, 728 insertions(+), 39 deletions(-) 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 */