diff --git a/litellm/__init__.py b/litellm/__init__.py index 8961de940a0..ae0fee11aeb 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -26,6 +26,7 @@ def _dev_env_hot_reload_enabled() -> bool: if os.getenv("LITELLM_MODE", "DEV") == "DEV": _dotenv.load_dotenv(override=_dev_env_hot_reload_enabled()) +from collections.abc import Sequence from typing import ( Any, Callable, @@ -217,6 +218,9 @@ add_user_information_to_llm_headers: Optional[bool] = ( overwrite_user_with_key_hash: bool = ( False # force the outgoing `user` param to the hashed api key, so providers see a stable, tamper-proof id ) +bedrock_request_metadata_fields: Optional[Sequence[str]] = ( + None # allow-list of `user_api_key_*` fields (+ `spend_logs_metadata`) sent as Bedrock `requestMetadata` +) store_audit_logs = False # Enterprise feature, allow users to see audit logs skip_system_message_in_guardrail: bool = False skip_tool_message_in_guardrail: bool = False diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 85918d40e12..e4937136108 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -39,6 +39,12 @@ from litellm.llms.anthropic.chat.transformation import ( AnthropicConfig, ) from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.bedrock.request_metadata import ( + bedrock_request_metadata_headers, + bedrock_request_metadata_is_owned, + merge_bedrock_invoke_headers, + resolve_bedrock_request_metadata, +) from litellm.types.llms.bedrock import * from litellm.types.llms.openai import ( AllMessageValues, @@ -1652,6 +1658,13 @@ class AmazonConverseConfig(BaseConfig): user_continue_message=litellm_params.pop("user_continue_message", None), ) + request_metadata: Final = resolve_bedrock_request_metadata( + litellm_params=litellm_params, caller_metadata=_data.get("requestMetadata") + ) + if bedrock_request_metadata_is_owned(): + _data.pop("requestMetadata", None) + if request_metadata is not None: + _data["requestMetadata"] = request_metadata data: Final[RequestObject] = {"messages": bedrock_messages, **_data} return data @@ -1705,6 +1718,13 @@ class AmazonConverseConfig(BaseConfig): user_continue_message=litellm_params.pop("user_continue_message", None), ) + request_metadata: Final = resolve_bedrock_request_metadata( + litellm_params=litellm_params, caller_metadata=_data.get("requestMetadata") + ) + if bedrock_request_metadata_is_owned(): + _data.pop("requestMetadata", None) + if request_metadata is not None: + _data["requestMetadata"] = request_metadata data: Final[RequestObject] = {"messages": bedrock_messages, **_data} return data @@ -2258,7 +2278,8 @@ class AmazonConverseConfig(BaseConfig): ) -> dict: if api_key: headers["Authorization"] = f"Bearer {api_key}" - return headers + owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params) + return merge_bedrock_invoke_headers(headers, (), metadata_headers, owned_names) def should_fake_stream( self, diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py index ddbb036df40..1671585be2d 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py @@ -13,6 +13,10 @@ import httpx from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import BedrockError +from litellm.llms.bedrock.request_metadata import ( + bedrock_request_metadata_headers, + merge_bedrock_invoke_headers, +) from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.passthrough.utils import CommonUtils from litellm.types.llms.openai import AllMessageValues @@ -169,9 +173,12 @@ class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM): """ Validate the environment and return headers. - For Bedrock, we don't need Bearer token auth since we use AWS SigV4. + For Bedrock, we don't need Bearer token auth since we use AWS SigV4. This path signs the + same ``/model/{id}/invoke`` endpoint as ``AmazonInvokeConfig``, so it owns the request + metadata header on the same terms rather than letting a caller supply it. """ - return headers + owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params) + return merge_bedrock_invoke_headers(headers, (), metadata_headers, owned_names) def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BedrockError: """Return the appropriate error class for Bedrock.""" diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index 430d0a92b51..76f91aa9115 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -20,6 +20,10 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.bedrock.chat.invoke_handler import make_call, make_sync_call from litellm.llms.bedrock.common_utils import BedrockError +from litellm.llms.bedrock.request_metadata import ( + bedrock_request_metadata_headers, + merge_bedrock_invoke_headers, +) from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -417,15 +421,13 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): api_base: str | None = None, ) -> dict: raw_guardrail_config: Final = optional_params.pop("guardrailConfig", None) - if raw_guardrail_config is None: - return headers - existing_header_names: Final = frozenset(name.lower() for name in headers) - guardrail_headers: Final = { - name: value - for name, value in _bedrock_invoke_guardrail_headers(raw_guardrail_config).items() - if name.lower() not in existing_header_names - } - return {**headers, **guardrail_headers} + guardrail_headers: Final = ( + () + if raw_guardrail_config is None + else tuple(_bedrock_invoke_guardrail_headers(raw_guardrail_config).items()) + ) + owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params) + return merge_bedrock_invoke_headers(headers, guardrail_headers, metadata_headers, owned_names) def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: return BedrockError(status_code=status_code, message=error_message) diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 372cf110f7c..ecdffe7d763 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -37,6 +37,10 @@ from litellm.llms.bedrock.common_utils import ( normalize_tool_input_schema_types_for_bedrock_invoke, pop_bedrock_invoke_output_config_format, ) +from litellm.llms.bedrock.request_metadata import ( + bedrock_request_metadata_headers, + merge_bedrock_invoke_headers, +) from litellm.types.llms.anthropic import ( ANTHROPIC_BETA_HEADER_VALUES, ANTHROPIC_TOOL_SEARCH_BETA_HEADER, @@ -89,7 +93,8 @@ class AmazonAnthropicClaudeMessagesConfig( api_key: str | None = None, api_base: str | None = None, ) -> tuple[dict, str | None]: - return headers, api_base + owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params) + return merge_bedrock_invoke_headers(headers, (), metadata_headers, owned_names), api_base def sign_request( self, diff --git a/litellm/llms/bedrock/request_metadata.py b/litellm/llms/bedrock/request_metadata.py new file mode 100644 index 00000000000..1f4e5886508 --- /dev/null +++ b/litellm/llms/bedrock/request_metadata.py @@ -0,0 +1,199 @@ +""" +Resolve AWS Bedrock ``requestMetadata`` from LiteLLM proxy identity and caller metadata. + +Bedrock attaches request metadata to CloudTrail records and to the dimension AWS Cost +Explorer groups on, so everything here is opt-in: nothing is forwarded unless the operator +sets ``litellm.bedrock_request_metadata_fields`` (``litellm_settings`` on the proxy). + +Two properties are load-bearing for that billing record and are asserted by the tests: +proxy identity is resolved first so it can never be evicted by caller-supplied pairs, and the +whole ``user_api_key_`` prefix is reserved so a caller cannot write a proxy-authoritative +looking key. Values that break Bedrock's constraints are dropped rather than sanitised or +rejected, because an operator flipping this setting on must not turn a working request into a +400 and a silently rewritten attribution key is worse than an absent one. +""" + +from __future__ import annotations + +import json +import re +from collections.abc import Mapping +from typing import Final + +import litellm + +BEDROCK_REQUEST_METADATA_HEADER: Final = "X-Amzn-Bedrock-Request-Metadata" +BEDROCK_REQUEST_METADATA_MAX_PAIRS: Final = 16 +BEDROCK_REQUEST_METADATA_IDENTITY_PREFIX: Final = "user_api_key_" +BEDROCK_REQUEST_METADATA_CLIENT_FIELD: Final = "spend_logs_metadata" + +_METADATA_PARAM_NAMES: Final[tuple[str, ...]] = ("metadata", "litellm_metadata") +_KEY_PATTERN: Final = re.compile(r"^[a-zA-Z0-9\s:_@$#=/+,.-]{1,256}$") +_VALUE_PATTERN: Final = re.compile(r"^[a-zA-Z0-9\s:_@$#=/+,.-]{0,256}$") +_OWNED_HEADER_NAMES: Final[frozenset[str]] = frozenset((BEDROCK_REQUEST_METADATA_HEADER.lower(),)) + + +def _is_forwardable(key: str, value: str) -> bool: + return _KEY_PATTERN.match(key) is not None and _VALUE_PATTERN.match(value) is not None + + +def _text_pairs(source: object) -> tuple[tuple[str, str], ...]: + if not isinstance(source, Mapping): + return () + return tuple((key, value) for key, value in source.items() if isinstance(key, str) and isinstance(value, str)) + + +def _allowed_fields() -> tuple[str, ...]: + """ + The operator allow-list, deduplicated so a field repeated in config cannot consume a second + reserved slot and shrink the client budget for nothing. First occurrence wins, which keeps + the operator's declared precedence intact. + """ + configured: Final[object] = litellm.bedrock_request_metadata_fields + if not isinstance(configured, (list, tuple)): + return () + fields: Final = tuple(str(field) for field in configured) + return tuple(field for index, field in enumerate(fields) if field not in fields[:index]) + + +def _metadata_sources(litellm_params: Mapping[str, object] | None) -> tuple[Mapping[str, object], ...]: + """``metadata`` on /v1/chat/completions, ``litellm_metadata`` on the LITELLM_METADATA_ROUTES.""" + if litellm_params is None: + return () + return tuple( + source + for name in _METADATA_PARAM_NAMES + for source in (litellm_params.get(name),) + if isinstance(source, Mapping) + ) + + +def _identity_pairs( + sources: tuple[Mapping[str, object], ...], + allowed_fields: tuple[str, ...], +) -> tuple[tuple[str, str], ...]: + return tuple( + (field, value) + for field in allowed_fields + if field.startswith(BEDROCK_REQUEST_METADATA_IDENTITY_PREFIX) + for value in (_first_text(sources, field),) + if value is not None and _is_forwardable(field, value) + )[:BEDROCK_REQUEST_METADATA_MAX_PAIRS] + + +def _first_text(sources: tuple[Mapping[str, object], ...], field: str) -> str | None: + return next((value for source in sources if isinstance(value := source.get(field), str)), None) + + +def _client_pairs( + sources: tuple[Mapping[str, object], ...], + allowed_fields: tuple[str, ...], + caller_metadata: object, + budget: int, +) -> tuple[tuple[str, str], ...]: + spend_logs_pairs: Final = ( + tuple(pair for source in sources for pair in _text_pairs(source.get(BEDROCK_REQUEST_METADATA_CLIENT_FIELD))) + if BEDROCK_REQUEST_METADATA_CLIENT_FIELD in allowed_fields + else () + ) + candidates: Final = tuple( + (key, value) + for key, value in (*_text_pairs(caller_metadata), *spend_logs_pairs) + if not key.startswith(BEDROCK_REQUEST_METADATA_IDENTITY_PREFIX) and _is_forwardable(key, value) + ) + return tuple( + pair + for index, pair in enumerate(candidates) + if pair[0] not in tuple(earlier for earlier, _ in candidates[:index]) + )[:budget] + + +def resolve_bedrock_request_metadata( + litellm_params: Mapping[str, object] | None, + caller_metadata: object = None, +) -> dict[str, str] | None: + """ + Resolve the ``requestMetadata`` pairs to send to Bedrock, or ``None`` when the feature is + off or nothing survives Bedrock's constraints. The result is a plain dict because it is + written straight onto the Converse body, which Bedrock types as ``dict[str, str]``. + + ``caller_metadata`` is any ``requestMetadata`` the caller passed explicitly. It has already + been validated (and rejected with a 400) by the Converse transformation, so it is only + filtered here for the reserved identity prefix and the remaining slot budget. + """ + allowed_fields: Final = _allowed_fields() + if not allowed_fields: + return None + sources: Final = _metadata_sources(litellm_params) + identity: Final = _identity_pairs(sources, allowed_fields) + client: Final = _client_pairs( + sources=sources, + allowed_fields=allowed_fields, + caller_metadata=caller_metadata, + budget=BEDROCK_REQUEST_METADATA_MAX_PAIRS - len(identity), + ) + resolved: Final = {key: value for key, value in (*identity, *client)} + return resolved or None + + +def bedrock_request_metadata_is_owned() -> bool: + """ + Whether the proxy OWNS the request-metadata field and header name for this request. + + Ownership follows the operator's opt-in alone, never whether anything resolved, because a + caller can suppress the resolver by omitting the allow-listed fields or by sending values + that all fail Bedrock's rules. Owned-but-empty has to mean "absent on the wire" rather than + "fall back to whatever the caller supplied", or the reserved-prefix guarantee is bypassable + by anyone who can make the resolver produce nothing. + """ + return bool(_allowed_fields()) + + +def bedrock_request_metadata_headers( + litellm_params: Mapping[str, object] | None, +) -> tuple[frozenset[str], tuple[tuple[str, str], ...]]: + """ + The signed ``X-Amzn-Bedrock-Request-Metadata`` header for the Invoke paths, which have no + body field for request metadata. + + Returns the header names the proxy OWNS and, separately, the pairs to send. Ownership is + reported whenever forwarding is enabled, including when nothing resolves, because a caller + can suppress the resolver (omit the allow-listed fields, or send values that all fail + Bedrock's rules) and an owned-but-empty result must still evict the caller's header rather + than fall back to it. + """ + if not bedrock_request_metadata_is_owned(): + return frozenset(), () + resolved: Final = resolve_bedrock_request_metadata(litellm_params) + if resolved is None: + return _OWNED_HEADER_NAMES, () + return _OWNED_HEADER_NAMES, ((BEDROCK_REQUEST_METADATA_HEADER, json.dumps(resolved, separators=(",", ":"))),) + + +def merge_bedrock_invoke_headers( + headers: dict[str, str], + caller_owned: tuple[tuple[str, str], ...], + proxy_owned: tuple[tuple[str, str], ...], + proxy_owned_names: frozenset[str], +) -> dict[str, str]: + """ + Merge the ``X-Amzn-*`` headers the Invoke paths derive from params. + + ``caller_owned`` (the guardrail headers) defers to a header the caller already set, which is + the long-standing behaviour for those. ``proxy_owned_names`` are dropped from the caller's + headers unconditionally and re-supplied only from ``proxy_owned``, because those names carry + proxy-authenticated identity into an AWS billing record that the caller must not be able to + write. Names are compared case-insensitively so a caller cannot leave a second spelling in + the dict and let the transport pick the winner. + """ + if not caller_owned and not proxy_owned and not proxy_owned_names: + return headers + existing_names: Final = frozenset(name.lower() for name in headers) + return { + name: value + for name, value in ( + *((n, v) for n, v in headers.items() if n.lower() not in proxy_owned_names), + *((n, v) for n, v in caller_owned if n.lower() not in existing_names), + *proxy_owned, + ) + } diff --git a/tests/test_litellm/llms/bedrock/test_request_metadata.py b/tests/test_litellm/llms/bedrock/test_request_metadata.py new file mode 100644 index 00000000000..ad14db5c85f --- /dev/null +++ b/tests/test_litellm/llms/bedrock/test_request_metadata.py @@ -0,0 +1,422 @@ +import asyncio +import json +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +import litellm +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig +from litellm.llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import ( + AmazonBedrockOpenAIConfig, +) +from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( + AmazonInvokeConfig, +) +from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import ( + AmazonAnthropicClaudeMessagesConfig, +) +from litellm.llms.bedrock.request_metadata import ( + BEDROCK_REQUEST_METADATA_HEADER, + BEDROCK_REQUEST_METADATA_MAX_PAIRS, + resolve_bedrock_request_metadata, +) + +MODEL = "anthropic.claude-3-5-sonnet-20240620-v1:0" +MESSAGES = [{"role": "user", "content": "hi"}] +ALL_FIELDS = [ + "user_api_key_alias", + "user_api_key_team_alias", + "user_api_key_user_email", + "spend_logs_metadata", +] +IDENTITY = {"user_api_key_alias": "prod-key", "user_api_key_team_alias": "platform"} + + +@pytest.fixture(autouse=True) +def reset_setting(): + previous = litellm.bedrock_request_metadata_fields + yield + litellm.bedrock_request_metadata_fields = previous + + +def litellm_params(metadata_key, **metadata): + return {metadata_key: dict(metadata)} + + +def converse_body(litellm_params_value, optional_params=None): + return AmazonConverseConfig()._transform_request( + model=MODEL, + messages=MESSAGES, + optional_params=dict(optional_params or {}), + litellm_params=dict(litellm_params_value), + ) + + +def converse_body_async(litellm_params_value, optional_params=None): + """The proxy serves completions through the async transform, so every rule asserted against + the sync body has to be asserted against this one too or half the product is untested.""" + return asyncio.run( + AmazonConverseConfig()._async_transform_request( + model=MODEL, + messages=MESSAGES, + optional_params=dict(optional_params or {}), + litellm_params=dict(litellm_params_value), + ) + ) + + +CONVERSE_DRIVERS = [converse_body, converse_body_async] + + +@pytest.mark.parametrize("setting", [None, []]) +def test_feature_off_by_default_leaves_body_and_headers_untouched(setting): + litellm.bedrock_request_metadata_fields = setting + params = litellm_params("metadata", spend_logs_metadata={"team": "x"}, **IDENTITY) + + assert "requestMetadata" not in converse_body(params) + assert BEDROCK_REQUEST_METADATA_HEADER not in AmazonInvokeConfig().validate_environment( + headers={}, model=MODEL, messages=MESSAGES, optional_params={}, litellm_params=dict(params) + ) + messages_headers, _ = AmazonAnthropicClaudeMessagesConfig().validate_anthropic_messages_environment( + headers={}, model=MODEL, messages=MESSAGES, optional_params={}, litellm_params=dict(params) + ) + assert BEDROCK_REQUEST_METADATA_HEADER not in messages_headers + + +@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) +def test_resolver_reads_both_metadata_variable_names(metadata_key): + """`/v1/chat/completions` populates `metadata`; the LITELLM_METADATA_ROUTES populate + `litellm_metadata`. Reading only one silently forwards nothing on the other route.""" + litellm.bedrock_request_metadata_fields = ALL_FIELDS + params = litellm_params(metadata_key, spend_logs_metadata={"cost_center": "cc-1"}, **IDENTITY) + + assert converse_body(params)["requestMetadata"] == {**IDENTITY, "cost_center": "cc-1"} + + +@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) +def test_invoke_messages_header_reads_both_metadata_variable_names(metadata_key): + litellm.bedrock_request_metadata_fields = ALL_FIELDS + params = litellm_params(metadata_key, **IDENTITY) + + headers, _ = AmazonAnthropicClaudeMessagesConfig().validate_anthropic_messages_environment( + headers={}, model=MODEL, messages=MESSAGES, optional_params={}, litellm_params=params + ) + + assert json.loads(headers[BEDROCK_REQUEST_METADATA_HEADER]) == IDENTITY + + +@pytest.mark.parametrize("reverse_client_keys", [False, True]) +@pytest.mark.parametrize("field_order", [ALL_FIELDS, list(reversed(ALL_FIELDS))]) +@pytest.mark.parametrize("client_source", ["spend_logs_metadata", "requestMetadata"]) +def test_identity_survives_a_caller_filling_every_slot(reverse_client_keys, field_order, client_source): + """A caller sending 16 keys of its own must not evict the identity the feature exists to + produce. Driven over every input ordering so the invariant is not an accident of one.""" + litellm.bedrock_request_metadata_fields = field_order + client_keys = [f"client_{index:02d}" for index in range(BEDROCK_REQUEST_METADATA_MAX_PAIRS)] + client_pairs = {key: "v" for key in (reversed(client_keys) if reverse_client_keys else client_keys)} + if client_source == "spend_logs_metadata": + params, optional_params = litellm_params("metadata", spend_logs_metadata=client_pairs, **IDENTITY), {} + else: + params, optional_params = litellm_params("metadata", **IDENTITY), {"requestMetadata": client_pairs} + + resolved = converse_body(params, optional_params)["requestMetadata"] + + assert len(resolved) == BEDROCK_REQUEST_METADATA_MAX_PAIRS + for key, value in IDENTITY.items(): + assert resolved[key] == value + assert len([key for key in resolved if key.startswith("client_")]) == ( + BEDROCK_REQUEST_METADATA_MAX_PAIRS - len(IDENTITY) + ) + + +@pytest.mark.parametrize( + "field_order", + [ + ["user_api_key_alias", "user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata"], + ["user_api_key_alias", "user_api_key_team_alias", "user_api_key_alias", "spend_logs_metadata"], + ["user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata", "user_api_key_team_alias"], + ], +) +def test_a_field_repeated_in_the_allow_list_does_not_consume_a_client_slot(field_order): + """An operator repeating a field in YAML must not inflate the reserved count and shrink the + client budget. Asserts the client keys that should have fitted actually reach the wire, since + asserting only that identity survives passes with or without the deduplication.""" + litellm.bedrock_request_metadata_fields = field_order + client_keys = [f"client_{index:02d}" for index in range(BEDROCK_REQUEST_METADATA_MAX_PAIRS - 1)] + params = litellm_params("metadata", spend_logs_metadata={key: "v" for key in client_keys}, **IDENTITY) + + resolved = converse_body(params)["requestMetadata"] + + expected_client_slots = BEDROCK_REQUEST_METADATA_MAX_PAIRS - len(IDENTITY) + assert resolved == {**IDENTITY, **{key: "v" for key in client_keys[:expected_client_slots]}} + assert len(resolved) == BEDROCK_REQUEST_METADATA_MAX_PAIRS + assert client_keys[expected_client_slots - 1] in resolved + + +@pytest.mark.parametrize("client_source", ["spend_logs_metadata", "requestMetadata"]) +@pytest.mark.parametrize( + "forged_key", + ["user_api_key_team_alias", "user_api_key_org_alias", "user_api_key_hash"], +) +def test_caller_cannot_forge_or_shadow_a_reserved_identity_key(forged_key, client_source): + """`user_api_key_org_alias` and `user_api_key_hash` are names the proxy does not set here, + so an exact-key reservation would let the forged value through under a name that reads as + proxy-authoritative in the AWS billing record.""" + litellm.bedrock_request_metadata_fields = ALL_FIELDS + forged = {forged_key: "attacker-controlled"} + if client_source == "spend_logs_metadata": + params, optional_params = litellm_params("metadata", spend_logs_metadata=forged, **IDENTITY), {} + else: + params, optional_params = litellm_params("metadata", **IDENTITY), {"requestMetadata": forged} + + resolved = converse_body(params, optional_params)["requestMetadata"] + + assert resolved == IDENTITY + assert "attacker-controlled" not in resolved.values() + + +def test_identity_violating_the_character_class_is_dropped_and_the_request_succeeds(): + """A team alias with an apostrophe must not turn a working request into a 400 the moment + an operator flips the setting on.""" + litellm.bedrock_request_metadata_fields = ALL_FIELDS + params = litellm_params( + "metadata", + user_api_key_alias="prod-key", + user_api_key_team_alias="O'Brien's team", + user_api_key_user_email="x" * 300, + ) + + body = converse_body(params) + + assert body["requestMetadata"] == {"user_api_key_alias": "prod-key"} + assert body["messages"] + + +def test_caller_supplied_violation_still_raises_bad_request(): + litellm.bedrock_request_metadata_fields = ALL_FIELDS + + with pytest.raises(litellm.exceptions.BadRequestError): + converse_body( + litellm_params("metadata", **IDENTITY), + {"requestMetadata": {"team": "O'Brien's team"}}, + ) + + +def test_non_string_and_absent_identity_values_are_dropped(): + litellm.bedrock_request_metadata_fields = ALL_FIELDS + ["user_api_key_spend"] + params = litellm_params("metadata", user_api_key_alias="prod-key", user_api_key_spend=1.25) + + assert converse_body(params)["requestMetadata"] == {"user_api_key_alias": "prod-key"} + + +def test_email_is_separately_opt_in(): + """PII crossing into CloudTrail only when the operator names the field.""" + identity_with_email = {**IDENTITY, "user_api_key_user_email": "owner@example.com"} + litellm.bedrock_request_metadata_fields = ["user_api_key_alias", "user_api_key_team_alias"] + assert ( + "user_api_key_user_email" + not in converse_body(litellm_params("metadata", **identity_with_email))["requestMetadata"] + ) + + litellm.bedrock_request_metadata_fields = ALL_FIELDS + assert converse_body(litellm_params("metadata", **identity_with_email))["requestMetadata"] == identity_with_email + + +def test_resolver_returns_none_when_nothing_survives(): + litellm.bedrock_request_metadata_fields = ALL_FIELDS + assert resolve_bedrock_request_metadata(litellm_params=None) is None + assert resolve_bedrock_request_metadata(litellm_params={"metadata": {"unrelated": "x"}}) is None + + +def test_invoke_header_is_json_encoded_and_signed(): + litellm.bedrock_request_metadata_fields = ALL_FIELDS + params = litellm_params("metadata", spend_logs_metadata={"cost_center": "cc-1"}, **IDENTITY) + + headers = AmazonInvokeConfig().validate_environment( + headers={"anthropic-version": "bedrock-2023-05-31"}, + model=MODEL, + messages=MESSAGES, + optional_params={}, + litellm_params=params, + ) + + assert json.loads(headers[BEDROCK_REQUEST_METADATA_HEADER]) == {**IDENTITY, "cost_center": "cc-1"} + signed = BaseAWSLLM()._filter_headers_for_aws_signature(headers) + assert BEDROCK_REQUEST_METADATA_HEADER in signed + assert "anthropic-version" not in signed + + +def test_a_caller_supplied_guardrail_header_still_wins(): + """The no-displace rule is deliberate for the guardrail headers and must survive the + request-metadata header becoming proxy-owned.""" + litellm.bedrock_request_metadata_fields = ALL_FIELDS + + headers = AmazonInvokeConfig().validate_environment( + headers={"X-Amzn-Bedrock-GuardrailIdentifier": "caller-set"}, + model=MODEL, + messages=MESSAGES, + optional_params={"guardrailConfig": {"guardrailIdentifier": "gid", "guardrailVersion": "DRAFT"}}, + litellm_params=litellm_params("metadata", **IDENTITY), + ) + + assert headers["X-Amzn-Bedrock-GuardrailIdentifier"] == "caller-set" + assert headers["X-Amzn-Bedrock-GuardrailVersion"] == "DRAFT" + + +FORGED = '{"user_api_key_alias":"FORGED-KEY","user_api_key_team_alias":"FORGED-TEAM"}' + + +def invoke_headers(caller_headers, params, optional_params=None): + return AmazonInvokeConfig().validate_environment( + headers=dict(caller_headers), + model=MODEL, + messages=MESSAGES, + optional_params=dict(optional_params or {}), + litellm_params=dict(params), + ) + + +def messages_headers(caller_headers, params): + resolved, _ = AmazonAnthropicClaudeMessagesConfig().validate_anthropic_messages_environment( + headers=dict(caller_headers), + model=MODEL, + messages=MESSAGES, + optional_params={}, + litellm_params=dict(params), + ) + return resolved + + +def openai_invoke_headers(caller_headers, params): + return AmazonBedrockOpenAIConfig().validate_environment( + headers=dict(caller_headers), + model=MODEL, + messages=MESSAGES, + optional_params={}, + litellm_params=dict(params), + ) + + +def converse_headers(caller_headers, params): + return AmazonConverseConfig().validate_environment( + headers=dict(caller_headers), + model=MODEL, + messages=MESSAGES, + optional_params={}, + litellm_params=dict(params), + ) + + +HEADER_DRIVERS = [invoke_headers, messages_headers, openai_invoke_headers, converse_headers] + + +def metadata_header_values(headers): + return [value for name, value in headers.items() if name.lower() == BEDROCK_REQUEST_METADATA_HEADER.lower()] + + +def test_converse_still_sets_the_bearer_authorization_header(): + """Converse owns the metadata header now, and that must not disturb the api_key path its + validate_environment existed for. Closing the forgery hole cannot break authentication.""" + litellm.bedrock_request_metadata_fields = ALL_FIELDS + + headers = AmazonConverseConfig().validate_environment( + headers={}, + model=MODEL, + messages=MESSAGES, + optional_params={}, + litellm_params=dict(litellm_params("metadata", **IDENTITY)), + api_key="sk-converse-bearer", + ) + + assert headers["Authorization"] == "Bearer sk-converse-bearer" + assert metadata_header_values(headers) == [json.dumps(IDENTITY, separators=(",", ":"))] + + +@pytest.mark.parametrize("driver", HEADER_DRIVERS) +@pytest.mark.parametrize( + "caller_header_name", + [BEDROCK_REQUEST_METADATA_HEADER, BEDROCK_REQUEST_METADATA_HEADER.lower(), "x-AMZN-bedrock-Request-METADATA"], +) +def test_a_caller_cannot_forge_the_request_metadata_header(driver, caller_header_name): + """`extra_headers` puts caller-supplied names into the same dict the proxy merges into, so a + deferring merge would sign the caller's forged identity into the AWS billing record. Every + spelling must lose, or a second variant is left for the transport to choose between.""" + litellm.bedrock_request_metadata_fields = ALL_FIELDS + + headers = driver({caller_header_name: FORGED}, litellm_params("metadata", **IDENTITY)) + + values = metadata_header_values(headers) + assert values == [json.dumps(IDENTITY, separators=(",", ":"))] + assert "FORGED" not in json.dumps(headers) + + +@pytest.mark.parametrize("driver", HEADER_DRIVERS) +def test_a_caller_cannot_forge_the_header_when_the_resolver_yields_nothing(driver): + """Forwarding enabled but nothing resolvable, which a caller can arrange by supplying values + that all fail Bedrock's rules. Owned-but-empty must mean no header on the wire, never a + fallback to the caller's.""" + litellm.bedrock_request_metadata_fields = ALL_FIELDS + unresolvable = litellm_params("metadata", user_api_key_alias="O'Brien's key", user_api_key_team_alias="x" * 300) + + headers = driver({BEDROCK_REQUEST_METADATA_HEADER: FORGED}, unresolvable) + + assert metadata_header_values(headers) == [] + assert "FORGED" not in json.dumps(headers) + + +@pytest.mark.parametrize("driver", CONVERSE_DRIVERS) +@pytest.mark.parametrize( + "forged_key", + ["user_api_key_team_alias", "user_api_key_org_alias", "user_api_key_hash"], +) +def test_a_caller_cannot_keep_reserved_body_keys_when_the_resolver_yields_nothing(forged_key, driver): + """The Converse body has the same fail-open shape as the header: with forwarding on and + nothing resolvable, leaving the caller's `requestMetadata` in place would keep their + reserved-prefix keys on the wire. Owned-but-empty must remove the field outright.""" + litellm.bedrock_request_metadata_fields = ALL_FIELDS + + body = driver(litellm_params("metadata"), {"requestMetadata": {forged_key: "FORGED"}}) + + assert "requestMetadata" not in body + assert "FORGED" not in json.dumps(body) + + +@pytest.mark.parametrize("driver", CONVERSE_DRIVERS) +def test_benign_caller_body_metadata_still_survives_when_no_identity_resolves(driver): + """Removing the field must be scoped to the reserved keys being the only thing left, not a + blanket drop of the caller's own attribution pairs.""" + litellm.bedrock_request_metadata_fields = ALL_FIELDS + + body = driver( + litellm_params("metadata"), + {"requestMetadata": {"cost_center": "cc-9", "user_api_key_team_alias": "FORGED"}}, + ) + + assert body["requestMetadata"] == {"cost_center": "cc-9"} + + +@pytest.mark.parametrize("driver", CONVERSE_DRIVERS) +def test_caller_body_metadata_is_left_alone_when_forwarding_is_off(driver): + """With the feature off the proxy does not own the field, so the pre-existing pass-through + behaviour for a caller-supplied `requestMetadata` must be unchanged.""" + litellm.bedrock_request_metadata_fields = None + caller_supplied = {"user_api_key_team_alias": "caller-set", "cost_center": "cc-9"} + + body = driver(litellm_params("metadata", **IDENTITY), {"requestMetadata": caller_supplied}) + + assert body["requestMetadata"] == caller_supplied + + +@pytest.mark.parametrize("driver", HEADER_DRIVERS) +def test_a_caller_header_is_left_alone_when_forwarding_is_off(driver): + """The proxy only claims the name when the operator turned forwarding on; with the feature + off this is an ordinary passthrough header and stripping it would be a regression.""" + litellm.bedrock_request_metadata_fields = None + + headers = driver({BEDROCK_REQUEST_METADATA_HEADER: FORGED}, litellm_params("metadata", **IDENTITY)) + + assert metadata_header_values(headers) == [FORGED]