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
This commit is contained in:
Dennis Henry 2026-05-05 21:36:22 -04:00 committed by Yuneng Jiang
parent 93d8375cbc
commit b36fb1dc19
No known key found for this signature in database
2 changed files with 153 additions and 6 deletions

View file

@ -151,8 +151,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
@ -1447,6 +1453,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,
@ -1497,12 +1552,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 = "",
@ -1542,12 +1597,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),

View file

@ -1551,6 +1551,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 (