diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 66ee920afd0..215db3a4ee4 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -66,7 +66,9 @@ from litellm.llms.bedrock.common_utils import ( ) from litellm.llms.bedrock.request_metadata import ( bedrock_request_metadata_headers, - bedrock_request_metadata_is_owned, + bedrock_request_metadata_is_rewritten, + bedrock_request_metadata_violation, + caller_request_metadata_violations_are_dropped, merge_bedrock_invoke_headers, resolve_bedrock_request_metadata, ) @@ -354,77 +356,15 @@ class AmazonConverseConfig(BaseConfig): and v is not None } - def _validate_request_metadata(self, metadata: dict) -> None: - """ - Validate requestMetadata according to AWS Bedrock Converse API constraints. - - Constraints: - - Maximum of 16 items - - Keys: 1-256 characters, pattern [a-zA-Z0-9\\s:_@$#=/+,-.]{1,256} - - Values: 0-256 characters, pattern [a-zA-Z0-9\\s:_@$#=/+,-.]{0,256} - """ - import re - - if not isinstance(metadata, dict): + def _validate_request_metadata(self, metadata: object) -> None: + violation: Final = bedrock_request_metadata_violation(metadata) + if violation is not None: raise litellm.exceptions.BadRequestError( - message="requestMetadata must be a dictionary", + message=violation.message, model="bedrock", llm_provider="bedrock", ) - if len(metadata) > 16: - raise litellm.exceptions.BadRequestError( - message="requestMetadata can contain a maximum of 16 items", - model="bedrock", - llm_provider="bedrock", - ) - - 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}$") - - for key, value in metadata.items(): - if not isinstance(key, str): - raise litellm.exceptions.BadRequestError( - message="requestMetadata keys must be strings", - model="bedrock", - llm_provider="bedrock", - ) - - if not isinstance(value, str): - raise litellm.exceptions.BadRequestError( - message="requestMetadata values must be strings", - model="bedrock", - llm_provider="bedrock", - ) - - if len(key) == 0 or len(key) > 256: - raise litellm.exceptions.BadRequestError( - message="requestMetadata key length must be 1-256 characters", - model="bedrock", - llm_provider="bedrock", - ) - - if len(value) > 256: - raise litellm.exceptions.BadRequestError( - message="requestMetadata value length must be 0-256 characters", - model="bedrock", - llm_provider="bedrock", - ) - - if not key_pattern.match(key): - raise litellm.exceptions.BadRequestError( - message=f"requestMetadata key '{key}' contains invalid characters. Allowed: [a-zA-Z0-9\\s:_@$#=/+,.-]", - model="bedrock", - llm_provider="bedrock", - ) - - if not value_pattern.match(value): - raise litellm.exceptions.BadRequestError( - message=f"requestMetadata value '{value}' contains invalid characters. Allowed: [a-zA-Z0-9\\s:_@$#=/+,.-]", - model="bedrock", - llm_provider="bedrock", - ) - @staticmethod def _is_openai_gpt_reasoning_model(model: str) -> bool: return re.search(r"openai\.gpt-\d", model) is not None @@ -1244,7 +1184,7 @@ class AmazonConverseConfig(BaseConfig): elif param == "context_management" and isinstance(value, (dict, list)): self._map_context_management_param(value, optional_params) if param == "requestMetadata": - self._map_request_metadata_param(value, optional_params) + self._map_request_metadata_param(value, optional_params, drop_params) if param == "service_tier" and isinstance(value, str): self._map_service_tier_param(value, optional_params) @@ -1277,9 +1217,10 @@ class AmazonConverseConfig(BaseConfig): return optional_params - def _map_request_metadata_param(self, value: object, optional_params: dict) -> None: + def _map_request_metadata_param(self, value: object, optional_params: dict, drop_params: bool) -> None: if value is not None and isinstance(value, dict): - self._validate_request_metadata(value) + if not caller_request_metadata_violations_are_dropped(drop_params): + self._validate_request_metadata(value) optional_params["requestMetadata"] = value def _map_context_management_param(self, value: dict | list, optional_params: dict) -> None: @@ -1737,7 +1678,7 @@ class AmazonConverseConfig(BaseConfig): # Extract requestMetadata before processing other parameters request_metadata: Final = inference_params.pop("requestMetadata", None) - if request_metadata is not None: + if request_metadata is not None and not caller_request_metadata_violations_are_dropped(drop_params): self._validate_request_metadata(request_metadata) output_config: OutputConfigBlock | None = inference_params.pop("outputConfig", None) @@ -2164,7 +2105,7 @@ class AmazonConverseConfig(BaseConfig): request_metadata: Final = resolve_bedrock_request_metadata( litellm_params=litellm_params, caller_metadata=_data.get("requestMetadata") ) - if bedrock_request_metadata_is_owned(): + if bedrock_request_metadata_is_rewritten(litellm_params.get("drop_params") is True): _data.pop("requestMetadata", None) if request_metadata is not None: _data["requestMetadata"] = request_metadata @@ -2244,7 +2185,7 @@ class AmazonConverseConfig(BaseConfig): request_metadata: Final = resolve_bedrock_request_metadata( litellm_params=litellm_params, caller_metadata=_data.get("requestMetadata") ) - if bedrock_request_metadata_is_owned(): + if bedrock_request_metadata_is_rewritten(litellm_params.get("drop_params") is True): _data.pop("requestMetadata", None) if request_metadata is not None: _data["requestMetadata"] = request_metadata diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py index b94435345d3..ce28e579dcf 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py @@ -95,7 +95,12 @@ class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig): model_info: Mapping[str, object] | None = None, ) -> dict: return AmazonConverseConfig.map_openai_params( - self, non_default_params, optional_params, model, drop_params, model_info + self, + {key: value for key, value in non_default_params.items() if key != "requestMetadata"}, + optional_params, + model, + drop_params, + model_info, ) def transform_request( @@ -110,7 +115,9 @@ class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig): self, model=model, messages=messages, - optional_params=_without_tool_config_injection_points(optional_params), + optional_params=_without_tool_config_injection_points( + {key: value for key, value in optional_params.items() if key != "requestMetadata"} + ), litellm_params=litellm_params, headers=headers, ) diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index 4d40e0dec85..61e084fe544 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -1021,7 +1021,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): converse_config: Final = AmazonConverseConfig() mapped_params = converse_config.map_openai_params( - non_default_params=optional_params, + non_default_params={k: v for k, v in optional_params.items() if k != "requestMetadata"}, optional_params={}, model=model, drop_params=False, diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py index fb0ecd5b825..36cd87aff06 100644 --- a/litellm/llms/bedrock/passthrough/transformation.py +++ b/litellm/llms/bedrock/passthrough/transformation.py @@ -10,6 +10,7 @@ from httpx import Response from pydantic import TypeAdapter from litellm._logging import verbose_logger +from litellm.constants import AWS_ERROR_TYPE_HEADER from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, PassthroughStreamCollector from litellm.types.llms.openai import ChatCompletionSystemMessage, ChatCompletionUserMessage @@ -17,6 +18,16 @@ from litellm.types.utils import ModelResponseStream from ..base_aws_llm import BaseAWSLLM from ..common_utils import BedrockError, BedrockEventStreamDecoderBase, BedrockModelInfo +from ..request_metadata import ( + BEDROCK_REQUEST_METADATA_HEADER, + BedrockRequestMetadataViolation, + bedrock_request_metadata_headers, + bedrock_request_metadata_is_owned, + bedrock_request_metadata_violation, + caller_request_metadata_violations_are_dropped, + merge_bedrock_invoke_headers, + resolve_bedrock_request_metadata, +) if TYPE_CHECKING: from botocore.eventstream import EventStreamMessage @@ -28,10 +39,87 @@ if TYPE_CHECKING: BEDROCK_STREAMING_ACTIONS: Final = frozenset({"invoke-with-response-stream", "converse-stream"}) +_BEDROCK_CONVERSE_ACTIONS: Final = frozenset({"converse", "converse-stream"}) +_BEDROCK_INVOKE_ACTIONS: Final = frozenset({"invoke", "invoke-with-response-stream"}) + + +def _bedrock_endpoint_action(endpoint: str) -> str: + return endpoint.partition("?")[0].rstrip("/").rsplit("/", 1)[-1] def is_bedrock_streaming_endpoint(endpoint: str) -> bool: - return endpoint.partition("?")[0].rstrip("/").rsplit("/", 1)[-1] in BEDROCK_STREAMING_ACTIONS + return _bedrock_endpoint_action(endpoint) in BEDROCK_STREAMING_ACTIONS + + +def _caller_request_metadata_violation( + action: str, body: Mapping[str, object] +) -> BedrockRequestMetadataViolation | None: + if action in _BEDROCK_INVOKE_ACTIONS and "requestMetadata" in body and body["requestMetadata"] is None: + return BedrockRequestMetadataViolation("type", "requestMetadata must be a dictionary") + return bedrock_request_metadata_violation(body.get("requestMetadata")) + + +def _aws_error_type(action: str, violation: BedrockRequestMetadataViolation) -> str: + if action in _BEDROCK_CONVERSE_ACTIONS and violation.kind == "type": + return "SerializationException" + return "ValidationException" + + +def _caller_request_metadata(action: str, body: Mapping[str, object], litellm_params: Mapping[str, object]) -> object: + violation: Final = ( + None + if caller_request_metadata_violations_are_dropped(litellm_params.get("drop_params") is True) + else _caller_request_metadata_violation(action, body) + ) + if violation is not None: + raise BedrockError( + status_code=400, + message=violation.message, + headers={AWS_ERROR_TYPE_HEADER: _aws_error_type(action, violation)}, + ) + return body.get("requestMetadata") + + +def _with_body_request_metadata( + action: str, body: Mapping[str, object], litellm_params: Mapping[str, object] +) -> Mapping[str, object]: + resolved: Final = resolve_bedrock_request_metadata( + litellm_params, caller_metadata=_caller_request_metadata(action, body, litellm_params) + ) + without_caller: Final = {key: value for key, value in body.items() if key != "requestMetadata"} + return without_caller if resolved is None else {**without_caller, "requestMetadata": resolved} + + +def _without_owned_request_metadata_header(headers: Mapping[str, str]) -> Mapping[str, str]: + if not bedrock_request_metadata_is_owned(): + return headers + return {name: value for name, value in headers.items() if name.lower() != BEDROCK_REQUEST_METADATA_HEADER.lower()} + + +def _with_request_metadata( + action: str | None, + headers: Mapping[str, str], + body: Mapping[str, object], + litellm_params: Mapping[str, object], +) -> tuple[Mapping[str, str], Mapping[str, object]]: + if action in _BEDROCK_CONVERSE_ACTIONS: + return headers, _with_body_request_metadata(action, body, litellm_params) + if action in _BEDROCK_INVOKE_ACTIONS and "requestMetadata" in body: + return _without_owned_request_metadata_header(headers), _with_body_request_metadata( + action, body, litellm_params + ) + if action in _BEDROCK_INVOKE_ACTIONS: + return _with_invoke_request_metadata(headers, body, litellm_params) + return headers, body + + +def _with_invoke_request_metadata( + headers: Mapping[str, str], + body: Mapping[str, object], + litellm_params: Mapping[str, object], +) -> tuple[Mapping[str, str], Mapping[str, object]]: + owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params) + return merge_bedrock_invoke_headers(dict(headers), (), metadata_headers, owned_names), body _TEXT_ONLY_DELTA_FIELDS: Final = frozenset({"content", "role"}) @@ -284,18 +372,24 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD def sign_request( self, - headers: dict, - litellm_params: dict, - request_data: dict | None, + headers: dict[str, str], + litellm_params: dict[str, object], + request_data: dict[str, object] | None, api_base: str, model: str | None = None, ) -> tuple[dict, bytes | None]: optional_params: Final = litellm_params.copy() + signed_headers, signed_body = _with_request_metadata( + action=_bedrock_endpoint_action(api_base) if bedrock_request_metadata_is_owned() else None, + headers=headers, + body=request_data or {}, + litellm_params=litellm_params, + ) return self._sign_request( service_name="bedrock", - headers=headers, + headers=dict(signed_headers), optional_params=optional_params, - request_data=request_data or {}, + request_data=dict(signed_body), api_base=api_base, model=model, api_key=self.get_bedrock_bearer_token(optional_params), diff --git a/litellm/llms/bedrock/request_metadata.py b/litellm/llms/bedrock/request_metadata.py index 1f4e5886508..71a2dc78578 100644 --- a/litellm/llms/bedrock/request_metadata.py +++ b/litellm/llms/bedrock/request_metadata.py @@ -8,17 +8,22 @@ sets ``litellm.bedrock_request_metadata_fields`` (``litellm_settings`` on the pr 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. +looking key. Identity and ``spend_logs_metadata`` 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. For the same reason an Invoke header that would pass Bedrock's 8,500-character limit +keeps only the leading pairs that fit. A caller's own malformed ``requestMetadata`` is rejected with +a 400 unless ``drop_params`` is set, in which case its bad pairs are dropped the same way, with or +without the setting. """ from __future__ import annotations import json import re -from collections.abc import Mapping -from typing import Final +from collections.abc import Mapping, Sized +from dataclasses import dataclass +from typing import Final, Literal import litellm @@ -26,15 +31,27 @@ 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" +BEDROCK_REQUEST_METADATA_HEADER_MAX_CHARS: Final = 8500 _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}$") +_ALLOWED_CHARACTERS: Final = r"[a-zA-Z0-9\s:_@$#=/+,.-]" +_KEY_PATTERN: Final = re.compile(rf"{_ALLOWED_CHARACTERS}{{1,256}}", re.ASCII) +_VALUE_PATTERN: Final = re.compile(rf"{_ALLOWED_CHARACTERS}{{0,256}}", re.ASCII) _OWNED_HEADER_NAMES: Final[frozenset[str]] = frozenset((BEDROCK_REQUEST_METADATA_HEADER.lower(),)) +@dataclass(frozen=True, slots=True) +class BedrockRequestMetadataViolation: + kind: Literal["type", "constraint"] + message: str + + def _is_forwardable(key: str, value: str) -> bool: - return _KEY_PATTERN.match(key) is not None and _VALUE_PATTERN.match(value) is not None + return _KEY_PATTERN.fullmatch(key) is not None and _VALUE_PATTERN.fullmatch(value) is not None + + +def _is_reserved(key: str) -> bool: + return key.strip().lower().startswith(BEDROCK_REQUEST_METADATA_IDENTITY_PREFIX) def _text_pairs(source: object) -> tuple[tuple[str, str], ...]: @@ -99,13 +116,16 @@ def _client_pairs( 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) + if not _is_reserved(key) 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] + first_index: Final = {key: index for index, (key, _) in reversed(tuple(enumerate(candidates)))} + return tuple(pair for index, pair in enumerate(candidates) if first_index[pair[0]] == index)[:budget] + + +def _forwardable_caller_pairs(caller_metadata: object) -> tuple[tuple[str, str], ...]: + return tuple(pair for pair in _text_pairs(caller_metadata) if _is_forwardable(*pair))[ + :BEDROCK_REQUEST_METADATA_MAX_PAIRS + ] def resolve_bedrock_request_metadata( @@ -113,17 +133,18 @@ def resolve_bedrock_request_metadata( 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]``. + Resolve the ``requestMetadata`` pairs to send to Bedrock, or ``None`` when 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. + ``caller_metadata`` is any ``requestMetadata`` the caller passed explicitly. Callers reject a + malformed one with a 400 unless ``drop_params`` is set, so here its pairs are only filtered for + Bedrock's rules, the reserved identity prefix, and the remaining slot budget. With the feature + off only Bedrock's rules apply, since no identity is written for a caller to forge. """ allowed_fields: Final = _allowed_fields() if not allowed_fields: - return None + return dict(_forwardable_caller_pairs(caller_metadata)) or None sources: Final = _metadata_sources(litellm_params) identity: Final = _identity_pairs(sources, allowed_fields) client: Final = _client_pairs( @@ -136,6 +157,36 @@ def resolve_bedrock_request_metadata( return resolved or None +def _pair_violation(key: str, value: str) -> BedrockRequestMetadataViolation | None: + if _KEY_PATTERN.fullmatch(key) is None: + return BedrockRequestMetadataViolation( + "constraint", f"requestMetadata key {key!r} must be 1-256 characters of {_ALLOWED_CHARACTERS}" + ) + if _VALUE_PATTERN.fullmatch(value) is None: + return BedrockRequestMetadataViolation( + "constraint", f"requestMetadata value for key {key!r} must be 0-256 characters of {_ALLOWED_CHARACTERS}" + ) + return None + + +def bedrock_request_metadata_violation(caller_metadata: object) -> BedrockRequestMetadataViolation | None: + """Why a caller's ``requestMetadata`` breaks Bedrock's rules, in the order Bedrock checks them: a + wrong JSON type before the pair count and character rules. ``None`` means absent or valid.""" + text_pairs: Final = _text_pairs(caller_metadata) + entry_count: Final = len(caller_metadata) if isinstance(caller_metadata, Sized) else 0 + if caller_metadata is None: + return None + if not isinstance(caller_metadata, Mapping): + return BedrockRequestMetadataViolation("type", "requestMetadata must be a dictionary") + if len(text_pairs) != entry_count: + return BedrockRequestMetadataViolation("type", "requestMetadata keys and values must be strings") + if len(text_pairs) > BEDROCK_REQUEST_METADATA_MAX_PAIRS: + return BedrockRequestMetadataViolation( + "constraint", f"requestMetadata can contain a maximum of {BEDROCK_REQUEST_METADATA_MAX_PAIRS} items" + ) + return next((violation for key, value in text_pairs if (violation := _pair_violation(key, value))), None) + + def bedrock_request_metadata_is_owned() -> bool: """ Whether the proxy OWNS the request-metadata field and header name for this request. @@ -149,25 +200,51 @@ def bedrock_request_metadata_is_owned() -> bool: return bool(_allowed_fields()) +def caller_request_metadata_violations_are_dropped(drop_params: bool) -> bool: + return drop_params or litellm.drop_params is True + + +def bedrock_request_metadata_is_rewritten(drop_params: bool) -> bool: + return bedrock_request_metadata_is_owned() or caller_request_metadata_violations_are_dropped(drop_params) + + +def _header_value(pairs: Mapping[str, str]) -> str: + return json.dumps(pairs, separators=(",", ":")) + + +def _fits_header(pairs: Mapping[str, str]) -> bool: + return len(_header_value(pairs)) <= BEDROCK_REQUEST_METADATA_HEADER_MAX_CHARS + + +def _longest_fitting_prefix(resolved: Mapping[str, str]) -> dict[str, str]: + pairs: Final = tuple(resolved.items()) + return next( + (dict(pairs[:count]) for count in range(len(pairs), 0, -1) if _fits_header(dict(pairs[:count]))), + {}, + ) + + 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. + The signed ``X-Amzn-Bedrock-Request-Metadata`` header for an Invoke request whose body carries + no ``requestMetadata`` of its own. 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. + than fall back to it. Only the leading pairs that fit Bedrock's header limit are sent, so + identity always survives and long ``spend_logs_metadata`` pairs drop off the end. """ if not bedrock_request_metadata_is_owned(): return frozenset(), () resolved: Final = resolve_bedrock_request_metadata(litellm_params) - if resolved is None: + fitting: Final = {} if resolved is None else _longest_fitting_prefix(resolved) + if not fitting: return _OWNED_HEADER_NAMES, () - return _OWNED_HEADER_NAMES, ((BEDROCK_REQUEST_METADATA_HEADER, json.dumps(resolved, separators=(",", ":"))),) + return _OWNED_HEADER_NAMES, ((BEDROCK_REQUEST_METADATA_HEADER, _header_value(fitting)),) def merge_bedrock_invoke_headers( diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 464cfee270a..3aef5724bdc 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -1087,6 +1087,16 @@ def _guardrail_refusal_status(error: Exception) -> int | None: return None +def _aws_error_type(error: Exception, handled: ProxyException) -> str | None: + """The AWS error type to answer the handled error with: ValidationException for a guardrail block, else the type + the error carried, as LiteLLM's own Bedrock request checks set it. None when the shared handler changed the + guardrail block's status.""" + refusal_status: Final = _guardrail_refusal_status(error) + if refusal_status is not None: + return "ValidationException" if str(refusal_status) == handled.code else None + return next((value for name, value in handled.headers.items() if name.lower() == AWS_ERROR_TYPE_HEADER), None) + + def _bedrock_error_response(error: Exception, handled: HTTPException) -> Response | None: """Bedrock's own error response, for an error Bedrock answered the passthrough request with an AWS error type.""" @@ -1111,9 +1121,9 @@ async def bedrock_passthrough_error_response( user_api_key_dict: UserAPIKeyAuth, proxy_logging_obj: ProxyLoggingType, ) -> Response: - """Run the shared error handling for a Bedrock passthrough request, then answer a guardrail block or an error - Bedrock returned the way Bedrock does, so an AWS SDK reads its code and message. Anything else is raised as the - shared handler raises it.""" + """Run the shared error handling for a Bedrock passthrough request, then answer a guardrail block, an error + Bedrock returned, or an error carrying an AWS error type the way Bedrock does, so an AWS SDK reads its code and + message. Anything else is raised as the shared handler raises it.""" try: await processor.handle_llm_api_exception( e=error, @@ -1126,9 +1136,9 @@ async def bedrock_passthrough_error_response( raise return response except ProxyException as handled: - refusal_status: Final = _guardrail_refusal_status(error) - if refusal_status is not None and str(refusal_status) == handled.code: - raise AwsErrorTypeProxyException(handled, "ValidationException") from handled + aws_error_type: Final = _aws_error_type(error, handled) + if aws_error_type is not None: + raise AwsErrorTypeProxyException(handled, aws_error_type) from handled raise diff --git a/tests/integration/providers/test_bedrock_passthrough_request_metadata_wire.py b/tests/integration/providers/test_bedrock_passthrough_request_metadata_wire.py new file mode 100644 index 00000000000..5fd12405ab2 --- /dev/null +++ b/tests/integration/providers/test_bedrock_passthrough_request_metadata_wire.py @@ -0,0 +1,316 @@ +import base64 +import json +import uuid +from urllib.parse import unquote +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import boto3 +import pytest +import yaml +from botocore.config import Config +from botocore.exceptions import ClientError +from integration._support.budget import case_timeout_seconds +from integration._support.client import Gateway +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.upstream import _aws_event_frame +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0" +_EVENT_STREAM: Final = "application/vnd.amazon.eventstream" +_METADATA_HEADER: Final = "x-amzn-bedrock-request-metadata" +_HEADER_MAX_CHARS: Final = 8500 +_LONG_SPEND_LOGS_METADATA: Final = tuple((f"{index:02d}".ljust(256, "k"), "\n" * 256) for index in range(12)) +_CONVERSE_REPLY: Final[Mapping[str, JsonValue]] = { + "output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 5, "outputTokens": 1, "totalTokens": 6}, + "metrics": {"latencyMs": 3}, +} +_INVOKE_REPLY: Final[Mapping[str, JsonValue]] = { + "id": "msg_bedrock", + "type": "message", + "role": "assistant", + "model": "claude", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 1}, +} +_CONVERSE_STREAM_FRAMES: Final = ( + _aws_event_frame("messageStart", {"role": "assistant"}, "sc", "u"), + _aws_event_frame("messageStop", {"stopReason": "end_turn"}, "sc", "u"), + _aws_event_frame("metadata", {"usage": {"inputTokens": 5, "outputTokens": 1, "totalTokens": 6}}, "sc", "u"), +) + + +def _bedrock_peer(request: Request) -> Reply: + if request.target.endswith("/converse-stream"): + return Reply(chunks=_CONVERSE_STREAM_FRAMES, content_type=_EVENT_STREAM) + if request.target.endswith("/invoke-with-response-stream"): + chunk: Final = {"bytes": base64.b64encode(json.dumps(_INVOKE_REPLY).encode()).decode()} + return Reply(chunks=(_aws_event_frame("chunk", chunk, "sc", "u"),), content_type=_EVENT_STREAM) + if request.target.endswith("/converse"): + return Reply(body=json.dumps(_CONVERSE_REPLY).encode()) + return Reply(body=json.dumps(_INVOKE_REPLY).encode()) + + +def _messages(marker: str) -> Sequence[JsonValue]: + return [{"role": "user", "content": [{"text": marker}]}] + + +def _invoke_body(marker: str) -> Mapping[str, JsonValue]: + return { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 16, + "messages": [{"role": "user", "content": marker}], + } + + +def _signed_header_names(request: Request) -> frozenset[str]: + return frozenset(request.headers["authorization"].split("SignedHeaders=")[1].split(",")[0].split(";")) + + +def _header_value(pairs: tuple[tuple[str, str], ...]) -> str: + return json.dumps(dict(pairs), separators=(",", ":")) + + +def _send_with_caller_metadata( + runtime: object, action: str, model: str, marker: str, caller_metadata: Mapping[str, JsonValue] +) -> None: + if action == "invoke": + getattr(runtime, "invoke_model")( + modelId=model, body=json.dumps({**_invoke_body(marker), "requestMetadata": caller_metadata}) + ) + return + getattr(runtime, action.replace("-", "_"))( + modelId=model, messages=_messages(marker), requestMetadata=caller_metadata + ) + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + wire: Wire + strict_model: str + dropping_model: str + chat_model: str + key: str + identity: Mapping[str, str] + long_key: str + long_identity: tuple[tuple[str, str], ...] + + def runtime(self, key: str, monkeypatch: pytest.MonkeyPatch): # noqa: ANN201 # boto3 clients are untyped + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", key) + return boto3.client( + "bedrock-runtime", + endpoint_url=f"{str(self.proxy.client.base_url).rstrip('/')}/bedrock", + region_name="us-east-1", + config=Config( + read_timeout=case_timeout_seconds(), + retries={"total_max_attempts": 1}, + parameter_validation=False, + ), + ) + + def received(self, marker: str) -> tuple[Request, ...]: + return tuple(request for request in self.wire.drain() if marker.encode() in request.body) + + +@pytest.fixture(scope="module") +def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + suffix: Final = uuid.uuid4().hex[:12] + strict_model: Final = f"bedrock-metadata-{suffix}" + dropping_model: Final = f"bedrock-metadata-dropping-{suffix}" + chat_model: Final = f"bedrock-metadata-chat-{suffix}" + directory: Final = tmp_path_factory.mktemp("bedrock-request-metadata") + base_config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + with wire_server(_bedrock_peer) as wire: + deployment: Final = { + "model": f"bedrock/{_MODEL_ID}", + "api_base": wire.url, + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", + "use_in_pass_through": True, + } + config: Final = { + **base_config, + "model_list": [ + *base_config.get("model_list", []), + {"model_name": strict_model, "litellm_params": deployment}, + {"model_name": dropping_model, "litellm_params": {**deployment, "drop_params": True}}, + { + "model_name": chat_model, + "litellm_params": {**deployment, "model": f"bedrock/converse/{_MODEL_ID}", "drop_params": True}, + }, + ], + "litellm_settings": { + **base_config.get("litellm_settings", {}), + "drop_params": False, + "bedrock_request_metadata_fields": [ + "user_api_key_alias", + "user_api_key_team_alias", + "spend_logs_metadata", + ], + }, + } + config_path: Final = directory / "bedrock-request-metadata.yaml" + config_path.write_text(yaml.safe_dump(config)) + with ( + owned_proxy(session_gateway, directory, {}, config=config_path, workers=2) as proxy, + proxy.scenario() as scenario, + ): + team_alias: Final = f"platform-{suffix}" + team_id: Final = scenario.team(team_alias=team_alias) + key_alias: Final = f"prod-key-{suffix}" + long_key_alias: Final = f"long-key-{suffix}" + yield Rig( + proxy=proxy, + wire=wire, + strict_model=strict_model, + dropping_model=dropping_model, + chat_model=chat_model, + key=scenario.key( + key_alias=key_alias, + team_id=team_id, + metadata={"spend_logs_metadata": {"agent": "triage-bot"}}, + ), + identity={ + "user_api_key_alias": key_alias, + "user_api_key_team_alias": team_alias, + "agent": "triage-bot", + }, + long_key=scenario.key( + key_alias=long_key_alias, + team_id=team_id, + metadata={"spend_logs_metadata": dict(_LONG_SPEND_LOGS_METADATA)}, + ), + long_identity=( + ("user_api_key_alias", long_key_alias), + ("user_api_key_team_alias", team_alias), + *_LONG_SPEND_LOGS_METADATA, + ), + ) + + +def test_boto3_passthrough_calls_carry_the_key_identity_in_signed_request_metadata( + rig: Rig, monkeypatch: pytest.MonkeyPatch +) -> None: + marker: Final = f"who-is-calling-{uuid.uuid4().hex}" + forged: Final = {"user_api_key_team_alias": "forged", "ticket": "T-42", "bad!": "x"} + runtime: Final = rig.runtime(rig.key, monkeypatch) + model: Final = rig.dropping_model + + converse: Final = runtime.converse(modelId=model, messages=_messages(marker), requestMetadata=forged) + converse_stream: Final = [ + next(iter(event)) + for event in runtime.converse_stream(modelId=model, messages=_messages(marker), requestMetadata=forged)[ + "stream" + ] + ] + invoke: Final = json.loads( + runtime.invoke_model(modelId=model, body=json.dumps(_invoke_body(marker)))["body"].read() + ) + invoke_stream: Final = [ + json.loads(event["chunk"]["bytes"]) + for event in runtime.invoke_model_with_response_stream(modelId=model, body=json.dumps(_invoke_body(marker)))[ + "body" + ] + ] + assert converse["output"] == _CONVERSE_REPLY["output"], converse + assert converse_stream == ["messageStart", "messageStop", "metadata"], converse_stream + assert invoke == _INVOKE_REPLY, invoke + assert invoke_stream == [_INVOKE_REPLY], invoke_stream + + received: Final = rig.received(marker) + assert [request.target for request in received] == [ + f"/model/{_MODEL_ID}/converse", + f"/model/{_MODEL_ID}/converse-stream", + f"/model/{_MODEL_ID}/invoke", + f"/model/{_MODEL_ID}/invoke-with-response-stream", + ], received + converse_requests: Final = received[:2] + invoke_requests: Final = received[2:] + assert [json.loads(request.body) for request in converse_requests] == [ + {"messages": _messages(marker), "requestMetadata": {**rig.identity, "ticket": "T-42"}}, + {"messages": _messages(marker), "requestMetadata": {**rig.identity, "ticket": "T-42"}}, + ], [request.body for request in converse_requests] + assert [_METADATA_HEADER in request.headers for request in converse_requests] == [False, False] + assert [json.loads(request.body) for request in invoke_requests] == [_invoke_body(marker), _invoke_body(marker)] + assert [json.loads(request.headers[_METADATA_HEADER]) for request in invoke_requests] == [ + rig.identity, + rig.identity, + ], [request.headers for request in invoke_requests] + assert all(_METADATA_HEADER in _signed_header_names(request) for request in invoke_requests), [ + request.headers["authorization"] for request in invoke_requests + ] + + +def test_invoke_metadata_too_long_for_the_header_signs_the_leading_pairs_that_fit( + rig: Rig, monkeypatch: pytest.MonkeyPatch +) -> None: + marker: Final = f"long-key-{uuid.uuid4().hex}" + runtime: Final = rig.runtime(rig.long_key, monkeypatch) + + invoke: Final = json.loads( + runtime.invoke_model(modelId=rig.strict_model, body=json.dumps(_invoke_body(marker)))["body"].read() + ) + + assert invoke == _INVOKE_REPLY, invoke + (request,) = rig.received(marker) + assert json.loads(request.body) == _invoke_body(marker) + header: Final = request.headers[_METADATA_HEADER] + sent_pairs: Final = len(json.loads(header)) + assert tuple(json.loads(header).items()) == rig.long_identity[:sent_pairs], header + assert 2 <= sent_pairs < len(rig.long_identity), header + assert len(header) <= _HEADER_MAX_CHARS < len(_header_value(rig.long_identity[: sent_pairs + 1])) + assert _METADATA_HEADER in _signed_header_names(request), request.headers["authorization"] + + +@pytest.mark.parametrize( + ("action", "caller_metadata", "error_code"), + [ + pytest.param("converse", {"ticket": 42}, "SerializationException", id="converse-number-value"), + pytest.param("converse-stream", {"bad!": "x"}, "ValidationException", id="converse-stream-bad-key"), + pytest.param("invoke", {"bad!": "x"}, "ValidationException", id="invoke-bad-key"), + ], +) +def test_malformed_caller_request_metadata_reaches_boto3_as_the_bedrock_error_without_an_upstream_call( + rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str, caller_metadata: Mapping[str, JsonValue], error_code: str +) -> None: + marker: Final = f"malformed-{uuid.uuid4().hex}" + runtime: Final = rig.runtime(rig.key, monkeypatch) + + with pytest.raises(ClientError) as raised: + _send_with_caller_metadata(runtime, action, rig.strict_model, marker, caller_metadata) + + assert raised.value.response["Error"]["Code"] == error_code, raised.value.response + assert raised.value.response["ResponseMetadata"]["HTTPStatusCode"] == 400, raised.value.response + assert rig.received(marker) == () + + +def test_chat_on_a_dropping_converse_deployment_sends_only_the_valid_caller_pairs(rig: Rig) -> None: + marker: Final = f"chat-{uuid.uuid4().hex}" + + response: Final = rig.proxy.client.post( + "/chat/completions", + json={ + "model": rig.chat_model, + "messages": [{"role": "user", "content": marker}], + "requestMetadata": {"ticket": "T-42", "bad!": "x"}, + }, + headers={"Authorization": f"Bearer {rig.key}"}, + timeout=case_timeout_seconds(), + ) + + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == "ok", response.text + (request,) = rig.received(marker) + assert unquote(request.target) == f"/model/{_MODEL_ID}/converse", request.target + assert json.loads(request.body)["requestMetadata"] == {**rig.identity, "ticket": "T-42"}, request.body diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_nova_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_nova_transformation.py index f0f0f9160fb..de4b135e6d0 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_nova_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_amazon_nova_transformation.py @@ -115,3 +115,22 @@ def test_tool_config_injection_point_is_neither_placed_nor_credited(local_model_ assert [tool["toolSpec"]["name"] for tool in request["toolConfig"]["tools"]] == ["f"] assert "cachePoint" not in json.dumps(request) assert GATEWAY_INJECTED_CACHE_METADATA_KEY not in metadata + + +@pytest.mark.parametrize("value", ["b c", "bad!"], ids=["nbsp", "outside-bedrock-charset"]) +def test_caller_request_metadata_is_ignored_because_invoke_model_has_no_field_for_it(local_model_cost_map, value): + caller_metadata = {"ticket": value} + optional_params = AmazonInvokeNovaConfig().map_openai_params( + non_default_params={"max_tokens": 8, "requestMetadata": caller_metadata}, + optional_params={}, + model=MODEL, + drop_params=False, + ) + + request = _transform_request( + messages=[{"role": "user", "content": "hi"}], + optional_params={**optional_params, "requestMetadata": caller_metadata}, + ) + + assert "requestMetadata" not in request + assert request["inferenceConfig"] == {"maxTokens": 8} diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index 0ed713f056e..6092934fc7a 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -4199,7 +4199,7 @@ def test_request_metadata_key_constraints(): # Test empty key invalid_metadata = {"": "value"} - with pytest.raises(Exception, match=r"(?i)key length|empty"): + with pytest.raises(Exception, match=r"key '' must be 1-256 characters"): config._transform_request( model="anthropic.claude-haiku-4-5-20251001-v1:0", messages=messages, @@ -4241,6 +4241,18 @@ def test_request_metadata_value_constraints(): ) +@pytest.mark.parametrize("metadata", [{"ticket": "a\u00a0b"}, {"a\u2003b": "ticket"}]) +def test_request_metadata_rejects_whitespace_outside_ascii(metadata): + with pytest.raises(litellm.exceptions.BadRequestError, match=r"must be [01]-256 characters of"): + AmazonConverseConfig()._transform_request( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + messages=[{"role": "user", "content": "Hello!"}], + optional_params={"requestMetadata": metadata}, + litellm_params={}, + headers={}, + ) + + def test_request_metadata_character_pattern(): """Test character pattern validation for keys and values.""" config = AmazonConverseConfig() diff --git a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py index 425beefac8a..b050bb7625b 100644 --- a/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py +++ b/tests/unit/llms/bedrock/files/test_bedrock_files_transformation.py @@ -1760,6 +1760,23 @@ class TestBedrockBatchNonChatEndpointRecords: assert "metadata" not in model_input assert "requestMetadata" not in model_input + def test_nova_record_request_metadata_stays_out_of_the_bedrock_body(self): + model_input = self._transform( + { + "custom_id": "4d", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": self.NOVA_MODEL, + "messages": [{"role": "user", "content": "hi"}], + "requestMetadata": {"ticket": "b\u00a0c"}, + }, + } + ) + + assert "requestMetadata" not in model_input + assert model_input["messages"] == [{"role": "user", "content": [{"text": "hi"}]}] + @pytest.mark.parametrize( "body", [ diff --git a/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py index 84d3160e609..da5c4699925 100644 --- a/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py +++ b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py @@ -17,6 +17,7 @@ import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.passthrough.transformation import PassthroughStreamCollector +from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.passthrough.transformation import ( BedrockPassthroughConfig, is_bedrock_streaming_endpoint, @@ -898,3 +899,293 @@ def test_stream_collector_logs_nothing_for_an_unrecognized_endpoint(): collector.add(_event_frame("messageStart", {"role": "assistant"})) assert collector.build_logged_response(_converse_stream_logging_obj()) is None + + +REQUEST_METADATA_HEADER = "x-amzn-bedrock-request-metadata" +REQUEST_METADATA_FIELDS = ["user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata"] +PROXY_IDENTITY = { + "user_api_key_alias": "prod-key", + "user_api_key_team_alias": "platform", + "spend_logs_metadata": {"agent": "triage-bot"}, +} +EXPECTED_REQUEST_METADATA = { + "user_api_key_alias": "prod-key", + "user_api_key_team_alias": "platform", + "agent": "triage-bot", +} +CONVERSE_BODY = {"messages": [{"role": "user", "content": [{"text": "hi"}]}]} +INVOKE_BODY = { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 8, + "messages": [{"role": "user", "content": "hi"}], +} + + +async def _sent_passthrough_request( + action: str, body: dict[str, object], identity: dict[str, object], **kwargs: object +) -> httpx.Request: + requests: Final[list[httpx.Request]] = [] + + def respond(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(200, json={"ok": True}, request=request) + + router: Final = litellm.Router( + model_list=[ + { + "model_name": "bedrock-pt", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "aws_access_key_id": "AKIAROUTERKEY", + "aws_secret_access_key": "router-secret", + "aws_region_name": "us-west-2", + }, + } + ] + ) + await router.allm_passthrough_route( + model="bedrock-pt", + method="POST", + endpoint=f"/model/bedrock-pt/{action}", + data=body, + client=AsyncHTTPHandler(transport=httpx.MockTransport(respond)), + litellm_metadata=identity, + **kwargs, + ) + return requests[0] + + +def _signed_header_names(request: httpx.Request) -> frozenset[str]: + return frozenset(request.headers["authorization"].split("SignedHeaders=")[1].split(",")[0].split(";")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("action", "query_params"), + [("converse", None), ("converse-stream", None), ("converse/", None), ("converse", {"trace": "1"})], +) +async def test_converse_passthrough_signs_proxy_identity_into_request_metadata( + action: str, query_params: dict[str, str] | None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + caller_body: Final = {**CONVERSE_BODY, "requestMetadata": {"user_api_key_team_alias": "forged", "ticket": "T-42"}} + + request: Final = await _sent_passthrough_request( + action, caller_body, PROXY_IDENTITY, request_query_params=query_params + ) + + assert json.loads(request.content) == { + **CONVERSE_BODY, + "requestMetadata": {**EXPECTED_REQUEST_METADATA, "ticket": "T-42"}, + } + assert caller_body["requestMetadata"] == {"user_api_key_team_alias": "forged", "ticket": "T-42"} + assert REQUEST_METADATA_HEADER not in request.headers + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +async def test_invoke_passthrough_sends_proxy_identity_as_a_signed_header( + action: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + + request: Final = await _sent_passthrough_request( + action, + INVOKE_BODY, + PROXY_IDENTITY, + request_headers={f"x-pass-{REQUEST_METADATA_HEADER}": '{"user_api_key_team_alias":"forged"}'}, + ) + + assert json.loads(request.headers[REQUEST_METADATA_HEADER]) == EXPECTED_REQUEST_METADATA + assert REQUEST_METADATA_HEADER in _signed_header_names(request) + assert json.loads(request.content) == INVOKE_BODY + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +async def test_invoke_passthrough_keeps_caller_body_request_metadata_in_the_body( + action: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + caller_body: Final = {**INVOKE_BODY, "requestMetadata": {"user_api_key_team_alias": "forged", "ticket": "T-42"}} + + request: Final = await _sent_passthrough_request( + action, + caller_body, + PROXY_IDENTITY, + request_headers={f"x-pass-{REQUEST_METADATA_HEADER.upper()}": '{"user_api_key_team_alias":"forged"}'}, + ) + + assert json.loads(request.content) == { + **INVOKE_BODY, + "requestMetadata": {**EXPECTED_REQUEST_METADATA, "ticket": "T-42"}, + } + assert REQUEST_METADATA_HEADER not in request.headers + + +@pytest.mark.asyncio +@pytest.mark.parametrize("setting", [None, []]) +async def test_passthrough_requests_are_untouched_when_forwarding_is_off( + setting: list[str] | None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", setting) + caller_body: Final = {**CONVERSE_BODY, "requestMetadata": {"user_api_key_team_alias": "forged"}} + + converse: Final = await _sent_passthrough_request("converse", caller_body, PROXY_IDENTITY) + invoke: Final = await _sent_passthrough_request("invoke", INVOKE_BODY, PROXY_IDENTITY) + + assert json.loads(converse.content) == caller_body + assert REQUEST_METADATA_HEADER not in invoke.headers + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("action", "body"), [("converse", CONVERSE_BODY), ("invoke", INVOKE_BODY)]) +async def test_passthrough_drops_reserved_caller_keys_when_nothing_resolves( + action: str, body: dict[str, object], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + caller_body: Final = {**body, "requestMetadata": {"user_api_key_team_alias": "forged"}} + + request: Final = await _sent_passthrough_request(action, caller_body, {}) + + assert json.loads(request.content) == body + assert REQUEST_METADATA_HEADER not in request.headers + + +MALFORMED_REQUEST_METADATA: Final = ( + pytest.param( + {"user_api_key_team_alias": "forged", "bad!": "x"}, "key 'bad!' must be", "ValidationException", id="bad-key" + ), + pytest.param({"ticket": "a\u00a0b"}, "value for key 'ticket' must be", "ValidationException", id="bad-value"), + pytest.param("ticket=T-42", "must be a dictionary", "SerializationException", id="not-a-dictionary"), + pytest.param({f"key_{index}": "v" for index in range(17)}, "maximum of 16 items", "ValidationException", id="17"), + pytest.param({"ticket": 42}, "must be strings", "SerializationException", id="non-string-value"), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["converse", "converse-stream"]) +@pytest.mark.parametrize(("caller_metadata", "violation", "converse_error_type"), MALFORMED_REQUEST_METADATA) +async def test_malformed_converse_request_metadata_is_rejected_with_the_error_type_bedrock_returns( + action: str, caller_metadata: object, violation: str, converse_error_type: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + + with pytest.raises(BedrockError, match=violation) as rejected: + await _sent_passthrough_request(action, {**CONVERSE_BODY, "requestMetadata": caller_metadata}, PROXY_IDENTITY) + + assert rejected.value.status_code == 400 + assert httpx.Headers(rejected.value.headers).get("x-amzn-errortype") == converse_error_type + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +@pytest.mark.parametrize(("caller_metadata", "violation", "converse_error_type"), MALFORMED_REQUEST_METADATA) +async def test_malformed_invoke_request_metadata_is_rejected_as_a_bedrock_validation_error( + action: str, caller_metadata: object, violation: str, converse_error_type: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + + with pytest.raises(BedrockError, match=violation) as rejected: + await _sent_passthrough_request(action, {**INVOKE_BODY, "requestMetadata": caller_metadata}, PROXY_IDENTITY) + + assert rejected.value.status_code == 400 + assert httpx.Headers(rejected.value.headers).get("x-amzn-errortype") == "ValidationException" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +async def test_invoke_null_request_metadata_is_rejected_like_bedrock_does( + action: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + + with pytest.raises(BedrockError, match="must be a dictionary") as rejected: + await _sent_passthrough_request(action, {**INVOKE_BODY, "requestMetadata": None}, PROXY_IDENTITY) + + assert rejected.value.status_code == 400 + assert httpx.Headers(rejected.value.headers).get("x-amzn-errortype") == "ValidationException" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["converse", "converse-stream"]) +async def test_converse_null_request_metadata_is_treated_as_absent_like_bedrock_does( + action: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + + request: Final = await _sent_passthrough_request(action, {**CONVERSE_BODY, "requestMetadata": None}, PROXY_IDENTITY) + + assert json.loads(request.content) == {**CONVERSE_BODY, "requestMetadata": EXPECTED_REQUEST_METADATA} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +async def test_invoke_metadata_too_long_for_the_header_sends_the_leading_pairs_that_fit( + action: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + long_pairs: Final = {f"{index:02d}".ljust(256, "k"): "\n" * 256 for index in range(12)} + resolved: Final = tuple( + {"user_api_key_alias": "prod-key", "user_api_key_team_alias": "platform", **long_pairs}.items() + ) + + request: Final = await _sent_passthrough_request( + action, + INVOKE_BODY, + {**PROXY_IDENTITY, "spend_logs_metadata": long_pairs}, + request_headers={f"x-pass-{REQUEST_METADATA_HEADER}": '{"user_api_key_team_alias":"forged"}'}, + ) + + header: Final = request.headers[REQUEST_METADATA_HEADER] + sent: Final = json.loads(header) + one_more_pair: Final = json.dumps(dict(resolved[: len(sent) + 1]), separators=(",", ":")) + assert sent == dict(resolved[: len(sent)]) + assert len(header) <= 8500 < len(one_more_pair) + assert json.loads(request.content) == INVOKE_BODY + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("action", "body"), [("converse", CONVERSE_BODY), ("invoke", INVOKE_BODY)]) +@pytest.mark.parametrize("drop_globally", [True, False]) +async def test_malformed_caller_request_metadata_is_dropped_when_params_are_dropped( + action: str, body: dict[str, object], drop_globally: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + monkeypatch.setattr(litellm, "drop_params", drop_globally) + caller_body: Final = {**body, "requestMetadata": {"user_api_key_team_alias": "forged", "bad!": "x"}} + + dropped: Final = await _sent_passthrough_request( + action, caller_body, PROXY_IDENTITY, **({} if drop_globally else {"drop_params": True}) + ) + + assert json.loads(dropped.content) == {**body, "requestMetadata": EXPECTED_REQUEST_METADATA} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("action", "body", "caller_metadata"), + [ + ("converse", CONVERSE_BODY, {"user_api_key_team_alias": "caller-team", "bad!": "x"}), + ("invoke", INVOKE_BODY, {"user_api_key_team_alias": "caller-team", "bad!": "x"}), + ("invoke", INVOKE_BODY, None), + ], +) +@pytest.mark.parametrize("drop_globally", [True, False]) +async def test_passthrough_body_is_sent_as_is_when_params_are_dropped_and_forwarding_is_off( + action: str, + body: dict[str, object], + caller_metadata: dict[str, str] | None, + drop_globally: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None) + monkeypatch.setattr(litellm, "drop_params", drop_globally) + caller_body: Final = {**body, "requestMetadata": caller_metadata} + + sent: Final = await _sent_passthrough_request( + action, caller_body, PROXY_IDENTITY, **({} if drop_globally else {"drop_params": True}) + ) + + assert json.loads(sent.content) == caller_body + assert REQUEST_METADATA_HEADER not in sent.headers diff --git a/tests/unit/llms/bedrock/test_request_metadata.py b/tests/unit/llms/bedrock/test_request_metadata.py index 5a14bbea9f2..6b74ad90cff 100644 --- a/tests/unit/llms/bedrock/test_request_metadata.py +++ b/tests/unit/llms/bedrock/test_request_metadata.py @@ -1,8 +1,10 @@ import asyncio import json +from collections.abc import Callable, Mapping +from typing import Final import pytest - +from pydantic import TypeAdapter import litellm from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM @@ -18,6 +20,7 @@ from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_tran ) from litellm.llms.bedrock.request_metadata import ( BEDROCK_REQUEST_METADATA_HEADER, + BEDROCK_REQUEST_METADATA_HEADER_MAX_CHARS, BEDROCK_REQUEST_METADATA_MAX_PAIRS, resolve_bedrock_request_metadata, ) @@ -128,6 +131,20 @@ def test_identity_survives_a_caller_filling_every_slot( ) +def test_a_key_sent_twice_keeps_the_caller_value_and_uses_one_slot(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS) + caller: Final = { + f"client_{index:02d}": "caller" for index in range(BEDROCK_REQUEST_METADATA_MAX_PAIRS - len(IDENTITY) - 1) + } + spend_logs: Final = {**{key: "spend" for key in caller}, "extra": "spend"} + + resolved: Final = resolve_bedrock_request_metadata( + litellm_params("metadata", spend_logs_metadata=spend_logs, **IDENTITY), caller_metadata=caller + ) + + assert resolved == {**IDENTITY, **caller, "extra": "spend"} + + @pytest.mark.parametrize( "field_order", [ @@ -158,7 +175,14 @@ def test_a_field_repeated_in_the_allow_list_does_not_consume_a_client_slot( @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"], + [ + "user_api_key_team_alias", + "user_api_key_org_alias", + "user_api_key_hash", + "USER_API_KEY_TEAM_ALIAS", + "User_Api_Key_Hash", + " user_api_key_user_id", + ], ) def test_caller_cannot_forge_or_shadow_a_reserved_identity_key( forged_key, @@ -181,14 +205,21 @@ def test_caller_cannot_forge_or_shadow_a_reserved_identity_key( assert "attacker-controlled" not in resolved.values() -def test_identity_violating_the_character_class_is_dropped_and_the_request_succeeds(monkeypatch: pytest.MonkeyPatch): +@pytest.mark.parametrize( + "team_alias", + ["O'Brien's team", "platform\u00a0team", "platform\u3000team", "x" * 256 + "\n"], + ids=["apostrophe", "no-break-space", "ideographic-space", "257-chars-ending-in-newline"], +) +def test_identity_violating_the_character_class_is_dropped_and_the_request_succeeds( + team_alias: str, monkeypatch: pytest.MonkeyPatch +) -> None: """A team alias with an apostrophe must not turn a working request into a 400 the moment - an operator flips the setting on.""" + an operator flips the setting on. Bedrock also rejects non-ASCII whitespace (checked live 2026-10-16).""" monkeypatch.setattr(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_team_alias=team_alias, user_api_key_user_email="x" * 300, ) @@ -198,13 +229,54 @@ def test_identity_violating_the_character_class_is_dropped_and_the_request_succe assert body["messages"] -def test_caller_supplied_violation_still_raises_bad_request(monkeypatch: pytest.MonkeyPatch): +@pytest.mark.parametrize( + "caller_metadata", + [{"team": "O'Brien's team"}, {"k" * 256 + "\n": "v"}, {"team": "v" * 256 + "\n"}], + ids=["bad-character", "257-char-key-ending-in-newline", "257-char-value-ending-in-newline"], +) +def test_caller_supplied_violation_still_raises_bad_request( + caller_metadata: Mapping[str, str], monkeypatch: pytest.MonkeyPatch +) -> None: monkeypatch.setattr(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"}}, + converse_body(litellm_params("metadata", **IDENTITY), {"requestMetadata": caller_metadata}) + + +@pytest.mark.parametrize("driver", CONVERSE_DRIVERS) +@pytest.mark.parametrize("drop_globally", [True, False]) +@pytest.mark.parametrize(("fields", "forwarded_identity"), [(ALL_FIELDS, IDENTITY), (None, {})]) +def test_caller_violations_are_dropped_like_the_passthrough_routes_when_params_are_dropped( + driver: Callable[[Mapping[str, object], Mapping[str, object]], Mapping[str, object]], + drop_globally: bool, + fields: list[str] | None, + forwarded_identity: Mapping[str, str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", fields) + monkeypatch.setattr(litellm, "drop_params", drop_globally) + params: Final = {**litellm_params("metadata", **IDENTITY), **({} if drop_globally else {"drop_params": True})} + + body: Final = driver(params, {"requestMetadata": {"team": "O'Brien's team", "ticket": "T-42"}}) + + assert body["requestMetadata"] == {**forwarded_identity, "ticket": "T-42"} + + +@pytest.mark.parametrize("fields", [ALL_FIELDS, None]) +def test_mapped_caller_violations_are_dropped_when_params_are_dropped( + fields: list[str] | None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", fields) + caller: Final = {"team": "O'Brien's team", "ticket": "T-42"} + + mapped: Final = AmazonConverseConfig().map_openai_params( + non_default_params={"requestMetadata": caller}, optional_params={}, model=MODEL, drop_params=True + ) + + assert mapped["requestMetadata"] == caller + with pytest.raises(litellm.exceptions.BadRequestError): + AmazonConverseConfig().map_openai_params( + non_default_params={"requestMetadata": caller}, optional_params={}, model=MODEL, drop_params=False ) @@ -320,6 +392,27 @@ def metadata_header_values(headers): return [value for name, value in headers.items() if name.lower() == BEDROCK_REQUEST_METADATA_HEADER.lower()] +@pytest.mark.parametrize("driver", HEADER_DRIVERS) +def test_the_metadata_header_keeps_the_leading_pairs_that_fit_bedrocks_header_limit( + driver: Callable[[Mapping[str, str], Mapping[str, object]], Mapping[str, str]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS) + long_pairs: Final = {f"{index:02d}".ljust(256, "k"): "\n" * 256 for index in range(12)} + params: Final = litellm_params("metadata", spend_logs_metadata=long_pairs, **IDENTITY) + resolved: Final = tuple((resolve_bedrock_request_metadata(params) or {}).items()) + + headers: Final = driver({}, params) + (header,) = (value for name, value in headers.items() if name.lower() == BEDROCK_REQUEST_METADATA_HEADER.lower()) + sent: Final = TypeAdapter(dict[str, str]).validate_json(header) + one_more_pair: Final = json.dumps(dict(resolved[: len(sent) + 1]), separators=(",", ":")) + + assert len(header) <= BEDROCK_REQUEST_METADATA_HEADER_MAX_CHARS + assert tuple(sent.items()) == resolved[: len(sent)] + assert IDENTITY.items() <= sent.items() + assert len(one_more_pair) > BEDROCK_REQUEST_METADATA_HEADER_MAX_CHARS + + def test_converse_still_sets_the_bearer_authorization_header(monkeypatch: pytest.MonkeyPatch): """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.""" diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 21170abaeb1..ed89cc37544 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -33,6 +33,7 @@ from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException, from litellm.proxy._types import AwsErrorTypeProxyException, ProxyException from litellm.proxy.utils import ProxyLogging from litellm.passthrough.main import raise_for_provider_status +from litellm.llms.bedrock.common_utils import BedrockError from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( @@ -2137,6 +2138,47 @@ class TestBedrockPassthroughAwsErrors: assert exc_info.value.aws_error_type == "ValidationException" assert (exc_info.value.error.message, exc_info.value.error.code) == (message, status_code) + + @pytest.mark.asyncio + async def test_bedrock_request_check_is_raised_with_the_aws_error_type_it_carries( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + error: Final = BedrockError( + status_code=400, + message="requestMetadata can contain a maximum of 16 items", + headers={"x-amzn-errortype": "ValidationException"}, + ) + + with pytest.raises(AwsErrorTypeProxyException) as exc_info: + await self._answer(error) + + assert exc_info.value.aws_error_type == "ValidationException" + assert "requestMetadata can contain a maximum of 16 items" in exc_info.value.error.message + assert exc_info.value.error.code == "400" + + @pytest.mark.asyncio + async def test_bedrock_error_without_an_aws_error_type_keeps_the_proxy_shape( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + + with pytest.raises(ProxyException) as exc_info: + await self._answer(BedrockError(status_code=400, message="bad request")) + + assert exc_info.value.code == "400" + + @pytest.mark.asyncio + async def test_error_with_a_non_ascii_header_value_keeps_the_proxy_shape( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + error: Final = HTTPException(status_code=403, detail="model not allowed", headers={"x-model-group": "modèle"}) + + with pytest.raises(ProxyException) as exc_info: + await self._answer(error) + + assert (exc_info.value.code, exc_info.value.headers.get("x-model-group")) == ("403", "modèle") @pytest.mark.asyncio @pytest.mark.parametrize( ("error", "message", "code"),