fix(proxy): leave room for callback headers on streamed pass-through responses

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(

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