fix(proxy): preserve typed headers without rebinding messages

Signed-off-by: Jaeyeon Kim <anencore94@gmail.com>
This commit is contained in:
Jaeyeon Kim 2026-11-01 22:42:05 +00:00
parent dd7d4eba77
commit eacacd4446

View file

@ -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)