From 73de892654f209ac7f5022b1bc7bf1e2f8ac674a Mon Sep 17 00:00:00 2001 From: Dennis Henry Date: Tue, 5 May 2026 21:36:22 -0400 Subject: [PATCH] fix: replace user api key auth with authorization or cookie for mcp server creation (#27190) * fix: replace user api key auth with authorization or cookie for mcp server creation * updated tests --- .../mcp_management_endpoints.py | 67 ++++++++++++-- .../test_mcp_management_endpoints.py | 92 +++++++++++++++++++ 2 files changed, 153 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index c62e40e90d..7bda0f87cc 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -152,8 +152,14 @@ if MCP_AVAILABLE: UserAPIKeyAuth, UserMCPManagementMode, ) - from litellm.proxy.auth.user_api_key_auth import user_api_key_auth - from litellm.proxy.common_utils.http_parsing_utils import _read_request_body + from litellm.proxy.auth.user_api_key_auth import ( + _user_api_key_auth_builder, + user_api_key_auth, + ) + from litellm.proxy.common_utils.http_parsing_utils import ( + _read_request_body, + populate_request_with_path_params, + ) from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view from litellm.proxy.management_helpers.utils import management_endpoint_wrapper from litellm.types.mcp import MCPCredentials @@ -1492,6 +1498,55 @@ if MCP_AVAILABLE: return _redact_mcp_credentials(temp_record) + async def _mcp_oauth_user_api_key_auth(request: Request) -> UserAPIKeyAuth: + """ + Auth dependency for MCP OAuth browser-navigation endpoints (/authorize, /token). + + Tries the Authorization header first. Falls back to decoding the UI + 'token' session cookie (set by SSO login) to extract the API key, which + allows browser-based OAuth redirects to work without an explicit + Authorization header. + """ + import jwt as _jwt + + from litellm.proxy.proxy_server import master_key + + auth_header = request.headers.get("Authorization", "") + api_key = auth_header # _get_bearer_token will strip "Bearer " prefix + + if not api_key: + token_cookie = request.cookies.get("token") + if token_cookie and master_key: + try: + decoded = _jwt.decode( + token_cookie, + master_key, + algorithms=["HS256"], + # UI session cookies may omit exp; don't require it. + options={"verify_exp": False}, + ) + if decoded.get("login_method") in ("sso", "username_password"): + cookie_key = decoded.get("key", "") + if cookie_key: + api_key = f"Bearer {cookie_key}" + except _jwt.InvalidTokenError: + pass + + request_data = await _read_request_body(request=request) + request_data = populate_request_with_path_params( + request_data=request_data, request=request + ) + + return await _user_api_key_auth_builder( + request=request, + api_key=api_key, + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data=request_data, + ) + async def _get_cached_temporary_mcp_server_or_404( server_id: str, user_api_key_dict: UserAPIKeyAuth, @@ -1542,12 +1597,12 @@ if MCP_AVAILABLE: @router.get( "/server/oauth/{server_id}/authorize", include_in_schema=False, - dependencies=[Depends(user_api_key_auth)], + dependencies=[Depends(_mcp_oauth_user_api_key_auth)], ) async def mcp_authorize( request: Request, server_id: str, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + user_api_key_dict: UserAPIKeyAuth = Depends(_mcp_oauth_user_api_key_auth), client_id: Optional[str] = None, redirect_uri: str = Query(...), state: str = "", @@ -1587,12 +1642,12 @@ if MCP_AVAILABLE: @router.post( "/server/oauth/{server_id}/token", include_in_schema=False, - dependencies=[Depends(user_api_key_auth)], + dependencies=[Depends(_mcp_oauth_user_api_key_auth)], ) async def mcp_token( request: Request, server_id: str, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + user_api_key_dict: UserAPIKeyAuth = Depends(_mcp_oauth_user_api_key_auth), grant_type: str = Form(...), code: Optional[str] = Form(None), redirect_uri: Optional[str] = Form(None), 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 24b8e1595b..f0909afcbf 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 @@ -1553,6 +1553,98 @@ class TestTemporaryMCPSessionEndpoints: assert "permission" in str(exc_info.value) + @pytest.mark.asyncio + async def test_mcp_oauth_user_api_key_auth_falls_back_to_token_cookie(self): + """ + When the Authorization header is absent but a valid 'token' cookie is + present (browser navigation), _mcp_oauth_user_api_key_auth should + decode the cookie JWT and authenticate via the API key stored in it. + """ + import jwt + + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _mcp_oauth_user_api_key_auth, + ) + + master_key = "test-master-key" + api_key_in_cookie = "sk-test-cookie-key" + token_cookie = jwt.encode( + { + "user_id": "user@example.com", + "key": api_key_in_cookie, + "user_role": "proxy_admin", + "login_method": "sso", + }, + master_key, + algorithm="HS256", + ) + + mock_request = MagicMock() + mock_request.headers = {} + mock_request.cookies = {"token": token_cookie} + + expected_auth = generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key=api_key_in_cookie + ) + fake_proxy_server = types.SimpleNamespace(master_key=master_key) + + with ( + patch.dict(sys.modules, {"litellm.proxy.proxy_server": fake_proxy_server}), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._user_api_key_auth_builder", + AsyncMock(return_value=expected_auth), + ) as auth_builder_mock, + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._read_request_body", + AsyncMock(return_value={}), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.populate_request_with_path_params", + side_effect=lambda request_data, request: request_data, + ), + ): + result = await _mcp_oauth_user_api_key_auth(mock_request) + + assert result is expected_auth + _, call_kwargs = auth_builder_mock.call_args + assert call_kwargs["api_key"] == f"Bearer {api_key_in_cookie}" + + @pytest.mark.asyncio + async def test_mcp_oauth_user_api_key_auth_uses_authorization_header_when_present( + self, + ): + """When Authorization header is present it takes priority over the cookie.""" + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _mcp_oauth_user_api_key_auth, + ) + + expected_auth = generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.PROXY_ADMIN + ) + mock_request = MagicMock() + mock_request.headers = {"Authorization": "Bearer sk-header-key"} + mock_request.cookies = {} + + with ( + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._user_api_key_auth_builder", + AsyncMock(return_value=expected_auth), + ) as auth_builder_mock, + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints._read_request_body", + AsyncMock(return_value={}), + ), + patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.populate_request_with_path_params", + side_effect=lambda request_data, request: request_data, + ), + ): + result = await _mcp_oauth_user_api_key_auth(mock_request) + + assert result is expected_auth + _, call_kwargs = auth_builder_mock.call_args + assert call_kwargs["api_key"] == "Bearer sk-header-key" + @pytest.mark.asyncio async def test_mcp_authorize_proxies_to_discoverable_endpoint(self): from litellm.proxy.management_endpoints.mcp_management_endpoints import (