From 9deefc0f766837822e192154290091d8b57bf5bd Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 20 Apr 2026 16:52:59 -0700 Subject: [PATCH 1/3] fix: align MCP broker endpoint access controls with existing auth patterns --- .../mcp_server/discoverable_endpoints.py | 8 +++++ .../mcp_management_endpoints.py | 30 +++++++++++++++---- .../test_mcp_management_endpoints.py | 8 ++--- 3 files changed, 37 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 65e2e3e983..792a9dace1 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -323,6 +323,14 @@ async def authorize_with_server( ) parsed = urlparse(redirect_uri) + if parsed.scheme not in ("http", "https"): + raise HTTPException( + status_code=400, + detail={ + "error": "invalid_redirect_uri", + "message": "redirect_uri must use http or https scheme", + }, + ) base_url = urlunparse(parsed._replace(query="")) request_base_url = get_request_base_url(request) encoded_state = encode_state_with_base_url( diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 56bbfe0300..9cb5ec95ab 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -19,6 +19,7 @@ import functools import importlib import json import os +from urllib.parse import urlparse from dataclasses import dataclass from datetime import datetime, timedelta, timezone from typing import Any, Dict, Iterable, List, Literal, Optional @@ -1336,7 +1337,9 @@ if MCP_AVAILABLE: return _redact_mcp_credentials(temp_record) - def _get_cached_temporary_mcp_server_or_404(server_id: str) -> MCPServer: + def _get_cached_temporary_mcp_server_or_404( + server_id: str, request: Optional[Request] = None + ) -> MCPServer: server = get_cached_temporary_mcp_server(server_id) if server is None: # Fall back to real DB/config server (e.g. for the user-side OAuth flow @@ -1344,10 +1347,14 @@ if MCP_AVAILABLE: from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) + from litellm.proxy.auth.ip_address_utils import IPAddressUtils + client_ip = IPAddressUtils.get_mcp_client_ip(request) if request else None server = global_mcp_server_manager.get_mcp_server_by_id( server_id - ) or global_mcp_server_manager.get_mcp_server_by_name(server_id) + ) or global_mcp_server_manager.get_mcp_server_by_name( + server_id, client_ip=client_ip + ) if server is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -1358,10 +1365,12 @@ if MCP_AVAILABLE: @router.get( "/server/oauth/{server_id}/authorize", include_in_schema=False, + dependencies=[Depends(user_api_key_auth)], ) async def mcp_authorize( request: Request, server_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), client_id: Optional[str] = None, redirect_uri: str = Query(...), state: str = "", @@ -1370,7 +1379,16 @@ if MCP_AVAILABLE: response_type: Optional[str] = None, scope: Optional[str] = None, ): - mcp_server = _get_cached_temporary_mcp_server_or_404(server_id) + parsed_redirect = urlparse(redirect_uri) + if parsed_redirect.scheme not in ("http", "https"): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ + "error": "invalid_redirect_uri", + "message": "redirect_uri must use http or https scheme", + }, + ) + mcp_server = _get_cached_temporary_mcp_server_or_404(server_id, request=request) # Use the server's stored client_id when the caller doesn't supply one resolved_client_id = mcp_server.client_id or client_id or "" if not resolved_client_id: @@ -1399,10 +1417,12 @@ if MCP_AVAILABLE: @router.post( "/server/oauth/{server_id}/token", include_in_schema=False, + dependencies=[Depends(user_api_key_auth)], ) async def mcp_token( request: Request, server_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), grant_type: str = Form(...), code: Optional[str] = Form(None), redirect_uri: Optional[str] = Form(None), @@ -1412,7 +1432,7 @@ if MCP_AVAILABLE: refresh_token: Optional[str] = Form(None), scope: Optional[str] = Form(None), ): - mcp_server = _get_cached_temporary_mcp_server_or_404(server_id) + mcp_server = _get_cached_temporary_mcp_server_or_404(server_id, request=request) resolved_client_id = mcp_server.client_id or client_id or "" if not resolved_client_id: raise HTTPException( @@ -1443,7 +1463,7 @@ if MCP_AVAILABLE: include_in_schema=False, ) async def mcp_register(request: Request, server_id: str): - mcp_server = _get_cached_temporary_mcp_server_or_404(server_id) + mcp_server = _get_cached_temporary_mcp_server_or_404(server_id, request=request) request_data = await _read_request_body(request=request) data: dict = {**request_data} 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 e8d31b4951..c1a1acb433 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 @@ -1486,7 +1486,7 @@ class TestTemporaryMCPSessionEndpoints: ) assert result is authorize_response - get_server.assert_called_once_with("server-1") + get_server.assert_called_once_with("server-1", request=request) authorize_mock.assert_awaited_once_with( request=request, mcp_server=server, @@ -1533,7 +1533,7 @@ class TestTemporaryMCPSessionEndpoints: ) assert result is exchange_response - get_server.assert_called_once_with("server-1") + get_server.assert_called_once_with("server-1", request=request) exchange_mock.assert_awaited_once_with( request=request, mcp_server=server, @@ -1581,7 +1581,7 @@ class TestTemporaryMCPSessionEndpoints: ) assert result is exchange_response - get_server.assert_called_once_with("server-1") + get_server.assert_called_once_with("server-1", request=request) exchange_mock.assert_awaited_once_with( request=request, mcp_server=server, @@ -1628,7 +1628,7 @@ class TestTemporaryMCPSessionEndpoints: result = await mcp_register(request=request, server_id="server-1") assert result is register_response - get_server.assert_called_once_with("server-1") + get_server.assert_called_once_with("server-1", request=request) read_body.assert_awaited_once_with(request=request) register_mock.assert_awaited_once_with( request=request, From 99f007f51d9961a3818b43ecb2bb74e6d69a2418 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 20 Apr 2026 16:59:54 -0700 Subject: [PATCH 2/3] refactor: consolidate redirect_uri scheme check into shared handler --- .../management_endpoints/mcp_management_endpoints.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 9cb5ec95ab..8e54af0f96 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -19,7 +19,6 @@ import functools import importlib import json import os -from urllib.parse import urlparse from dataclasses import dataclass from datetime import datetime, timedelta, timezone from typing import Any, Dict, Iterable, List, Literal, Optional @@ -1379,15 +1378,6 @@ if MCP_AVAILABLE: response_type: Optional[str] = None, scope: Optional[str] = None, ): - parsed_redirect = urlparse(redirect_uri) - if parsed_redirect.scheme not in ("http", "https"): - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail={ - "error": "invalid_redirect_uri", - "message": "redirect_uri must use http or https scheme", - }, - ) mcp_server = _get_cached_temporary_mcp_server_or_404(server_id, request=request) # Use the server's stored client_id when the caller doesn't supply one resolved_client_id = mcp_server.client_id or client_id or "" From b6de470ce97d0e7c2054d1640535ce86a4393aa8 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 20 Apr 2026 17:00:34 -0700 Subject: [PATCH 3/3] fix: add access control to register endpoint to match authorize and token --- .../proxy/management_endpoints/mcp_management_endpoints.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index 8e54af0f96..f18e699045 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -1451,8 +1451,13 @@ if MCP_AVAILABLE: @router.post( "/server/oauth/{server_id}/register", include_in_schema=False, + dependencies=[Depends(user_api_key_auth)], ) - async def mcp_register(request: Request, server_id: str): + async def mcp_register( + request: Request, + server_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + ): mcp_server = _get_cached_temporary_mcp_server_or_404(server_id, request=request) request_data = await _read_request_body(request=request) data: dict = {**request_data}