diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 23c1202ea07..af99e0a88ac 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -104,6 +104,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 ( + GUARDRAIL_SCAN_IDS_METADATA_KEY, + GUARDRAIL_SCAN_METADATA_METADATA_KEY, + TRUSTED_PILLAR_RESPONSE_HEADERS_METADATA_KEY, + get_logging_caching_headers, # pyright: ignore[reportUnknownVariableType] # untyped request_data param +) 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 @@ -176,14 +182,63 @@ pass_through_endpoint_logging: Final = PassThroughEndpointLogging() _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]) _KEPT_OUT_OF_LITELLM_PARAMS: Final = _METADATA_KEYS | frozenset(wire_names(ProxyRequestState)) _PRICING_MODEL_OVERRIDE_PARAMS: Final = frozenset(("base_model",)) +# The metadata fields get_logging_caching_headers turns into response headers. +_GUARDRAIL_HEADER_METADATA_FIELDS: Final = frozenset( + ( + "applied_guardrails", + "applied_policies", + "policy_sources", + "semantic-similarity", + GUARDRAIL_SCAN_IDS_METADATA_KEY, + GUARDRAIL_SCAN_METADATA_METADATA_KEY, + TRUSTED_PILLAR_RESPONSE_HEADERS_METADATA_KEY, + "pillar_response_headers", + "pillar_flagged", + ) +) 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_header_metadata( + data: dict[str, object], # mutable-ok: strips the request body in place, like strip_client_guardrail_information +) -> None: + """Drop the metadata fields that become guardrail response headers, so only the guardrails that ran set them.""" + 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 _merged_metadata_entry(prior: object, value: object) -> object: + """``value``, or both lists combined without repeats when ``prior`` and ``value`` are lists.""" + if not isinstance(prior, list) or not isinstance(value, list): + return value + prior_entries: Final = _GUARDRAIL_ENTRIES.validate_python(prior) + return [*prior_entries, *(e for e in _GUARDRAIL_ENTRIES.validate_python(value) if e not in prior_entries)] + + +def _guardrail_response_headers(*request_data: object) -> Mapping[str, str]: + """The ``x-litellm-applied-guardrails`` family of response headers for what guardrails recorded in these dicts. + + Reads the metadata of each dict and of its ``litellm_params``. The pre-call and post-call hooks record on + different dicts, so list entries such as the applied guardrail names are combined. + """ + merged: Final = dict[str, object]() + for data in request_data: + fields = as_str_mapping(data) or {} + for scope in (fields, as_str_mapping(fields.get("litellm_params")) or {}): + for metadata_key in sorted(_METADATA_KEYS): + for key, value in (as_str_mapping(scope.get(metadata_key)) or {}).items(): + merged[key] = _merged_metadata_entry(merged.get(key), value) + return _RESPONSE_HEADERS.validate_python(get_logging_caching_headers({"metadata": merged}) or {}) + + class _RegisteredPassThroughRoute(TypedDict): endpoint_id: ReadOnly[str] path: ReadOnly[str | None] @@ -410,7 +465,10 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): headers: httpx.Headers, litellm_call_id: str | None = None, custom_headers: Mapping[str, str] | None = None, + guardrail_request_data: Sequence[object] = (), ) -> dict: + """The upstream response headers to relay, plus the call id, ``custom_headers`` and the guardrail headers + for what guardrails recorded in ``guardrail_request_data``.""" # Exclude headers that uvicorn writes itself (server, date) and # encoding/length headers that don't survive re-serialization. # If we forward the upstream's Server header, uvicorn adds its @@ -435,6 +493,7 @@ 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) + return_headers.update(_guardrail_response_headers(*guardrail_request_data)) return return_headers @@ -1343,6 +1402,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_header_metadata(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) @@ -1706,6 +1766,7 @@ async def pass_through_request( _response_headers = HttpPassThroughEndpointHelpers.get_response_headers( headers=relay_response.headers, litellm_call_id=litellm_call_id, + guardrail_request_data=(_parsed_body, kwargs), # pyright: ignore[reportUnknownArgumentType] # untyped request dicts ) callback_headers = await proxy_logging_obj.post_call_response_headers_hook( data=_parsed_body or {}, @@ -1808,6 +1869,7 @@ async def pass_through_request( _response_headers = HttpPassThroughEndpointHelpers.get_response_headers( headers=detected_relay_response.headers, litellm_call_id=litellm_call_id, + guardrail_request_data=(_parsed_body, kwargs), # pyright: ignore[reportUnknownArgumentType] # untyped request dicts ) callback_headers = await proxy_logging_obj.post_call_response_headers_hook( data=_parsed_body or {}, @@ -1884,6 +1946,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 ), ) @@ -2052,6 +2115,7 @@ async def pass_through_request( response_headers: Final = HttpPassThroughEndpointHelpers.get_response_headers( headers=response.headers, custom_headers=custom_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) @@ -2110,7 +2174,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, + **_guardrail_response_headers(_parsed_body, kwargs, post_call_guardrail_data), # pyright: ignore[reportUnknownArgumentType] # untyped request dicts + }, ) except Exception as e: custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers( @@ -2120,6 +2187,7 @@ async def pass_through_request( cache_key=None, api_base=str(url._uri_reference) if url else None, ) + custom_headers.update(_guardrail_response_headers(_parsed_body, kwargs, post_call_guardrail_data)) # pyright: ignore[reportUnknownArgumentType, reportUnknownMemberType] # untyped request and header dicts if CustomGuardrail._is_guardrail_intervention(e): verbose_proxy_logger.warning( "pass_through_endpoint: request blocked by guardrail - %s", 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..c3337083e78 --- /dev/null +++ b/tests/integration/authorization/test_passthrough_applied_guardrails_header.py @@ -0,0 +1,164 @@ +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" +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" +_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 _openai_reply(request: Request) -> Reply: + streaming: Final = _JSON.validate_json(request.body).get("stream") is True + return chat_reply("chatcmpl-passthrough-header", "gpt-5.6", "Hello", stream=streaming) + + +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)) + 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())) + path: Final = directory / "passthrough-applied-guardrails-header.yaml" + path.write_text( + yaml.safe_dump( + { + **config, + "guardrails": [ + { + "guardrail_name": name, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": mode, + "default_on": False, + "api_base": vendor.url, + "api_key": "synthetic-guardrail-key", + }, + } + for name, mode in ((PRE_CALL, "pre_call"), (POST_CALL, "post_call")) + ], + } + ) + ) + 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}" + + +def test_client_cannot_claim_a_guardrail_ran(rig: HeaderRig) -> None: + gateway: Final = rig.owned.gateway + forged: Final = { + **_route_body(OPENAI_ROUTE), + "metadata": {"applied_guardrails": ["forged-guardrail"], "guardrail_scan_ids": ["forged-scan"]}, + } + 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 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 294ad45d795..51d4de3ef48 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 @@ -51,6 +51,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( SafeRouteAdder, _build_passthrough_failure_request_payload, _registered_pass_through_routes, + _strip_client_guardrail_header_metadata, _truncate_upstream_error_body, _update_metadata_with_tags_in_header, _with_trace_context, @@ -11472,3 +11473,37 @@ 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"]}}, + None, + ), + ) + + assert headers["x-litellm-applied-guardrails"] == "pre,post" + assert headers["x-litellm-guardrail-scan-id"] == "scan-1" + + 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_header_metadata(body) + + assert body == {"metadata": {"tags": ["team-a"]}, "litellm_metadata": {}}