scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

feat(bedrock): thread aws_session_tags into STS AssumeRole

Browse source 
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>
This commit is contained in:
ryan-crabbe-berri 2026-09-09 12:37:36 -07:00
parent e8140eb269
commit 7b43977460
24 changed files with 728 additions and 39 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/litellm_core_utils/get_litellm_params.py
View file

@ -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",
}

108
litellm/llms/bedrock/base_aws_llm.py
View file

@ -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': <str>, 'Value': <str>} 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)

3
litellm/llms/bedrock/batches/handler.py
View file

@ -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(

2
litellm/llms/bedrock/chat/converse_handler.py
View file

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

2
litellm/llms/bedrock/common_utils.py
View file

@ -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

2
litellm/llms/bedrock/embed/embedding.py
View file

@ -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

2
litellm/llms/sagemaker/chat/handler.py
View file

@ -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

2
litellm/llms/sagemaker/completion/handler.py
View file

@ -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

1
litellm/proxy/auth/auth_utils.py
View file

@ -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

5
litellm/types/llms/bedrock.py
View file

@ -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.

4
litellm/types/router.py
View file

@ -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

1
tests/llm_translation/test_bedrock_dynamic_auth_params_unit_tests.py
View file

@ -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):

52
tests/test_litellm/llms/bedrock/batches/test_handler.py
View file

@ -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"]

1
tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_moonshot_transformation.py
View file

@ -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"}],
}

4
tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py
View file

@ -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 (

54
tests/test_litellm/llms/bedrock/chat/test_bedrock_converse_handler.py
View file

@ -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"]

57
tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py
View file

@ -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

279
tests/test_litellm/llms/bedrock/test_base_aws_llm.py
View file

@ -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

56
tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py
View file

@ -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) #
# --------------------------------------------------------------------------- #

42
tests/test_litellm/llms/sagemaker/test_sagemaker_chat_handler.py
View file

@ -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

48
tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py
View file

@ -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

4
tests/test_litellm/proxy/auth/test_auth_utils.py
View file

@ -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):

24
tests/test_litellm/types/test_router.py
View file

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

13
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -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 */