Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_migration_job_prisma_boot_cost

litellm/llms/bedrock/common_utils.py

@@ -1487,6 +1487,7 @@ class CommonBatchFilesUtils:
             aws_role_name=optional_params.get("aws_role_name"),
             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"),
         )
 
         # Prepare the request data

litellm/llms/bedrock/files/handler.py

@@ -113,6 +113,7 @@ class BedrockFilesHandler(BaseAWSLLM):
             aws_role_name=optional_params.get("aws_role_name"),
             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"),
         )
 
         # Create S3 client

litellm/llms/bedrock/files/transformation.py

@@ -146,6 +146,7 @@ class _BedrockS3RequestParams(BaseModel):
     aws_role_name: str | None = None
     aws_web_identity_token: str | None = None
     aws_sts_endpoint: str | None = None
+    aws_external_id: str | None = None
     s3_region_name: str | None = None
     s3_endpoint_url: str | None = None
 
@@ -1029,6 +1030,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
             aws_role_name=optional_params.get("aws_role_name"),
             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"),
         )
 
         # Calculate SHA256 hash of the content (REQUIRED for S3)
@@ -1290,6 +1292,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
             aws_role_name=request_params.aws_role_name,
             aws_web_identity_token=request_params.aws_web_identity_token,
             aws_sts_endpoint=request_params.aws_sts_endpoint,
+            aws_external_id=request_params.aws_external_id,
         )
 
         empty_body_hash: Final = hashlib.sha256(b"").hexdigest()

litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py

@@ -32,6 +32,9 @@ from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS
 from litellm.exceptions import ModifyResponseException
 from litellm.integrations.custom_guardrail import CustomGuardrail
 from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys
+from litellm.litellm_core_utils.litellm_logging import (
+    _get_masked_values,  # pyright: ignore[reportPrivateUsage]  # the shared header-masking helper has no public name
+)
 from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import bedrock_guardrail_cost
 from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler
 from litellm.llms.base_llm.guardrail_translation.utils import (
@@ -1170,11 +1173,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
             aws_region_name=aws_region_name,
             api_key=api_key,
         )
+        headers_dict: Final = dict(prepared_request.headers)  # mutable-ok: the masking helper requires a dict
         verbose_proxy_logger.debug(
             "Bedrock AI request body: %s, url %s, headers: %s",
             bedrock_request_data,
             prepared_request.url,
-            prepared_request.headers,
+            _get_masked_values(headers_dict),
         )
 
         httpx_response: Final = await self._sign_and_post(

tests/test_litellm/llms/bedrock/files/test_bedrock_files_handler.py

@@ -204,3 +204,69 @@ def test_should_forward_trusted_model_credentials_to_retrieve_provider_config():
     assert response is mock_response
     litellm_params = mock_retrieve_file.call_args.kwargs["litellm_params"]
     assert litellm_params["_litellm_internal_model_credentials"] is trusted_credentials
+
+
+@pytest.mark.asyncio
+async def test_afile_content_assumes_role_with_external_id(monkeypatch):
+    """A trust policy requiring sts:ExternalId must be satisfied by the deployment's aws_external_id."""
+    import datetime
+
+    import boto3
+    from botocore.exceptions import ClientError
+
+    monkeypatch.delenv("AWS_EXTERNAL_ID", raising=False)
+
+    class FakeSTSClient:
+        def get_caller_identity(self):
+            return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"}
+
+        def assume_role(self, **params):
+            if params.get("ExternalId") != "external-id-files-download":
+                raise ClientError(
+                    {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:AssumeRole"}},
+                    "AssumeRole",
+                )
+            return {
+                "Credentials": {
+                    "AccessKeyId": "ASIAFILESDOWNLOADROLE",
+                    "SecretAccessKey": "assumed-secret",
+                    "SessionToken": "assumed-session-token",
+                    "Expiration": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30),
+                }
+            }
+
+    class FakeS3Body:
+        def read(self):
+            return b'{"custom_id": "req-1"}'
+
+    class FakeS3Client:
+        def get_object(self, Bucket, Key):
+            return {"Body": FakeS3Body()}
+
+    def fake_boto3_client(service_name, **kwargs):
+        if service_name == "sts":
+            return FakeSTSClient()
+        return FakeS3Client()
+
+    optional_params = {
+        "_litellm_internal_model_credentials": MappingProxyType({"s3_bucket_name": "safe-bucket"}),
+        "aws_region_name": "us-east-1",
+        "aws_access_key_id": "AKIAFILESDOWNLOADCALLER",
+        "aws_secret_access_key": "pod-caller-secret",
+        "aws_role_name": "arn:aws:iam::999999999999:role/litellm-files-download-role",
+        "aws_session_name": "litellm-files-download-session",
+        "aws_external_id": "external-id-files-download",
+    }
+
+    with patch.object(boto3, "client", side_effect=fake_boto3_client) as mock_boto3_client:
+        response = await BedrockFilesHandler().afile_content(
+            file_content_request={"file_id": "s3://safe-bucket/litellm-bedrock-files-model-id-abc.jsonl"},
+            optional_params=optional_params,
+            timeout=10.0,
+            max_retries=None,
+        )
+
+    s3_client_kwargs = next(call.kwargs for call in mock_boto3_client.call_args_list if call.args[0] == "s3")
+    assert s3_client_kwargs["aws_access_key_id"] == "ASIAFILESDOWNLOADROLE"
+    assert s3_client_kwargs["aws_session_token"] == "assumed-session-token"
+    assert response.content == b'{"custom_id": "req-1"}'

tests/test_litellm/llms/bedrock/files/test_bedrock_files_transformation.py

@@ -2404,3 +2404,111 @@ class TestBedrockFilesS3SignatureEncoding:
             body=None,
             headers=litellm_params[S3_SIGNED_GET_HEADERS_PARAM],
         )
