From 71975a744dd733cb539e5522928e9ba0d753403f Mon Sep 17 00:00:00 2001 From: Kentaro Takiguchi Date: Tue, 3 Nov 2026 09:19:34 +0000 Subject: [PATCH 1/5] fix(google_genai): forward proxy client headers in generate_content endpoints Merge proxy-forwarded `headers` (populated by `forward_client_headers_to_llm_api` or deployment `litellm_params`) into `extra_headers` across sync/async `generate_content` and `generate_content_stream` entrypoints for both native Gemini/Vertex AI handlers and non-Gemini completion adapters. Explicit `extra_headers` take case-insensitive precedence on header-name conflicts. --- litellm/google_genai/main.py | 48 ++- .../google_genai/test_google_genai_main.py | 373 +++++++++++++++++- 2 files changed, 413 insertions(+), 8 deletions(-) diff --git a/litellm/google_genai/main.py b/litellm/google_genai/main.py index 40c9c7bc36c..1e4aa7d1224 100644 --- a/litellm/google_genai/main.py +++ b/litellm/google_genai/main.py @@ -1,10 +1,11 @@ import asyncio import contextvars +from collections.abc import Mapping from functools import partial from typing import TYPE_CHECKING, Any, ClassVar, Final import httpx -from pydantic import ConfigDict, TypeAdapter +from pydantic import ConfigDict, TypeAdapter, ValidationError import litellm from litellm.constants import request_timeout @@ -40,6 +41,7 @@ else: # Initialize any necessary instances or variables here base_llm_http_handler = BaseLLMHTTPHandler() _GENERATE_CONTENT_RESULT: Final = TypeAdapter(object) +_FORWARDED_HEADERS: Final = TypeAdapter[Mapping[str, object]](Mapping[str, object]) ################################################# @@ -236,6 +238,35 @@ def _merge_native_request_fields( return {**native_request_fields, **(extra_body or {})} +def _validated_client_headers(client_headers: object) -> Mapping[str, object] | None: + if client_headers is None: + return None + try: + return _FORWARDED_HEADERS.validate_python(client_headers) + except ValidationError: + return None + + +def _merge_client_forwarded_headers( + extra_headers: Mapping[str, object] | None, + client_headers: object, +) -> dict[str, object] | None: # mutable-ok: downstream handlers expect dict | None + """ + Merge proxy-forwarded ``headers`` (set when ``forward_client_headers_to_llm_api`` + is enabled or on deployment ``litellm_params``) into ``extra_headers``. + + ``extra_headers`` wins on case-insensitive header-name conflicts. + """ + forwarded: Final = _validated_client_headers(client_headers) + if not forwarded: + return dict(extra_headers) if extra_headers is not None else None + explicit_names: Final = frozenset(name.lower() for name in (extra_headers or ())) + return { + **{name: value for name, value in forwarded.items() if name.lower() not in explicit_names}, + **(extra_headers or {}), + } + + @client async def agenerate_content( model: str, @@ -349,6 +380,7 @@ def generate_content( # Extract systemInstruction from kwargs to pass to handler system_instruction: Final = kwargs.get("systemInstruction") or kwargs.get("system_instruction") + merged_extra_headers: Final = _merge_client_forwarded_headers(extra_headers, kwargs.get("headers")) # Check if we should use the adapter (when provider config is None) if setup_result.generate_content_provider_config is None: @@ -360,7 +392,7 @@ def generate_content( tools=tools, _is_async=_is_async, litellm_params=setup_result.litellm_params, - extra_headers=extra_headers, + extra_headers=merged_extra_headers, **kwargs, ) @@ -375,7 +407,7 @@ def generate_content( custom_llm_provider=setup_result.custom_llm_provider, litellm_params=setup_result.litellm_params, logging_obj=setup_result.litellm_logging_obj, - extra_headers=extra_headers, + extra_headers=merged_extra_headers, extra_body=_merge_native_request_fields(setup_result.native_request_fields, extra_body), timeout=timeout or request_timeout, _is_async=_is_async, @@ -442,6 +474,7 @@ async def agenerate_content_stream( # Extract systemInstruction from kwargs to pass to handler system_instruction: Final = kwargs.get("systemInstruction") or kwargs.get("system_instruction") + merged_extra_headers: Final = _merge_client_forwarded_headers(extra_headers, kwargs.get("headers")) # Check if we should use the adapter (when provider config is None) if setup_result.generate_content_provider_config is None: @@ -456,7 +489,7 @@ async def agenerate_content_stream( litellm_params=setup_result.litellm_params, tools=tools, stream=True, - extra_headers=extra_headers, + extra_headers=merged_extra_headers, **kwargs, ) @@ -472,7 +505,7 @@ async def agenerate_content_stream( custom_llm_provider=setup_result.custom_llm_provider, litellm_params=setup_result.litellm_params, logging_obj=setup_result.litellm_logging_obj, - extra_headers=extra_headers, + extra_headers=merged_extra_headers, extra_body=_merge_native_request_fields(setup_result.native_request_fields, extra_body), timeout=timeout or request_timeout, _is_async=True, @@ -534,6 +567,7 @@ def generate_content_stream( # Extract systemInstruction from kwargs to pass to handler system_instruction: Final = kwargs.get("systemInstruction") or kwargs.get("system_instruction") + merged_extra_headers: Final = _merge_client_forwarded_headers(extra_headers, kwargs.get("headers")) # Check if we should use the adapter (when provider config is None) if setup_result.generate_content_provider_config is None: @@ -548,7 +582,7 @@ def generate_content_stream( _is_async=_is_async, litellm_params=setup_result.litellm_params, stream=True, - extra_headers=extra_headers, + extra_headers=merged_extra_headers, **kwargs, ) @@ -563,7 +597,7 @@ def generate_content_stream( custom_llm_provider=setup_result.custom_llm_provider, litellm_params=setup_result.litellm_params, logging_obj=setup_result.litellm_logging_obj, - extra_headers=extra_headers, + extra_headers=merged_extra_headers, extra_body=_merge_native_request_fields(setup_result.native_request_fields, extra_body), timeout=timeout or request_timeout, _is_async=_is_async, diff --git a/tests/unit/google_genai/test_google_genai_main.py b/tests/unit/google_genai/test_google_genai_main.py index 5cac61e727e..8f909cfc381 100644 --- a/tests/unit/google_genai/test_google_genai_main.py +++ b/tests/unit/google_genai/test_google_genai_main.py @@ -4,7 +4,7 @@ Test to verify the Google GenAI generate_content adapter functionality """ import json -from collections.abc import Iterator +from collections.abc import Iterator, Mapping from typing import Final import pytest @@ -382,3 +382,374 @@ async def test_vertex_generate_content_spend_follows_the_deployment_location( recorded: Final = await _logged_vertex_generate_content_cost(location, stream=stream) assert recorded == pytest.approx(expected), (location, stream) + + +class _WireCapture: + def __init__(self, url: str) -> None: + self.url: Final = url + self._items: Final[list[tuple[str, Mapping[str, str]]]] = [] # mutable-ok: test wire log + + def record(self, path: str, headers: Mapping[str, str]) -> None: + self._items.append((path, headers)) + + @property + def requests(self) -> tuple[tuple[str, Mapping[str, str]], ...]: + return tuple(self._items) + + +def _mock_wire_reply(path: str, stream_requested: bool) -> tuple[bytes, str]: + if "streamGenerateContent" in path: + return ( + b'data: {"candidates":[{"content":{"role":"model","parts":[{"text":"pong"}]},' + b'"finishReason":"STOP"}],' + b'"usageMetadata":{"promptTokenCount":4,"candidatesTokenCount":1,"totalTokenCount":5}}\n\n', + "text/event-stream", + ) + if ":generateContent" in path: + return ( + b'{"candidates":[{"content":{"role":"model","parts":[{"text":"pong"}]},' + b'"finishReason":"STOP"}],' + b'"usageMetadata":{"promptTokenCount":4,"candidatesTokenCount":1,"totalTokenCount":5}}', + "application/json", + ) + if stream_requested: + return ( + b'data: {"id":"chatcmpl-1","object":"chat.completion.chunk","created":1,' + b'"model":"gpt-4o-mini","choices":[{"index":0,' + b'"delta":{"role":"assistant","content":"pong"},"finish_reason":"stop"}]}\n\n' + b"data: [DONE]\n\n", + "text/event-stream", + ) + return ( + b'{"id":"chatcmpl-1","object":"chat.completion","created":1,"model":"gpt-4o-mini",' + b'"choices":[{"index":0,"message":{"role":"assistant","content":"pong"},' + b'"finish_reason":"stop"}],' + b'"usage":{"prompt_tokens":4,"completion_tokens":1,"total_tokens":5}}', + "application/json", + ) + + +@pytest.fixture +def header_capture_wire() -> Iterator[_WireCapture]: + import threading + from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + + holder: Final[list[_WireCapture]] = [] # mutable-ok: single-slot fixture bootstrap reference + + class _Handler(BaseHTTPRequestHandler): + def do_POST(self) -> None: + length: Final = int(self.headers.get("Content-Length", "0")) + raw_body: Final = self.rfile.read(length) + lowered_headers: Final = {k.lower(): v for k, v in self.headers.items()} + holder[0].record(self.path, lowered_headers) + + payload: Final = json.loads(raw_body.decode("utf-8")) if raw_body else {} + reply: Final = _mock_wire_reply(self.path, bool(payload.get("stream"))) + body: Final = reply[0] + content_type: Final = reply[1] + + self.send_response(200) + self.send_header("Content-Type", content_type) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, format: str, *args: object) -> None: + return + + server: Final = ThreadingHTTPServer(("127.0.0.1", 0), _Handler) + host: Final = server.server_address[0] + port: Final = server.server_address[1] + host_str: Final = host.decode("ascii") if isinstance(host, bytes) else host + capture: Final = _WireCapture(f"http://{host_str}:{port}") + holder.append(capture) + thread: Final = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield capture + finally: + server.shutdown() + server.server_close() + thread.join(timeout=5) + + +def _extract_chunk_texts(payload: object) -> tuple[str, ...]: + from pydantic import TypeAdapter + + obj_map: Final = TypeAdapter[Mapping[str, object]](Mapping[str, object]) + map_seq: Final = TypeAdapter[tuple[Mapping[str, object], ...]](tuple[Mapping[str, object], ...]) + parsed: Final = obj_map.validate_python(payload.model_dump() if hasattr(payload, "model_dump") else payload) + candidates: Final = map_seq.validate_python(parsed.get("candidates", ())) + if not candidates: + return () + content: Final = obj_map.validate_python(candidates[0].get("content", {})) + parts: Final = map_seq.validate_python(content.get("parts", ())) + return tuple(str(part["text"]) for part in parts if isinstance(part.get("text"), str) and part["text"]) + + +def _sse_payload_texts(line: str) -> tuple[str, ...]: + if not line.startswith("data: "): + return () + payload: Final = line[6:].strip() + if not payload or payload == "[DONE]": + return () + return _extract_chunk_texts(json.loads(payload)) + + +def _chunk_sse_texts(raw: bytes) -> tuple[str, ...]: + from itertools import chain + + return tuple(chain.from_iterable(_sse_payload_texts(line) for line in raw.decode("utf-8").splitlines())) + + +def _extract_sse_texts(chunks: tuple[bytes, ...]) -> tuple[str, ...]: + from itertools import chain + + return tuple(chain.from_iterable(_chunk_sse_texts(raw) for raw in chunks)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "entrypoint", + ["sync", "async", "sync_stream", "async_stream"], +) +async def test_native_generate_content_forwards_client_headers( + header_capture_wire: _WireCapture, + entrypoint: str, +) -> None: + from collections.abc import AsyncIterator + + from pydantic import TypeAdapter + + from litellm.google_genai.main import ( + agenerate_content, + agenerate_content_stream, + generate_content, + generate_content_stream, + ) + + contents: Final = [{"role": "user", "parts": [{"text": "ping"}]}] + forwarded: Final = {"x-team-id": "platform", "x-trace-tag": f"native-{entrypoint}"} + + if entrypoint == "sync": + resp: Final = generate_content( + model="gemini/gemini-2.0-flash", + contents=contents, + api_base=header_capture_wire.url, + api_key="gemini-secret", + headers=forwarded, + ) + assert _extract_chunk_texts(resp) == ("pong",) + elif entrypoint == "async": + async_resp: Final = await agenerate_content( + model="gemini/gemini-2.0-flash", + contents=contents, + api_base=header_capture_wire.url, + api_key="gemini-secret", + headers=forwarded, + ) + assert _extract_chunk_texts(async_resp) == ("pong",) + elif entrypoint == "sync_stream": + sync_stream: Final = generate_content_stream( + model="gemini/gemini-2.0-flash", + contents=contents, + api_base=header_capture_wire.url, + api_key="gemini-secret", + headers=forwarded, + ) + assert isinstance(sync_stream, Iterator) + sync_chunks: Final = tuple(TypeAdapter(bytes).validate_python(chunk) for chunk in sync_stream) + assert _extract_sse_texts(sync_chunks) == ("pong",) + else: + async_stream: Final = await agenerate_content_stream( + model="gemini/gemini-2.0-flash", + contents=contents, + api_base=header_capture_wire.url, + api_key="gemini-secret", + headers=forwarded, + ) + assert isinstance(async_stream, AsyncIterator) + async_chunks: Final = tuple([TypeAdapter(bytes).validate_python(chunk) async for chunk in async_stream]) + assert _extract_sse_texts(async_chunks) == ("pong",) + + assert len(header_capture_wire.requests) == 1 + wire_headers: Final = header_capture_wire.requests[0][1] + assert wire_headers["x-goog-api-key"] == "gemini-secret" + assert wire_headers["x-team-id"] == "platform" + assert wire_headers["x-trace-tag"] == f"native-{entrypoint}" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "entrypoint", + ["sync", "async", "sync_stream", "async_stream"], +) +async def test_adapter_generate_content_forwards_client_headers( + header_capture_wire: _WireCapture, + entrypoint: str, +) -> None: + from collections.abc import AsyncIterator + + from pydantic import TypeAdapter + + from litellm.google_genai.main import ( + agenerate_content, + agenerate_content_stream, + generate_content, + generate_content_stream, + ) + + contents: Final = [{"role": "user", "parts": [{"text": "ping"}]}] + forwarded: Final = { + "x-team-id": "platform", + "anthropic-beta": "context-1m-2025-08-07", + "x-trace-tag": f"adapter-{entrypoint}", + } + + if entrypoint == "sync": + resp: Final = generate_content( + model="openai/gpt-4o-mini", + contents=contents, + api_base=f"{header_capture_wire.url}/v1", + api_key="openai-secret", + headers=forwarded, + ) + assert _extract_chunk_texts(resp) == ("pong",) + elif entrypoint == "async": + async_resp: Final = await agenerate_content( + model="openai/gpt-4o-mini", + contents=contents, + api_base=f"{header_capture_wire.url}/v1", + api_key="openai-secret", + headers=forwarded, + ) + assert _extract_chunk_texts(async_resp) == ("pong",) + elif entrypoint == "sync_stream": + sync_stream: Final = generate_content_stream( + model="openai/gpt-4o-mini", + contents=contents, + api_base=f"{header_capture_wire.url}/v1", + api_key="openai-secret", + headers=forwarded, + ) + assert isinstance(sync_stream, AsyncIterator) + sync_chunks: Final = tuple([TypeAdapter(bytes).validate_python(chunk) async for chunk in sync_stream]) + assert _extract_sse_texts(sync_chunks) == ("pong",) + else: + async_stream: Final = await agenerate_content_stream( + model="openai/gpt-4o-mini", + contents=contents, + api_base=f"{header_capture_wire.url}/v1", + api_key="openai-secret", + headers=forwarded, + ) + assert isinstance(async_stream, AsyncIterator) + async_chunks: Final = tuple([TypeAdapter(bytes).validate_python(chunk) async for chunk in async_stream]) + assert _extract_sse_texts(async_chunks) == ("pong",) + + assert len(header_capture_wire.requests) == 1 + req_path: Final = header_capture_wire.requests[0][0] + wire_headers: Final = header_capture_wire.requests[0][1] + assert req_path == "/v1/chat/completions" + assert wire_headers["authorization"] == "Bearer openai-secret" + assert wire_headers["x-team-id"] == "platform" + assert wire_headers["anthropic-beta"] == "context-1m-2025-08-07" + assert wire_headers["x-trace-tag"] == f"adapter-{entrypoint}" + + +@pytest.mark.parametrize("target", ["native", "adapter"]) +def test_extra_headers_override_forwarded_client_headers_case_insensitively( + header_capture_wire: _WireCapture, + target: str, +) -> None: + from litellm.google_genai.main import _merge_client_forwarded_headers, generate_content + + forwarded: Final = {"x-custom": "from-proxy", "x-only-proxy": "proxy-val"} + explicit: Final = {"X-Custom": "from-caller", "X-Only-Caller": "caller-val"} + + resp: Final = generate_content( + model="gemini/gemini-2.0-flash" if target == "native" else "openai/gpt-4o-mini", + contents=[{"role": "user", "parts": [{"text": "ping"}]}], + api_base=header_capture_wire.url if target == "native" else f"{header_capture_wire.url}/v1", + api_key="test-secret", + headers=forwarded, + extra_headers=explicit, + ) + + assert _extract_chunk_texts(resp) == ("pong",) + assert forwarded == {"x-custom": "from-proxy", "x-only-proxy": "proxy-val"} + assert explicit == {"X-Custom": "from-caller", "X-Only-Caller": "caller-val"} + assert len(header_capture_wire.requests) == 1 + wire_headers: Final = header_capture_wire.requests[0][1] + assert wire_headers["x-custom"] == "from-caller" + assert wire_headers["x-only-proxy"] == "proxy-val" + assert wire_headers["x-only-caller"] == "caller-val" + + assert _merge_client_forwarded_headers(None, None) is None + assert _merge_client_forwarded_headers(None, {}) is None + assert _merge_client_forwarded_headers({}, None) == {} + source_extra: Final = {"x-b": "2"} + merged_extra: Final = _merge_client_forwarded_headers(source_extra, None) + assert merged_extra == {"x-b": "2"} and merged_extra is not source_extra + source_headers: Final = {"x-a": "1"} + merged_only: Final = _merge_client_forwarded_headers(None, source_headers) + assert merged_only == {"x-a": "1"} and merged_only is not source_headers + assert _merge_client_forwarded_headers({}, {"x-a": "1"}) == {"x-a": "1"} + assert _merge_client_forwarded_headers({"X-Custom": "from-caller"}, {"x-custom": "from-proxy"}) == { + "X-Custom": "from-caller" + } + assert _merge_client_forwarded_headers({"x-b": "2"}, "not-a-mapping") == {"x-b": "2"} + assert _merge_client_forwarded_headers(None, "not-a-mapping") is None + + +@pytest.mark.asyncio +async def test_proxy_forwarded_headers_reach_generate_content_without_leaking_proxy_auth( + header_capture_wire: _WireCapture, +) -> None: + from starlette.datastructures import Headers + + from litellm.google_genai.main import agenerate_content + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, clean_headers + + incoming_headers: Final = Headers( + { + "x-goog-api-key": "sk-litellm-virtual-key", + "authorization": "Bearer sk-litellm-virtual-key", + "x-cost-center": "core-platform", + "x-request-org": "acme", + "user-agent": "google-genai-sdk/1.0", + } + ) + cleaned: Final = clean_headers( + incoming_headers, + litellm_key_header_name=None, + forward_llm_provider_auth_headers=False, + authenticated_with_header="x-goog-api-key", + ) + request_data: Final[Mapping[str, object]] = { + "model": "gemini/gemini-2.0-flash", + "contents": [{"role": "user", "parts": [{"text": "ping"}]}], + } + backend_data: Final = LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call( + headers=cleaned, + request_data=request_data, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-litellm-virtual-key"), + general_settings={"forward_client_headers_to_llm_api": True}, + ) + + resp: Final = await agenerate_content( + model="gemini/gemini-2.0-flash", + contents=request_data["contents"], + api_base=header_capture_wire.url, + api_key="real-upstream-gemini-key", + headers=backend_data.get("headers"), + ) + + assert _extract_chunk_texts(resp) == ("pong",) + assert len(header_capture_wire.requests) == 1 + wire_headers: Final = header_capture_wire.requests[0][1] + assert wire_headers["x-goog-api-key"] == "real-upstream-gemini-key" + assert wire_headers["x-cost-center"] == "core-platform" + assert wire_headers["x-request-org"] == "acme" + assert "authorization" not in wire_headers From 8fbc8f7602e493b8bedd57a397dbfa3a712d7c2c Mon Sep 17 00:00:00 2001 From: Kentaro Takiguchi Date: Tue, 3 Nov 2026 10:27:20 +0000 Subject: [PATCH 2/5] test(google_genai): cover router deployment headers and log ignored forwarded headers - Add a Router test proving deployment `litellm_params.headers` reach the upstream wire for native and adapter generate_content routes - Consolidate the native/adapter entrypoint tests into one parametrized test and table-drive the merge helper edge cases - Narrow forwarded headers to str-to-str mappings and emit a debug log when a malformed `headers` value is ignored --- litellm/google_genai/main.py | 8 +- .../google_genai/test_google_genai_main.py | 335 ++++++++---------- 2 files changed, 157 insertions(+), 186 deletions(-) diff --git a/litellm/google_genai/main.py b/litellm/google_genai/main.py index 1e4aa7d1224..fd28c1d02b1 100644 --- a/litellm/google_genai/main.py +++ b/litellm/google_genai/main.py @@ -8,6 +8,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter, ValidationError import litellm +from litellm._logging import verbose_logger from litellm.constants import request_timeout # Import the adapter for fallback to completion format @@ -41,7 +42,7 @@ else: # Initialize any necessary instances or variables here base_llm_http_handler = BaseLLMHTTPHandler() _GENERATE_CONTENT_RESULT: Final = TypeAdapter(object) -_FORWARDED_HEADERS: Final = TypeAdapter[Mapping[str, object]](Mapping[str, object]) +_FORWARDED_HEADERS: Final = TypeAdapter[Mapping[str, str]](Mapping[str, str]) ################################################# @@ -238,12 +239,13 @@ def _merge_native_request_fields( return {**native_request_fields, **(extra_body or {})} -def _validated_client_headers(client_headers: object) -> Mapping[str, object] | None: +def _validated_client_headers(client_headers: object) -> Mapping[str, str] | None: if client_headers is None: return None try: return _FORWARDED_HEADERS.validate_python(client_headers) - except ValidationError: + except ValidationError as exc: + verbose_logger.debug("Ignoring forwarded `headers` that are not a str-to-str mapping: %s", exc) return None diff --git a/tests/unit/google_genai/test_google_genai_main.py b/tests/unit/google_genai/test_google_genai_main.py index 8f909cfc381..f58c3dda0c1 100644 --- a/tests/unit/google_genai/test_google_genai_main.py +++ b/tests/unit/google_genai/test_google_genai_main.py @@ -431,6 +431,12 @@ def _mock_wire_reply(path: str, stream_requested: bool) -> tuple[bytes, str]: @pytest.fixture def header_capture_wire() -> Iterator[_WireCapture]: + """ + Loopback upstream that records outbound wire headers. + + A real socket (allowed by ``tests/unit/conftest.py`` for loopback hosts) is used instead of + ``respx`` because async calls go through litellm's aiohttp transport, which ``respx`` cannot intercept. + """ import threading from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer @@ -440,14 +446,9 @@ def header_capture_wire() -> Iterator[_WireCapture]: def do_POST(self) -> None: length: Final = int(self.headers.get("Content-Length", "0")) raw_body: Final = self.rfile.read(length) - lowered_headers: Final = {k.lower(): v for k, v in self.headers.items()} - holder[0].record(self.path, lowered_headers) - + holder[0].record(self.path, {k.lower(): v for k, v in self.headers.items()}) payload: Final = json.loads(raw_body.decode("utf-8")) if raw_body else {} - reply: Final = _mock_wire_reply(self.path, bool(payload.get("stream"))) - body: Final = reply[0] - content_type: Final = reply[1] - + body, content_type = _mock_wire_reply(self.path, bool(payload.get("stream"))) self.send_response(200) self.send_header("Content-Type", content_type) self.send_header("Content-Length", str(len(body))) @@ -458,10 +459,7 @@ def header_capture_wire() -> Iterator[_WireCapture]: return server: Final = ThreadingHTTPServer(("127.0.0.1", 0), _Handler) - host: Final = server.server_address[0] - port: Final = server.server_address[1] - host_str: Final = host.decode("ascii") if isinstance(host, bytes) else host - capture: Final = _WireCapture(f"http://{host_str}:{port}") + capture: Final = _WireCapture(f"http://127.0.0.1:{server.server_address[1]}") holder.append(capture) thread: Final = threading.Thread(target=server.serve_forever, daemon=True) thread.start() @@ -496,27 +494,30 @@ def _sse_payload_texts(line: str) -> tuple[str, ...]: return _extract_chunk_texts(json.loads(payload)) -def _chunk_sse_texts(raw: bytes) -> tuple[str, ...]: - from itertools import chain - - return tuple(chain.from_iterable(_sse_payload_texts(line) for line in raw.decode("utf-8").splitlines())) - - def _extract_sse_texts(chunks: tuple[bytes, ...]) -> tuple[str, ...]: from itertools import chain - return tuple(chain.from_iterable(_chunk_sse_texts(raw) for raw in chunks)) + lines: Final = chain.from_iterable(raw.decode("utf-8").splitlines() for raw in chunks) + return tuple(chain.from_iterable(_sse_payload_texts(line) for line in lines)) -@pytest.mark.asyncio -@pytest.mark.parametrize( - "entrypoint", - ["sync", "async", "sync_stream", "async_stream"], -) -async def test_native_generate_content_forwards_client_headers( - header_capture_wire: _WireCapture, +_UPSTREAM_KEY: Final = "upstream-secret" + + +def _target_route(target: str, wire: _WireCapture) -> tuple[str, str]: + if target == "native": + return "gemini/gemini-2.0-flash", wire.url + return "openai/gpt-4o-mini", f"{wire.url}/v1" + + +async def _invoke_generate_content( entrypoint: str, -) -> None: + *, + model: str, + api_base: str, + headers: Mapping[str, str], + extra_headers: Mapping[str, str], +) -> tuple[str, ...]: from collections.abc import AsyncIterator from pydantic import TypeAdapter @@ -529,177 +530,152 @@ async def test_native_generate_content_forwards_client_headers( ) contents: Final = [{"role": "user", "parts": [{"text": "ping"}]}] - forwarded: Final = {"x-team-id": "platform", "x-trace-tag": f"native-{entrypoint}"} - if entrypoint == "sync": - resp: Final = generate_content( - model="gemini/gemini-2.0-flash", - contents=contents, - api_base=header_capture_wire.url, - api_key="gemini-secret", - headers=forwarded, + return _extract_chunk_texts( + generate_content( + model=model, + contents=contents, + api_base=api_base, + api_key=_UPSTREAM_KEY, + headers=headers, + extra_headers=dict(extra_headers), + ) ) - assert _extract_chunk_texts(resp) == ("pong",) - elif entrypoint == "async": - async_resp: Final = await agenerate_content( - model="gemini/gemini-2.0-flash", - contents=contents, - api_base=header_capture_wire.url, - api_key="gemini-secret", - headers=forwarded, + if entrypoint == "async": + return _extract_chunk_texts( + await agenerate_content( + model=model, + contents=contents, + api_base=api_base, + api_key=_UPSTREAM_KEY, + headers=headers, + extra_headers=dict(extra_headers), + ) ) - assert _extract_chunk_texts(async_resp) == ("pong",) - elif entrypoint == "sync_stream": - sync_stream: Final = generate_content_stream( - model="gemini/gemini-2.0-flash", + stream: Final = ( + generate_content_stream( + model=model, contents=contents, - api_base=header_capture_wire.url, - api_key="gemini-secret", - headers=forwarded, + api_base=api_base, + api_key=_UPSTREAM_KEY, + headers=headers, + extra_headers=dict(extra_headers), ) - assert isinstance(sync_stream, Iterator) - sync_chunks: Final = tuple(TypeAdapter(bytes).validate_python(chunk) for chunk in sync_stream) - assert _extract_sse_texts(sync_chunks) == ("pong",) - else: - async_stream: Final = await agenerate_content_stream( - model="gemini/gemini-2.0-flash", + if entrypoint == "sync_stream" + else await agenerate_content_stream( + model=model, contents=contents, - api_base=header_capture_wire.url, - api_key="gemini-secret", - headers=forwarded, + api_base=api_base, + api_key=_UPSTREAM_KEY, + headers=headers, + extra_headers=dict(extra_headers), ) - assert isinstance(async_stream, AsyncIterator) - async_chunks: Final = tuple([TypeAdapter(bytes).validate_python(chunk) async for chunk in async_stream]) - assert _extract_sse_texts(async_chunks) == ("pong",) - - assert len(header_capture_wire.requests) == 1 - wire_headers: Final = header_capture_wire.requests[0][1] - assert wire_headers["x-goog-api-key"] == "gemini-secret" - assert wire_headers["x-team-id"] == "platform" - assert wire_headers["x-trace-tag"] == f"native-{entrypoint}" + ) + as_bytes: Final = TypeAdapter(bytes) + # native sync streams implement both protocols; adapter sync streams are async-only + if entrypoint == "sync_stream" and isinstance(stream, Iterator): + return _extract_sse_texts(tuple(as_bytes.validate_python(chunk) for chunk in stream)) + assert isinstance(stream, AsyncIterator) + return _extract_sse_texts(tuple([as_bytes.validate_python(chunk) async for chunk in stream])) @pytest.mark.asyncio -@pytest.mark.parametrize( - "entrypoint", - ["sync", "async", "sync_stream", "async_stream"], -) -async def test_adapter_generate_content_forwards_client_headers( +@pytest.mark.parametrize("entrypoint", ["sync", "async", "sync_stream", "async_stream"]) +@pytest.mark.parametrize("target", ["native", "adapter"]) +async def test_generate_content_forwards_client_headers_behind_explicit_extra_headers( header_capture_wire: _WireCapture, + target: str, entrypoint: str, ) -> None: - from collections.abc import AsyncIterator + model, api_base = _target_route(target, header_capture_wire) + forwarded: Final = {"x-team-id": "platform", "x-custom": "from-proxy", "anthropic-beta": "context-1m-2025-08-07"} + explicit: Final = {"X-Custom": "from-caller"} - from pydantic import TypeAdapter - - from litellm.google_genai.main import ( - agenerate_content, - agenerate_content_stream, - generate_content, - generate_content_stream, + texts: Final = await _invoke_generate_content( + entrypoint, model=model, api_base=api_base, headers=forwarded, extra_headers=explicit ) - contents: Final = [{"role": "user", "parts": [{"text": "ping"}]}] - forwarded: Final = { - "x-team-id": "platform", - "anthropic-beta": "context-1m-2025-08-07", - "x-trace-tag": f"adapter-{entrypoint}", - } - - if entrypoint == "sync": - resp: Final = generate_content( - model="openai/gpt-4o-mini", - contents=contents, - api_base=f"{header_capture_wire.url}/v1", - api_key="openai-secret", - headers=forwarded, - ) - assert _extract_chunk_texts(resp) == ("pong",) - elif entrypoint == "async": - async_resp: Final = await agenerate_content( - model="openai/gpt-4o-mini", - contents=contents, - api_base=f"{header_capture_wire.url}/v1", - api_key="openai-secret", - headers=forwarded, - ) - assert _extract_chunk_texts(async_resp) == ("pong",) - elif entrypoint == "sync_stream": - sync_stream: Final = generate_content_stream( - model="openai/gpt-4o-mini", - contents=contents, - api_base=f"{header_capture_wire.url}/v1", - api_key="openai-secret", - headers=forwarded, - ) - assert isinstance(sync_stream, AsyncIterator) - sync_chunks: Final = tuple([TypeAdapter(bytes).validate_python(chunk) async for chunk in sync_stream]) - assert _extract_sse_texts(sync_chunks) == ("pong",) - else: - async_stream: Final = await agenerate_content_stream( - model="openai/gpt-4o-mini", - contents=contents, - api_base=f"{header_capture_wire.url}/v1", - api_key="openai-secret", - headers=forwarded, - ) - assert isinstance(async_stream, AsyncIterator) - async_chunks: Final = tuple([TypeAdapter(bytes).validate_python(chunk) async for chunk in async_stream]) - assert _extract_sse_texts(async_chunks) == ("pong",) - + assert texts == ("pong",) + assert forwarded == {"x-team-id": "platform", "x-custom": "from-proxy", "anthropic-beta": "context-1m-2025-08-07"} + assert explicit == {"X-Custom": "from-caller"} assert len(header_capture_wire.requests) == 1 - req_path: Final = header_capture_wire.requests[0][0] wire_headers: Final = header_capture_wire.requests[0][1] - assert req_path == "/v1/chat/completions" - assert wire_headers["authorization"] == "Bearer openai-secret" assert wire_headers["x-team-id"] == "platform" assert wire_headers["anthropic-beta"] == "context-1m-2025-08-07" - assert wire_headers["x-trace-tag"] == f"adapter-{entrypoint}" + assert wire_headers["x-custom"] == "from-caller" +@pytest.mark.parametrize( + "extra_headers, client_headers, expected", + [ + (None, None, None), + (None, {}, None), + ({}, None, {}), + ({"x-b": "2"}, None, {"x-b": "2"}), + (None, {"x-a": "1"}, {"x-a": "1"}), + ({}, {"x-a": "1"}, {"x-a": "1"}), + ({"X-Custom": "from-caller"}, {"x-custom": "from-proxy", "x-a": "1"}, {"x-a": "1", "X-Custom": "from-caller"}), + ({"x-b": "2"}, "not-a-mapping", {"x-b": "2"}), + (None, "not-a-mapping", None), + ({"x-b": "2"}, {"x-a": 1}, {"x-b": "2"}), + ], + ids=[ + "nothing", + "empty-forwarded", + "empty-explicit", + "explicit-only", + "forwarded-only", + "empty-explicit-with-forwarded", + "case-insensitive-precedence", + "non-mapping-forwarded-kept-explicit", + "non-mapping-forwarded", + "non-str-value-forwarded", + ], +) +def test_merge_client_forwarded_headers( + extra_headers: Mapping[str, str] | None, + client_headers: object, + expected: Mapping[str, str] | None, +) -> None: + from litellm.google_genai.main import _merge_client_forwarded_headers + + merged: Final = _merge_client_forwarded_headers(extra_headers, client_headers) + + assert merged == expected + assert merged is None or (merged is not extra_headers and merged is not client_headers) + + +@pytest.mark.asyncio @pytest.mark.parametrize("target", ["native", "adapter"]) -def test_extra_headers_override_forwarded_client_headers_case_insensitively( +async def test_router_deployment_headers_reach_generate_content( header_capture_wire: _WireCapture, target: str, ) -> None: - from litellm.google_genai.main import _merge_client_forwarded_headers, generate_content + from litellm import Router - forwarded: Final = {"x-custom": "from-proxy", "x-only-proxy": "proxy-val"} - explicit: Final = {"X-Custom": "from-caller", "X-Only-Caller": "caller-val"} + model, api_base = _target_route(target, header_capture_wire) + router: Final = Router( + model_list=[ + { + "model_name": "genai-deployment", + "litellm_params": { + "model": model, + "api_base": api_base, + "api_key": _UPSTREAM_KEY, + "headers": {"x-deployment-tag": "blue"}, + }, + } + ] + ) - resp: Final = generate_content( - model="gemini/gemini-2.0-flash" if target == "native" else "openai/gpt-4o-mini", + resp: Final = await router.agenerate_content( + model="genai-deployment", contents=[{"role": "user", "parts": [{"text": "ping"}]}], - api_base=header_capture_wire.url if target == "native" else f"{header_capture_wire.url}/v1", - api_key="test-secret", - headers=forwarded, - extra_headers=explicit, ) assert _extract_chunk_texts(resp) == ("pong",) - assert forwarded == {"x-custom": "from-proxy", "x-only-proxy": "proxy-val"} - assert explicit == {"X-Custom": "from-caller", "X-Only-Caller": "caller-val"} assert len(header_capture_wire.requests) == 1 - wire_headers: Final = header_capture_wire.requests[0][1] - assert wire_headers["x-custom"] == "from-caller" - assert wire_headers["x-only-proxy"] == "proxy-val" - assert wire_headers["x-only-caller"] == "caller-val" - - assert _merge_client_forwarded_headers(None, None) is None - assert _merge_client_forwarded_headers(None, {}) is None - assert _merge_client_forwarded_headers({}, None) == {} - source_extra: Final = {"x-b": "2"} - merged_extra: Final = _merge_client_forwarded_headers(source_extra, None) - assert merged_extra == {"x-b": "2"} and merged_extra is not source_extra - source_headers: Final = {"x-a": "1"} - merged_only: Final = _merge_client_forwarded_headers(None, source_headers) - assert merged_only == {"x-a": "1"} and merged_only is not source_headers - assert _merge_client_forwarded_headers({}, {"x-a": "1"}) == {"x-a": "1"} - assert _merge_client_forwarded_headers({"X-Custom": "from-caller"}, {"x-custom": "from-proxy"}) == { - "X-Custom": "from-caller" - } - assert _merge_client_forwarded_headers({"x-b": "2"}, "not-a-mapping") == {"x-b": "2"} - assert _merge_client_forwarded_headers(None, "not-a-mapping") is None + assert header_capture_wire.requests[0][1]["x-deployment-tag"] == "blue" @pytest.mark.asyncio @@ -712,44 +688,37 @@ async def test_proxy_forwarded_headers_reach_generate_content_without_leaking_pr from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, clean_headers - incoming_headers: Final = Headers( - { - "x-goog-api-key": "sk-litellm-virtual-key", - "authorization": "Bearer sk-litellm-virtual-key", - "x-cost-center": "core-platform", - "x-request-org": "acme", - "user-agent": "google-genai-sdk/1.0", - } - ) cleaned: Final = clean_headers( - incoming_headers, + Headers( + { + "x-goog-api-key": "sk-litellm-virtual-key", + "authorization": "Bearer sk-litellm-virtual-key", + "x-cost-center": "core-platform", + "user-agent": "google-genai-sdk/1.0", + } + ), litellm_key_header_name=None, forward_llm_provider_auth_headers=False, authenticated_with_header="x-goog-api-key", ) - request_data: Final[Mapping[str, object]] = { - "model": "gemini/gemini-2.0-flash", - "contents": [{"role": "user", "parts": [{"text": "ping"}]}], - } backend_data: Final = LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call( headers=cleaned, - request_data=request_data, + request_data={"model": "gemini/gemini-2.0-flash"}, user_api_key_dict=UserAPIKeyAuth(api_key="sk-litellm-virtual-key"), general_settings={"forward_client_headers_to_llm_api": True}, ) resp: Final = await agenerate_content( model="gemini/gemini-2.0-flash", - contents=request_data["contents"], + contents=[{"role": "user", "parts": [{"text": "ping"}]}], api_base=header_capture_wire.url, - api_key="real-upstream-gemini-key", + api_key=_UPSTREAM_KEY, headers=backend_data.get("headers"), ) assert _extract_chunk_texts(resp) == ("pong",) assert len(header_capture_wire.requests) == 1 wire_headers: Final = header_capture_wire.requests[0][1] - assert wire_headers["x-goog-api-key"] == "real-upstream-gemini-key" + assert wire_headers["x-goog-api-key"] == _UPSTREAM_KEY assert wire_headers["x-cost-center"] == "core-platform" - assert wire_headers["x-request-org"] == "acme" assert "authorization" not in wire_headers From 8ed52b127dfb5f8b63e878d206c691b65017d795 Mon Sep 17 00:00:00 2001 From: Kentaro Takiguchi Date: Tue, 3 Nov 2026 12:13:50 +0000 Subject: [PATCH 3/5] refactor(google_genai): drop helper docstring and align proxy auth test Remove the redundant docstring on `_merge_client_forwarded_headers` and pass `litellm_auth_header_name(incoming_headers)` into `clean_headers` so the proxy header test exercises the same `authenticated_with_header` lookup used in request setup --- litellm/google_genai/main.py | 6 ----- .../google_genai/test_google_genai_main.py | 25 +++++++++++-------- 2 files changed, 15 insertions(+), 16 deletions(-) diff --git a/litellm/google_genai/main.py b/litellm/google_genai/main.py index fd28c1d02b1..628aade7eba 100644 --- a/litellm/google_genai/main.py +++ b/litellm/google_genai/main.py @@ -253,12 +253,6 @@ def _merge_client_forwarded_headers( extra_headers: Mapping[str, object] | None, client_headers: object, ) -> dict[str, object] | None: # mutable-ok: downstream handlers expect dict | None - """ - Merge proxy-forwarded ``headers`` (set when ``forward_client_headers_to_llm_api`` - is enabled or on deployment ``litellm_params``) into ``extra_headers``. - - ``extra_headers`` wins on case-insensitive header-name conflicts. - """ forwarded: Final = _validated_client_headers(client_headers) if not forwarded: return dict(extra_headers) if extra_headers is not None else None diff --git a/tests/unit/google_genai/test_google_genai_main.py b/tests/unit/google_genai/test_google_genai_main.py index f58c3dda0c1..a30dbd6e501 100644 --- a/tests/unit/google_genai/test_google_genai_main.py +++ b/tests/unit/google_genai/test_google_genai_main.py @@ -686,20 +686,25 @@ async def test_proxy_forwarded_headers_reach_generate_content_without_leaking_pr from litellm.google_genai.main import agenerate_content from litellm.proxy._types import UserAPIKeyAuth - from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, clean_headers + from litellm.proxy.litellm_pre_call_utils import ( + LiteLLMProxyRequestSetup, + clean_headers, + litellm_auth_header_name, + ) + incoming_headers: Final = Headers( + { + "x-goog-api-key": "sk-litellm-virtual-key", + "authorization": "Bearer sk-litellm-virtual-key", + "x-cost-center": "core-platform", + "user-agent": "google-genai-sdk/1.0", + } + ) cleaned: Final = clean_headers( - Headers( - { - "x-goog-api-key": "sk-litellm-virtual-key", - "authorization": "Bearer sk-litellm-virtual-key", - "x-cost-center": "core-platform", - "user-agent": "google-genai-sdk/1.0", - } - ), + incoming_headers, litellm_key_header_name=None, forward_llm_provider_auth_headers=False, - authenticated_with_header="x-goog-api-key", + authenticated_with_header=litellm_auth_header_name(incoming_headers), ) backend_data: Final = LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call( headers=cleaned, From a7c1ddb99eeecbadff1026d348202ecb01cec445 Mon Sep 17 00:00:00 2001 From: Kentaro Takiguchi Date: Wed, 4 Nov 2026 01:31:52 +0000 Subject: [PATCH 4/5] fix(proxy): strip the header used for proxy auth when forwarding provider auth headers With forward_llm_provider_auth_headers on, clean_headers kept every provider auth header except x-api-key, so a virtual key sent as x-goog-api-key reached the provider. generate_content now forwards client headers, which made this overwrite the deployment key litellm_auth_header_name now follows get_api_key's header precedence, honors litellm_key_header_name, and skips an empty authorization header. clean_headers drops whichever header that resolves to --- litellm/proxy/litellm_pre_call_utils.py | 41 ++++-- .../google_genai/test_google_genai_main.py | 60 ++++++++ .../unit/proxy/test_litellm_pre_call_utils.py | 134 ++++++++++++++++++ 3 files changed, 225 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 4444d6d25f9..19c41f2cfc1 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1498,12 +1498,30 @@ def _tenant_service_name(user_api_key_dict: UserAPIKeyAuth) -> str | None: ) -def litellm_auth_header_name(headers: Headers) -> str: - if "x-litellm-api-key" in headers: - return "x-litellm-api-key" - if "authorization" in headers: - return "authorization" - return "x-api-key" +_PROXY_AUTH_HEADER_PRECEDENCE: Final = tuple( + header.value.lower() + for header in ( + SpecialHeaders.custom_litellm_api_key, + SpecialHeaders.openai_authorization, + SpecialHeaders.azure_authorization, + SpecialHeaders.anthropic_authorization, + SpecialHeaders.google_ai_studio_authorization, + SpecialHeaders.azure_apim_authorization, + ) +) + + +def litellm_auth_header_name(headers: Headers, litellm_key_header_name: str | None = None) -> str: + if litellm_key_header_name is not None and litellm_key_header_name in headers: + return litellm_key_header_name.lower() + return next( + ( + name + for name in _PROXY_AUTH_HEADER_PRECEDENCE + if name in headers and (name != "authorization" or headers[name]) + ), + "x-api-key", + ) def forwarded_anthropic_oauth_token(headers: Headers) -> str | None: @@ -1560,6 +1578,8 @@ def clean_headers( elif forward_llm_provider_auth_headers and header_lower in _SPECIAL_HEADERS_CACHE: if litellm_key_lower and header_lower == litellm_key_lower: continue + if authenticated_with_header is not None and header_lower == authenticated_with_header.lower(): + continue if header_lower == "authorization": continue # Never forward x-litellm-api-key (it's for proxy auth only) @@ -2669,13 +2689,14 @@ async def add_litellm_data_to_request( forward_llm_auth = general_settings.get("forward_llm_provider_auth_headers", False) if not forward_llm_auth: forward_llm_auth = getattr(litellm, "forward_llm_provider_auth_headers", False) + litellm_key_header_name: Final = ( + general_settings.get("litellm_key_header_name") if general_settings is not None else None + ) _headers: Final[dict[str, str]] = clean_headers( request.headers, - litellm_key_header_name=( - general_settings.get("litellm_key_header_name") if general_settings is not None else None - ), + litellm_key_header_name=litellm_key_header_name, forward_llm_provider_auth_headers=forward_llm_auth, - authenticated_with_header=litellm_auth_header_name(request.headers), + authenticated_with_header=litellm_auth_header_name(request.headers, litellm_key_header_name), ) _strip_untrusted_request_header_controls( _headers, diff --git a/tests/unit/google_genai/test_google_genai_main.py b/tests/unit/google_genai/test_google_genai_main.py index a30dbd6e501..0f59daee1b6 100644 --- a/tests/unit/google_genai/test_google_genai_main.py +++ b/tests/unit/google_genai/test_google_genai_main.py @@ -727,3 +727,63 @@ async def test_proxy_forwarded_headers_reach_generate_content_without_leaking_pr assert wire_headers["x-goog-api-key"] == _UPSTREAM_KEY assert wire_headers["x-cost-center"] == "core-platform" assert "authorization" not in wire_headers + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("incoming", "expected_upstream_key"), + [ + ({"x-goog-api-key": "sk-litellm-virtual-key", "x-cost-center": "core-platform"}, _UPSTREAM_KEY), + ( + { + "x-litellm-api-key": "sk-litellm-virtual-key", + "x-goog-api-key": "AIza-caller-key", + "x-cost-center": "core-platform", + }, + "AIza-caller-key", + ), + ], + ids=["proxy-key-in-goog-header", "caller-key-beside-proxy-key"], +) +async def test_provider_auth_forwarding_never_sends_proxy_key_to_generate_content( + header_capture_wire: _WireCapture, + incoming: Mapping[str, str], + expected_upstream_key: str, +) -> None: + from starlette.datastructures import Headers + + from litellm.google_genai.main import agenerate_content + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.litellm_pre_call_utils import ( + LiteLLMProxyRequestSetup, + clean_headers, + litellm_auth_header_name, + ) + + incoming_headers: Final = Headers(incoming) + cleaned: Final = clean_headers( + incoming_headers, + forward_llm_provider_auth_headers=True, + authenticated_with_header=litellm_auth_header_name(incoming_headers), + ) + backend_data: Final = LiteLLMProxyRequestSetup.add_litellm_data_for_backend_llm_call( + headers=cleaned, + request_data={"model": "gemini/gemini-2.0-flash"}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-litellm-virtual-key"), + general_settings={"forward_client_headers_to_llm_api": True}, + ) + + resp: Final = await agenerate_content( + model="gemini/gemini-2.0-flash", + contents=[{"role": "user", "parts": [{"text": "ping"}]}], + api_base=header_capture_wire.url, + api_key=_UPSTREAM_KEY, + headers=backend_data.get("headers"), + ) + + assert _extract_chunk_texts(resp) == ("pong",) + assert len(header_capture_wire.requests) == 1 + wire_headers: Final = header_capture_wire.requests[0][1] + assert wire_headers["x-goog-api-key"] == expected_upstream_key + assert wire_headers["x-cost-center"] == "core-platform" + assert "x-litellm-api-key" not in wire_headers diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 91c87c63067..074aa0c34e8 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -4,6 +4,7 @@ import copy import json import os import time +from collections.abc import Mapping from datetime import datetime, timezone from types import SimpleNamespace from typing import Final @@ -58,6 +59,7 @@ from litellm.proxy.litellm_pre_call_utils import ( forwarded_anthropic_oauth_token, get_dynamic_logging_metadata, get_metadata_variable_name, + litellm_auth_header_name, move_guardrails_to_metadata, ) from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload @@ -6447,6 +6449,84 @@ def test_clean_headers_strips_x_api_key_when_byok_enabled_but_x_api_key_was_auth assert "x-api-key" not in result +@pytest.mark.parametrize( + ("headers", "litellm_key_header_name", "expected"), + [ + ({"x-litellm-api-key": "sk-proxy", "authorization": "Bearer sk-other"}, None, "x-litellm-api-key"), + ({"authorization": "Bearer sk-proxy", "x-goog-api-key": "AIza-caller"}, None, "authorization"), + ({"authorization": "", "x-goog-api-key": "sk-proxy"}, None, "x-goog-api-key"), + ({"api-key": "sk-proxy", "x-api-key": "sk-ant-caller"}, None, "api-key"), + ({"x-api-key": "sk-proxy", "x-goog-api-key": "AIza-caller"}, None, "x-api-key"), + ({"x-goog-api-key": "sk-proxy"}, None, "x-goog-api-key"), + ({"ocp-apim-subscription-key": "sk-proxy"}, None, "ocp-apim-subscription-key"), + ({"x-custom-key": "sk-proxy", "authorization": "Bearer sk-other"}, "X-Custom-Key", "x-custom-key"), + ({"authorization": "Bearer sk-proxy"}, "x-custom-key", "authorization"), + ({"content-type": "application/json"}, None, "x-api-key"), + ], + ids=[ + "litellm-key-beats-authorization", + "authorization-beats-goog", + "empty-authorization-skipped", + "azure-beats-anthropic", + "anthropic-beats-goog", + "goog", + "azure-apim", + "custom-header-wins-when-present", + "custom-header-absent", + "no-auth-header", + ], +) +def test_litellm_auth_header_name_follows_user_api_key_auth_precedence( + headers: Mapping[str, str], litellm_key_header_name: str | None, expected: str +) -> None: + assert litellm_auth_header_name(Headers(headers), litellm_key_header_name) == expected + + +@pytest.mark.parametrize( + ("headers", "expected"), + [ + ( + {"x-goog-api-key": "sk-proxy-virtual-key", "x-cost-center": "core-platform"}, + {"x-cost-center": "core-platform"}, + ), + ( + {"x-litellm-api-key": "sk-proxy-virtual-key", "x-goog-api-key": "AIza-caller-key"}, + {"x-goog-api-key": "AIza-caller-key"}, + ), + ( + {"api-key": "sk-proxy-virtual-key", "x-goog-api-key": "AIza-caller-key"}, + {"x-goog-api-key": "AIza-caller-key"}, + ), + ], + ids=["goog-was-proxy-auth", "goog-beside-litellm-key", "goog-beside-azure-proxy-key"], +) +def test_clean_headers_byok_strips_only_the_header_used_for_proxy_auth( + headers: Mapping[str, str], expected: Mapping[str, str] +) -> None: + incoming: Final = Headers(headers) + + result: Final = clean_headers( + incoming, + forward_llm_provider_auth_headers=True, + authenticated_with_header=litellm_auth_header_name(incoming), + ) + + assert result == expected + + +def test_clean_headers_byok_strips_custom_litellm_key_header_named_x_api_key() -> None: + incoming: Final = Headers({"x-api-key": "sk-proxy-virtual-key", "authorization": "Bearer sk-ant-api03-caller"}) + + result: Final = clean_headers( + incoming, + litellm_key_header_name="x-api-key", + forward_llm_provider_auth_headers=True, + authenticated_with_header=litellm_auth_header_name(incoming, "x-api-key"), + ) + + assert "x-api-key" not in result + + # --------------------------------------------------------------------------- # Team guardrail + global policy regression tests # --------------------------------------------------------------------------- @@ -7965,6 +8045,60 @@ async def test_add_litellm_data_to_request_keeps_every_forwarded_credential_out_ assert value not in logged +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("request_headers", "general_settings", "expected_forwarded_auth"), + [ + ( + {"x-goog-api-key": "sk-virtual-key", "x-cost-center": "core-platform"}, + FORWARD_PROVIDER_AUTH_HEADERS, + {}, + ), + ( + { + "x-litellm-api-key": "Bearer sk-virtual-key", + "x-goog-api-key": "AIza-caller-key", + "x-cost-center": "core-platform", + }, + FORWARD_PROVIDER_AUTH_HEADERS, + {"x-goog-api-key": "AIza-caller-key"}, + ), + ( + {"x-api-key": "sk-virtual-key", "authorization": "Bearer sk-other", "x-cost-center": "core-platform"}, + {**FORWARD_PROVIDER_AUTH_HEADERS, "litellm_key_header_name": "x-api-key"}, + {}, + ), + ], + ids=["goog-was-proxy-auth", "goog-beside-litellm-key", "custom-litellm-key-header"], +) +async def test_add_litellm_data_to_request_never_forwards_the_header_used_for_proxy_auth( + request_headers: Mapping[str, str], + general_settings: Mapping[str, object], + expected_forwarded_auth: Mapping[str, str], +) -> None: + request_mock: Final = _make_request_mock( + "/v1/chat/completions", {"Content-Type": "application/json", **request_headers} + ) + + updated: Final = await add_litellm_data_to_request( + data={"model": "gemini-2.0-flash", "messages": [{"role": "user", "content": "hello"}]}, + request=request_mock, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings=dict(general_settings), + version="test-version", + ) + + forwarded_auth: Final = { + name: value + for name, value in updated["headers"].items() + if name.lower() in {"x-goog-api-key", "x-api-key", "x-litellm-api-key"} + } + assert forwarded_auth == expected_forwarded_auth + assert updated["headers"]["x-cost-center"] == "core-platform" + assert "sk-virtual-key" not in str(updated.get("api_key")) + + _SERVICE_ACCOUNT_TOKEN_REFERENCE: Final = "oidc/file//var/run/secrets/kubernetes.io/serviceaccount/token" From af2cf4d70606e09f36074182f0532d2419a14912 Mon Sep 17 00:00:00 2001 From: Kentaro Takiguchi Date: Wed, 4 Nov 2026 05:23:59 +0000 Subject: [PATCH 5/5] fix(proxy): strip duplicate and mixed-case proxy auth credentials in clean_headers Google GenAI SDK clients attach x-goog-api-key alongside Authorization or x-litellm-api-key, so stripping only the winning proxy auth header name still leaked the proxy virtual key upstream when forward_llm_provider_auth_headers and forward_client_headers_to_llm_api were both enabled. Empty higher-precedence headers and mixed-case Mapping keys also bypassed litellm_auth_header_name or overwrote upstream credentials with empty values Normalize non-empty headers case-insensitively in litellm_auth_header_name and unify clean_headers provider-auth filtering so empty provider auth headers and any header carrying the bearer-stripped proxy credential are stripped while separately supplied caller provider keys stay forwarded --- litellm/proxy/litellm_pre_call_utils.py | 78 +++++----- .../google_genai/test_google_genai_main.py | 41 +++++- .../unit/proxy/test_litellm_pre_call_utils.py | 138 ++++++++++++++++-- 3 files changed, 205 insertions(+), 52 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 19c41f2cfc1..cc0686a0636 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1511,17 +1511,16 @@ _PROXY_AUTH_HEADER_PRECEDENCE: Final = tuple( ) -def litellm_auth_header_name(headers: Headers, litellm_key_header_name: str | None = None) -> str: - if litellm_key_header_name is not None and litellm_key_header_name in headers: +def _strip_bearer(value: str) -> str: + stripped: Final = value.strip() + return stripped[7:].strip() if stripped[:7].lower() == "bearer " else stripped + + +def litellm_auth_header_name(headers: Mapping[str, str], litellm_key_header_name: str | None = None) -> str: + present: Final = {name.lower() for name, value in headers.items() if value} + if litellm_key_header_name is not None and litellm_key_header_name.lower() in present: return litellm_key_header_name.lower() - return next( - ( - name - for name in _PROXY_AUTH_HEADER_PRECEDENCE - if name in headers and (name != "authorization" or headers[name]) - ), - "x-api-key", - ) + return next((name for name in _PROXY_AUTH_HEADER_PRECEDENCE if name in present), "x-api-key") def forwarded_anthropic_oauth_token(headers: Headers) -> str | None: @@ -1539,11 +1538,11 @@ def forwarded_anthropic_oauth_token(headers: Headers) -> str | None: def clean_headers( - headers: Headers, + headers: Mapping[str, str], litellm_key_header_name: str | None = None, forward_llm_provider_auth_headers: bool = False, authenticated_with_header: str | None = None, -) -> dict[str, str]: +) -> dict[str, str]: # mutable-ok: downstream callers mutate returned headers dict """ Removes litellm api key from headers @@ -1560,38 +1559,35 @@ def clean_headers( from litellm.llms.anthropic.common_utils import is_anthropic_oauth_key from litellm.llms.chatgpt.common_utils import is_chatgpt_oauth_key - clean_headers: Final = {} litellm_key_lower: Final = litellm_key_header_name.lower() if litellm_key_header_name is not None else None - for header, value in headers.items(): - header_lower = header.lower() + auth_header_lower: Final = authenticated_with_header.lower() if authenticated_with_header is not None else None + custom_excluded: Final = frozenset({litellm_key_lower, auth_header_lower}) + provider_excluded: Final = frozenset({"authorization", "x-litellm-api-key", *custom_excluded}) + proxy_auth_credential: Final = ( + next( + ( + credential + for name, raw in headers.items() + if name.lower() == auth_header_lower and (credential := _strip_bearer(raw)) + ), + None, + ) + if auth_header_lower is not None + else None + ) + def _keep(header: str, value: str) -> bool: + header_lower: Final = header.lower() if header_lower == "authorization" and (is_anthropic_oauth_key(value) or is_chatgpt_oauth_key(value)): - if authenticated_with_header is None or authenticated_with_header.lower() != "authorization": - clean_headers[header] = value - continue - # Special handling for x-api-key: forward it based on authenticated_with_header - elif header_lower == "x-api-key": - if forward_llm_provider_auth_headers and ( - authenticated_with_header is None or authenticated_with_header.lower() != "x-api-key" - ): - clean_headers[header] = value - elif forward_llm_provider_auth_headers and header_lower in _SPECIAL_HEADERS_CACHE: - if litellm_key_lower and header_lower == litellm_key_lower: - continue - if authenticated_with_header is not None and header_lower == authenticated_with_header.lower(): - continue - if header_lower == "authorization": - continue - # Never forward x-litellm-api-key (it's for proxy auth only) - if header_lower == "x-litellm-api-key": - continue - clean_headers[header] = value - # Check if header should be excluded: either in special headers cache or matches custom litellm key - elif header_lower not in _SPECIAL_HEADERS_CACHE and ( - litellm_key_lower is None or header_lower != litellm_key_lower - ): - clean_headers[header] = value - return clean_headers + return auth_header_lower != "authorization" and _strip_bearer(value) != proxy_auth_credential + if forward_llm_provider_auth_headers and header_lower in _SPECIAL_HEADERS_CACHE: + if header_lower in provider_excluded: + return False + credential: Final = _strip_bearer(value) + return bool(credential) and credential != proxy_auth_credential + return header_lower not in _SPECIAL_HEADERS_CACHE and header_lower not in custom_excluded + + return {header: value for header, value in headers.items() if _keep(header, value)} def _request_fields(value: object) -> Mapping[str, object]: diff --git a/tests/unit/google_genai/test_google_genai_main.py b/tests/unit/google_genai/test_google_genai_main.py index 0f59daee1b6..a80e8b21515 100644 --- a/tests/unit/google_genai/test_google_genai_main.py +++ b/tests/unit/google_genai/test_google_genai_main.py @@ -734,6 +734,30 @@ async def test_proxy_forwarded_headers_reach_generate_content_without_leaking_pr ("incoming", "expected_upstream_key"), [ ({"x-goog-api-key": "sk-litellm-virtual-key", "x-cost-center": "core-platform"}, _UPSTREAM_KEY), + ( + { + "authorization": "Bearer sk-litellm-virtual-key", + "x-goog-api-key": "sk-litellm-virtual-key", + "x-cost-center": "core-platform", + }, + _UPSTREAM_KEY, + ), + ( + { + "x-api-key": "", + "x-goog-api-key": "sk-litellm-virtual-key", + "x-cost-center": "core-platform", + }, + _UPSTREAM_KEY, + ), + ( + { + "authorization": "Bearer sk-litellm-virtual-key", + "x-goog-api-key": "", + "x-cost-center": "core-platform", + }, + _UPSTREAM_KEY, + ), ( { "x-litellm-api-key": "sk-litellm-virtual-key", @@ -742,8 +766,23 @@ async def test_proxy_forwarded_headers_reach_generate_content_without_leaking_pr }, "AIza-caller-key", ), + ( + { + "authorization": "Bearer sk-litellm-virtual-key", + "x-goog-api-key": "AIza-caller-key", + "x-cost-center": "core-platform", + }, + "AIza-caller-key", + ), + ], + ids=[ + "proxy-key-in-goog-header", + "sdk-dual-bearer-and-goog-proxy-key", + "empty-higher-precedence-header", + "empty-goog-header-beside-bearer-key", + "caller-key-beside-litellm-key", + "caller-key-beside-bearer-key", ], - ids=["proxy-key-in-goog-header", "caller-key-beside-proxy-key"], ) async def test_provider_auth_forwarding_never_sends_proxy_key_to_generate_content( header_capture_wire: _WireCapture, diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 074aa0c34e8..14e92cc343e 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -6453,25 +6453,33 @@ def test_clean_headers_strips_x_api_key_when_byok_enabled_but_x_api_key_was_auth ("headers", "litellm_key_header_name", "expected"), [ ({"x-litellm-api-key": "sk-proxy", "authorization": "Bearer sk-other"}, None, "x-litellm-api-key"), + ({"x-litellm-api-key": "", "x-goog-api-key": "sk-proxy"}, None, "x-goog-api-key"), ({"authorization": "Bearer sk-proxy", "x-goog-api-key": "AIza-caller"}, None, "authorization"), ({"authorization": "", "x-goog-api-key": "sk-proxy"}, None, "x-goog-api-key"), ({"api-key": "sk-proxy", "x-api-key": "sk-ant-caller"}, None, "api-key"), + ({"x-api-key": "", "x-goog-api-key": "sk-proxy"}, None, "x-goog-api-key"), ({"x-api-key": "sk-proxy", "x-goog-api-key": "AIza-caller"}, None, "x-api-key"), ({"x-goog-api-key": "sk-proxy"}, None, "x-goog-api-key"), + ({"X-Goog-Api-Key": "sk-proxy"}, None, "x-goog-api-key"), ({"ocp-apim-subscription-key": "sk-proxy"}, None, "ocp-apim-subscription-key"), - ({"x-custom-key": "sk-proxy", "authorization": "Bearer sk-other"}, "X-Custom-Key", "x-custom-key"), + ({"X-Custom-Key": "sk-proxy", "authorization": "Bearer sk-other"}, "x-custom-key", "x-custom-key"), + ({"x-custom-key": "", "x-goog-api-key": "sk-proxy"}, "x-custom-key", "x-goog-api-key"), ({"authorization": "Bearer sk-proxy"}, "x-custom-key", "authorization"), ({"content-type": "application/json"}, None, "x-api-key"), ], ids=[ "litellm-key-beats-authorization", + "empty-litellm-key-skipped", "authorization-beats-goog", "empty-authorization-skipped", "azure-beats-anthropic", + "empty-anthropic-skipped", "anthropic-beats-goog", "goog", + "goog-mixed-case-dict", "azure-apim", - "custom-header-wins-when-present", + "custom-header-mixed-case-dict", + "empty-custom-header-skipped", "custom-header-absent", "no-auth-header", ], @@ -6479,6 +6487,7 @@ def test_clean_headers_strips_x_api_key_when_byok_enabled_but_x_api_key_was_auth def test_litellm_auth_header_name_follows_user_api_key_auth_precedence( headers: Mapping[str, str], litellm_key_header_name: str | None, expected: str ) -> None: + assert litellm_auth_header_name(headers, litellm_key_header_name) == expected assert litellm_auth_header_name(Headers(headers), litellm_key_header_name) == expected @@ -6489,29 +6498,88 @@ def test_litellm_auth_header_name_follows_user_api_key_auth_precedence( {"x-goog-api-key": "sk-proxy-virtual-key", "x-cost-center": "core-platform"}, {"x-cost-center": "core-platform"}, ), + ( + {"X-Goog-Api-Key": "sk-proxy-virtual-key", "x-cost-center": "core-platform"}, + {"x-cost-center": "core-platform"}, + ), + ( + { + "authorization": "Bearer sk-proxy-virtual-key", + "x-goog-api-key": "sk-proxy-virtual-key", + "x-cost-center": "core-platform", + }, + {"x-cost-center": "core-platform"}, + ), + ( + { + "x-litellm-api-key": "sk-proxy-virtual-key", + "x-goog-api-key": "sk-proxy-virtual-key", + "x-cost-center": "core-platform", + }, + {"x-cost-center": "core-platform"}, + ), + ( + { + "x-api-key": "", + "x-goog-api-key": "sk-proxy-virtual-key", + "x-cost-center": "core-platform", + }, + {"x-cost-center": "core-platform"}, + ), ( {"x-litellm-api-key": "sk-proxy-virtual-key", "x-goog-api-key": "AIza-caller-key"}, {"x-goog-api-key": "AIza-caller-key"}, ), + ( + {"authorization": "Bearer sk-proxy-virtual-key", "x-goog-api-key": "AIza-caller-key"}, + {"x-goog-api-key": "AIza-caller-key"}, + ), ( {"api-key": "sk-proxy-virtual-key", "x-goog-api-key": "AIza-caller-key"}, {"x-goog-api-key": "AIza-caller-key"}, ), ], - ids=["goog-was-proxy-auth", "goog-beside-litellm-key", "goog-beside-azure-proxy-key"], + ids=[ + "goog-was-proxy-auth", + "goog-mixed-case-was-proxy-auth", + "duplicate-proxy-key-in-bearer-and-goog", + "duplicate-proxy-key-in-litellm-and-goog", + "empty-provider-header-and-goog-proxy-auth", + "goog-beside-litellm-key", + "goog-beside-bearer-proxy-key", + "goog-beside-azure-proxy-key", + ], ) def test_clean_headers_byok_strips_only_the_header_used_for_proxy_auth( headers: Mapping[str, str], expected: Mapping[str, str] ) -> None: + assert ( + clean_headers( + headers, + forward_llm_provider_auth_headers=True, + authenticated_with_header=litellm_auth_header_name(headers), + ) + == expected + ) incoming: Final = Headers(headers) - - result: Final = clean_headers( - incoming, - forward_llm_provider_auth_headers=True, - authenticated_with_header=litellm_auth_header_name(incoming), + assert ( + clean_headers( + incoming, + forward_llm_provider_auth_headers=True, + authenticated_with_header=litellm_auth_header_name(incoming), + ) + == expected ) - assert result == expected + +def test_clean_headers_strips_authenticated_custom_header_without_litellm_key_name() -> None: + cleaned: Final = clean_headers( + {"X-Custom-Proxy-Key": "sk-proxy-virtual-key", "x-cost-center": "core-platform"}, + forward_llm_provider_auth_headers=True, + authenticated_with_header="x-custom-proxy-key", + ) + + assert cleaned == {"x-cost-center": "core-platform"} def test_clean_headers_byok_strips_custom_litellm_key_header_named_x_api_key() -> None: @@ -8054,6 +8122,38 @@ async def test_add_litellm_data_to_request_keeps_every_forwarded_credential_out_ FORWARD_PROVIDER_AUTH_HEADERS, {}, ), + ( + {"X-Goog-Api-Key": "sk-virtual-key", "x-cost-center": "core-platform"}, + FORWARD_PROVIDER_AUTH_HEADERS, + {}, + ), + ( + { + "authorization": "Bearer sk-virtual-key", + "x-goog-api-key": "sk-virtual-key", + "x-cost-center": "core-platform", + }, + FORWARD_PROVIDER_AUTH_HEADERS, + {}, + ), + ( + { + "Authorization": "Bearer sk-virtual-key", + "X-Goog-Api-Key": "sk-virtual-key", + "x-cost-center": "core-platform", + }, + FORWARD_PROVIDER_AUTH_HEADERS, + {}, + ), + ( + { + "x-api-key": "", + "x-goog-api-key": "sk-virtual-key", + "x-cost-center": "core-platform", + }, + FORWARD_PROVIDER_AUTH_HEADERS, + {}, + ), ( { "x-litellm-api-key": "Bearer sk-virtual-key", @@ -8068,8 +8168,26 @@ async def test_add_litellm_data_to_request_keeps_every_forwarded_credential_out_ {**FORWARD_PROVIDER_AUTH_HEADERS, "litellm_key_header_name": "x-api-key"}, {}, ), + ( + { + "x-custom-proxy-key": "sk-virtual-key", + "x-goog-api-key": "AIza-caller-key", + "x-cost-center": "core-platform", + }, + {**FORWARD_PROVIDER_AUTH_HEADERS, "litellm_key_header_name": "x-custom-proxy-key"}, + {"x-goog-api-key": "AIza-caller-key"}, + ), + ], + ids=[ + "goog-was-proxy-auth", + "goog-mixed-case-was-proxy-auth", + "duplicate-proxy-key-in-bearer-and-goog", + "duplicate-proxy-key-mixed-case", + "empty-higher-precedence-header", + "goog-beside-litellm-key", + "custom-litellm-key-header", + "custom-litellm-key-header-preserves-goog-byok", ], - ids=["goog-was-proxy-auth", "goog-beside-litellm-key", "custom-litellm-key-header"], ) async def test_add_litellm_data_to_request_never_forwards_the_header_used_for_proxy_auth( request_headers: Mapping[str, str],