fix(proxy): keep the parsed key when the custom key header is absent (#44469)

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-10-10 11:48:04 -07:00 • committed by GitHub
parent 09558baca6
commit b9c49aeb35
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 93 additions and 12 deletions

View file

@ -1625,7 +1625,7 @@ class MCPRequestHandler:
)
litellm_api_key: Final = (
headers.get(custom_key_header_name)
if custom_key_header_name is not None
if custom_key_header_name is not None and custom_key_header_name in headers
else next(
(header_value for header_name in admission_header_names if (header_value := headers.get(header_name))),
None,

View file

@ -1601,7 +1601,7 @@ async def user_api_key_auth_builder(
request=request,
)
custom_litellm_key_header_name: Final = general_settings.get("litellm_key_header_name")
if custom_litellm_key_header_name is not None:
if custom_litellm_key_header_name is not None and custom_litellm_key_header_name in request.headers:
api_key = get_api_key_from_custom_header(
request=request,
custom_litellm_key_header_name=custom_litellm_key_header_name,

View file

@ -9813,7 +9813,7 @@ class TestSessionBearerEgressScrub:
{},
)
async def test_caller_admission_credential_does_not_fall_back_when_custom_header_is_absent(self) -> None:
async def test_caller_admission_credential_falls_back_to_standard_headers_when_custom_header_is_absent(self) -> None:
caller_value: Final = "Bearer sk-caller-admission-key-123"
headers: Final = Headers(
{
@ -9837,16 +9837,12 @@ class TestSessionBearerEgressScrub:
admitted_credential=admitted_credential,
)
assert admitted_credential == "sk-caller-admission-key-123"
assert result == (
{"Authorization": caller_value},
{
"x-litellm-api-key": caller_value,
"authorization": caller_value,
"x-mcp-auth": caller_value,
"x-mcp-echo_srv-authorization": caller_value,
},
caller_value,
{"echo_srv": {"Authorization": caller_value}},
None,
{"x-litellm-api-key": caller_value},
None,
{},
)
async def test_caller_admission_credential_uses_master_key_alias_as_admission_gate(self) -> None:

View file

@ -1176,6 +1176,91 @@ async def test_x_litellm_api_key():
assert valid_token.token != hash_token(master_key)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("headers", "api_key", "custom_litellm_key_header"),
[
({"Authorization": "Bearer sk-lit9211-master"}, "Bearer sk-lit9211-master", None),
({"x-litellm-api-key": "Bearer sk-lit9211-master"}, "", "Bearer sk-lit9211-master"),
(
{"X-Custom-Key": "Bearer sk-lit9211-master", "Authorization": "Bearer sk-wrong"},
"Bearer sk-wrong",
None,
),
],
)
async def test_auth_custom_key_header_precedence_and_fallback(
headers: dict[str, str],
api_key: str,
custom_litellm_key_header: str | None,
monkeypatch: pytest.MonkeyPatch,
):
from fastapi import Request
from starlette.datastructures import URL
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
from litellm.proxy import proxy_server
monkeypatch.setattr(proxy_server, "master_key", "sk-lit9211-master")
monkeypatch.setattr(proxy_server, "general_settings", {"litellm_key_header_name": "X-Custom-Key"})
request_headers = [(name.lower().encode(), value.encode()) for name, value in headers.items()]
request = Request(
scope={
"type": "http",
"method": "POST",
"path": "/chat/completions",
"headers": request_headers,
},
receive=AsyncMock(return_value={"type": "http.request", "body": b"", "more_body": False}),
)
request._url = URL(url="/chat/completions")
valid_token = await user_api_key_auth(
request=request,
api_key=api_key,
custom_litellm_key_header=custom_litellm_key_header,
)
assert valid_token.user_role == LitellmUserRoles.PROXY_ADMIN
assert valid_token.api_key == LITELLM_PROXY_MASTER_KEY_ALIAS
@pytest.mark.asyncio
@pytest.mark.parametrize("custom_header_value", [b"Bearer sk-wrong", b""])
async def test_auth_rejects_wrong_or_empty_present_custom_header(
custom_header_value: bytes, monkeypatch: pytest.MonkeyPatch
):
from fastapi import Request
from starlette.datastructures import URL
from litellm.proxy import proxy_server
from litellm.proxy._types import ProxyException
monkeypatch.setattr(proxy_server, "master_key", "sk-lit9211-master")
monkeypatch.setattr(proxy_server, "general_settings", {"litellm_key_header_name": "X-Custom-Key"})
monkeypatch.setattr(proxy_server, "prisma_client", MagicMock(get_data=AsyncMock(return_value=None)))
request = Request(
scope={
"type": "http",
"method": "POST",
"path": "/chat/completions",
"headers": [
(b"x-custom-key", custom_header_value),
(b"authorization", b"Bearer sk-lit9211-master"),
],
},
receive=AsyncMock(return_value={"type": "http.request", "body": b"", "more_body": False}),
)
request._url = URL(url="/chat/completions")
with pytest.raises(ProxyException) as exc_info:
await user_api_key_auth(request=request, api_key="Bearer sk-lit9211-master")
assert exc_info.value.code == "401"
@pytest.mark.asyncio
async def test_user_api_key_from_query_param():
"""Ensure user_api_key_auth reads API key from `key` query parameter."""