diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index c62e40e90d4..eb42cd4912d 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -1492,6 +1492,53 @@ 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"], + ) + 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, 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 24b8e1595b4..01283461de0 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 @@ -2,7 +2,7 @@ import os import sys import types import json -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from types import SimpleNamespace from typing import List, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -1553,6 +1553,151 @@ 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-with-at-least-32-bytes" + 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_rejects_expired_token_cookie(self): + """Expired UI session cookies must not authenticate MCP OAuth redirects.""" + import jwt + + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + _mcp_oauth_user_api_key_auth, + ) + + master_key = "test-master-key-with-at-least-32-bytes" + api_key_in_cookie = "sk-expired-cookie-key" + token_cookie = jwt.encode( + { + "user_id": "user@example.com", + "key": api_key_in_cookie, + "user_role": "proxy_admin", + "login_method": "sso", + "exp": datetime.now(timezone.utc) - timedelta(minutes=5), + }, + 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="" + ) + 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"] == "" + + @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 (