From 0c07a613a7a14c8795894de9ed9134214f40a25d Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 3 Nov 2026 04:53:59 -0800 Subject: [PATCH] fix(proxy): leave room for callback headers on streamed pass-through responses --- .../pass_through_endpoints.py | 33 +++++++++++- ...t_passthrough_applied_guardrails_header.py | 53 ++++++++++++++++++- 2 files changed, 83 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index f202e2bd407..a03b436d47c 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -234,6 +234,15 @@ _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: @@ -314,6 +323,20 @@ def _guardrail_headers_with_room( ) +def _with_callback_headers( + response_headers: Mapping[str, str], callback_headers: Mapping[str, str] +) -> dict[str, str]: # mutable-ok: the starlette response takes a plain dict + """``response_headers`` with a callback's headers added, less a guardrail header they then leave no room for""" + merged: Final = {**response_headers, **callback_headers} + guardrail_headers: Final = { + name: value + for name, value in response_headers.items() + if name in _GUARDRAIL_HEADER_NAMES and name not in callback_headers + } + kept: Final = _guardrail_headers_with_room(guardrail_headers, merged) + return {name: value for name, value in merged.items() if name not in guardrail_headers or name in kept} + + def _recorded_guardrail_headers( parsed_body: object = None, kwargs: object = None, post_call_data: object = None ) -> Mapping[str, str]: @@ -2129,7 +2152,10 @@ async def pass_through_request( request_headers=dict(redact_credential_headers(request.headers)), ) if callback_headers: - _response_headers.update(callback_headers) + _response_headers = _with_callback_headers( # rebind-ok: replaces this branch's headers with the room kept + _response_headers, # pyright: ignore[reportUnknownArgumentType] # untyped header dict + callback_headers, + ) return StreamingResponse( wrap_passthrough_sse_bytes_with_keepalive_pings( @@ -2234,7 +2260,10 @@ async def pass_through_request( request_headers=dict(redact_credential_headers(request.headers)), ) if callback_headers: - _response_headers.update(callback_headers) + _response_headers = _with_callback_headers( # rebind-ok: replaces this branch's headers with the room kept + _response_headers, # pyright: ignore[reportUnknownArgumentType] # untyped header dict + callback_headers, + ) return StreamingResponse( wrap_passthrough_sse_bytes_with_keepalive_pings( diff --git a/tests/integration/authorization/test_passthrough_applied_guardrails_header.py b/tests/integration/authorization/test_passthrough_applied_guardrails_header.py index df6fefd26a4..43ee159095f 100644 --- a/tests/integration/authorization/test_passthrough_applied_guardrails_header.py +++ b/tests/integration/authorization/test_passthrough_applied_guardrails_header.py @@ -28,6 +28,23 @@ HEADER: Final = "x-litellm-applied-guardrails" OPENAI_ROUTE: Final = "/openai/v1/chat/completions" ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages" GEMINI_ROUTE: Final = "/gemini/v1beta/models/gemini-2.5-flash:generateContent" +CALLBACK_HEADERS_MODEL: Final = "passthrough-header-callback-headers" +UNREQUESTED_STREAM_MODEL: Final = "passthrough-header-callback-headers-unrequested-stream" +CALLBACK_HEADER: Final = "x-callback-padding" +CALLBACK_HEADER_VALUE: Final = "c" * 12_000 +_CALLBACK_HEADERS_MODULE: Final = f""" +from litellm.integrations.custom_logger import CustomLogger + + +class LargeHeadersHook(CustomLogger): + async def async_post_call_response_headers_hook(self, data, user_api_key_dict, response, *args, **kwargs): + if str(data.get("model")).startswith("{CALLBACK_HEADERS_MODEL}"): + return {{"{CALLBACK_HEADER}": "{CALLBACK_HEADER_VALUE}"}} + return None + + +proxy_handler_instance = LargeHeadersHook() +""" _CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" _JSON: Final = TypeAdapter(dict[str, JsonValue]) @@ -61,7 +78,12 @@ def _openai_reply(request: Request) -> 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) + return chat_reply( + "chatcmpl-passthrough-header", + "gpt-5.6", + "Hello", + stream=payload.get("stream") is True or payload.get("model") == UNREQUESTED_STREAM_MODEL, + ) def _anthropic_reply(_request: Request) -> Reply: @@ -100,11 +122,16 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I } ) config: Final = _JSON.validate_python(yaml.safe_load(_CONFIG.read_text())) + (directory / "passthrough_large_headers_hook.py").write_text(_CALLBACK_HEADERS_MODULE) path: Final = directory / "passthrough-applied-guardrails-header.yaml" path.write_text( yaml.safe_dump( { **config, + "litellm_settings": { + **_JSON.validate_python(config.get("litellm_settings", {})), + "callbacks": ["passthrough_large_headers_hook.proxy_handler_instance"], + }, "guardrails": [ { "guardrail_name": name, @@ -294,6 +321,30 @@ def test_header_value_with_a_line_break_is_left_out(rig: HeaderRig, suffix: str) assert "x-injected" not in response.headers +@pytest.mark.parametrize( + ("model", "stream"), + [(CALLBACK_HEADERS_MODEL, True), (UNREQUESTED_STREAM_MODEL, False)], + ids=["requested-stream", "unrequested-stream"], +) +@pytest.mark.parametrize(("names", "kept"), [(300, False), (100, True)], ids=["no-room", "room"]) +def test_a_callback_header_on_a_stream_leaves_room_or_costs_the_guardrail_header( + rig: HeaderRig, model: str, stream: bool, names: int, kept: bool +) -> None: + gateway: Final = rig.owned.gateway + body: Final[dict[str, JsonValue]] = { + **_route_body(OPENAI_ROUTE), + "model": model, + "stream": stream, + "guardrails": [f"caller-named-{index}" for index in range(names)], + } + with gateway.scenario() as scenario: + response: Final = _send(gateway, scenario, OPENAI_ROUTE, body) + + assert response.status_code == 200, response.text + assert response.headers.get(CALLBACK_HEADER) == CALLBACK_HEADER_VALUE + assert (HEADER in response.headers) is kept + + @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