From ef3b24fc14ed2ebfbad7340c4f6fe014d75851ec Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 3 Nov 2026 04:19:18 -0800 Subject: [PATCH] fix(proxy): keep guardrail headers from pushing a response past a client's header limit --- .../pass_through_endpoints.py | 39 +++++++++++++++---- .../test_pass_through_endpoints.py | 37 ++++++++++++++++++ 2 files changed, 68 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 62568420f37..b73089db6c8 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -230,6 +230,10 @@ _JOINED_HEADER_METADATA_FIELDS: Final = frozenset( _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: @@ -285,9 +289,29 @@ def _is_sendable_header_value(value: str) -> bool: ) -def _with_header_in_budget(kept: Mapping[str, str], header: tuple[str, str]) -> Mapping[str, str]: +def _with_header_in_budget( + kept: Mapping[str, str], header: tuple[str, str], budget: int = _GUARDRAIL_HEADERS_MAX_BYTES +) -> 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 + return MappingProxyType({**kept, header[0]: header[1]}) if size <= budget else kept + + +def _guardrail_headers_with_room( + guardrail_headers: Mapping[str, str], response_headers: Mapping[str, str] +) -> Mapping[str, str]: + """``guardrail_headers``, without one the response's other headers leave no room for, so adding them can't cost a + client the response""" + others: Final = sum( + len(name) + len(value) + _HEADER_FRAMING_BYTES + for name, value in response_headers.items() + if name not in guardrail_headers + ) + no_headers: Final[Mapping[str, str]] = MappingProxyType({}) + return reduce( + partial(_with_header_in_budget, budget=_RESPONSE_HEADERS_MAX_BYTES - others), + guardrail_headers.items(), + no_headers, + ) def _recorded_guardrail_headers( @@ -595,13 +619,12 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): } return_headers.update(sanitized_custom_headers) guardrail_headers: Final = _guardrail_response_headers(*guardrail_request_data) - return_headers.update( - reduce( - _with_header_in_budget, - (header for header in earlier_guardrail_headers.items() if header[0] not in guardrail_headers), - guardrail_headers, - ) + 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 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 6cf010e67d0..9a0ab4d3d99 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 @@ -13851,6 +13851,43 @@ class TestGuardrailResponseHeaders: 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_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" + def test_guardrail_header_past_the_size_budget_is_left_out(self): names: Final = [f"g{index}" for index in range(2_000)]