From b854823bb981a48343cdc28395f1a2d5b78c127d Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Tue, 5 May 2026 09:18:42 +0530 Subject: [PATCH] Fix failing tests --- litellm/proxy/auth/user_api_key_auth.py | 74 +++++++++++++++---- .../mcp_management_endpoints.py | 23 ++++-- .../test_mcp_management_endpoints.py | 24 +++++- 3 files changed, 97 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 9d3c06e641f..98be18162a2 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2030,24 +2030,18 @@ def _should_skip_budget_checks( return False -@tracer.wrap() -async def user_api_key_auth( +async def run_user_api_key_auth_pipeline( request: Request, - api_key: str = fastapi.Security(api_key_header), - azure_api_key_header: str = fastapi.Security(azure_api_key_header), - anthropic_api_key_header: Optional[str] = fastapi.Security( - anthropic_api_key_header - ), - google_ai_studio_api_key_header: Optional[str] = fastapi.Security( - google_ai_studio_api_key_header - ), - azure_apim_header: Optional[str] = fastapi.Security(azure_apim_header), - custom_litellm_key_header: Optional[str] = fastapi.Security( - custom_litellm_key_header - ), + api_key: str, + azure_api_key_header: str, + anthropic_api_key_header: Optional[str], + google_ai_studio_api_key_header: Optional[str], + azure_apim_header: Optional[str], + custom_litellm_key_header: Optional[str], ) -> UserAPIKeyAuth: """ - Parent function to authenticate user api key / jwt token. + Shared implementation for ``user_api_key_auth`` and for call sites that must + run the same auth pipeline without FastAPI ``Security()`` injection. """ request_data = await _read_request_body(request=request) @@ -2104,6 +2098,56 @@ async def user_api_key_auth( return user_api_key_auth_obj +async def user_api_key_auth_from_request_headers(request: Request) -> UserAPIKeyAuth: + """ + Run the same auth as ``Depends(user_api_key_auth)`` using headers on ``request``. + + Used when a route cannot use the FastAPI dependency (e.g. MCP OAuth broker + ``/authorize`` / ``/token`` resolving optional ``Authorization``). + """ + h = request.headers + return await run_user_api_key_auth_pipeline( + request=request, + api_key=h.get("authorization") or "", + azure_api_key_header=h.get("api-key") or "", + anthropic_api_key_header=h.get("x-api-key"), + google_ai_studio_api_key_header=h.get("x-goog-api-key"), + azure_apim_header=h.get("ocp-apim-subscription-key"), + custom_litellm_key_header=h.get("x-litellm-api-key"), + ) + + +@tracer.wrap() +async def user_api_key_auth( + request: Request, + api_key: str = fastapi.Security(api_key_header), + azure_api_key_header: str = fastapi.Security(azure_api_key_header), + anthropic_api_key_header: Optional[str] = fastapi.Security( + anthropic_api_key_header + ), + google_ai_studio_api_key_header: Optional[str] = fastapi.Security( + google_ai_studio_api_key_header + ), + azure_apim_header: Optional[str] = fastapi.Security(azure_apim_header), + custom_litellm_key_header: Optional[str] = fastapi.Security( + custom_litellm_key_header + ), +) -> UserAPIKeyAuth: + """ + Parent function to authenticate user api key / jwt token. + """ + + return await run_user_api_key_auth_pipeline( + request=request, + api_key=api_key, + azure_api_key_header=azure_api_key_header, + anthropic_api_key_header=anthropic_api_key_header, + google_ai_studio_api_key_header=google_ai_studio_api_key_header, + azure_apim_header=azure_apim_header, + custom_litellm_key_header=custom_litellm_key_header, + ) + + async def _return_user_api_key_auth_obj( user_obj: Optional[LiteLLM_UserTable], api_key: str, diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 9f174287826..af742ce5f0b 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -1501,13 +1501,24 @@ if MCP_AVAILABLE: When absent, unauthenticated access is still allowed for **temp-cache** servers only (browser OAuth). When present, global-registry access follows admin / allowlist rules via ``_get_cached_temporary_mcp_server_or_404``. + + Only non-empty **string** ``Authorization`` values trigger a full auth + pipeline import (tests and mocks may attach MagicMock headers). """ - authorization = ( - request.headers.get("authorization") - or request.headers.get("Authorization") - or "" - ).strip() - if not authorization: + try: + headers = request.headers + except Exception: + return None + raw: object = None + for key in ("authorization", "Authorization"): + try: + candidate = headers.get(key) + except Exception: + candidate = None + if isinstance(candidate, str) and candidate.strip(): + raw = candidate + break + if not isinstance(raw, str) or not raw.strip(): return None from litellm.proxy.auth.user_api_key_auth import ( user_api_key_auth_from_request_headers, diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 21fef1703a8..cd54fac39cf 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -1563,6 +1563,10 @@ class TestTemporaryMCPSessionEndpoints: server = generate_mock_mcp_server_config_record(server_id="server-1") authorize_response = MagicMock() with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._try_resolve_mcp_oauth_broker_user", + AsyncMock(return_value=None), + ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", return_value=server, @@ -1585,7 +1589,9 @@ class TestTemporaryMCPSessionEndpoints: ) assert result is authorize_response - get_server.assert_awaited_once_with("server-1", request=request) + get_server.assert_awaited_once_with( + "server-1", user_api_key_dict=None, request=request + ) authorize_mock.assert_awaited_once_with( request=request, mcp_server=server, @@ -1609,6 +1615,10 @@ class TestTemporaryMCPSessionEndpoints: exchange_response = {"access_token": "token"} with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._try_resolve_mcp_oauth_broker_user", + AsyncMock(return_value=None), + ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", return_value=server, @@ -1632,7 +1642,9 @@ class TestTemporaryMCPSessionEndpoints: ) assert result is exchange_response - get_server.assert_awaited_once_with("server-1", request=request) + get_server.assert_awaited_once_with( + "server-1", user_api_key_dict=None, request=request + ) exchange_mock.assert_awaited_once_with( request=request, mcp_server=server, @@ -1657,6 +1669,10 @@ class TestTemporaryMCPSessionEndpoints: exchange_response = {"access_token": "new-token", "refresh_token": "new-rt"} with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._try_resolve_mcp_oauth_broker_user", + AsyncMock(return_value=None), + ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints._get_cached_temporary_mcp_server_or_404", return_value=server, @@ -1680,7 +1696,9 @@ class TestTemporaryMCPSessionEndpoints: ) assert result is exchange_response - get_server.assert_awaited_once_with("server-1", request=request) + get_server.assert_awaited_once_with( + "server-1", user_api_key_dict=None, request=request + ) exchange_mock.assert_awaited_once_with( request=request, mcp_server=server,