scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-10 08:28:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 108

fix(proxy): keep an upstream's guardrail headers beside a streamed callback header

Browse source 
Only a guardrail header this proxy recorded gives way to a callback header that leaves it no room, so a stream
whose proxy recorded nothing passes the upstream's headers through unchanged, as the buffered path does.
This commit is contained in:
yucheng 2026-11-05 01:20:49 -08:00
parent 392c26895b
commit 7c488f9f44
2 changed files with 42 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -251,15 +251,6 @@ _GUARDRAIL_HEADERS_MAX_BYTES: Final = 8192
_RESPONSE_HEADERS_MAX_BYTES: Final = 15360
# The ": " and the line ending around each header on the wire.
_HEADER_FRAMING_BYTES: Final = 4
_GUARDRAIL_HEADER_NAMES: Final = frozenset(
(
"x-litellm-applied-guardrails",
"x-litellm-applied-policies",
"x-litellm-guardrail-scan-id",
"x-litellm-guardrail-scan-metadata",
"x-litellm-policy-sources",
)
)
def _is_mergeable_client_metadata(key: str, value: object) -> bool:
@ -339,12 +330,14 @@ def _guardrail_headers_with_room(
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 they then leave no room for"""
"""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 response_headers.items()
if name in _GUARDRAIL_HEADER_NAMES and name not in callback_headers
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)
@ -2238,10 +2231,11 @@ async def pass_through_request(
# 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,
guardrail_request_data=(_parsed_body, kwargs), # pyright: ignore[reportUnknownArgumentType] # untyped request dicts
earlier_guardrail_headers=streamed_guardrail_headers,
)
relay_response: Final = await _log_passthrough_upstream_failure(
response=response,
@ -2267,7 +2261,7 @@ async def pass_through_request(
request_headers=dict(redact_credential_headers(request.headers)),
)
if callback_headers:
_add_callback_headers(_response_headers, callback_headers) # pyright: ignore[reportUnknownArgumentType] # untyped header dict
_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(
@ -2343,10 +2337,11 @@ async def pass_through_request(
# 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,
guardrail_request_data=(_parsed_body, kwargs), # pyright: ignore[reportUnknownArgumentType] # untyped request dicts
earlier_guardrail_headers=detected_stream_guardrail_headers,
)
detected_relay_response: Final = await _log_passthrough_upstream_failure(
response=response,
@ -2372,7 +2367,7 @@ async def pass_through_request(
request_headers=dict(redact_credential_headers(request.headers)),
)
if callback_headers:
_add_callback_headers(_response_headers, callback_headers) # pyright: ignore[reportUnknownArgumentType] # untyped header dict
_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(

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

@ -68,6 +68,7 @@ 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,
@ -14087,6 +14088,36 @@ class TestGuardrailResponseHeaders:
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)]