From f51ca91711d3302c3d7717c0abae61a923d0dbe5 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 31 Oct 2026 11:59:07 -0700 Subject: [PATCH] fix(proxy): keep guardrail headers on upstream errors, and merge mapping metadata from both guardrail passes --- .../pass_through_endpoints.py | 59 +++++++++++---- ...t_passthrough_applied_guardrails_header.py | 74 ++++++++++++++++++- .../test_pass_through_endpoints.py | 12 +++ 3 files changed, 127 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 7f5cdbc6777..8f907faf8d3 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -226,7 +226,12 @@ def _strip_client_guardrail_fields( def _merged_metadata_entry(prior: object, value: object) -> object: - """``value``, or both combined without repeats when ``prior`` and ``value`` are lists or tuples.""" + """``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)) @@ -1274,6 +1279,26 @@ 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.""" + guardrail_headers: Final = _guardrail_response_headers(parsed_body, kwargs) + 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, @@ -1823,6 +1848,13 @@ 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. + _response_headers = HttpPassThroughEndpointHelpers.get_response_headers( + headers=response.headers, + litellm_call_id=litellm_call_id, + guardrail_request_data=(_parsed_body, kwargs), # pyright: ignore[reportUnknownArgumentType] # untyped request dicts + ) relay_response: Final = await _log_passthrough_upstream_failure( response=response, user_api_key_dict=user_api_key_dict, @@ -1840,11 +1872,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, - 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 {}, user_api_key_dict=user_api_key_dict, @@ -1926,6 +1953,13 @@ 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. + _response_headers = HttpPassThroughEndpointHelpers.get_response_headers( + headers=response.headers, + litellm_call_id=litellm_call_id, + guardrail_request_data=(_parsed_body, kwargs), # pyright: ignore[reportUnknownArgumentType] # untyped request dicts + ) detected_relay_response: Final = await _log_passthrough_upstream_failure( response=response, user_api_key_dict=user_api_key_dict, @@ -1943,11 +1977,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, - 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 {}, user_api_key_dict=user_api_key_dict, @@ -2047,11 +2076,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 ) if response.status_code < 400 and response_body is not None and post_call_guardrails: @@ -2191,7 +2222,7 @@ async def pass_through_request( response_headers: Final = HttpPassThroughEndpointHelpers.get_response_headers( headers=response.headers, - custom_headers=custom_headers, + custom_headers={**custom_headers, **upstream_guardrail_headers}, guardrail_request_data=(_parsed_body, kwargs, post_call_guardrail_data), # pyright: ignore[reportUnknownArgumentType] # untyped request dicts ) emitted_call_id: Final = ( diff --git a/tests/integration/authorization/test_passthrough_applied_guardrails_header.py b/tests/integration/authorization/test_passthrough_applied_guardrails_header.py index 5a6656ca999..b6b40547d8b 100644 --- a/tests/integration/authorization/test_passthrough_applied_guardrails_header.py +++ b/tests/integration/authorization/test_passthrough_applied_guardrails_header.py @@ -19,6 +19,8 @@ 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" HEADER: Final = "x-litellm-applied-guardrails" OPENAI_ROUTE: Final = "/openai/v1/chat/completions" ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages" @@ -41,9 +43,18 @@ 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 _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) + 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) def _anthropic_reply(_request: Request) -> Reply: @@ -72,6 +83,7 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I 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)) upstreams: Final = MappingProxyType( { OPENAI_ROUTE: stack.enter_context(wire_server(_openai_reply)), @@ -92,11 +104,16 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I "guardrail": "generic_guardrail_api", "mode": mode, "default_on": False, - "api_base": vendor.url, + "api_base": api_base, "api_key": "synthetic-guardrail-key", }, } - for name, mode in ((PRE_CALL, "pre_call"), (POST_CALL, "post_call"), (NON_LATIN1, "pre_call")) + 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), + ) ], } ) @@ -174,3 +191,52 @@ def test_guardrail_name_a_header_cannot_carry_is_left_out(rig: HeaderRig) -> Non 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 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 6f3ff0b5ed4..e2a318732dd 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 @@ -11980,6 +11980,18 @@ class TestGuardrailResponseHeaders: 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"]}}