diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 0a29a71cd06..3b0f9ff79bf 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -205,6 +205,8 @@ _GUARDRAIL_HEADER_METADATA_FIELDS: Final = frozenset( ) # h11's header field-value grammar: no NUL, line breaks or other whitespace beyond spaces and tabs between characters. _SENDABLE_HEADER_VALUE: Final = re.compile(rb"(?:[^\x00\s]+(?:[ \t]+[^\x00\s]+)*)?") +# 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 def _is_mergeable_client_metadata(key: str, value: object) -> bool: @@ -244,6 +246,11 @@ def _is_sendable_header_value(value: str) -> bool: ) +def _with_header_in_budget(kept: Mapping[str, str], header: tuple[str, str]) -> 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 + + def _guardrail_response_headers( parsed_body: object = None, kwargs: object = None, post_call_data: object = None ) -> Mapping[str, str]: @@ -251,7 +258,8 @@ def _guardrail_response_headers( Pre-call guardrails record on the request body's metadata, which ``init_kwargs_for_pass_through_endpoint`` moves to ``kwargs["litellm_params"]``, and post-call guardrails record on ``post_call_data``, so list entries such as the - applied guardrail names are combined. A ``litellm_params`` object sent in the request body is never read. + applied guardrail names are combined. A ``litellm_params`` object sent in the request body is never read. A header + that would take these headers past ``_GUARDRAIL_HEADERS_MAX_BYTES`` is left out. """ scopes: Final = ( as_str_mapping(parsed_body) or {}, @@ -270,7 +278,9 @@ def _guardrail_response_headers( for field, entries in groupby(sorted(recorded, key=lambda entry: entry[0]), key=lambda entry: entry[0]) } headers: Final = _RESPONSE_HEADERS.validate_python(get_logging_caching_headers({"metadata": merged}) or {}) - return MappingProxyType({name: value for name, value in headers.items() if _is_sendable_header_value(value)}) + sendable: Final = ((name, value) for name, value in headers.items() if _is_sendable_header_value(value)) + no_headers: Final[Mapping[str, str]] = MappingProxyType({}) + return reduce(_with_header_in_budget, sendable, no_headers) class _RegisteredPassThroughRoute(TypedDict): diff --git a/tests/integration/authorization/test_passthrough_applied_guardrails_header.py b/tests/integration/authorization/test_passthrough_applied_guardrails_header.py index b4c05f6a1ef..df6fefd26a4 100644 --- a/tests/integration/authorization/test_passthrough_applied_guardrails_header.py +++ b/tests/integration/authorization/test_passthrough_applied_guardrails_header.py @@ -292,3 +292,20 @@ def test_header_value_with_a_line_break_is_left_out(rig: HeaderRig, suffix: str) assert response.headers.get("x-litellm-applied-policies") == TAG_POLICY assert "x-litellm-policy-sources" not in response.headers assert "x-injected" not in response.headers + + +@pytest.mark.parametrize("stream", [False, True]) +def test_response_to_a_body_naming_thousands_of_guardrails_stays_readable(rig: HeaderRig, stream: bool) -> None: + gateway: Final = rig.owned.gateway + body: Final[dict[str, JsonValue]] = { + **_route_body(OPENAI_ROUTE), + "stream": stream, + "guardrails": [f"caller-named-{index}" for index in range(10_000)], + "metadata": {"tags": [f"{POLICY_TAG_PREFIX}prod"]}, + } + with gateway.scenario() as scenario: + response: Final = _send(gateway, scenario, OPENAI_ROUTE, body) + + assert response.status_code == 200, response.text + assert HEADER not in response.headers + assert response.headers.get("x-litellm-applied-policies") == TAG_POLICY 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 1d9a3be5e2f..5db94e7bcea 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 @@ -12125,3 +12125,18 @@ class TestGuardrailResponseHeaders: assert headers == {"x-litellm-applied-policies": "tag-policy"} assert h11.Response(status_code=200, headers=[*headers.items()]).status_code == 200 + + def test_guardrail_headers_past_the_size_budget_are_left_out_so_a_client_can_read_the_response(self): + names: Final = [f"caller-named-{index}" for index in range(10_000)] + headers: Final = _guardrail_response_headers( + {}, {"litellm_params": {"metadata": {"applied_guardrails": names, "applied_policies": ["tag-policy"]}}} + ) + server: Final = h11.Connection(h11.SERVER) + server.receive_data(b"GET / HTTP/1.1\r\nhost: proxy\r\n\r\n") + server.next_event() + client: Final = h11.Connection(h11.CLIENT) + client.send(h11.Request(method="GET", target="/", headers=[("host", "proxy")])) + client.receive_data(server.send(h11.Response(status_code=200, headers=[*headers.items()])) or b"") + + assert headers == {"x-litellm-applied-policies": "tag-policy"} + assert isinstance(client.next_event(), h11.Response)