diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 281d2e8d3e1..2dbc3b190cc 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -19,7 +19,7 @@ from collections.abc import ( from dataclasses import dataclass from datetime import datetime from functools import partial, reduce -from itertools import chain, count, groupby, takewhile +from itertools import chain, count, groupby, product, takewhile from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, TypedDict, TypeVar, cast from urllib.parse import parse_qs, urlencode, urlparse @@ -124,7 +124,12 @@ from litellm.proxy.common_request_processing import ( open_sse_before_first_byte, resolve_litellm_call_id, ) -from litellm.proxy.common_utils.callback_utils import strip_callback_config +from litellm.proxy.common_utils.callback_utils import ( + GUARDRAIL_SCAN_IDS_METADATA_KEY, + GUARDRAIL_SCAN_METADATA_METADATA_KEY, + get_logging_caching_headers, # pyright: ignore[reportUnknownVariableType] # untyped request_data param + strip_callback_config, +) from litellm.proxy.common_utils.error_body_call_id import JSON_OBJECT, error_body_call_id, with_call_id from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports _read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -204,6 +209,7 @@ _blocked_response_cost_logging: Final = PassThroughEndpointLogging(log_dispatch= _METADATA_KEYS: Final = frozenset(("litellm_metadata", "metadata")) _DECODED_JSON: Final = TypeAdapter(object) _GUARDRAIL_ENTRIES: Final = TypeAdapter(list[object]) +_RESPONSE_HEADERS: Final = TypeAdapter(Mapping[str, str]) _FIELDS_GUARDRAILS_READ_BY_SHAPE: Final = ( ("guardrails", (list, dict)), ("guardrail_config", (dict,)), @@ -224,12 +230,172 @@ _POST_CALL_SELECTION_FIELDS: Final = ( ) _KEPT_OUT_OF_LITELLM_PARAMS: Final = _METADATA_KEYS | frozenset(wire_names(ProxyRequestState)) _PRICING_MODEL_OVERRIDE_PARAMS: Final = frozenset(("base_model",)) +_GUARDRAIL_HEADER_METADATA_FIELDS: Final = frozenset( + ( + "applied_guardrails", + "applied_policies", + "policy_sources", + GUARDRAIL_SCAN_IDS_METADATA_KEY, + GUARDRAIL_SCAN_METADATA_METADATA_KEY, + ) +) +_JOINED_HEADER_METADATA_FIELDS: Final = frozenset( + ("applied_guardrails", "applied_policies", GUARDRAIL_SCAN_IDS_METADATA_KEY) +) +# RFC 9110's field-value, the one both h11 and Node's HTTP parser accept: no control characters, and only spaces and +# tabs between visible characters. +_SENDABLE_HEADER_VALUE: Final = re.compile(rb"(?:[\x21-\x7e\x80-\xff]+(?:[ \t]+[\x21-\x7e\x80-\xff]+)*)?") +# Half of the 16 KB response header limit in Node.js, the smallest common client limit, leaving room for the others. +_GUARDRAIL_HEADERS_MAX_BYTES: Final = 8192 +# Node.js's limit for a whole response, less 1 KB for the status line and what the server adds. +_RESPONSE_HEADERS_MAX_BYTES: Final = 15360 +# The ": " and the line ending around each header on the wire. +_HEADER_FRAMING_BYTES: Final = 4 def _is_mergeable_client_metadata(key: str, value: object) -> bool: return not key.startswith("user_api_key_") and (key != "model_info" or isinstance(value, Mapping)) +def _strip_client_guardrail_fields( + data: dict[str, object], # mutable-ok: strips the request body in place, like strip_client_guardrail_information +) -> None: + """Drop the guardrail logging and the metadata fields that become guardrail response headers from a client body, + so only the guardrails that ran set them.""" + strip_client_guardrail_information(data) + for metadata in (as_str_mapping(data.get(key)) for key in sorted(_METADATA_KEYS)): + if isinstance(metadata, dict): + for field in _GUARDRAIL_HEADER_METADATA_FIELDS & metadata.keys(): + metadata.pop(field) + + +def _is_name_list(value: object) -> bool: + return isinstance(value, list | tuple) and all( + isinstance(entry, str) for entry in _GUARDRAIL_ENTRIES.validate_python(value) + ) + + +def _joinable_header_metadata(recorded: Mapping[str, object]) -> Mapping[str, object]: + """``recorded`` without a value ``get_logging_caching_headers`` can't join into a header, which a custom guardrail + can record in the name lists.""" + return { + field: value + for field, value in recorded.items() + if field not in _JOINED_HEADER_METADATA_FIELDS or _is_name_list(value) + } + + +def _metadata_entry_identity(entry: object) -> object: + return entry if isinstance(entry, str) else (json.dumps(entry, sort_keys=True, default=str),) + + +def _merged_metadata_entry(prior: object, value: object) -> object: + """``value``, or both combined: lists and tuples without repeats, and mappings with ``value``'s keys winning, as + ``add_policy_sources_to_metadata`` combines policy sources on one request.""" + prior_mapping: Final = as_str_mapping(prior) + value_mapping: Final = as_str_mapping(value) + if prior_mapping is not None and value_mapping is not None: + return {**prior_mapping, **value_mapping} + if not isinstance(prior, list | tuple) or not isinstance(value, list | tuple): + return value + entries: Final = (*_GUARDRAIL_ENTRIES.validate_python(prior), *_GUARDRAIL_ENTRIES.validate_python(value)) + return tuple({_metadata_entry_identity(entry): entry for entry in entries}.values()) + + +def _is_sendable_header_value(value: str) -> bool: + return all(ord(character) < 256 for character in value) and bool( + _SENDABLE_HEADER_VALUE.fullmatch(value.encode("latin-1")) + ) + + +def _with_header_in_budget(kept: Mapping[str, str], header: tuple[str, str]) -> Mapping[str, str]: + size: Final = sum(len(name) + len(value) for name, value in (*kept.items(), header)) + return MappingProxyType({**kept, header[0]: header[1]}) if size <= _GUARDRAIL_HEADERS_MAX_BYTES else kept + + +def _with_header_in_response_limit(headers: Mapping[str, str], header: tuple[str, str]) -> Mapping[str, str]: + merged: Final = {**headers, header[0]: header[1]} + size: Final = sum(len(name) + len(value) + _HEADER_FRAMING_BYTES for name, value in merged.items()) + return merged if size <= _RESPONSE_HEADERS_MAX_BYTES else headers + + +def _guardrail_headers_with_room( + guardrail_headers: Mapping[str, str], response_headers: Mapping[str, str] +) -> Mapping[str, str]: + """``response_headers`` with each of ``guardrail_headers`` that fits, so adding them can't cost a client the + response. One that doesn't fit leaves the header of that name as it was.""" + return reduce(_with_header_in_response_limit, guardrail_headers.items(), response_headers) + + +def _add_callback_headers( + response_headers: dict[str, str], # mutable-ok: adds to the response's headers in place, like the dict's update + callback_headers: Mapping[str, str], + own_guardrail_headers: Mapping[str, str], +) -> None: + """Add a callback's headers to ``response_headers``, less a guardrail header this proxy set, ``own_guardrail_headers``, + that they then leave no room for. The upstream's guardrail headers are kept as they came.""" + guardrail_headers: Final = { + name: value + for name, value in own_guardrail_headers.items() + if response_headers.get(name) == value and name not in callback_headers + } + for name in guardrail_headers: + response_headers.pop(name) + response_headers.update(callback_headers) + response_headers.update(_guardrail_headers_with_room(guardrail_headers, response_headers)) + + +def _recorded_guardrail_headers( + parsed_body: object = None, kwargs: object = None, post_call_data: object = None +) -> Mapping[str, str]: + """The ``x-litellm-applied-guardrails`` family of response headers for what guardrails recorded on this request. + + Pre-call guardrails record on the request body's metadata, which ``init_kwargs_for_pass_through_endpoint`` moves to + ``kwargs["litellm_params"]``, and post-call guardrails record on ``post_call_data``, so list entries such as the + applied guardrail names are combined. A ``litellm_params`` object sent in the request body is never read. A header + that would take these headers past ``_GUARDRAIL_HEADERS_MAX_BYTES``, or whose recorded value can't be joined, is + left out. + """ + scopes: Final = ( + as_str_mapping(parsed_body) or {}, + as_str_mapping((as_str_mapping(kwargs) or {}).get("litellm_params")) or {}, + as_str_mapping(post_call_data) or {}, + ) + metadatas: Final = tuple( + as_str_mapping(scope.get(metadata_key)) or {} for scope, metadata_key in product(scopes, sorted(_METADATA_KEYS)) + ) + recorded: Final = chain.from_iterable( + ((field, metadata[field]) for field in sorted(_GUARDRAIL_HEADER_METADATA_FIELDS & metadata.keys())) + for metadata in metadatas + ) + merged: Final = { + field: reduce(_merged_metadata_entry, (value for _, value in entries)) + for field, entries in groupby(sorted(recorded, key=lambda entry: entry[0]), key=lambda entry: entry[0]) + } + headers: Final = _RESPONSE_HEADERS.validate_python( + get_logging_caching_headers({"metadata": _joinable_header_metadata(merged)}) or {} + ) + sendable: Final = ( + (name, value) + for name, value in headers.items() + if len(value) <= _GUARDRAIL_HEADERS_MAX_BYTES and _is_sendable_header_value(value) + ) + no_headers: Final[Mapping[str, str]] = MappingProxyType({}) + return reduce(_with_header_in_budget, sendable, no_headers) + + +def _guardrail_response_headers( + parsed_body: object = None, kwargs: object = None, post_call_data: object = None +) -> Mapping[str, str]: + try: + return _recorded_guardrail_headers(parsed_body, kwargs, post_call_data) + except (TypeError, ValueError, RecursionError) as error: + verbose_proxy_logger.warning( + "Leaving out guardrail response headers: a recorded value can't be encoded (%s)", type(error).__name__ + ) + return MappingProxyType({}) + + def _holds_string_headers(value: object) -> bool: headers: Final = as_str_mapping(value) return headers is not None and all(isinstance(header, str) for header in headers.values()) @@ -469,6 +635,8 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): headers: httpx.Headers, litellm_call_id: str | None = None, custom_headers: Mapping[str, str] | None = None, + guardrail_request_data: Sequence[object] = (), + earlier_guardrail_headers: Mapping[str, str] = MappingProxyType({}), ) -> dict: # Exclude headers that uvicorn writes itself (server, date) and # encoding/length headers that don't survive re-serialization. @@ -494,6 +662,13 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): key: value for key, value in custom_headers.items() if key.lower() not in excluded_headers } return_headers.update(sanitized_custom_headers) + guardrail_headers: Final = _guardrail_response_headers(*guardrail_request_data) + combined_guardrail_headers: Final = reduce( + _with_header_in_budget, + (header for header in earlier_guardrail_headers.items() if header[0] not in guardrail_headers), + guardrail_headers, + ) + return_headers.update(_guardrail_headers_with_room(combined_guardrail_headers, return_headers)) return return_headers @@ -1482,6 +1657,29 @@ async def _log_passthrough_upstream_failure( return relay_response +async def _log_passthrough_upstream_failure_keeping_guardrail_headers( + response: httpx.Response, + user_api_key_dict: UserAPIKeyAuth, + request_payload: dict[str, object], # mutable-ok: the failure hooks lift fields onto request_data in place + logging_obj: LiteLLMLoggingObj, + parsed_body: Mapping[str, object] | None, + kwargs: Mapping[str, object] | None, +) -> Mapping[str, str]: + """Log an upstream error like ``_log_passthrough_upstream_failure`` and return the guardrail response headers, + read before the failure hooks rebuild ``kwargs``' metadata without the guardrail fields. A successful response + has no failure hooks to read around, so it returns none.""" + guardrail_headers: Final[Mapping[str, str]] = ( + _guardrail_response_headers(parsed_body, kwargs) if response.status_code >= 400 else {} + ) + await _log_passthrough_upstream_failure( + response=response, + user_api_key_dict=user_api_key_dict, + request_payload=request_payload, + logging_obj=logging_obj, + ) + return guardrail_headers + + async def _relay_reporting_failures( stream: AsyncGenerator[bytes, None], upstream_status: int, @@ -1681,7 +1879,7 @@ async def pass_through_request( ) client_body: Final = TypeAdapter(dict[str, object]).validate_python(_parsed_body) - strip_client_guardrail_information(client_body) + _strip_client_guardrail_fields(client_body) if not _key_or_team_allows_client_pricing_override(user_api_key_dict): _strip_client_pricing_overrides(client_body) eu_jev_client_body: Final = _eu_jev_client_body(custom_llm_provider, client_body) @@ -2031,6 +2229,14 @@ async def pass_through_request( ) apply_upstream_reported_operation(logging_obj, response.headers) + # Built before the failure hooks rebuild kwargs' metadata without the guardrail fields. The relay + # response differs only in body-framing headers, which get_response_headers drops anyway. + streamed_guardrail_headers: Final = _guardrail_response_headers(_parsed_body, kwargs) # pyright: ignore[reportUnknownArgumentType] # untyped request dicts + _response_headers = HttpPassThroughEndpointHelpers.get_response_headers( + headers=response.headers, + litellm_call_id=litellm_call_id, + earlier_guardrail_headers=streamed_guardrail_headers, + ) relay_response: Final = await _log_passthrough_upstream_failure( response=response, user_api_key_dict=user_api_key_dict, @@ -2048,10 +2254,6 @@ async def pass_through_request( ) # Call response headers hook for streaming pass-through - _response_headers = HttpPassThroughEndpointHelpers.get_response_headers( - headers=relay_response.headers, - litellm_call_id=litellm_call_id, - ) callback_headers = await proxy_logging_obj.post_call_response_headers_hook( data=_parsed_body or {}, user_api_key_dict=user_api_key_dict, @@ -2059,7 +2261,7 @@ async def pass_through_request( request_headers=dict(redact_credential_headers(request.headers)), ) if callback_headers: - _response_headers.update(callback_headers) + _add_callback_headers(_response_headers, callback_headers, streamed_guardrail_headers) # pyright: ignore[reportUnknownArgumentType] # untyped header dict return StreamingResponse( wrap_passthrough_sse_bytes_with_keepalive_pings( @@ -2133,6 +2335,14 @@ async def pass_through_request( logging_obj.stream = True logging_obj.model_call_details["stream"] = True + # Built before the failure hooks rebuild kwargs' metadata without the guardrail fields. The relay + # response differs only in body-framing headers, which get_response_headers drops anyway. + detected_stream_guardrail_headers: Final = _guardrail_response_headers(_parsed_body, kwargs) # pyright: ignore[reportUnknownArgumentType] # untyped request dicts + _response_headers = HttpPassThroughEndpointHelpers.get_response_headers( + headers=response.headers, + litellm_call_id=litellm_call_id, + earlier_guardrail_headers=detected_stream_guardrail_headers, + ) detected_relay_response: Final = await _log_passthrough_upstream_failure( response=response, user_api_key_dict=user_api_key_dict, @@ -2150,10 +2360,6 @@ async def pass_through_request( ) # Call response headers hook for detected streaming pass-through - _response_headers = HttpPassThroughEndpointHelpers.get_response_headers( - headers=detected_relay_response.headers, - litellm_call_id=litellm_call_id, - ) callback_headers = await proxy_logging_obj.post_call_response_headers_hook( data=_parsed_body or {}, user_api_key_dict=user_api_key_dict, @@ -2161,7 +2367,7 @@ async def pass_through_request( request_headers=dict(redact_credential_headers(request.headers)), ) if callback_headers: - _response_headers.update(callback_headers) + _add_callback_headers(_response_headers, callback_headers, detected_stream_guardrail_headers) # pyright: ignore[reportUnknownArgumentType] # untyped header dict return StreamingResponse( wrap_passthrough_sse_bytes_with_keepalive_pings( @@ -2229,6 +2435,7 @@ async def pass_through_request( headers=HttpPassThroughEndpointHelpers.get_response_headers( headers=response.headers, custom_headers=relay_custom_headers, + guardrail_request_data=(_parsed_body, kwargs), # pyright: ignore[reportUnknownArgumentType] # untyped request dicts ), ) @@ -2252,11 +2459,13 @@ async def pass_through_request( upstream_model=upstream_model, ) failure_request_payload["response_body"] = response_body - await _log_passthrough_upstream_failure( + upstream_guardrail_headers: Final = await _log_passthrough_upstream_failure_keeping_guardrail_headers( response=response, user_api_key_dict=user_api_key_dict, - request_payload=failure_request_payload, + request_payload=failure_request_payload, # pyright: ignore[reportUnknownArgumentType] # untyped request dicts logging_obj=logging_obj, + parsed_body=_parsed_body, # pyright: ignore[reportUnknownArgumentType] # untyped request dicts + kwargs=kwargs, # pyright: ignore[reportUnknownArgumentType] # untyped request dicts ) policy_pipelines: Final = _post_call_selection_metadata( @@ -2412,6 +2621,8 @@ async def pass_through_request( response_headers: Final = HttpPassThroughEndpointHelpers.get_response_headers( headers=response.headers, custom_headers=custom_headers, + earlier_guardrail_headers=upstream_guardrail_headers, + guardrail_request_data=(_parsed_body, kwargs, post_call_guardrail_data), # pyright: ignore[reportUnknownArgumentType] # untyped request dicts ) emitted_call_id: Final = ( JSON_OBJECT.validate_python(response_headers).get(LITELLM_CALL_ID_HEADER) @@ -2440,6 +2651,7 @@ async def pass_through_request( e.guardrail_name, str(e.message or "")[:200], ) + modified_response_headers: Final = _guardrail_response_headers(_parsed_body, kwargs, post_call_guardrail_data) # pyright: ignore[reportUnknownArgumentType] # untyped request dicts try: await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, @@ -2470,7 +2682,10 @@ async def pass_through_request( content=json.dumps(error_body), status_code=200, media_type="application/json", - headers={"x-litellm-call-id": litellm_call_id}, + headers={ + "x-litellm-call-id": litellm_call_id, + **modified_response_headers, + }, ) except Exception as e: custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( @@ -2480,6 +2695,8 @@ async def pass_through_request( cache_key=None, api_base=str(url._uri_reference) if url else None, ) + guardrail_headers: Final = _guardrail_response_headers(_parsed_body, kwargs, post_call_guardrail_data) # pyright: ignore[reportUnknownArgumentType] # untyped request dicts + custom_headers.update(guardrail_headers) # pyright: ignore[reportUnknownMemberType] # untyped header dict if CustomGuardrail._is_guardrail_intervention(e): verbose_proxy_logger.warning( "pass_through_endpoint: request blocked by guardrail - %s", @@ -2516,6 +2733,7 @@ async def pass_through_request( ######################################################### if isinstance(e, ProxyException): + e.headers = {**e.headers, **guardrail_headers} raise if isinstance(e, HTTPException): raise ProxyException( diff --git a/tests/integration/authorization/test_passthrough_applied_guardrails_header.py b/tests/integration/authorization/test_passthrough_applied_guardrails_header.py new file mode 100644 index 00000000000..43ee159095f --- /dev/null +++ b/tests/integration/authorization/test_passthrough_applied_guardrails_header.py @@ -0,0 +1,362 @@ +import json +from collections.abc import Iterator, Mapping +from contextlib import ExitStack +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.anthropic_sse import message_json +from integration._support.client import Gateway, Scenario +from integration._support.openai_wire import chat_reply +from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +PRE_CALL: Final = "passthrough-header-pre" +POST_CALL: Final = "passthrough-header-post" +NON_LATIN1: Final = "passthrough-header-pre\u2014em" +BLOCKER: Final = "passthrough-header-block" +MISSING_MODEL: Final = "passthrough-header-missing-model" +PILLAR: Final = "passthrough-header-pillar" +TAG_POLICY: Final = "passthrough-header-tag-policy" +POLICY_TAG_PREFIX: Final = "passthrough-header-tag-" +HEADER: Final = "x-litellm-applied-guardrails" +OPENAI_ROUTE: Final = "/openai/v1/chat/completions" +ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages" +GEMINI_ROUTE: Final = "/gemini/v1beta/models/gemini-2.5-flash:generateContent" +CALLBACK_HEADERS_MODEL: Final = "passthrough-header-callback-headers" +UNREQUESTED_STREAM_MODEL: Final = "passthrough-header-callback-headers-unrequested-stream" +CALLBACK_HEADER: Final = "x-callback-padding" +CALLBACK_HEADER_VALUE: Final = "c" * 12_000 +_CALLBACK_HEADERS_MODULE: Final = f""" +from litellm.integrations.custom_logger import CustomLogger + + +class LargeHeadersHook(CustomLogger): + async def async_post_call_response_headers_hook(self, data, user_api_key_dict, response, *args, **kwargs): + if str(data.get("model")).startswith("{CALLBACK_HEADERS_MODEL}"): + return {{"{CALLBACK_HEADER}": "{CALLBACK_HEADER_VALUE}"}} + return None + + +proxy_handler_instance = LargeHeadersHook() +""" +_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" +_JSON: Final = TypeAdapter(dict[str, JsonValue]) + + +def _route_body(route: str) -> dict[str, JsonValue]: + match route: + case "/anthropic/v1/messages": + return {"model": "claude-sonnet-5-5", "max_tokens": 16, "messages": [{"role": "user", "content": "Hi"}]} + case "/gemini/v1beta/models/gemini-2.5-flash:generateContent": + return {"contents": [{"role": "user", "parts": [{"text": "Hi"}]}]} + case _: + return {"model": "gpt-5.6", "messages": [{"role": "user", "content": "Hi"}]} + + +def _allow(_request: Request) -> Reply: + return Reply(body=json.dumps({"action": "NONE"}).encode()) + + +def _block(_request: Request) -> Reply: + return Reply(body=json.dumps({"action": "BLOCKED", "blocked_reason": "synthetic block"}).encode()) + + +def _pillar_allow(_request: Request) -> Reply: + return Reply(body=json.dumps({"flagged": False, "session_id": "pillar-vendor-session"}).encode()) + + +def _openai_reply(request: Request) -> Reply: + payload: Final = _JSON.validate_json(request.body) + if payload.get("model") == MISSING_MODEL: + return Reply( + status=404, + body=json.dumps({"error": {"message": "model not found", "type": "invalid_request_error"}}).encode(), + ) + return chat_reply( + "chatcmpl-passthrough-header", + "gpt-5.6", + "Hello", + stream=payload.get("stream") is True or payload.get("model") == UNREQUESTED_STREAM_MODEL, + ) + + +def _anthropic_reply(_request: Request) -> Reply: + return Reply(body=message_json("msg_passthrough_header", "claude-sonnet-5-5", "Hello")) + + +def _gemini_reply(_request: Request) -> Reply: + return Reply( + body=json.dumps( + { + "candidates": [{"content": {"role": "model", "parts": [{"text": "Hello"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, + } + ).encode() + ) + + +@dataclass(frozen=True, slots=True) +class HeaderRig: + owned: OwnedProxy + upstreams: Mapping[str, Wire] + + +@pytest.fixture(scope="module") +def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[HeaderRig]: + directory: Final = tmp_path_factory.mktemp("passthrough-applied-guardrails-header") + with ExitStack() as stack: + vendor: Final = stack.enter_context(wire_server(_allow)) + blocking_vendor: Final = stack.enter_context(wire_server(_block)) + pillar_vendor: Final = stack.enter_context(wire_server(_pillar_allow)) + upstreams: Final = MappingProxyType( + { + OPENAI_ROUTE: stack.enter_context(wire_server(_openai_reply)), + ANTHROPIC_ROUTE: stack.enter_context(wire_server(_anthropic_reply)), + GEMINI_ROUTE: stack.enter_context(wire_server(_gemini_reply)), + } + ) + config: Final = _JSON.validate_python(yaml.safe_load(_CONFIG.read_text())) + (directory / "passthrough_large_headers_hook.py").write_text(_CALLBACK_HEADERS_MODULE) + path: Final = directory / "passthrough-applied-guardrails-header.yaml" + path.write_text( + yaml.safe_dump( + { + **config, + "litellm_settings": { + **_JSON.validate_python(config.get("litellm_settings", {})), + "callbacks": ["passthrough_large_headers_hook.proxy_handler_instance"], + }, + "guardrails": [ + { + "guardrail_name": name, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": mode, + "default_on": False, + "api_base": api_base, + "api_key": "synthetic-guardrail-key", + }, + } + for name, mode, api_base in ( + (PRE_CALL, "pre_call", vendor.url), + (POST_CALL, "post_call", vendor.url), + (NON_LATIN1, "pre_call", vendor.url), + (BLOCKER, "pre_call", blocking_vendor.url), + ) + ] + + [ + { + "guardrail_name": PILLAR, + "litellm_params": { + "guardrail": "pillar", + "mode": "pre_call", + "default_on": False, + "api_base": pillar_vendor.url, + "api_key": "synthetic-pillar-key", + }, + } + ], + "policies": {TAG_POLICY: {"guardrails": {"add": [PRE_CALL]}}}, + "policy_attachments": [{"policy": TAG_POLICY, "tags": [f"{POLICY_TAG_PREFIX}*"]}], + } + ) + ) + environment: Final = { + "OPENAI_API_BASE": upstreams[OPENAI_ROUTE].url, + "OPENAI_API_KEY": "synthetic-openai-key", + "ANTHROPIC_API_BASE": upstreams[ANTHROPIC_ROUTE].url, + "ANTHROPIC_API_KEY": "synthetic-anthropic-key", + "GEMINI_API_BASE": upstreams[GEMINI_ROUTE].url, + "GEMINI_API_KEY": "synthetic-gemini-key", + } + owned: Final = stack.enter_context(owned_proxy_process(session_gateway, directory, environment, config=path)) + yield HeaderRig(owned, upstreams) + + +def _send( + gateway: Gateway, scenario: Scenario, route: str, body: Mapping[str, JsonValue], *guardrails: str +) -> httpx.Response: + key: Final = scenario.key(metadata={"guardrails": list(guardrails)}) if guardrails else scenario.key() + return gateway.request("POST", route, body, key=key, headers={"x-goog-api-key": key}) + + +@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE]) +def test_response_names_the_guardrail_that_ran(rig: HeaderRig, route: str) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + response: Final = _send(gateway, scenario, route, _route_body(route), PRE_CALL) + + assert response.status_code == 200, response.text + assert response.headers.get(HEADER) == PRE_CALL + + +def test_streamed_response_names_the_guardrail_that_ran(rig: HeaderRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + response: Final = _send( + gateway, scenario, OPENAI_ROUTE, {**_route_body(OPENAI_ROUTE), "stream": True}, PRE_CALL + ) + + assert response.status_code == 200, response.text + assert response.headers.get("content-type", "").startswith("text/event-stream") + assert response.headers.get(HEADER) == PRE_CALL + + +def test_response_names_pre_and_post_call_guardrails(rig: HeaderRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + response: Final = _send(gateway, scenario, OPENAI_ROUTE, _route_body(OPENAI_ROUTE), PRE_CALL, POST_CALL) + + assert response.status_code == 200, response.text + assert response.headers.get(HEADER) == f"{PRE_CALL},{POST_CALL}" + + +@pytest.mark.parametrize("field", ["metadata", "litellm_params"]) +def test_client_cannot_claim_a_guardrail_ran(rig: HeaderRig, field: str) -> None: + gateway: Final = rig.owned.gateway + claim: Final[dict[str, JsonValue]] = { + "applied_guardrails": ["forged-guardrail"], + "guardrail_scan_ids": ["forged-scan"], + } + forged: Final = {**_route_body(OPENAI_ROUTE), field: claim if field == "metadata" else {"metadata": claim}} + with gateway.scenario() as scenario: + response: Final = _send(gateway, scenario, OPENAI_ROUTE, forged) + + assert response.status_code == 200, response.text + assert HEADER not in response.headers + assert "x-litellm-guardrail-scan-id" not in response.headers + + +def test_guardrail_name_a_header_cannot_carry_is_left_out(rig: HeaderRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + response: Final = _send(gateway, scenario, OPENAI_ROUTE, _route_body(OPENAI_ROUTE), NON_LATIN1) + + assert response.status_code == 200, response.text + assert HEADER not in response.headers + + +@pytest.mark.parametrize("stream", [False, True]) +def test_upstream_error_names_the_guardrail_that_ran(rig: HeaderRig, stream: bool) -> None: + gateway: Final = rig.owned.gateway + body: Final = {**_route_body(OPENAI_ROUTE), "model": MISSING_MODEL, "stream": stream} + with gateway.scenario() as scenario: + response: Final = _send(gateway, scenario, OPENAI_ROUTE, body, PRE_CALL) + + assert response.status_code == 404, response.text + assert response.headers.get(HEADER) == PRE_CALL + + +def test_blocked_request_names_the_blocking_guardrail(rig: HeaderRig) -> None: + gateway: Final = rig.owned.gateway + upstream: Final = rig.upstreams[OPENAI_ROUTE] + upstream.drain() + with gateway.scenario() as scenario: + response: Final = _send(gateway, scenario, OPENAI_ROUTE, _route_body(OPENAI_ROUTE), PRE_CALL, BLOCKER) + + assert response.status_code == 400, response.text + assert response.headers.get(HEADER) == f"{PRE_CALL},{BLOCKER}" + assert upstream.drain() == () + + +@pytest.mark.parametrize("field", ["metadata", "litellm_metadata"]) +def test_client_claims_beside_a_real_guardrail_are_dropped(rig: HeaderRig, field: str) -> None: + gateway: Final = rig.owned.gateway + claim: Final[dict[str, JsonValue]] = { + "applied_guardrails": ["forged-guardrail"], + "guardrail_scan_ids": ["forged-scan"], + } + with gateway.scenario() as scenario: + response: Final = _send(gateway, scenario, OPENAI_ROUTE, {**_route_body(OPENAI_ROUTE), field: claim}, PRE_CALL) + + assert response.status_code == 200, response.text + assert response.headers.get(HEADER) == PRE_CALL + assert "forged-scan" not in response.headers.get("x-litellm-guardrail-scan-id", "") + + +@pytest.mark.parametrize("claim", ["none", 7, ""]) +def test_non_list_client_claim_does_not_fail_a_guarded_request(rig: HeaderRig, claim: JsonValue) -> None: + gateway: Final = rig.owned.gateway + body: Final = {**_route_body(OPENAI_ROUTE), "metadata": {"applied_guardrails": claim}} + with gateway.scenario() as scenario: + response: Final = _send(gateway, scenario, OPENAI_ROUTE, body, PRE_CALL) + + assert response.status_code == 200, response.text + assert response.headers.get(HEADER) == PRE_CALL + + +def test_pillar_headers_stay_off_pass_through_responses(rig: HeaderRig) -> None: + gateway: Final = rig.owned.gateway + forged: Final[dict[str, JsonValue]] = {"pillar_session_id_response": "client-session", "pillar_scanners": {"x": 1}} + with gateway.scenario() as scenario: + response: Final = _send( + gateway, scenario, OPENAI_ROUTE, {**_route_body(OPENAI_ROUTE), "metadata": forged}, PILLAR + ) + + assert response.status_code == 200, response.text + assert response.headers.get(HEADER) == PILLAR + assert [name for name, _ in response.headers.multi_items() if name.startswith("x-pillar-")] == [] + + +@pytest.mark.parametrize("suffix", ["team\n", "x\r\nX-Injected: yes"]) +def test_header_value_with_a_line_break_is_left_out(rig: HeaderRig, suffix: str) -> None: + gateway: Final = rig.owned.gateway + body: Final[dict[str, JsonValue]] = { + **_route_body(OPENAI_ROUTE), + "metadata": {"tags": [f"{POLICY_TAG_PREFIX}{suffix}"]}, + } + with gateway.scenario() as scenario: + response: Final = _send(gateway, scenario, OPENAI_ROUTE, body) + + assert response.status_code == 200, response.text + assert response.headers.get(HEADER) == PRE_CALL + assert response.headers.get("x-litellm-applied-policies") == TAG_POLICY + assert "x-litellm-policy-sources" not in response.headers + assert "x-injected" not in response.headers + + +@pytest.mark.parametrize( + ("model", "stream"), + [(CALLBACK_HEADERS_MODEL, True), (UNREQUESTED_STREAM_MODEL, False)], + ids=["requested-stream", "unrequested-stream"], +) +@pytest.mark.parametrize(("names", "kept"), [(300, False), (100, True)], ids=["no-room", "room"]) +def test_a_callback_header_on_a_stream_leaves_room_or_costs_the_guardrail_header( + rig: HeaderRig, model: str, stream: bool, names: int, kept: bool +) -> None: + gateway: Final = rig.owned.gateway + body: Final[dict[str, JsonValue]] = { + **_route_body(OPENAI_ROUTE), + "model": model, + "stream": stream, + "guardrails": [f"caller-named-{index}" for index in range(names)], + } + with gateway.scenario() as scenario: + response: Final = _send(gateway, scenario, OPENAI_ROUTE, body) + + assert response.status_code == 200, response.text + assert response.headers.get(CALLBACK_HEADER) == CALLBACK_HEADER_VALUE + assert (HEADER in response.headers) is kept + + +@pytest.mark.parametrize("stream", [False, True]) +def test_response_to_a_body_naming_thousands_of_guardrails_stays_readable(rig: HeaderRig, stream: bool) -> None: + gateway: Final = rig.owned.gateway + body: Final[dict[str, JsonValue]] = { + **_route_body(OPENAI_ROUTE), + "stream": stream, + "guardrails": [f"caller-named-{index}" for index in range(10_000)], + "metadata": {"tags": [f"{POLICY_TAG_PREFIX}prod"]}, + } + with gateway.scenario() as scenario: + response: Final = _send(gateway, scenario, OPENAI_ROUTE, body) + + assert response.status_code == 200, response.text + assert HEADER not in response.headers + assert response.headers.get("x-litellm-applied-policies") == TAG_POLICY diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index b3d711c3e9d..80d64543449 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -717,6 +717,20 @@ def test_post_call_policy_pipeline_decides_what_a_flagged_step_does(rig: AuthGua assert response.status_code == 200, response.text assert len(rig.key_post_call_vendor.drain()) == 1 + assert KEY_POST_CALL in response.headers["x-litellm-applied-guardrails"].split(",") + + +def test_response_a_policy_pipeline_modifies_names_the_step_that_ran(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key() + _drain(rig) + + response: Final = _send(gateway, OPENAI_ROUTE, _tagged_body(OPENAI_ROUTE, PIPELINE_MODIFIES), key) + + assert response.status_code == 200, response.text + assert response.json()["error"]["type"] == "content_filter" + assert KEY_POST_CALL in response.headers["x-litellm-applied-guardrails"].split(",") def test_pre_call_policy_pipeline_relays_the_provider_bytes(rig: AuthGuardrailRig) -> None: diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index d1782b1fe21..ceabde021c1 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -11,13 +11,14 @@ import zlib from collections.abc import AsyncIterator, Callable, Mapping, Sequence from contextlib import ExitStack, contextmanager from dataclasses import dataclass -from datetime import datetime +from datetime import datetime, timezone from io import BytesIO from types import MappingProxyType, ModuleType, SimpleNamespace from typing import Final, Optional, cast from unittest.mock import AsyncMock, MagicMock, call, patch from uuid import UUID +import h11 import httpx import pytest import respx @@ -67,6 +68,10 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( WebsocketFrameModelVerdict, _build_passthrough_failure_request_payload, _registered_pass_through_routes, + _add_callback_headers, + _guardrail_response_headers, + _log_passthrough_upstream_failure_keeping_guardrail_headers, + _strip_client_guardrail_fields, _truncate_upstream_error_body, _update_metadata_with_tags_in_header, _with_trace_context, @@ -505,6 +510,61 @@ async def test_pass_through_request_preserves_proxy_exception_status(): assert exc.value.param == "limit" +class _RecordsThenRaisesProxyException(CustomGuardrail): + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: str, + ) -> None: + data["metadata"] = {"applied_guardrails": ["passed-guard"]} # mutable-ok: a guardrail records on the live body + raise ProxyException(message="blocked by guardrail", type="None", param="None", code=400) + + +@pytest.mark.asyncio +async def test_guardrail_proxy_exception_carries_the_guardrails_that_ran(): + guardrail: Final = _RecordsThenRaisesProxyException( + guardrail_name="passed-guard", event_hook=GuardrailEventHooks.pre_call, default_on=True + ) + mock_request: Final = MagicMock(spec=Request) + mock_request.method = "POST" + mock_request.body = AsyncMock(return_value=b'{"model": "gpt-5.6", "messages": []}') + mock_request.headers = Headers({"content-type": "application/json"}) + mock_request.query_params = QueryParams({}) + + with ( + pytest.MonkeyPatch.context() as litellm_patch, + patch( # test-quality-ok: proxy_logging_obj is a proxy_server module global read inside pass_through_request + "litellm.proxy.proxy_server.proxy_logging_obj", ProxyLogging(user_api_key_cache=DualCache()) + ), + ): + litellm_patch.setattr(litellm, "callbacks", [guardrail]) + with pytest.raises(ProxyException) as exc: + await pass_through_request( + request=mock_request, + target="http://test.com/v1/chat/completions", + custom_headers={}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + assert exc.value.headers["x-litellm-applied-guardrails"] == "passed-guard" + + +@pytest.mark.asyncio +async def test_successful_response_keeps_no_guardrail_headers_from_before_the_post_call_pass(): + headers: Final = await _log_passthrough_upstream_failure_keeping_guardrail_headers( + response=httpx.Response(status_code=200), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + request_payload={}, + logging_obj=MagicMock(), + parsed_body={"metadata": {"applied_guardrails": ["pre-call-guard"]}}, + kwargs=None, + ) + + assert headers == {} + + def test_is_langfuse_route(): """ Test that the is_langfuse_route method correctly identifies Langfuse routes @@ -13229,13 +13289,23 @@ async def test_passthrough_header_and_callback_share_finalized_cost( client, cleanup = _inject_fake_passthrough_client(httpx.MockTransport(upstream), timeout=None) try: - request: Final = Request({ - "type": "http", "method": "GET", "path": "/passthrough-relay/results", - "headers": [], "query_string": b"", "scheme": "http", "server": ("test", 80), - }) + request: Final = Request( + { + "type": "http", + "method": "GET", + "path": "/passthrough-relay/results", + "headers": [], + "query_string": b"", + "scheme": "http", + "server": ("test", 80), + } + ) result: Final = await pass_through_request( - request=request, target="http://cost-upstream.invalid/results", custom_headers={}, - user_api_key_dict=UserAPIKeyAuth(api_key="hashed"), cost_per_request=flat_cost, + request=request, + target="http://cost-upstream.invalid/results", + custom_headers={}, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed"), + cost_per_request=flat_cost, ) payload: Final = await asyncio.wait_for(observed, timeout=10) await GLOBAL_LOGGING_WORKER.flush() @@ -13790,9 +13860,7 @@ class TestInheritAuthGuardrails: async def test_request_body_guardrails_are_added_to_the_key_guardrails(self): auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) - pre_call, _ = await self._run( - self._chat_body(guardrails=["body-guard"]), inherit=True, user_api_key_dict=auth - ) + pre_call, _ = await self._run(self._chat_body(guardrails=["body-guard"]), inherit=True, user_api_key_dict=auth) assert pre_call[0]["metadata"]["guardrails"] == ["key-guard", "body-guard"] assert "guardrails" not in pre_call[0] @@ -13844,3 +13912,301 @@ class TestInheritAuthGuardrails: assert "guardrails" not in pre_call[0].get("metadata", {}) post_call.assert_not_awaited() + + +class TestGuardrailResponseHeaders: + def test_pre_and_post_call_guardrails_are_combined_without_repeats(self): + headers: Final = HttpPassThroughEndpointHelpers.get_response_headers( + headers=httpx.Headers({"content-type": "application/json"}), + guardrail_request_data=( + {"model": "gpt-5.6"}, + {"litellm_params": {"metadata": {"applied_guardrails": ["pre"], "guardrail_scan_ids": ["scan-1"]}}}, + {"metadata": {"applied_guardrails": ["pre", "post"]}}, + ), + ) + + assert headers["x-litellm-applied-guardrails"] == "pre,post" + assert headers["x-litellm-guardrail-scan-id"] == "scan-1" + + def test_scan_ids_recorded_as_tuples_before_and_after_the_call_are_both_kept(self): + headers: Final = HttpPassThroughEndpointHelpers.get_response_headers( + headers=httpx.Headers({"content-type": "application/json"}), + guardrail_request_data=( + {}, + {"litellm_params": {"metadata": {"guardrail_scan_ids": ("pre-scan",)}}}, + {"metadata": {"guardrail_scan_ids": ("post-scan",)}}, + ), + ) + + assert headers["x-litellm-guardrail-scan-id"] == "pre-scan,post-scan" + + def test_policy_sources_recorded_before_and_after_the_call_are_both_kept(self): + headers: Final = HttpPassThroughEndpointHelpers.get_response_headers( + headers=httpx.Headers({"content-type": "application/json"}), + guardrail_request_data=( + {}, + {"litellm_params": {"metadata": {"policy_sources": {"pre-policy": "key", "shared": "team"}}}}, + {"metadata": {"policy_sources": {"post-policy": "tag", "shared": "key"}}}, + ), + ) + + assert headers["x-litellm-policy-sources"] == "pre-policy=key; shared=key; post-policy=tag" + + def test_a_litellm_params_object_in_the_request_body_is_never_read(self): + forged: Final = {"metadata": {"applied_guardrails": ["forged"], "guardrail_scan_ids": ["fake"]}} + + headers: Final = HttpPassThroughEndpointHelpers.get_response_headers( + headers=httpx.Headers({"content-type": "application/json"}), + guardrail_request_data=( + {"model": "gpt-5.6", "litellm_params": forged}, + {"litellm_params": {"metadata": {}}}, + {"model": "gpt-5.6", "litellm_params": forged, "metadata": {"guardrails": ["post"]}}, + ), + ) + + assert headers == {"content-type": "application/json"} + + def test_guardrail_name_outside_latin1_is_left_out_of_the_headers(self): + headers: Final = HttpPassThroughEndpointHelpers.get_response_headers( + headers=httpx.Headers({"content-type": "application/json"}), + guardrail_request_data=( + {}, + { + "litellm_params": { + "metadata": {"applied_guardrails": ["pii\u2014guard"], "guardrail_scan_ids": ["s1"]} + } + }, + ), + ) + + assert headers == {"content-type": "application/json", "x-litellm-guardrail-scan-id": "s1"} + assert Response(content=b"{}", headers=headers).status_code == 200 + + def test_metadata_outside_the_header_fields_is_never_read(self): + class _Unreadable: + __hash__ = object.__hash__ + + def __eq__(self, other: object) -> bool: + raise AssertionError("non-header metadata was read") + + headers: Final = HttpPassThroughEndpointHelpers.get_response_headers( + headers=httpx.Headers({"content-type": "application/json"}), + guardrail_request_data=( + {"metadata": {"unused": [_Unreadable()]}}, + {"litellm_params": {"metadata": {"unused": [_Unreadable()]}}}, + ), + ) + + assert headers == {"content-type": "application/json"} + + def test_no_guardrail_data_adds_no_headers(self): + headers: Final = HttpPassThroughEndpointHelpers.get_response_headers( + headers=httpx.Headers({"content-type": "application/json"}), + guardrail_request_data=({"metadata": {"tags": ["team-a"]}},), + ) + + assert headers == {"content-type": "application/json"} + + def test_client_sent_guardrail_header_fields_are_stripped(self): + body: Final = { + "metadata": {"applied_guardrails": ["forged"], "pillar_flagged": True, "tags": ["team-a"]}, + "litellm_metadata": {"guardrail_scan_ids": ["forged-scan"]}, + } + + _strip_client_guardrail_fields(body) + + assert body == {"metadata": {"pillar_flagged": True, "tags": ["team-a"]}, "litellm_metadata": {}} + + def test_pillar_fields_recorded_on_the_request_add_no_pillar_headers(self): + headers: Final = _guardrail_response_headers( + {"metadata": {"pillar_session_id_response": "client-session"}}, + { + "litellm_params": { + "metadata": { + "applied_guardrails": ["pillar"], + "pillar_flagged": False, + "pillar_response_headers": {"x-pillar-flagged": "false", "x-pillar-session-id": "vendor"}, + "_pillar_response_headers_trusted": True, + } + } + }, + ) + + assert headers == {"x-litellm-applied-guardrails": "pillar"} + + @pytest.mark.parametrize("tag", ["prod-team\n", "prod-x\r\nX-Injected: yes", "prod\x00"]) + def test_header_value_an_http_response_cannot_carry_is_left_out(self, tag: str): + headers: Final = _guardrail_response_headers( + {}, + { + "litellm_params": { + "metadata": {"applied_policies": ["tag-policy"], "policy_sources": {"tag-policy": f"tag:{tag}"}} + } + }, + ) + + assert headers == {"x-litellm-applied-policies": "tag-policy"} + assert h11.Response(status_code=200, headers=[*headers.items()]).status_code == 200 + + @pytest.mark.parametrize( + ("upstream_bytes", "names", "kept"), + [(12000, 500, False), (12000, 1, True), (8000, 500, True), (15340, 1, False)], + ids=["large-and-large", "large-and-small", "both-within-the-limit", "upstream-at-the-limit"], + ) + def test_guardrail_headers_leave_the_whole_response_within_a_clients_limit(self, upstream_bytes, names, kept): + upstream: Final = "u" * upstream_bytes + headers: Final = HttpPassThroughEndpointHelpers.get_response_headers( + headers=httpx.Headers({"x-upstream": upstream}), + guardrail_request_data=( + {}, + {"litellm_params": {"metadata": {"applied_guardrails": [f"guardrail-{i:03d}" for i in range(names)]}}}, + ), + ) + + assert headers["x-upstream"] == upstream + assert ("x-litellm-applied-guardrails" in headers) is kept + assert sum(len(name) + len(value) for name, value in headers.items()) <= 15360 or not kept + + def test_many_small_upstream_headers_count_with_their_framing(self): + upstream: Final = httpx.Headers({f"x-{index:04d}": "v" for index in range(1_500)}) + headers: Final = HttpPassThroughEndpointHelpers.get_response_headers( + headers=upstream, + guardrail_request_data=({}, {"litellm_params": {"metadata": {"applied_guardrails": ["guard"]}}}), + ) + + assert "x-litellm-applied-guardrails" not in headers + assert len(headers) == 1_500 + + def test_an_upstream_guardrail_header_a_larger_one_cannot_replace_still_counts(self): + names: Final = [f"guardrail-{index:03d}" for index in range(280)] + upstream: Final = {"x-upstream": "u" * 10_000, "x-litellm-applied-policies": "p" * 4_000} + headers: Final = HttpPassThroughEndpointHelpers.get_response_headers( + headers=httpx.Headers(upstream), + guardrail_request_data=( + {}, + {"litellm_params": {"metadata": {"applied_guardrails": names, "applied_policies": names}}}, + ), + ) + + assert sum(len(name) + len(value) + 4 for name, value in headers.items()) <= 15360 + assert headers["x-upstream"] == upstream["x-upstream"] + + def test_an_upstream_guardrail_header_the_proxy_replaces_takes_no_room(self): + headers: Final = HttpPassThroughEndpointHelpers.get_response_headers( + headers=httpx.Headers({"x-litellm-applied-guardrails": "u" * 15_400}), + guardrail_request_data=({}, {"litellm_params": {"metadata": {"applied_guardrails": ["guard"]}}}), + ) + + assert headers["x-litellm-applied-guardrails"] == "guard" + + @staticmethod + def _upstream_with_guardrail_headers() -> dict[str, str]: + return { + "x-litellm-applied-guardrails": "upstream-guard", + "x-litellm-guardrail-scan-id": "upstream-scan", + "x-upstream-0": "o" * 8_000, + "x-upstream-1": "o" * 8_000, + } + + def test_upstream_guardrail_headers_survive_a_callback_header_when_the_proxy_recorded_nothing(self): + headers: Final = self._upstream_with_guardrail_headers() + + _add_callback_headers(headers, {"x-callback": "c"}, {}) + + assert headers == {**self._upstream_with_guardrail_headers(), "x-callback": "c"} + + def test_the_proxys_guardrail_header_gives_way_to_a_callback_header_with_no_room(self): + headers: Final = {"x-litellm-applied-guardrails": "guard", "x-upstream": "o" * 15_000} + + _add_callback_headers(headers, {"x-callback": "c" * 300}, {"x-litellm-applied-guardrails": "guard"}) + + assert headers == {"x-upstream": "o" * 15_000, "x-callback": "c" * 300} + + def test_the_proxys_guardrail_header_stays_beside_a_callback_header_with_room(self): + headers: Final = {"x-litellm-applied-guardrails": "guard", "x-upstream": "o"} + + _add_callback_headers(headers, {"x-callback": "c"}, {"x-litellm-applied-guardrails": "guard"}) + + assert headers == {"x-litellm-applied-guardrails": "guard", "x-upstream": "o", "x-callback": "c"} + + def test_guardrail_header_past_the_size_budget_is_left_out(self): + names: Final = [f"g{index}" for index in range(2_000)] + + assert _guardrail_response_headers({"metadata": {"applied_guardrails": names}}) == {} + + @pytest.mark.parametrize("value", [7, "name", {"a": 1}, ["ok", 7], [None]]) + @pytest.mark.parametrize("field", ["applied_guardrails", "applied_policies", "guardrail_scan_ids"]) + def test_recorded_value_that_is_not_a_list_of_names_costs_the_header(self, field, value): + assert _guardrail_response_headers(post_call_data={"metadata": {field: value}}) == {} + + def test_bad_post_call_value_costs_the_header_of_the_names_recorded_before_it(self): + assert ( + _guardrail_response_headers( + {"metadata": {"applied_guardrails": ["pre-call-guard"]}}, + post_call_data={"metadata": {"applied_guardrails": 7}}, + ) + == {} + ) + + @pytest.mark.parametrize( + "entry", + [{"s1"}, {"at": datetime(2026, 1, 1, tzinfo=timezone.utc)}, object()], + ids=["set", "datetime", "object"], + ) + def test_scan_metadata_json_cannot_encode_costs_the_headers_not_the_response(self, entry): + assert _guardrail_response_headers(post_call_data={"metadata": {"guardrail_scan_metadata": [entry]}}) == {} + + @pytest.mark.parametrize("field", ["applied_guardrails", "guardrail_scan_ids", "guardrail_scan_metadata"]) + def test_mixed_key_entry_recorded_before_and_after_the_call_costs_the_headers_not_the_response(self, field): + entries: Final = [{1: "a", "b": 2}] + + assert ( + _guardrail_response_headers({"metadata": {field: entries}}, post_call_data={"metadata": {field: entries}}) + == {} + ) + + @pytest.mark.parametrize("control", ["\x01", "\x1f", "\x7f", "\x00", "\n"]) + def test_guardrail_name_with_a_control_character_stays_out_of_the_header(self, control): + assert _guardrail_response_headers({"metadata": {"applied_guardrails": [f"a{control}b"]}}) == {} + + def test_guardrail_name_with_a_latin1_character_and_a_space_stays_in_the_header(self): + assert _guardrail_response_headers({"metadata": {"applied_guardrails": ["caf\xe9 guard"]}}) == { + "x-litellm-applied-guardrails": "caf\xe9 guard" + } + + def test_guardrail_headers_stay_in_budget_across_the_earlier_and_the_final_set(self): + final_value: Final = "n" * 8000 + response_headers: Final = HttpPassThroughEndpointHelpers.get_response_headers( + httpx.Headers(), + guardrail_request_data=({"metadata": {"applied_guardrails": [final_value]}},), + earlier_guardrail_headers={"x-litellm-policy-sources": "s" * 7900, "x-litellm-applied-guardrails": "old"}, + ) + + assert response_headers == {"x-litellm-applied-guardrails": final_value} + + def test_earlier_guardrail_header_within_the_budget_is_kept(self): + response_headers: Final = HttpPassThroughEndpointHelpers.get_response_headers( + httpx.Headers(), + guardrail_request_data=({"metadata": {"applied_guardrails": ["post-call"]}},), + earlier_guardrail_headers={"x-litellm-policy-sources": "p=q"}, + ) + + assert response_headers == { + "x-litellm-applied-guardrails": "post-call", + "x-litellm-policy-sources": "p=q", + } + + def test_guardrail_headers_past_the_size_budget_are_left_out_so_a_client_can_read_the_response(self): + names: Final = [f"caller-named-{index}" for index in range(10_000)] + headers: Final = _guardrail_response_headers( + {}, {"litellm_params": {"metadata": {"applied_guardrails": names, "applied_policies": ["tag-policy"]}}} + ) + server: Final = h11.Connection(h11.SERVER) + server.receive_data(b"GET / HTTP/1.1\r\nhost: proxy\r\n\r\n") + server.next_event() + client: Final = h11.Connection(h11.CLIENT) + client.send(h11.Request(method="GET", target="/", headers=[("host", "proxy")])) + client.receive_data(server.send(h11.Response(status_code=200, headers=[*headers.items()])) or b"") + + assert headers == {"x-litellm-applied-policies": "tag-policy"} + assert isinstance(client.next_event(), h11.Response)