tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-13 21:06:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add MCP servers header to the scope of header (#12266)

Browse Source 
* add mcp header param

* rename class

* fix tests

* fix tests

* added ruff fixes

* fix ruff checks

* added mypy type checking
This commit is contained in:
Jugal D. Bhatt
2025-07-03 08:52:14 +05:30
committed by GitHub
parent eef0623443
commit 5fe01b9a2b
6 changed files with 230 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py
+10 -4
View File
@@ -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
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
+36 -16
View File
@@ -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)
)
#########################################################
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
+2 -2
View File
@@ -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(
litellm/proxy/_experimental/mcp_server/server.py
+11 -10
View File
@@ -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
litellm/proxy/_types.py
+1
View File
@@ -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):
tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
+170 -20
View File
@@ -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"]