diff --git a/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py b/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py index b04fc3a0a4..3c62b9e4ea 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py +++ b/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py @@ -1,15 +1,21 @@ -from typing import Optional +from typing import Optional, List from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser from litellm.proxy._types import UserAPIKeyAuth -class LiteLLMAuthenticatedUser(AuthenticatedUser): +class MCPAuthenticatedUser(AuthenticatedUser): """ - Wrapper class to make UserAPIKeyAuth compatible with MCP's AuthenticatedUser + Wrapper class to make LiteLLM's authentication and configuration compatible with MCP's AuthenticatedUser. + + This class handles: + 1. User API key authentication information + 2. MCP authentication header + 3. MCP server configuration """ - def __init__(self, user_api_key_auth: UserAPIKeyAuth, mcp_auth_header: Optional[str] = None): + def __init__(self, user_api_key_auth: UserAPIKeyAuth, mcp_auth_header: Optional[str] = None, mcp_servers: Optional[List[str]] = None): self.user_api_key_auth = user_api_key_auth self.mcp_auth_header = mcp_auth_header + self.mcp_servers = mcp_servers diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 177fbafa5e..a9782db994 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -1,4 +1,5 @@ from typing import List, Optional, Tuple +import json from starlette.datastructures import Headers from starlette.requests import Request @@ -8,12 +9,14 @@ from litellm._logging import verbose_logger from litellm.proxy._types import LiteLLM_TeamTable, SpecialHeaders, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth - -class UserAPIKeyAuthMCP: +class MCPRequestHandler: """ - Class to handle Authentication for MCP requests + Class to handle MCP request processing, including: + 1. Authentication via LiteLLM API keys + 2. MCP server configuration and routing + 3. Header extraction and validation - Utilizes the main `user_api_key_auth` function to validate the request + Utilizes the main `user_api_key_auth` function to validate authentication """ LITELLM_API_KEY_HEADER_NAME_PRIMARY = SpecialHeaders.custom_litellm_api_key.value @@ -22,10 +25,16 @@ class UserAPIKeyAuthMCP: # This is the header to use if you want LiteLLM to use this header for authenticating to the MCP server LITELLM_MCP_AUTH_HEADER_NAME = SpecialHeaders.mcp_auth.value + LITELLM_MCP_SERVERS_HEADER_NAME = SpecialHeaders.mcp_servers.value + @staticmethod - async def user_api_key_auth_mcp(scope: Scope) -> Tuple[UserAPIKeyAuth, Optional[str]]: + async def process_mcp_request(scope: Scope) -> Tuple[UserAPIKeyAuth, Optional[str], Optional[List[str]]]: """ - Validate and extract headers from the ASGI scope for MCP requests. + Process and validate MCP request headers from the ASGI scope. + This includes: + 1. Extracting and validating authentication headers + 2. Processing MCP server configuration + 3. Handling MCP-specific headers Args: scope: ASGI scope containing request information @@ -33,15 +42,26 @@ class UserAPIKeyAuthMCP: Returns: UserAPIKeyAuth containing validated authentication information mcp_auth_header: Optional[str] MCP auth header to be passed to the MCP server + mcp_servers: Optional[List[str]] List of MCP servers to use Raises: HTTPException: If headers are invalid or missing required headers """ - headers = UserAPIKeyAuthMCP._safe_get_headers_from_scope(scope) + headers = MCPRequestHandler._safe_get_headers_from_scope(scope) litellm_api_key = ( - UserAPIKeyAuthMCP.get_litellm_api_key_from_headers(headers) or "" + MCPRequestHandler.get_litellm_api_key_from_headers(headers) or "" ) - mcp_auth_header = headers.get(UserAPIKeyAuthMCP.LITELLM_MCP_AUTH_HEADER_NAME) + mcp_auth_header = headers.get(MCPRequestHandler.LITELLM_MCP_AUTH_HEADER_NAME) + mcp_servers_header = headers.get(MCPRequestHandler.LITELLM_MCP_SERVERS_HEADER_NAME) + mcp_servers = None + if mcp_servers_header: + try: + mcp_servers = json.loads(mcp_servers_header) + if not isinstance(mcp_servers, list): + mcp_servers = None + except (json.JSONDecodeError, TypeError, ValueError) as e: + verbose_logger.debug(f"Error parsing mcp_servers header: {e}") + mcp_servers = None # Create a proper Request object with mock body method to avoid ASGI receive channel issues request = Request(scope=scope) @@ -57,7 +77,7 @@ class UserAPIKeyAuthMCP: api_key=litellm_api_key, request=request ) - return validated_user_api_key_auth, mcp_auth_header + return validated_user_api_key_auth, mcp_auth_header, mcp_servers @staticmethod def get_litellm_api_key_from_headers(headers: Headers) -> Optional[str]: @@ -71,12 +91,12 @@ class UserAPIKeyAuthMCP: headers: Starlette Headers object that handles case insensitivity """ # Headers object handles case insensitivity automatically - api_key = headers.get(UserAPIKeyAuthMCP.LITELLM_API_KEY_HEADER_NAME_PRIMARY) + api_key = headers.get(MCPRequestHandler.LITELLM_API_KEY_HEADER_NAME_PRIMARY) if api_key: return api_key auth_header = headers.get( - UserAPIKeyAuthMCP.LITELLM_API_KEY_HEADER_NAME_SECONDARY + MCPRequestHandler.LITELLM_API_KEY_HEADER_NAME_SECONDARY ) if auth_header: return auth_header @@ -101,7 +121,7 @@ class UserAPIKeyAuthMCP: for name, value in raw_headers } return Headers(headers_dict) - except Exception as e: + except (UnicodeDecodeError, AttributeError, TypeError) as e: verbose_logger.exception(f"Error getting headers from scope: {e}") # Return empty Headers object with empty dict return Headers({}) @@ -111,16 +131,16 @@ class UserAPIKeyAuthMCP: user_api_key_auth: Optional[UserAPIKeyAuth] = None, ) -> List[str]: """ - Apply least privilege + Get list of allowed MCP servers for the given user/key based on permissions """ from typing import List allowed_mcp_servers: List[str] = [] allowed_mcp_servers_for_key = ( - await UserAPIKeyAuthMCP._get_allowed_mcp_servers_for_key(user_api_key_auth) + await MCPRequestHandler._get_allowed_mcp_servers_for_key(user_api_key_auth) ) allowed_mcp_servers_for_team = ( - await UserAPIKeyAuthMCP._get_allowed_mcp_servers_for_team(user_api_key_auth) + await MCPRequestHandler._get_allowed_mcp_servers_for_team(user_api_key_auth) ) ######################################################### diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index c993e781e6..ed25889d3a 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -18,7 +18,7 @@ from mcp.types import Tool as MCPTool from litellm._logging import verbose_logger from litellm.experimental_mcp_client.client import MCPClient from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( - UserAPIKeyAuthMCP, + MCPRequestHandler, ) from litellm.proxy._types import ( LiteLLM_MCPServerTable, @@ -143,7 +143,7 @@ class MCPServerManager: """ Get the allowed MCP Servers for the user """ - allowed_mcp_servers = await UserAPIKeyAuthMCP.get_allowed_mcp_servers( + allowed_mcp_servers = await MCPRequestHandler.get_allowed_mcp_servers( user_api_key_auth ) verbose_logger.debug( diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 83f922a223..7414890299 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -14,13 +14,17 @@ from litellm._logging import verbose_logger from litellm.constants import MCP_TOOL_NAME_PREFIX from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( - UserAPIKeyAuthMCP, + MCPRequestHandler, ) from litellm.proxy._types import UserAPIKeyAuth from litellm.types.mcp_server.mcp_server_manager import MCPInfo from litellm.types.utils import StandardLoggingMCPToolCall from litellm.utils import client +from litellm.proxy._experimental.mcp_server.auth.litellm_auth_handler import ( + MCPAuthenticatedUser, +) + LITELLM_MCP_SERVER_NAME = "litellm-mcp-server" LITELLM_MCP_SERVER_VERSION = "1.0.0" LITELLM_MCP_SERVER_DESCRIPTION = "MCP Server for LiteLLM" @@ -55,9 +59,6 @@ if MCP_AVAILABLE: from mcp.types import TextContent as MCPTextContent from mcp.types import Tool as MCPTool - from litellm.proxy._experimental.mcp_server.auth.litellm_auth_handler import ( - LiteLLMAuthenticatedUser, - ) from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, ) @@ -350,8 +351,8 @@ if MCP_AVAILABLE: """Handle MCP requests through StreamableHTTP.""" try: # Validate headers and log request info - user_api_key_auth, mcp_auth_header = ( - await UserAPIKeyAuthMCP.user_api_key_auth_mcp(scope) + user_api_key_auth, mcp_auth_header, mcp_servers = ( + await MCPRequestHandler.process_mcp_request(scope) ) # Set the auth context variable for easy access in MCP functions set_auth_context( @@ -374,8 +375,8 @@ if MCP_AVAILABLE: """Handle MCP requests through SSE.""" try: # Validate headers and log request info - user_api_key_auth, mcp_auth_header = ( - await UserAPIKeyAuthMCP.user_api_key_auth_mcp(scope) + user_api_key_auth, mcp_auth_header, mcp_servers = ( + await MCPRequestHandler.process_mcp_request(scope) ) # Set the auth context variable for easy access in MCP functions set_auth_context( @@ -429,7 +430,7 @@ if MCP_AVAILABLE: user_api_key_auth: UserAPIKeyAuth object mcp_auth_header: MCP auth header to be passed to the MCP server """ - auth_user = LiteLLMAuthenticatedUser( + auth_user = MCPAuthenticatedUser( user_api_key_auth=user_api_key_auth, mcp_auth_header=mcp_auth_header, ) @@ -443,7 +444,7 @@ if MCP_AVAILABLE: Tuple[Optional[UserAPIKeyAuth], Optional[str]]: UserAPIKeyAuth object and MCP auth header """ auth_user = auth_context_var.get() - if auth_user and isinstance(auth_user, LiteLLMAuthenticatedUser): + if auth_user and isinstance(auth_user, MCPAuthenticatedUser): return auth_user.user_api_key_auth, auth_user.mcp_auth_header return None, None diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 98156a5ad6..0698bfc3f9 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2695,6 +2695,7 @@ class SpecialHeaders(enum.Enum): azure_apim_authorization = "Ocp-Apim-Subscription-Key" custom_litellm_api_key = "x-litellm-api-key" mcp_auth = "x-mcp-auth" + mcp_servers = "x-mcp-servers" class LitellmDataForBackendLLMCall(TypedDict, total=False): diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 8458f9a45a..2712815e3e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -13,13 +13,14 @@ sys.path.insert( ) # Adds the parent directory to the system path from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( - UserAPIKeyAuthMCP, + MCPRequestHandler, ) -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import UserAPIKeyAuth, SpecialHeaders +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @pytest.mark.asyncio -class TestUserAPIKeyAuthMCP: +class TestMCPRequestHandler: @pytest.mark.parametrize( "user_api_key_auth,object_permission_id,prisma_client_available,db_result,expected_result", @@ -91,7 +92,7 @@ class TestUserAPIKeyAuthMCP: with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): # Call the method - result = await UserAPIKeyAuthMCP._get_allowed_mcp_servers_for_key( + result = await MCPRequestHandler._get_allowed_mcp_servers_for_key( user_api_key_auth ) @@ -170,8 +171,8 @@ class TestUserAPIKeyAuthMCP: ), ], ) - async def test_user_api_key_auth_mcp(self, headers, expected_api_key, expected_mcp_auth_header): - """Test user_api_key_auth_mcp method with various header scenarios""" + async def test_process_mcp_request(self, headers, expected_api_key, expected_mcp_auth_header): + """Test process_mcp_request method with various header scenarios""" # Create ASGI scope with headers scope = { @@ -181,28 +182,33 @@ class TestUserAPIKeyAuthMCP: "headers": headers, } - # Mock the user_api_key_auth function - mock_auth_result = UserAPIKeyAuth( - api_key=expected_api_key, - user_id="test-user-id", - team_id="test-team-id", - ) + # Create an async mock for user_api_key_auth + async def mock_user_api_key_auth(api_key, request): + return UserAPIKeyAuth( + token="e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855" if api_key else None, + api_key=api_key, + user_id="test-user-id" if api_key else None, + team_id="test-team-id" if api_key else None, + user_role=None, + request_route=None + ) with patch( - "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth" - ) as mock_user_api_key_auth: - mock_user_api_key_auth.return_value = mock_auth_result - + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth", + side_effect=mock_user_api_key_auth, + ) as mock_auth: # Call the method - auth_result, mcp_auth_header = await UserAPIKeyAuthMCP.user_api_key_auth_mcp(scope) + auth_result, mcp_auth_header, mcp_servers = await MCPRequestHandler.process_mcp_request(scope) # Assert the results - assert auth_result == mock_auth_result + assert auth_result.api_key == expected_api_key + assert auth_result.user_id == ("test-user-id" if expected_api_key else None) + assert auth_result.team_id == ("test-team-id" if expected_api_key else None) assert mcp_auth_header == expected_mcp_auth_header # Verify user_api_key_auth was called with correct parameters - mock_user_api_key_auth.assert_called_once() - call_args = mock_user_api_key_auth.call_args + mock_auth.assert_called_once() + call_args = mock_auth.call_args # Check that api_key parameter is correct assert call_args.kwargs["api_key"] == expected_api_key @@ -213,3 +219,147 @@ class TestUserAPIKeyAuthMCP: # Verify the request has the correct scope assert request_param.scope == scope + + @pytest.mark.parametrize( + "headers,expected_result", + [ + # Test case 1: All headers present + ( + [ + (b"x-litellm-api-key", b"test-api-key"), + (b"x-mcp-auth", b"test-mcp-auth"), + (b"x-mcp-servers", b'["server1", "server2"]'), + ], + { + "api_key": "test-api-key", + "mcp_auth": "test-mcp-auth", + "mcp_servers": ["server1", "server2"], + } + ), + # Test case 2: Only API key present + ( + [(b"x-litellm-api-key", b"test-api-key")], + { + "api_key": "test-api-key", + "mcp_auth": None, + "mcp_servers": None, + } + ), + # Test case 3: Invalid JSON in mcp_servers + ( + [ + (b"x-litellm-api-key", b"test-api-key"), + (b"x-mcp-servers", b'invalid-json'), + ], + { + "api_key": "test-api-key", + "mcp_auth": None, + "mcp_servers": None, + } + ), + # Test case 4: mcp_servers not a list + ( + [ + (b"x-litellm-api-key", b"test-api-key"), + (b"x-mcp-servers", b'{"key": "value"}'), + ], + { + "api_key": "test-api-key", + "mcp_auth": None, + "mcp_servers": None, + } + ), + # Test case 5: Empty mcp_servers list + ( + [ + (b"x-litellm-api-key", b"test-api-key"), + (b"x-mcp-servers", b'[]'), + ], + { + "api_key": "test-api-key", + "mcp_auth": None, + "mcp_servers": [], + } + ), + # Test case 6: Using Authorization header instead of x-litellm-api-key + ( + [ + (b"authorization", b"Bearer test-api-key"), + (b"x-mcp-servers", b'["server1"]'), + ], + { + "api_key": "Bearer test-api-key", + "mcp_auth": None, + "mcp_servers": ["server1"], + } + ), + # Test case 7: Case insensitive header names + ( + [ + (b"X-LITELLM-API-KEY", b"test-api-key"), + (b"X-MCP-AUTH", b"test-mcp-auth"), + (b"X-MCP-SERVERS", b'["server1"]'), + ], + { + "api_key": "test-api-key", + "mcp_auth": "test-mcp-auth", + "mcp_servers": ["server1"], + } + ), + ] + ) + async def test_header_extraction(self, headers, expected_result): + """Test header extraction and processing from ASGI scope""" + + # Create ASGI scope with headers + scope = { + "type": "http", + "method": "POST", + "path": "/test", + "headers": headers, + } + + # Get headers using the internal method + extracted_headers = MCPRequestHandler._safe_get_headers_from_scope(scope) + + # Verify API key extraction + api_key = MCPRequestHandler.get_litellm_api_key_from_headers(extracted_headers) + assert api_key == expected_result["api_key"] + + # Verify MCP auth header + mcp_auth = extracted_headers.get(SpecialHeaders.mcp_auth.value) + assert mcp_auth == expected_result["mcp_auth"] + + # Verify MCP servers + mcp_servers_header = extracted_headers.get(SpecialHeaders.mcp_servers.value) + if mcp_servers_header: + try: + mcp_servers = json.loads(mcp_servers_header) + if not isinstance(mcp_servers, list): + mcp_servers = None + except: + mcp_servers = None + else: + mcp_servers = None + + assert mcp_servers == expected_result["mcp_servers"] + + # Test the full process_mcp_request method + mock_auth_result = UserAPIKeyAuth( + api_key=expected_result["api_key"], + user_id="test-user-id", + team_id="test-team-id", + ) + + with patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth" + ) as mock_user_api_key_auth: + mock_user_api_key_auth.return_value = mock_auth_result + + # Call the method + auth_result, mcp_auth_header, mcp_servers_result = await MCPRequestHandler.process_mcp_request(scope) + + # Assert the results + assert auth_result == mock_auth_result + assert mcp_auth_header == expected_result["mcp_auth"] + assert mcp_servers_result == expected_result["mcp_servers"]