From dd7d4eba779c48f2dd4b9a8bcc5c80b8f0c6aad8 Mon Sep 17 00:00:00 2001 From: Jaeyeon Kim Date: Sun, 1 Nov 2026 21:56:48 +0000 Subject: [PATCH 1/2] feat(proxy): optionally omit LiteLLM response headers across HTTP paths Signed-off-by: Jaeyeon Kim --- litellm/__init__.py | 1 + .../litellm_response_headers_middleware.py | 29 ++++++++++++++++ litellm/proxy/proxy_server.py | 3 ++ ...est_litellm_response_headers_middleware.py | 33 +++++++++++++++++++ 4 files changed, 66 insertions(+) create mode 100644 litellm/proxy/middleware/litellm_response_headers_middleware.py create mode 100644 tests/unit/proxy/middleware/test_litellm_response_headers_middleware.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 0e5101b49c3..cad8973f0b3 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -425,6 +425,7 @@ autorouter_presets_url: str = os.getenv( "https://raw.githubusercontent.com/BerriAI/litellm/main/litellm/proxy/public_endpoints/autorouter_presets.json", ) suppress_debug_info: bool = False +include_litellm_headers: bool = True dynamodb_table_name: Optional[str] = None s3_callback_params: Optional[Dict] = None s3_audit_callback_params: Optional[Dict] = None diff --git a/litellm/proxy/middleware/litellm_response_headers_middleware.py b/litellm/proxy/middleware/litellm_response_headers_middleware.py new file mode 100644 index 00000000000..620c63f8715 --- /dev/null +++ b/litellm/proxy/middleware/litellm_response_headers_middleware.py @@ -0,0 +1,29 @@ +"""Apply the proxy's opt-out to all HTTP response headers, including streams.""" + +from starlette.types import ASGIApp, Message, Receive, Scope, Send + +import litellm + + +class LiteLLMResponseHeadersMiddleware: + def __init__(self, app: ASGIApp) -> None: + self.app = app + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if scope["type"] != "http": + await self.app(scope, receive, send) + return + + async def send_response(message: Message) -> None: + if message["type"] == "http.response.start" and litellm.include_litellm_headers is False: + message = { + **message, + "headers": [ + (name, value) + for name, value in message.get("headers", []) + if not name.lower().startswith(b"x-litellm-") + ], + } + await send(message) + + await self.app(scope, receive, send_response) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c5fcbfb2564..450ff33146c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -828,6 +828,7 @@ from litellm.proxy.middleware.billable_request_metrics_middleware import ( from litellm.proxy.middleware.budget_reservation_release_middleware import ( BudgetReservationReleaseMiddleware, ) +from litellm.proxy.middleware.litellm_response_headers_middleware import LiteLLMResponseHeadersMiddleware from litellm.proxy.middleware.redis_request_batch_middleware import RedisRequestBatchMiddleware from litellm.proxy.plugin_routes import ( register_plugins_from_config, @@ -21907,6 +21908,8 @@ app.add_middleware( get_settings=lambda: get_admission_control_settings(general_settings), state=admission_control_state, ) +app.add_middleware(LiteLLMResponseHeadersMiddleware) + # Added last on purpose - last-added is outermost, and the client-visible URL # prefix must be resolved into scope["root_path"] before any inner middleware # or the router inspects the path. Only added when SERVER_ROOT_PATHS is diff --git a/tests/unit/proxy/middleware/test_litellm_response_headers_middleware.py b/tests/unit/proxy/middleware/test_litellm_response_headers_middleware.py new file mode 100644 index 00000000000..dbc37849ad5 --- /dev/null +++ b/tests/unit/proxy/middleware/test_litellm_response_headers_middleware.py @@ -0,0 +1,33 @@ +import pytest +from fastapi import FastAPI +from fastapi.responses import StreamingResponse +from fastapi.testclient import TestClient + +import litellm +from litellm.proxy.middleware.litellm_response_headers_middleware import LiteLLMResponseHeadersMiddleware + + +@pytest.mark.parametrize("include", [True, False]) +@pytest.mark.parametrize("status_code", [200, 400]) +@pytest.mark.parametrize("streaming", [True, False]) +def test_header_opt_out_preserves_body_and_other_headers(monkeypatch, include, status_code, streaming): + from starlette.responses import Response + + monkeypatch.setattr(litellm, "include_litellm_headers", include) + app = FastAPI() + app.add_middleware(LiteLLMResponseHeadersMiddleware) + + @app.get("/reply") + def reply(): + headers = {"X-LiteLLM-Call-ID": "call", "x-litellm-priority": "1", "retry-after": "2"} + if streaming: + return StreamingResponse(iter([b"hello", b" world"]), headers=headers, status_code=status_code) + return Response(b"hello world", headers=headers, status_code=status_code) + + with TestClient(app) as client: + response = client.get("/reply") + assert response.status_code == status_code + assert response.content == b"hello world" + assert response.headers["retry-after"] == "2" + actual = {k: v for k, v in response.headers.items() if k.startswith("x-litellm-")} + assert actual == ({"x-litellm-call-id": "call", "x-litellm-priority": "1"} if include else {}) From eacacd4446bcd3b5886e2b683d80c68f71270fff Mon Sep 17 00:00:00 2001 From: Jaeyeon Kim Date: Sun, 1 Nov 2026 22:42:05 +0000 Subject: [PATCH 2/2] fix(proxy): preserve typed headers without rebinding messages Signed-off-by: Jaeyeon Kim --- .../litellm_response_headers_middleware.py | 25 +++++++++++-------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/middleware/litellm_response_headers_middleware.py b/litellm/proxy/middleware/litellm_response_headers_middleware.py index 620c63f8715..879e5e58c5f 100644 --- a/litellm/proxy/middleware/litellm_response_headers_middleware.py +++ b/litellm/proxy/middleware/litellm_response_headers_middleware.py @@ -1,5 +1,8 @@ """Apply the proxy's opt-out to all HTTP response headers, including streams.""" +from typing import Final + +from starlette.datastructures import Headers from starlette.types import ASGIApp, Message, Receive, Scope, Send import litellm @@ -15,15 +18,17 @@ class LiteLLMResponseHeadersMiddleware: return async def send_response(message: Message) -> None: - if message["type"] == "http.response.start" and litellm.include_litellm_headers is False: - message = { - **message, - "headers": [ - (name, value) - for name, value in message.get("headers", []) - if not name.lower().startswith(b"x-litellm-") - ], - } - await send(message) + if message["type"] != "http.response.start" or litellm.include_litellm_headers is not False: + await send(message) + return + response: Final[Message] = { + **message, + "headers": [ + (name, value) + for name, value in Headers(scope=message).raw + if not name.lower().startswith(b"x-litellm-") + ], + } + await send(response) await self.app(scope, receive, send_response)