+
+
+def test_sign_s3_request_assumes_role_with_external_id(monkeypatch):
+    """A trust policy requiring sts:ExternalId must be satisfied when signing the S3 upload request."""
+    import datetime
+    from unittest.mock import patch
+
+    import boto3
+    from botocore.exceptions import ClientError
+
+    from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
+
+    monkeypatch.delenv("AWS_EXTERNAL_ID", raising=False)
+
+    class FakeSTSClient:
+        def get_caller_identity(self):
+            return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"}
+
+        def assume_role(self, **params):
+            if params.get("ExternalId") != "external-id-files-put":
+                raise ClientError(
+                    {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:AssumeRole"}},
+                    "AssumeRole",
+                )
+            return {
+                "Credentials": {
+                    "AccessKeyId": "ASIAFILESPUTROLE",
+                    "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": "AKIAFILESPUTCALLER",
+        "aws_secret_access_key": "pod-caller-secret",
+        "aws_role_name": "arn:aws:iam::999999999999:role/litellm-files-put-role",
+        "aws_session_name": "litellm-files-put-session",
+        "aws_external_id": "external-id-files-put",
+    }
+
+    with patch.object(boto3, "client", return_value=FakeSTSClient()):
+        signed_headers, _signed_body = BedrockFilesConfig()._sign_s3_request(
+            content='{"custom_id": "req-1"}',
+            api_base="https://s3.us-east-1.amazonaws.com/safe-bucket/litellm-bedrock-files-model-id-abc.jsonl",
+            optional_params=optional_params,
+        )
+
+    authorization = {key.lower(): value for key, value in signed_headers.items()}["authorization"]
+    assert "ASIAFILESPUTROLE" in authorization
+
+
+def test_sign_s3_get_request_assumes_role_with_external_id(monkeypatch):
+    """A trust policy requiring sts:ExternalId must be satisfied when signing the S3 download request."""
+    import datetime
+    from unittest.mock import patch
+
+    import boto3
+    from botocore.exceptions import ClientError
+
+    from litellm.llms.bedrock.files.transformation import (
+        BedrockFilesConfig,
+        _BedrockS3RequestParams,
+    )
+
+    monkeypatch.delenv("AWS_EXTERNAL_ID", raising=False)
+
+    class FakeSTSClient:
+        def get_caller_identity(self):
+            return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"}
+
+        def assume_role(self, **params):
+            if params.get("ExternalId") != "external-id-files-get":
+                raise ClientError(
+                    {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:AssumeRole"}},
+                    "AssumeRole",
+                )
+            return {
+                "Credentials": {
+                    "AccessKeyId": "ASIAFILESGETROLE",
+                    "SecretAccessKey": "assumed-secret",
+                    "SessionToken": "assumed-session-token",
+                    "Expiration": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30),
+                }
+            }
+
+    request_params = _BedrockS3RequestParams.model_validate(
+        {
+            "aws_region_name": "us-east-1",
+            "aws_access_key_id": "AKIAFILESGETCALLER",
+            "aws_secret_access_key": "pod-caller-secret",
+            "aws_role_name": "arn:aws:iam::999999999999:role/litellm-files-get-role",
+            "aws_session_name": "litellm-files-get-session",
+            "aws_external_id": "external-id-files-get",
+        }
+    )
+    assert request_params.aws_external_id == "external-id-files-get"
+
+    with patch.object(boto3, "client", return_value=FakeSTSClient()):
+        signed_headers = BedrockFilesConfig()._sign_s3_get_request(
+            api_base="https://s3.us-east-1.amazonaws.com/safe-bucket/litellm-bedrock-files-model-id-abc.jsonl",
+            aws_region_name="us-east-1",
+            request_params=request_params,
+        )
+
+    authorization = {key.lower(): value for key, value in signed_headers.items()}["authorization"]
+    assert "ASIAFILESGETROLE" in authorization

tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py

@@ -520,3 +520,56 @@ def test_merge_bedrock_aws_request_params_keeps_caller_credentials_without_stati
     assert merged["aws_secret_access_key"] == "caller-secret"
     assert merged["aws_session_token"] == "caller-token"
     assert merged["aws_region_name"] == "us-west-2"
+
+
+def test_sign_aws_request_assumes_role_with_external_id(monkeypatch):
+    """A trust policy requiring sts:ExternalId must be satisfied when signing batch API requests."""
+    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_EXTERNAL_ID", raising=False)
+
+    class FakeSTSClient:
+        def get_caller_identity(self):
+            return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"}
+
+        def assume_role(self, **params):
+            if params.get("ExternalId") != "external-id-batch-sign":
+                raise ClientError(
+                    {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:AssumeRole"}},
+                    "AssumeRole",
+                )
+            return {
+                "Credentials": {
+                    "AccessKeyId": "ASIABATCHSIGNROLE",
+                    "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_external_id": "external-id-batch-sign",
+    }
+
+    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 "ASIABATCHSIGNROLE" in authorization
+    assert signed_data == b'{"jobName": "litellm-batch-job"}'

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py

@@ -5590,3 +5590,52 @@ async def test_streaming_end_of_stream_block_emits_error_frame_instead_of_trunca
     assert payload["error"]["message"] == "Violated guardrail policy"
     assert payload["error"]["code"] == "400"
     assert payload["error"]["provider_specific_fields"]["guardrailIdentifier"] == "test-guardrail"
+
+
+@pytest.mark.asyncio
+async def test_apply_guardrail_debug_log_masks_signed_request_headers():
+    import logging
+
+    from litellm._logging import verbose_proxy_logger
+
+    session_token = "FakeSessionTokenValueThatMustNeverAppearInLogs1234567890"
+    guardrail = BedrockGuardrail(
+        guardrailIdentifier="test-guardrail",
+        guardrailVersion="DRAFT",
+        aws_access_key_id="ASIAFAKEACCESSKEYID1",
+        aws_secret_access_key="fakeSecretAccessKeyForSigning",
+        aws_session_token=session_token,
+        aws_region_name="us-east-1",
+    )
+
+    mock_response = MagicMock()
+    mock_response.status_code = 200
+    mock_response.json.return_value = {"action": "NONE", "outputs": []}
+
+    captured_records: list[logging.LogRecord] = []
+
+    class _RecordingHandler(logging.Handler):
+        def emit(self, record: logging.LogRecord) -> None:
+            captured_records.append(record)
+
+    handler = _RecordingHandler(level=logging.DEBUG)
+    previous_level = verbose_proxy_logger.level
+    verbose_proxy_logger.addHandler(handler)
+    verbose_proxy_logger.setLevel(logging.DEBUG)
+    try:
+        with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post:
+            mock_post.return_value = mock_response
+            await guardrail.make_bedrock_api_request(
+                source="INPUT",
+                messages=[{"role": "user", "content": "hello"}],
+                request_data={},
+            )
+    finally:
+        verbose_proxy_logger.removeHandler(handler)
+        verbose_proxy_logger.setLevel(previous_level)
+
+    rendered_messages = [record.getMessage() for record in captured_records]
+    header_lines = [message for message in rendered_messages if "headers:" in message]
+    assert header_lines, "expected the signed-request debug line to be logged"
+    assert any("X-Amz-Security-Token" in message for message in header_lines)
+    assert all(session_token not in message for message in rendered_messages)

tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py

@@ -917,7 +917,9 @@ async def test_bedrock_guardrail_make_api_request_passes_api_key():
             "Content-Type": "application/json",
             "Authorization": "Bearer test-api-key-789",
         }
-        mock_request_instance.prepare.return_value = Mock()
+        mock_request_instance.prepare.return_value = Mock(
+            headers=mock_request_instance.headers
+        )
         mock_aws_request.return_value = mock_request_instance
 
         await guardrail_hook.make_bedrock_api_request(