scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

fix(proxy): keep guardrail headers on upstream errors, and merge mapping metadata from both guardrail passes

Browse source
This commit is contained in:
yucheng 2026-10-31 11:59:07 -07:00
parent 76d4f5e028
commit f51ca91711
3 changed files with 127 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

59
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -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 = (

74
tests/integration/authorization/test_passthrough_applied_guardrails_header.py
View file

@ -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

12
tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py
View file

@ -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"]}}