diff --git a/litellm/google_genai/main.py b/litellm/google_genai/main.py index 40c9c7bc36c..628aade7eba 100644 --- a/litellm/google_genai/main.py +++ b/litellm/google_genai/main.py @@ -1,12 +1,14 @@ 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._logging import verbose_logger from litellm.constants import request_timeout # Import the adapter for fallback to completion format @@ -40,6 +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, str]](Mapping[str, str]) ################################################# @@ -236,6 +239,30 @@ def _merge_native_request_fields( return {**native_request_fields, **(extra_body or {})} +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 as exc: + verbose_logger.debug("Ignoring forwarded `headers` that are not a str-to-str mapping: %s", exc) + 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 + 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 +376,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 +388,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 +403,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 +470,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 +485,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 +501,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 +563,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 +578,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 +593,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/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 4444d6d25f9..cc0686a0636 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1498,12 +1498,29 @@ 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 _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 present), "x-api-key") def forwarded_anthropic_oauth_token(headers: Headers) -> str | None: @@ -1521,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 @@ -1542,36 +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 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]: @@ -2669,13 +2685,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 5cac61e727e..a80e8b21515 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,447 @@ 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]: + """ + 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 + + 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) + 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 {} + 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))) + 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) + 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() + 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 _extract_sse_texts(chunks: tuple[bytes, ...]) -> tuple[str, ...]: + from itertools import chain + + 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)) + + +_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, + *, + 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 + + from litellm.google_genai.main import ( + agenerate_content, + agenerate_content_stream, + generate_content, + generate_content_stream, + ) + + contents: Final = [{"role": "user", "parts": [{"text": "ping"}]}] + if entrypoint == "sync": + 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), + ) + ) + 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), + ) + ) + stream: Final = ( + generate_content_stream( + model=model, + contents=contents, + api_base=api_base, + api_key=_UPSTREAM_KEY, + headers=headers, + extra_headers=dict(extra_headers), + ) + if entrypoint == "sync_stream" + else await agenerate_content_stream( + model=model, + contents=contents, + api_base=api_base, + api_key=_UPSTREAM_KEY, + headers=headers, + extra_headers=dict(extra_headers), + ) + ) + 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"]) +@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: + 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"} + + texts: Final = await _invoke_generate_content( + entrypoint, model=model, api_base=api_base, headers=forwarded, extra_headers=explicit + ) + + 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 + wire_headers: Final = header_capture_wire.requests[0][1] + assert wire_headers["x-team-id"] == "platform" + assert wire_headers["anthropic-beta"] == "context-1m-2025-08-07" + 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"]) +async def test_router_deployment_headers_reach_generate_content( + header_capture_wire: _WireCapture, + target: str, +) -> None: + from litellm import Router + + 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 = await router.agenerate_content( + model="genai-deployment", + contents=[{"role": "user", "parts": [{"text": "ping"}]}], + ) + + assert _extract_chunk_texts(resp) == ("pong",) + assert len(header_capture_wire.requests) == 1 + assert header_capture_wire.requests[0][1]["x-deployment-tag"] == "blue" + + +@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, + 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( + incoming_headers, + litellm_key_header_name=None, + forward_llm_provider_auth_headers=False, + 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"] == _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), + ( + { + "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", + "x-goog-api-key": "AIza-caller-key", + "x-cost-center": "core-platform", + }, + "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", + ], +) +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 42e828f267e..cf2a88ac9d4 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,152 @@ 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"), + ({"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": "", "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-mixed-case-dict", + "empty-custom-header-skipped", + "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, litellm_key_header_name) == expected + 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-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-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) + assert ( + clean_headers( + incoming, + forward_llm_provider_auth_headers=True, + authenticated_with_header=litellm_auth_header_name(incoming), + ) + == 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: + 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 +8113,110 @@ 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-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", + "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"}, + {}, + ), + ( + { + "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", + ], +) +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"