From 31f293a9fc60a5ef7ff8c40ebfe4eab3fc5d11f2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 17 Jul 2026 13:49:02 -0400 Subject: [PATCH 1/3] feat(bedrock): forward bedrock_tags to CreateModelInvocationJob for batch jobs --- .../llms/bedrock/batches/transformation.py | 18 ++++ litellm/types/llms/bedrock.py | 7 +- .../bedrock/batches/test_transformation.py | 87 +++++++++++++++++++ 3 files changed, 111 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 4fcf7cf91cb..8648d6586e8 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -4,6 +4,7 @@ import time from typing import Any, Dict, List, Literal, Optional, Union, cast from httpx import Headers, Response +from pydantic import TypeAdapter, ValidationError from litellm.litellm_core_utils.cloud_storage_security import ( BEDROCK_MANAGED_S3_BATCH_PREFIX, @@ -19,6 +20,7 @@ from litellm.types.llms.bedrock import ( BedrockOutputDataConfig, BedrockS3InputDataConfig, BedrockS3OutputDataConfig, + BedrockTag, ) from litellm.types.llms.openai import ( AllMessageValues, @@ -38,6 +40,18 @@ _S3_BATCH_FILE_UUID_SUFFIX_PATTERN = re.compile( r"-[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}\.jsonl$" ) +_BEDROCK_TAGS_ADAPTER: TypeAdapter[list[BedrockTag]] = TypeAdapter(list[BedrockTag]) + + +def _validate_bedrock_tags(raw_tags: object) -> list[BedrockTag]: + try: + return _BEDROCK_TAGS_ADAPTER.validate_python(raw_tags, strict=True) + except ValidationError as e: + raise ValueError( + "Invalid 'bedrock_tags' value. Expected a list of {'key': , 'value': } dicts, " + f"e.g. [{{'key': 'team', 'value': 'genai'}}]. Got: {raw_tags!r}" + ) from e + class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): """ @@ -201,6 +215,10 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): "roleArn": role_arn, } + bedrock_tags = litellm_params.get("bedrock_tags") or optional_params.get("bedrock_tags") + if bedrock_tags is not None: + bedrock_request["tags"] = _validate_bedrock_tags(bedrock_tags) + # Add optional parameters if provided completion_window = create_batch_data.get("completion_window") if completion_window: diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index bdf6b8fefed..d9f8229dbed 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -985,6 +985,11 @@ class BedrockOutputDataConfig(TypedDict): s3OutputDataConfig: BedrockS3OutputDataConfig +class BedrockTag(TypedDict): + key: str + value: str + + class BedrockCreateBatchRequest(TypedDict, total=False): """ Request structure for creating a Bedrock batch inference job. @@ -999,7 +1004,7 @@ class BedrockCreateBatchRequest(TypedDict, total=False): outputDataConfig: BedrockOutputDataConfig timeoutDurationInHours: Optional[int] clientRequestToken: Optional[str] - tags: Optional[List[dict]] + tags: Optional[List[BedrockTag]] BedrockBatchJobStatus = Literal["Submitted", "InProgress", "Completed", "Failed", "Stopping", "Stopped"] diff --git a/tests/test_litellm/llms/bedrock/batches/test_transformation.py b/tests/test_litellm/llms/bedrock/batches/test_transformation.py index d1ad5943ae6..b38d271e210 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_transformation.py +++ b/tests/test_litellm/llms/bedrock/batches/test_transformation.py @@ -258,6 +258,93 @@ def test_create_request_no_timeout_for_non_24h_window(config): assert "timeoutDurationInHours" not in mock_sign.call_args.kwargs["data"] +def test_create_request_forwards_bedrock_tags_from_litellm_params(config): + tags = [ + {"key": "application", "value": "genai-proxy"}, + {"key": "team", "value": "ml-platform"}, + ] + with patch.object( + config.common_utils, + "generate_unique_job_name", + return_value="litellm-batch-1", + ), patch.object(config.common_utils, "sign_aws_request") as mock_sign: + mock_sign.return_value = ({}, b"{}") + config.transform_create_batch_request( + model="m", + create_batch_data={"input_file_id": "s3://b/in.jsonl"}, + optional_params={}, + litellm_params={ + "aws_batch_role_arn": "arn:aws:iam::1:role/r", + "bedrock_tags": tags, + }, + ) + assert mock_sign.call_args.kwargs["data"]["tags"] == tags + + +def test_create_request_forwards_bedrock_tags_from_optional_params(config): + tags = [{"key": "env", "value": "prod"}] + with patch.object( + config.common_utils, + "generate_unique_job_name", + return_value="litellm-batch-1", + ), patch.object(config.common_utils, "sign_aws_request") as mock_sign: + mock_sign.return_value = ({}, b"{}") + config.transform_create_batch_request( + model="m", + create_batch_data={"input_file_id": "s3://b/in.jsonl"}, + optional_params={"bedrock_tags": tags}, + litellm_params={"aws_batch_role_arn": "arn:aws:iam::1:role/r"}, + ) + assert mock_sign.call_args.kwargs["data"]["tags"] == tags + + +def test_create_request_omits_tags_when_bedrock_tags_absent(config): + with patch.object( + config.common_utils, + "generate_unique_job_name", + return_value="litellm-batch-1", + ), patch.object(config.common_utils, "sign_aws_request") as mock_sign: + mock_sign.return_value = ({}, b"{}") + config.transform_create_batch_request( + model="m", + create_batch_data={"input_file_id": "s3://b/in.jsonl"}, + optional_params={}, + litellm_params={"aws_batch_role_arn": "arn:aws:iam::1:role/r"}, + ) + assert "tags" not in mock_sign.call_args.kwargs["data"] + + +@pytest.mark.parametrize( + "bad_tags", + [ + ["application=genai-proxy"], + [{"key": "application"}], + [{"value": "genai-proxy"}], + [{"key": "application", "value": 42}], + {"key": "application", "value": "genai-proxy"}, + "application=genai-proxy", + ], +) +def test_create_request_rejects_malformed_bedrock_tags(config, bad_tags): + with patch.object( + config.common_utils, + "generate_unique_job_name", + return_value="litellm-batch-1", + ), patch.object(config.common_utils, "sign_aws_request") as mock_sign: + mock_sign.return_value = ({}, b"{}") + with pytest.raises(ValueError, match="Invalid 'bedrock_tags' value"): + config.transform_create_batch_request( + model="m", + create_batch_data={"input_file_id": "s3://b/in.jsonl"}, + optional_params={}, + litellm_params={ + "aws_batch_role_arn": "arn:aws:iam::1:role/r", + "bedrock_tags": bad_tags, + }, + ) + mock_sign.assert_not_called() + + # --------------------------------------------------------------------------- # # transform_create_batch_response - status mapping + LiteLLMBatch shape # --------------------------------------------------------------------------- # From f5dc1a30107d681e119256ded35cddeb19931ff9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 20 Jul 2026 14:26:35 -0700 Subject: [PATCH 2/3] test(router): prove request-level bedrock_tags override deployment-level tags for acreate_batch --- tests/test_litellm/test_router.py | 60 +++++++++++++++++++++++++++++++ 1 file changed, 60 insertions(+) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index c2c98c8869c..a13b3759865 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -5430,3 +5430,63 @@ class TestRouterRequestTimeoutPropagation: ) == 60 ) + + +@pytest.mark.asyncio +async def test_acreate_batch_request_bedrock_tags_override_deployment_tags(): + import httpx + + from litellm.llms.bedrock.common_utils import CommonBatchFilesUtils + + deployment_tags = [{"key": "application", "value": "config-level"}] + request_tags = [{"key": "application", "value": "request-level"}] + router = litellm.Router( + model_list=[ + { + "model_name": "bedrock-batch-model", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-sonnet-5", + "aws_batch_role_arn": "arn:aws:iam::123:role/batch-role", + "aws_region_name": "us-west-2", + "bedrock_tags": deployment_tags, + }, + } + ] + ) + + def fake_response(): + return httpx.Response( + status_code=200, + json={ + "jobArn": "arn:aws:bedrock:us-west-2:123:model-invocation-job/abc1234567", + "status": "Submitted", + }, + ) + + mock_client = MagicMock() + mock_client.post = AsyncMock(side_effect=lambda *args, **kwargs: fake_response()) + + with patch.object( + CommonBatchFilesUtils, + "sign_aws_request", + return_value=({"Authorization": "signed"}, b"{}"), + ) as mock_sign, patch( + "litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client", + return_value=mock_client, + ): + await router.acreate_batch( + model="bedrock-batch-model", + input_file_id="s3://bucket/input.jsonl", + endpoint="/v1/chat/completions", + completion_window="24h", + ) + assert mock_sign.call_args.kwargs["data"]["tags"] == deployment_tags + + await router.acreate_batch( + model="bedrock-batch-model", + input_file_id="s3://bucket/input.jsonl", + endpoint="/v1/chat/completions", + completion_window="24h", + bedrock_tags=request_tags, + ) + assert mock_sign.call_args.kwargs["data"]["tags"] == request_tags From c2bd8699be1072592b0552c42467bc7427a46d37 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 20 Jul 2026 15:00:41 -0700 Subject: [PATCH 3/3] fix(proxy): require admin opt-in for request-body bedrock_tags Caller-supplied bedrock_tags land as AWS resource tags under the proxy's AWS identity, letting an authenticated caller forge ownership or cost-allocation labels. Add bedrock_tags to _BANNED_REQUEST_BODY_PARAMS so per-request tags need general_settings.allow_client_side_credentials or configurable_clientside_auth_params on the deployment, matching the aws_bedrock_project_id precedent. Deployment-level bedrock_tags in litellm_params are unaffected. Also stop an explicit empty bedrock_tags list in litellm_params from falling through to optional_params --- .../llms/bedrock/batches/transformation.py | 3 +- litellm/proxy/auth/auth_utils.py | 1 + .../bedrock/batches/test_transformation.py | 19 +++++ .../proxy/auth/test_auth_utils.py | 85 +++++++++++++++++++ 4 files changed, 107 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 8648d6586e8..a4ff1c78467 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -215,7 +215,8 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): "roleArn": role_arn, } - bedrock_tags = litellm_params.get("bedrock_tags") or optional_params.get("bedrock_tags") + config_bedrock_tags = litellm_params.get("bedrock_tags") + bedrock_tags = config_bedrock_tags if config_bedrock_tags is not None else optional_params.get("bedrock_tags") if bedrock_tags is not None: bedrock_request["tags"] = _validate_bedrock_tags(bedrock_tags) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 38900260c98..293bb74e211 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -273,6 +273,7 @@ _BANNED_REQUEST_BODY_PARAMS: Tuple[str, ...] = ( # re-route the request's retention and accounting to any project # reachable with the deployment's shared AWS credentials. "aws_bedrock_project_id", + "bedrock_tags", # Provider-specific endpoint overrides that flow into the outbound # request via ``optional_params``. Same threat as ``api_base``: # ``s3_endpoint_url`` redirects Bedrock file uploads to attacker diff --git a/tests/test_litellm/llms/bedrock/batches/test_transformation.py b/tests/test_litellm/llms/bedrock/batches/test_transformation.py index b38d271e210..3681daffe5e 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_transformation.py +++ b/tests/test_litellm/llms/bedrock/batches/test_transformation.py @@ -298,6 +298,25 @@ def test_create_request_forwards_bedrock_tags_from_optional_params(config): assert mock_sign.call_args.kwargs["data"]["tags"] == tags +def test_create_request_empty_litellm_params_tags_do_not_fall_through(config): + with patch.object( + config.common_utils, + "generate_unique_job_name", + return_value="litellm-batch-1", + ), patch.object(config.common_utils, "sign_aws_request") as mock_sign: + mock_sign.return_value = ({}, b"{}") + config.transform_create_batch_request( + model="m", + create_batch_data={"input_file_id": "s3://b/in.jsonl"}, + optional_params={"bedrock_tags": [{"key": "env", "value": "prod"}]}, + litellm_params={ + "aws_batch_role_arn": "arn:aws:iam::1:role/r", + "bedrock_tags": [], + }, + ) + assert mock_sign.call_args.kwargs["data"]["tags"] == [] + + def test_create_request_omits_tags_when_bedrock_tags_absent(config): with patch.object( config.common_utils, diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index b5d8727f7e6..72bd215b9be 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -1944,6 +1944,91 @@ class TestIsRequestBodySafeBlocksRivaUseSsl: ) +class TestIsRequestBodySafeBlocksBedrockTags: + """``bedrock_tags`` lands as AWS resource tags on Bedrock batch jobs + created with the proxy's AWS identity, so a caller-supplied value can + forge ownership or cost-allocation labels; like + ``aws_bedrock_project_id`` it is blocked without an admin opt-in.""" + + def test_bedrock_tags_in_request_body_is_rejected(self): + with pytest.raises(ValueError, match="bedrock_tags"): + is_request_body_safe( + request_body={ + "model": "bedrock-batch-opus", + "bedrock_tags": [{"key": "application", "value": "genai-proxy"}], + }, + general_settings={}, + llm_router=None, + model="bedrock-batch-opus", + ) + + def test_admin_opt_in_proxy_wide_allows_bedrock_tags(self): + assert ( + is_request_body_safe( + request_body={ + "model": "bedrock-batch-opus", + "bedrock_tags": [{"key": "application", "value": "genai-proxy"}], + }, + general_settings={"allow_client_side_credentials": True}, + llm_router=None, + model="bedrock-batch-opus", + ) + is True + ) + + def test_admin_opt_in_per_deployment_allows_bedrock_tags(self): + from litellm import Router + + router = Router( + model_list=[ + { + "model_name": "bedrock-batch-opus", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-opus-4-7", + "configurable_clientside_auth_params": ["bedrock_tags"], + }, + } + ] + ) + assert ( + is_request_body_safe( + request_body={ + "model": "bedrock-batch-opus", + "bedrock_tags": [{"key": "application", "value": "genai-proxy"}], + }, + general_settings={}, + llm_router=router, + model="bedrock-batch-opus", + ) + is True + ) + + def test_per_deployment_opt_in_for_other_param_still_rejects_bedrock_tags(self): + from litellm import Router + + router = Router( + model_list=[ + { + "model_name": "bedrock-batch-opus", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-opus-4-7", + "configurable_clientside_auth_params": ["api_base"], + }, + } + ] + ) + with pytest.raises(ValueError, match="bedrock_tags"): + is_request_body_safe( + request_body={ + "model": "bedrock-batch-opus", + "bedrock_tags": [{"key": "application", "value": "genai-proxy"}], + }, + general_settings={}, + llm_router=router, + model="bedrock-batch-opus", + ) + + # ── is_request_body_safe nested-config recursion (VERIA-6) ────────────────────