mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-18 04:18:15 +00:00
Merge pull request #26142 from BerriAI/litellm_mcp_broker_endpoint_auth
[Fix] MCP broker OAuth endpoint access controls
This commit is contained in:
@@ -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(
|
||||
|
||||
@@ -1336,7 +1336,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 +1346,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 +1364,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 +1378,7 @@ if MCP_AVAILABLE:
|
||||
response_type: Optional[str] = None,
|
||||
scope: Optional[str] = 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)
|
||||
# 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 +1407,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 +1422,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(
|
||||
@@ -1441,9 +1451,14 @@ 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):
|
||||
mcp_server = _get_cached_temporary_mcp_server_or_404(server_id)
|
||||
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}
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user