diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 3b8751fe2c6..9de2c7b9d23 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -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( 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 0f507c80fb1..c98983598f4 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 @@ -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)]