mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
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:
parent
93d8375cbc
commit
b36fb1dc19
2 changed files with 153 additions and 6 deletions
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue