mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-25 01:02:15 +00:00
fix(mcp): support OAuth browser auth
This commit is contained in:
gym-cmd
parent
d0cc1a49f5
commit
0daeeab884
2 changed files with 193 additions and 1 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue