tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-11 19:48:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #14964 from BerriAI/litellm_dev_09_24_2025_p1

Browse Source 
feat: initial commit for v2 oauth flow
This commit is contained in:
Krish Dholakia
2025-09-27 11:22:31 -07:00
committed by GitHub
parent 55572ce3cd 3e3e3837a1
commit e27ac3b411
18 changed files with 787 additions and 152 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/test-linting.yml
+18 -1
View File
@@ -11,6 +11,9 @@ jobs:
steps:
- uses: actions/checkout@v4
with:
fetch-depth: 0
clean: true
- name: Set up Python
uses: actions/setup-python@v4
@@ -20,6 +23,11 @@ jobs:
- name: Install Poetry
uses: snok/install-poetry@v1
- name: Clean Python cache
run: |
find . -type d -name "__pycache__" -exec rm -rf {} + || true
find . -name "*.pyc" -delete || true
- name: Install dependencies
run: |
poetry install --with dev
@@ -31,6 +39,15 @@ jobs:
poetry run black .
cd ..
- name: Debug - Check file state
run: |
echo "Current branch:"
git branch --show-current
echo "Last 3 commits:"
git log --oneline -3
echo "File content around line 43:"
head -50 litellm/litellm_core_utils/custom_logger_registry.py | tail -10
- name: Run Ruff linting
run: |
cd litellm
@@ -44,7 +61,7 @@ jobs:
- name: Run MyPy type checking
run: |
cd litellm
poetry run mypy . --ignore-missing-imports
poetry run mypy . --ignore-missing-imports --disable-error-code=var-annotated
cd ..
- name: Check for circular imports
.github/workflows/test-litellm.yml
+1 -1
View File
@@ -40,4 +40,4 @@ jobs:
cd ..
- name: Run tests
run: |
poetry run pytest tests/test_litellm -x -vv -n 4
poetry run pytest tests/test_litellm --tb=short -vv --maxfail=10 -n 4
docs/my-website/sidebars.js
+1
View File
@@ -584,6 +584,7 @@ const sidebars = {
"budget_manager",
"caching/all_caches",
"completion/token_usage",
"sdk_custom_pricing",
"embedding/async_embedding",
"embedding/moderation",
"migration",
litellm/experimental_mcp_client/client.py
+12 -5
View File
@@ -1,10 +1,11 @@
"""
LiteLLM Proxy uses this MCP Client to connnect to other MCP servers.
"""
import asyncio
import base64
from datetime import timedelta
from typing import List, Optional
from typing import Dict, List, Optional
from mcp import ClientSession, StdioServerParameters
from mcp.client.sse import sse_client
@@ -46,6 +47,7 @@ class MCPClient:
auth_value: Optional[str] = None,
timeout: float = 60.0,
stdio_config: Optional[MCPStdioConfig] = None,
extra_headers: Optional[Dict[str, str]] = None,
):
self.server_url: str = server_url
self.transport_type: MCPTransport = transport_type
@@ -59,7 +61,7 @@ class MCPClient:
self._session_ctx = None
self._task: Optional[asyncio.Task] = None
self.stdio_config: Optional[MCPStdioConfig] = stdio_config
self.extra_headers: Optional[Dict[str, str]] = extra_headers
# handle the basic auth value if provided
if auth_value:
self.update_auth_value(auth_value)
@@ -115,6 +117,9 @@ class MCPClient:
await self._session.initialize()
else: # http
headers = self._get_auth_headers()
verbose_logger.debug(
"litellm headers for streamablehttp_client: ", headers
)
self._transport_ctx = streamablehttp_client(
url=self.server_url,
timeout=timedelta(seconds=self.timeout),
@@ -186,9 +191,7 @@ class MCPClient:
def _get_auth_headers(self) -> dict:
"""Generate authentication headers based on auth type."""
headers = {
"MCP-Protocol-Version": "2025-06-18"
}
headers = {"MCP-Protocol-Version": "2025-06-18"}
if self._mcp_auth_value:
if self.auth_type == MCPAuth.bearer_token:
@@ -200,6 +203,10 @@ class MCPClient:
elif self.auth_type == MCPAuth.authorization:
headers["Authorization"] = self._mcp_auth_value
# update the headers with the extra headers
if self.extra_headers:
headers.update(self.extra_headers)
return headers
async def list_tools(self) -> List[MCPTool]:
litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py
+13 -3
View File
@@ -1,4 +1,4 @@
from typing import List, Optional, Dict
from typing import Dict, List, Optional
from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser
@@ -8,17 +8,27 @@ from litellm.proxy._types import UserAPIKeyAuth
class MCPAuthenticatedUser(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 (deprecated)
3. MCP server configuration (can include access groups)
4. Server-specific authentication headers
5. OAuth2 headers
"""
def __init__(self, user_api_key_auth: UserAPIKeyAuth, mcp_auth_header: Optional[str] = None, mcp_servers: Optional[List[str]] = None, mcp_server_auth_headers: Optional[Dict[str, str]] = None, mcp_protocol_version: Optional[str] = None):
def __init__(
self,
user_api_key_auth: UserAPIKeyAuth,
mcp_auth_header: Optional[str] = None,
mcp_servers: Optional[List[str]] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
mcp_protocol_version: Optional[str] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
):
self.user_api_key_auth = user_api_key_auth
self.mcp_auth_header = mcp_auth_header
self.mcp_servers = mcp_servers
self.mcp_server_auth_headers = mcp_server_auth_headers or {}
self.mcp_protocol_version = mcp_protocol_version
self.oauth2_headers = oauth2_headers
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
+36 -13
View File
@@ -1,4 +1,4 @@
from typing import List, Optional, Tuple, Dict, Set
from typing import Dict, List, Optional, Set, Tuple
from starlette.datastructures import Headers
from starlette.requests import Request
@@ -36,7 +36,11 @@ class MCPRequestHandler:
async def process_mcp_request(
scope: Scope,
) -> Tuple[
UserAPIKeyAuth, Optional[str], Optional[List[str]], Optional[Dict[str, str]]
UserAPIKeyAuth,
Optional[str],
Optional[List[str]],
Optional[Dict[str, str]],
Optional[Dict[str, str]],
]:
"""
Process and validate MCP request headers from the ASGI scope.
@@ -44,6 +48,7 @@ class MCPRequestHandler:
1. Extracting and validating authentication headers
2. Processing MCP server configuration
3. Handling MCP-specific headers
4. Handling oauth2 headers
Args:
scope: ASGI scope containing request information
@@ -70,6 +75,9 @@ class MCPRequestHandler:
MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers)
)
# Get the oauth2 headers
oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers)
# Parse MCP servers from header
mcp_servers_header = headers.get(
MCPRequestHandler.LITELLM_MCP_SERVERS_HEADER_NAME
@@ -96,14 +104,18 @@ class MCPRequestHandler:
return b"{}"
request.body = mock_body # type: ignore
validated_user_api_key_auth = await user_api_key_auth(
api_key=litellm_api_key, request=request
)
if ".well-known" in str(request.url): # public routes
validated_user_api_key_auth = UserAPIKeyAuth()
else:
validated_user_api_key_auth = await user_api_key_auth(
api_key=litellm_api_key, request=request
)
return (
validated_user_api_key_auth,
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
)
@staticmethod
@@ -174,6 +186,17 @@ class MCPRequestHandler:
return server_auth_headers
@staticmethod
def _get_oauth2_headers_from_headers(headers: Headers) -> Dict[str, str]:
"""
Get the oauth2 headers from the request headers.
"""
oauth2_headers = {}
for header_name, header_value in headers.items():
if header_name.lower().startswith("authorization"):
oauth2_headers["Authorization"] = header_value
return oauth2_headers
@staticmethod
def _get_mcp_client_side_auth_header_name() -> str:
"""
@@ -359,10 +382,10 @@ class MCPRequestHandler:
return []
try:
team_obj: Optional[
LiteLLM_TeamTable
] = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": user_api_key_auth.team_id},
team_obj: Optional[LiteLLM_TeamTable] = (
await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": user_api_key_auth.team_id},
)
)
if team_obj is None:
verbose_logger.debug("team_obj is None")
@@ -535,10 +558,10 @@ class MCPRequestHandler:
verbose_logger.debug("prisma_client is None")
return []
team_obj: Optional[
LiteLLM_TeamTable
] = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": user_api_key_auth.team_id},
team_obj: Optional[LiteLLM_TeamTable] = (
await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": user_api_key_auth.team_id},
)
)
if team_obj is None:
verbose_logger.debug("team_obj is None")
litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py
+252
View File
@@ -0,0 +1,252 @@
import json
from typing import Optional
from urllib.parse import urlencode, urlparse, urlunparse
from fastapi import APIRouter, Form, HTTPException, Request
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
decrypt_value_helper,
encrypt_value_helper,
)
router = APIRouter(
tags=["mcp"],
)
def encode_state_with_base_url(base_url: str, original_state: str) -> str:
"""
Encode the base_url and original state using encryption.
Args:
base_url: The base URL to encode
original_state: The original state parameter
Returns:
An encrypted string that encodes both values
"""
state_data = {"base_url": base_url, "original_state": original_state}
state_json = json.dumps(state_data, sort_keys=True)
encrypted_state = encrypt_value_helper(state_json)
return encrypted_state
def decode_state_hash(encrypted_state: str) -> tuple[str, str]:
"""
Decode an encrypted state to retrieve the base_url and original state.
Args:
encrypted_state: The encrypted string to decode
Returns:
A tuple of (base_url, original_state)
Raises:
Exception: If decryption fails or data is malformed
"""
decrypted_json = decrypt_value_helper(encrypted_state, "oauth_state")
if decrypted_json is None:
raise ValueError("Failed to decrypt state parameter")
state_data = json.loads(decrypted_json)
return state_data["base_url"], state_data["original_state"]
@router.get("/{mcp_server_name}/authorize")
@router.get("/authorize")
async def authorize(
request: Request,
client_id: str,
redirect_uri: str,
state: str = "",
mcp_server_name: Optional[str] = None,
):
# Redirect to real GitHub OAuth
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
mcp_server = global_mcp_server_manager.get_mcp_server_by_name(client_id)
if mcp_server is None:
raise HTTPException(status_code=404, detail="MCP server not found")
if mcp_server.auth_type != "oauth2":
raise HTTPException(status_code=400, detail="MCP server is not OAuth2")
if mcp_server.client_id is None:
raise HTTPException(status_code=400, detail="MCP server client id is not set")
if mcp_server.authorization_url is None:
raise HTTPException(
status_code=400, detail="MCP server authorization url is not set"
)
if mcp_server.scopes is None:
raise HTTPException(status_code=400, detail="MCP server scopes is not set")
# Parse it to remove any existing query
parsed = urlparse(redirect_uri)
base_url = urlunparse(parsed._replace(query=""))
request_base_url = str(request.base_url).rstrip("/")
# Encode the base_url and original state in a unique hash
encoded_state = encode_state_with_base_url(base_url, state)
params = {
"client_id": mcp_server.client_id,
"redirect_uri": f"{request_base_url}/callback",
"scope": " ".join(mcp_server.scopes),
"state": encoded_state,
}
return RedirectResponse(f"{mcp_server.authorization_url}?{urlencode(params)}")
@router.post("/token")
async def token_endpoint(
request: Request,
grant_type: str = Form(...),
code: str = Form(None),
redirect_uri: str = Form(None),
client_id: str = Form(...),
client_secret: str = Form(...),
):
"""
Accept the authorization code from Claude and exchange it for GitHub token.
Forward the GitHub token back to Claude in standard OAuth format.
1. Call the token endpoint
2. Store the user's PAT in the db - and generate a LiteLLM virtual key
2. Return the token
3. Return a virtual key in this response
"""
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
mcp_server = global_mcp_server_manager.get_mcp_server_by_name(client_id)
if mcp_server is None:
raise HTTPException(status_code=404, detail="MCP server not found")
if grant_type != "authorization_code":
raise HTTPException(status_code=400, detail="Unsupported grant_type")
if mcp_server.token_url is None:
raise HTTPException(status_code=400, detail="MCP server token url is not set")
proxy_base_url = str(request.base_url).rstrip("/")
# Exchange code for real GitHub token
async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.Oauth2Check)
response = await async_client.post(
mcp_server.token_url,
headers={"Accept": "application/json"},
data={
"client_id": mcp_server.client_id,
"client_secret": mcp_server.client_secret,
"code": code,
"redirect_uri": f"{proxy_base_url}/callback",
},
)
response.raise_for_status()
github_token = response.json()["access_token"]
# Return to Claude in expected OAuth 2 format
### return a virtual key in this response
return JSONResponse(
{"access_token": github_token, "token_type": "Bearer", "expires_in": 3600}
)
@router.get("/callback")
async def callback(code: str, state: str):
try:
# Decode the state hash to get base_url and original state
base_url, original_state = decode_state_hash(state)
# Exchange code for token with GitHub
params = {"code": code, "state": original_state}
# Forward token to Claude ephemeral endpoint
complete_returned_url = f"{base_url}?{urlencode(params)}"
return RedirectResponse(url=complete_returned_url, status_code=302)
except Exception:
# fallback if state hash not found
return HTMLResponse(
"<html><body>Authentication incomplete. You can close this window.</body></html>"
)
# ------------------------------
# Optional .well-known endpoints for MCP + OAuth discovery
# ------------------------------
@router.get("/.well-known/oauth-protected-resource/{mcp_server_name}/mcp")
@router.get("/.well-known/oauth-protected-resource")
async def oauth_protected_resource_mcp(
request: Request, mcp_server_name: Optional[str] = None
):
request_base_url = str(request.base_url).rstrip("/")
return {
"authorization_servers": [
(
f"{request_base_url}/{mcp_server_name}"
if mcp_server_name
else f"{request_base_url}"
)
],
"resource": (
f"{request_base_url}/{mcp_server_name}/mcp"
if mcp_server_name
else f"{request_base_url}/mcp"
), # this is what Claude will call
}
@router.get("/.well-known/oauth-authorization-server/{mcp_server_name}")
@router.get("/.well-known/oauth-authorization-server")
async def oauth_authorization_server_mcp(
request: Request, mcp_server_name: Optional[str] = None
):
request_base_url = str(request.base_url).rstrip("/")
return {
"issuer": request_base_url, # point to your proxy
"authorization_endpoint": f"{request_base_url}/authorize",
"token_endpoint": f"{request_base_url}/token",
"response_types_supported": ["code"],
"grant_types_supported": ["authorization_code"],
"code_challenge_methods_supported": ["S256"],
"token_endpoint_auth_methods_supported": ["client_secret_post"],
# Claude expects a registration endpoint, even if we just fake it
"registration_endpoint": f"{request_base_url}/{mcp_server_name}/register",
}
# Alias for standard OpenID discovery
@router.get("/.well-known/openid-configuration")
async def openid_configuration(request: Request):
return await oauth_authorization_server_mcp(request)
@router.get("/.well-known/oauth-authorization-server/{mcp_server_name}/mcp")
@router.get("/.well-known/oauth-authorization-server")
async def oauth_authorization_server_root(
request: Request, mcp_server_name: Optional[str] = None
):
return await oauth_authorization_server_mcp(request, mcp_server_name)
@router.post("/{mcp_server_name}/register")
@router.post("/register")
async def register_client(request: Request, mcp_server_name: Optional[str] = None):
request_base_url = str(request.base_url).rstrip("/")
# return fixed GitHub client credentials
return {
"client_id": mcp_server_name or "dummy_client",
"client_secret": "dummy",
"redirect_uris": [f"{request_base_url}/mcp/callback"],
}
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
+111 -61
View File
@@ -39,7 +39,7 @@ from litellm.proxy._types import (
UserAPIKeyAuth,
)
from litellm.proxy.utils import ProxyLogging
from litellm.types.mcp import MCPStdioConfig
from litellm.types.mcp import MCPAuth, MCPStdioConfig
from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer
@@ -199,6 +199,12 @@ class MCPServerManager:
command=server_config.get("command", None) or "",
args=server_config.get("args", None) or [],
env=server_config.get("env", None) or {},
# oauth specific fields
client_id=server_config.get("client_id", None),
client_secret=server_config.get("client_secret", None),
scopes=server_config.get("scopes", None),
authorization_url=server_config.get("authorization_url", None),
token_url=server_config.get("token_url", None),
# TODO: utility fn the default values
transport=server_config.get("transport", MCPTransport.http),
auth_type=server_config.get("auth_type", None),
@@ -376,6 +382,7 @@ class MCPServerManager:
self,
server: MCPServer,
mcp_auth_header: Optional[str] = None,
extra_headers: Optional[Dict[str, str]] = None,
) -> MCPClient:
"""
Create an MCPClient instance for the given server.
@@ -405,6 +412,7 @@ class MCPServerManager:
auth_value=mcp_auth_header or server.authentication_token,
timeout=60.0,
stdio_config=stdio_config,
extra_headers=extra_headers,
)
else:
# For HTTP/SSE transports
@@ -415,12 +423,14 @@ class MCPServerManager:
auth_type=server.auth_type,
auth_value=mcp_auth_header or server.authentication_token,
timeout=60.0,
extra_headers=extra_headers,
)
async def _get_tools_from_server(
self,
server: MCPServer,
mcp_auth_header: Optional[str] = None,
extra_headers: Optional[Dict[str, str]] = None,
) -> List[MCPTool]:
"""
Helper method to get tools from a single MCP server with prefixed names.
@@ -441,6 +451,7 @@ class MCPServerManager:
client = self._create_mcp_client(
server=server,
mcp_auth_header=mcp_auth_header,
extra_headers=extra_headers,
)
tools = await self._fetch_tools_with_timeout(client, server.name)
@@ -550,6 +561,77 @@ class MCPServerManager:
)
return prefixed_tools
async def pre_call_tool_check(
self,
name: str,
arguments: Dict[str, Any],
server_name_from_prefix: str,
user_api_key_auth: Optional[UserAPIKeyAuth],
proxy_logging_obj: ProxyLogging,
):
pre_hook_kwargs = {
"name": name,
"arguments": arguments,
"server_name": server_name_from_prefix,
"user_api_key_auth": user_api_key_auth,
"user_api_key_user_id": (
getattr(user_api_key_auth, "user_id", None)
if user_api_key_auth
else None
),
"user_api_key_team_id": (
getattr(user_api_key_auth, "team_id", None)
if user_api_key_auth
else None
),
"user_api_key_end_user_id": (
getattr(user_api_key_auth, "end_user_id", None)
if user_api_key_auth
else None
),
"user_api_key_hash": (
getattr(user_api_key_auth, "api_key_hash", None)
if user_api_key_auth
else None
),
}
# Create MCP request object for processing
mcp_request_obj = proxy_logging_obj._create_mcp_request_object_from_kwargs(
pre_hook_kwargs
)
# Convert to LLM format for existing guardrail compatibility
synthetic_llm_data = proxy_logging_obj._convert_mcp_to_llm_format(
mcp_request_obj, pre_hook_kwargs
)
try:
# Use standard pre_call_hook with call_type="mcp_call"
modified_data = await proxy_logging_obj.pre_call_hook(
user_api_key_dict=user_api_key_auth, # type: ignore
data=synthetic_llm_data,
call_type="mcp_call", # type: ignore
)
if modified_data:
# Convert response back to MCP format and apply modifications
modified_kwargs = (
proxy_logging_obj._convert_mcp_hook_response_to_kwargs(
modified_data, pre_hook_kwargs
)
)
if modified_kwargs.get("arguments") != arguments:
arguments = modified_kwargs["arguments"]
except (
BlockedPiiEntityError,
GuardrailRaisedException,
HTTPException,
) as e:
# Re-raise guardrail exceptions to properly fail the MCP call
verbose_logger.error(f"Guardrail blocked MCP tool call pre call: {str(e)}")
raise e
async def call_tool(
self,
name: str,
@@ -558,6 +640,7 @@ class MCPServerManager:
mcp_auth_header: Optional[str] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
proxy_logging_obj: Optional[ProxyLogging] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
) -> CallToolResult:
"""
Call a tool with the given name and arguments (handles prefixed tool names)
@@ -602,65 +685,14 @@ class MCPServerManager:
# Using standard pre_call_hook with call_type="mcp_call"
#########################################################
if proxy_logging_obj:
pre_hook_kwargs = {
"name": name,
"arguments": arguments,
"server_name": server_name_from_prefix,
"user_api_key_auth": user_api_key_auth,
"user_api_key_user_id": getattr(user_api_key_auth, "user_id", None)
if user_api_key_auth
else None,
"user_api_key_team_id": getattr(user_api_key_auth, "team_id", None)
if user_api_key_auth
else None,
"user_api_key_end_user_id": getattr(
user_api_key_auth, "end_user_id", None
)
if user_api_key_auth
else None,
"user_api_key_hash": getattr(user_api_key_auth, "api_key_hash", None)
if user_api_key_auth
else None,
}
# Create MCP request object for processing
mcp_request_obj = proxy_logging_obj._create_mcp_request_object_from_kwargs(
pre_hook_kwargs
await self.pre_call_tool_check(
name=original_tool_name,
arguments=arguments,
server_name_from_prefix=server_name_from_prefix,
user_api_key_auth=user_api_key_auth,
proxy_logging_obj=proxy_logging_obj,
)
# Convert to LLM format for existing guardrail compatibility
synthetic_llm_data = proxy_logging_obj._convert_mcp_to_llm_format(
mcp_request_obj, pre_hook_kwargs
)
try:
# Use standard pre_call_hook with call_type="mcp_call"
modified_data = await proxy_logging_obj.pre_call_hook(
user_api_key_dict=user_api_key_auth, # type: ignore
data=synthetic_llm_data,
call_type="mcp_call", # type: ignore
)
if modified_data:
# Convert response back to MCP format and apply modifications
modified_kwargs = (
proxy_logging_obj._convert_mcp_hook_response_to_kwargs(
modified_data, pre_hook_kwargs
)
)
if modified_kwargs.get("arguments") != arguments:
arguments = modified_kwargs["arguments"]
except (
BlockedPiiEntityError,
GuardrailRaisedException,
HTTPException,
) as e:
# Re-raise guardrail exceptions to properly fail the MCP call
verbose_logger.error(
f"Guardrail blocked MCP tool call pre call: {str(e)}"
)
raise e
# Get server-specific auth header if available
server_auth_header = None
if mcp_server_auth_headers and mcp_server.alias:
@@ -672,9 +704,15 @@ class MCPServerManager:
if server_auth_header is None:
server_auth_header = mcp_auth_header
# oauth2 headers
extra_headers: Optional[Dict[str, str]] = None
if mcp_server.auth_type == MCPAuth.oauth2:
extra_headers = oauth2_headers
client = self._create_mcp_client(
server=mcp_server,
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
)
async with client:
@@ -834,6 +872,16 @@ class MCPServerManager:
return server
return None
def get_mcp_server_by_name(self, server_name: str) -> Optional[MCPServer]:
"""
Get the MCP Server from the server name
"""
registry = self.get_registry()
for server in registry.values():
if server.server_name == server_name:
return server
return None
def _generate_stable_server_id(
self,
server_name: str,
@@ -1023,9 +1071,11 @@ class MCPServerManager:
auth_type=_server_config.auth_type,
created_at=datetime.datetime.now(),
updated_at=datetime.datetime.now(),
description=_server_config.mcp_info.get("description")
if _server_config.mcp_info
else None,
description=(
_server_config.mcp_info.get("description")
if _server_config.mcp_info
else None
),
mcp_info=_server_config.mcp_info,
mcp_access_groups=_server_config.access_groups or [],
# Stdio-specific fields
litellm/proxy/_experimental/mcp_server/rest_endpoints.py
+3 -3
View File
@@ -177,9 +177,9 @@ if MCP_AVAILABLE:
return {
"tools": list_tools_result,
"error": "partial_failure" if error_message else None,
"message": error_message
if error_message
else "Successfully retrieved tools",
"message": (
error_message if error_message else "Successfully retrieved tools"
),
}
except Exception as e:
litellm/proxy/_experimental/mcp_server/server.py
+65 -25
View File
@@ -22,6 +22,7 @@ from litellm.proxy._experimental.mcp_server.utils import (
LITELLM_MCP_SERVER_VERSION,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.mcp import MCPAuth
from litellm.types.mcp_server.mcp_server_manager import MCPInfo, MCPServer
from litellm.types.utils import StandardLoggingMCPToolCall
from litellm.utils import client
@@ -178,6 +179,7 @@ if MCP_AVAILABLE:
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
) = get_auth_context()
verbose_logger.debug(
f"MCP list_tools - User API Key Auth from context: {user_api_key_auth}"
@@ -195,6 +197,7 @@ if MCP_AVAILABLE:
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
)
verbose_logger.info(
f"MCP list_tools - Successfully returned {len(tools)} tools"
@@ -235,6 +238,7 @@ if MCP_AVAILABLE:
mcp_auth_header,
_,
mcp_server_auth_headers,
oauth2_headers,
) = get_auth_context()
verbose_logger.debug(
@@ -266,6 +270,7 @@ if MCP_AVAILABLE:
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
**data, # for logging
)
except BlockedPiiEntityError as e:
@@ -357,6 +362,7 @@ if MCP_AVAILABLE:
mcp_auth_header: Optional[str],
mcp_servers: Optional[List[str]],
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
) -> List[MCPTool]:
"""
Helper method to fetch tools from MCP servers based on server filtering criteria.
@@ -365,7 +371,8 @@ if MCP_AVAILABLE:
user_api_key_auth: User authentication info for access control
mcp_auth_header: Optional auth header for MCP server (deprecated)
mcp_servers: Optional list of server names/aliases to filter by
mcp_server_auth_headers: Optional dict of server-specific auth headers {server_alias: auth_value}
mcp_server_auth_headers: Optional dict of server-specific auth headers
oauth2_headers: Optional dict of oauth2 headers
Returns:
List[MCPTool]: Combined list of tools from filtered servers
@@ -398,6 +405,10 @@ if MCP_AVAILABLE:
elif mcp_server_auth_headers and server.server_name is not None:
server_auth_header = mcp_server_auth_headers.get(server.server_name)
extra_headers: Optional[Dict[str, str]] = None
if server.auth_type == MCPAuth.oauth2:
extra_headers = oauth2_headers
# Fall back to deprecated mcp_auth_header if no server-specific header found
if server_auth_header is None:
server_auth_header = mcp_auth_header
@@ -406,6 +417,7 @@ if MCP_AVAILABLE:
tools = await global_mcp_server_manager._get_tools_from_server(
server=server,
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
)
all_tools.extend(tools)
verbose_logger.debug(
@@ -427,6 +439,7 @@ if MCP_AVAILABLE:
mcp_auth_header: Optional[str] = None,
mcp_servers: Optional[List[str]] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
) -> List[MCPTool]:
"""
List all available MCP tools.
@@ -450,6 +463,7 @@ if MCP_AVAILABLE:
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
)
verbose_logger.debug(
f"Successfully fetched {len(managed_tools)} tools from managed MCP servers"
@@ -492,6 +506,7 @@ if MCP_AVAILABLE:
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
mcp_auth_header: Optional[str] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
**kwargs: Any,
) -> List[Union[TextContent, ImageContent, EmbeddedResource]]:
"""
@@ -519,16 +534,16 @@ if MCP_AVAILABLE:
"litellm_logging_obj", None
)
if litellm_logging_obj:
litellm_logging_obj.model_call_details[
"mcp_tool_call_metadata"
] = standard_logging_mcp_tool_call
litellm_logging_obj.model_call_details["mcp_tool_call_metadata"] = (
standard_logging_mcp_tool_call
)
litellm_logging_obj.model = f"MCP: {name}"
# Try managed server tool first (pass the full prefixed name)
# Primary and recommended way to use MCP servers
#########################################################
mcp_server: Optional[
MCPServer
] = global_mcp_server_manager._get_mcp_server_from_tool_name(name)
mcp_server: Optional[MCPServer] = (
global_mcp_server_manager._get_mcp_server_from_tool_name(name)
)
if mcp_server:
standard_logging_mcp_tool_call["mcp_server_cost_info"] = (
mcp_server.mcp_info or {}
@@ -539,6 +554,7 @@ if MCP_AVAILABLE:
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
litellm_logging_obj=litellm_logging_obj,
)
@@ -591,6 +607,7 @@ if MCP_AVAILABLE:
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
mcp_auth_header: Optional[str] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
litellm_logging_obj: Optional[Any] = None,
) -> List[Union[TextContent, ImageContent, EmbeddedResource]]:
"""Handle tool execution for managed server tools"""
@@ -603,6 +620,7 @@ if MCP_AVAILABLE:
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
proxy_logging_obj=proxy_logging_obj,
)
verbose_logger.debug("CALL TOOL RESULT: %s", call_tool_result)
@@ -638,26 +656,32 @@ if MCP_AVAILABLE:
mcp_path_match = re.match(r"^/mcp/([^?#]+)(?:\?.*)?(?:#.*)?$", path)
if mcp_path_match:
servers_and_path = mcp_path_match.group(1)
if servers_and_path:
# Check if it contains commas (comma-separated servers)
if ',' in servers_and_path:
if "," in servers_and_path:
# For comma-separated, look for a path at the end
# Common patterns: /tools, /chat/completions, etc.
path_match = re.search(r'/([^/,]+(?:/[^/,]+)*)$', servers_and_path)
path_match = re.search(r"/([^/,]+(?:/[^/,]+)*)$", servers_and_path)
if path_match:
# Path found at the end, remove it from servers
path_part = '/' + path_match.group(1)
servers_part = servers_and_path[:-len(path_part)]
mcp_servers_from_path = [s.strip() for s in servers_part.split(',') if s.strip()]
path_part = "/" + path_match.group(1)
servers_part = servers_and_path[: -len(path_part)]
mcp_servers_from_path = [
s.strip() for s in servers_part.split(",") if s.strip()
]
else:
# No path, just comma-separated servers
mcp_servers_from_path = [s.strip() for s in servers_and_path.split(',') if s.strip()]
mcp_servers_from_path = [
s.strip() for s in servers_and_path.split(",") if s.strip()
]
else:
# Single server case - use regex approach for server/path separation
# This handles cases like "custom_solutions/user_123/chat/completions"
# where we want to extract "custom_solutions/user_123" as the server name
single_server_match = re.match(r"^([^/]+(?:/[^/]+)?)(?:/.*)?$", servers_and_path)
single_server_match = re.match(
r"^([^/]+(?:/[^/]+)?)(?:/.*)?$", servers_and_path
)
if single_server_match:
server_name = single_server_match.group(1)
mcp_servers_from_path = [server_name]
@@ -677,6 +701,7 @@ if MCP_AVAILABLE:
mcp_auth_header,
_,
mcp_server_auth_headers,
oauth2_headers,
) = await MCPRequestHandler.process_mcp_request(scope)
mcp_servers = mcp_servers_from_path
else:
@@ -685,8 +710,15 @@ if MCP_AVAILABLE:
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
) = await MCPRequestHandler.process_mcp_request(scope)
return user_api_key_auth, mcp_auth_header, mcp_servers, mcp_server_auth_headers
return (
user_api_key_auth,
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
)
async def handle_streamable_http_mcp(
scope: Scope, receive: Receive, send: Send
@@ -699,6 +731,7 @@ if MCP_AVAILABLE:
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
) = await extract_mcp_auth_context(scope, path)
verbose_logger.debug(
f"MCP request mcp_servers (header/path): {mcp_servers}"
@@ -712,6 +745,7 @@ if MCP_AVAILABLE:
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
)
# Ensure session managers are initialized
@@ -750,6 +784,7 @@ if MCP_AVAILABLE:
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
) = await extract_mcp_auth_context(scope, path)
verbose_logger.debug(
f"MCP request mcp_servers (header/path): {mcp_servers}"
@@ -762,6 +797,7 @@ if MCP_AVAILABLE:
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
)
if not _SESSION_MANAGERS_INITIALIZED:
@@ -809,6 +845,8 @@ if MCP_AVAILABLE:
# Mount the MCP handlers
app.mount("/", handle_streamable_http_mcp)
app.mount("/mcp", handle_streamable_http_mcp)
app.mount("/{mcp_server_name}/mcp", handle_streamable_http_mcp)
app.mount("/sse", handle_sse_mcp)
app.add_middleware(AuthContextMiddleware)
@@ -821,6 +859,7 @@ if MCP_AVAILABLE:
mcp_auth_header: Optional[str] = None,
mcp_servers: Optional[List[str]] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
) -> None:
"""
Set the UserAPIKeyAuth in the auth context variable.
@@ -836,17 +875,17 @@ if MCP_AVAILABLE:
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
)
auth_context_var.set(auth_user)
def get_auth_context() -> (
Tuple[
Optional[UserAPIKeyAuth],
Optional[str],
Optional[List[str]],
Optional[Dict[str, str]],
]
):
def get_auth_context() -> Tuple[
Optional[UserAPIKeyAuth],
Optional[str],
Optional[List[str]],
Optional[Dict[str, str]],
Optional[Dict[str, str]],
]:
"""
Get the UserAPIKeyAuth from the auth context variable.
@@ -861,8 +900,9 @@ if MCP_AVAILABLE:
auth_user.mcp_auth_header,
auth_user.mcp_servers,
auth_user.mcp_server_auth_headers,
auth_user.oauth2_headers,
)
return None, None, None, None
return None, None, None, None, None
########################################################
############ End of Auth Context Functions #############
litellm/proxy/_new_secret_config.yaml
+13
View File
@@ -15,3 +15,16 @@ model_list:
model: hosted_vllm/whisper-v3
api_base: "https://webhook.site/2f385e05-00aa-402b-86d1-efc9261471a5"
api_key: dummy
mcp_servers:
github_mcp:
url: "https://api.githubcopilot.com/mcp"
auth_type: oauth2
authorization_url: https://github.com/login/oauth/authorize
token_url: https://github.com/login/oauth/access_token
client_id: os.environ/GITHUB_OAUTH_CLIENT_ID
client_secret: os.environ/GITHUB_OAUTH_CLIENT_SECRET
scopes: ["public_repo", "user:email"]
# allowed_tools: ["list_tools"]
# disallowed_tools: ["repo_delete"]
litellm/proxy/auth/user_api_key_auth.py
+19 -6
View File
@@ -390,7 +390,6 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
pass_through_endpoints: Optional[List[dict]] = general_settings.get(
"pass_through_endpoints", None
)
passed_in_key: Optional[str] = None
## CHECK IF X-LITELM-API-KEY IS PASSED IN - supercedes Authorization header
api_key, passed_in_key = get_api_key(
custom_litellm_key_header=custom_litellm_key_header,
@@ -502,7 +501,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
end_user_object = result["end_user_object"]
org_id = result["org_id"]
token = result["token"]
team_membership: Optional[LiteLLM_TeamMembership] = result.get("team_membership", None)
team_membership: Optional[LiteLLM_TeamMembership] = result.get(
"team_membership", None
)
global_proxy_spend = await get_global_proxy_spend(
litellm_proxy_admin_name=litellm_proxy_admin_name,
@@ -537,10 +538,22 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
org_id=org_id,
parent_otel_span=parent_otel_span,
end_user_id=end_user_id,
user_tpm_limit=user_object.tpm_limit if user_object is not None else None,
user_rpm_limit=user_object.rpm_limit if user_object is not None else None,
team_member_rpm_limit=team_membership.safe_get_team_member_rpm_limit() if team_membership is not None else None,
team_member_tpm_limit=team_membership.safe_get_team_member_tpm_limit() if team_membership is not None else None,
user_tpm_limit=(
user_object.tpm_limit if user_object is not None else None
),
user_rpm_limit=(
user_object.rpm_limit if user_object is not None else None
),
team_member_rpm_limit=(
team_membership.safe_get_team_member_rpm_limit()
if team_membership is not None
else None
),
team_member_tpm_limit=(
team_membership.safe_get_team_member_tpm_limit()
if team_membership is not None
else None
),
)
# run through common checks
_ = await common_checks(
litellm/proxy/proxy_server.py
+76 -6
View File
@@ -151,6 +151,9 @@ from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.proxy._experimental.mcp_server.discoverable_endpoints import (
router as mcp_discoverable_endpoints_router,
)
from litellm.proxy._experimental.mcp_server.rest_endpoints import (
router as mcp_rest_endpoints_router,
)
@@ -249,9 +252,7 @@ from litellm.proxy.management_endpoints.customer_endpoints import (
from litellm.proxy.management_endpoints.internal_user_endpoints import (
router as internal_user_router,
)
from litellm.proxy.management_endpoints.internal_user_endpoints import (
user_update,
)
from litellm.proxy.management_endpoints.internal_user_endpoints import user_update
from litellm.proxy.management_endpoints.key_management_endpoints import (
delete_verification_tokens,
duration_in_seconds,
@@ -298,9 +299,7 @@ from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMi
from litellm.proxy.openai_files_endpoints.files_endpoints import (
router as openai_files_router,
)
from litellm.proxy.openai_files_endpoints.files_endpoints import (
set_files_config,
)
from litellm.proxy.openai_files_endpoints.files_endpoints import set_files_config
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
passthrough_endpoint_router,
)
@@ -9512,5 +9511,76 @@ app.include_router(ui_discovery_endpoints_router)
########################################################
# MCP Server
########################################################
# Dynamic MCP server routes - handle /{mcp_server_name}/mcp
@app.api_route(
"/{mcp_server_name}/mcp",
methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"],
)
async def dynamic_mcp_route(mcp_server_name: str, request: Request):
"""Handle dynamic MCP server routes like /github_mcp/mcp"""
try:
# Validate that the MCP server exists
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
from litellm.types.mcp import MCPAuth
mcp_server = global_mcp_server_manager.get_mcp_server_by_name(mcp_server_name)
if mcp_server is None:
raise HTTPException(
status_code=404, detail=f"MCP server '{mcp_server_name}' not found"
)
# Create a new scope with the correct path format that the MCP handler expects
# Transform /{mcp_server_name}/mcp to /mcp/{mcp_server_name}
scope = dict(request.scope)
scope["path"] = f"/mcp/{mcp_server_name}"
# Import the MCP handler
from litellm.proxy._experimental.mcp_server.server import (
handle_streamable_http_mcp,
)
# Create a custom send function to capture the response
response_started = False
response_body = b""
response_status = 200
response_headers = []
async def custom_send(message):
nonlocal response_started, response_body, response_status, response_headers
if message["type"] == "http.response.start":
response_started = True
response_status = message["status"]
response_headers = message.get("headers", [])
elif message["type"] == "http.response.body":
response_body += message.get("body", b"")
# Call the existing MCP handler
await handle_streamable_http_mcp(
scope, receive=request.receive, send=custom_send
)
# Return the response
from starlette.responses import Response
headers_dict = {k.decode(): v.decode() for k, v in response_headers}
return Response(
content=response_body,
status_code=response_status,
headers=headers_dict,
media_type=headers_dict.get("content-type", "application/json"),
)
except Exception as e:
verbose_proxy_logger.error(
f"Error handling dynamic MCP route for {mcp_server_name}: {str(e)}"
)
raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
app.mount(path=BASE_MCP_ROUTE, app=mcp_app)
app.include_router(mcp_rest_endpoints_router)
app.include_router(mcp_discoverable_endpoints_router)
litellm/types/mcp.py
+16 -4
View File
@@ -15,6 +15,7 @@ else:
MCPImageContent = Any
MCPTextContent = Any
class MCPTransport(str, enum.Enum):
sse = "sse"
http = "http"
@@ -26,17 +27,21 @@ class MCPSpecVersion(str, enum.Enum):
mar_2025 = "2025-03-26"
jun_2025 = "2025-06-18"
class MCPAuth(str, enum.Enum):
none = "none"
api_key = "api_key"
bearer_token = "bearer_token"
basic = "basic"
authorization = "authorization"
oauth2 = "oauth2"
# MCP Literals
MCPTransportType = Literal[MCPTransport.sse, MCPTransport.http, MCPTransport.stdio]
MCPSpecVersionType = Literal[MCPSpecVersion.nov_2024, MCPSpecVersion.mar_2025, MCPSpecVersion.jun_2025]
MCPSpecVersionType = Literal[
MCPSpecVersion.nov_2024, MCPSpecVersion.mar_2025, MCPSpecVersion.jun_2025
]
MCPAuthType = Optional[
Literal[
MCPAuth.none,
@@ -44,11 +49,11 @@ MCPAuthType = Optional[
MCPAuth.bearer_token,
MCPAuth.basic,
MCPAuth.authorization,
MCPAuth.oauth2,
]
]
class MCPServerCostInfo(TypedDict, total=False):
default_cost_per_query: Optional[float]
"""
@@ -82,6 +87,7 @@ class MCPPreCallRequestObject(BaseModel):
"""
Pydantic object used for MCP pre_call_hook request validation and modification
"""
tool_name: str
arguments: Dict[str, Any]
server_name: Optional[str] = None
@@ -93,6 +99,7 @@ class MCPPreCallResponseObject(BaseModel):
"""
Pydantic object used for MCP pre_call_hook response
"""
should_proceed: bool = True
modified_arguments: Optional[Dict[str, Any]] = None
error_message: Optional[str] = None
@@ -103,6 +110,7 @@ class MCPDuringCallRequestObject(BaseModel):
"""
Pydantic object used for MCP during_call_hook request
"""
tool_name: str
arguments: Dict[str, Any]
server_name: Optional[str] = None
@@ -114,6 +122,7 @@ class MCPDuringCallResponseObject(BaseModel):
"""
Pydantic object used for MCP during_call_hook response
"""
should_continue: bool = True
error_message: Optional[str] = None
hidden_params: HiddenParams = HiddenParams()
@@ -123,5 +132,8 @@ class MCPPostCallResponseObject(BaseModel):
"""
Pydantic object used for MCP post_call_hook response
"""
mcp_tool_call_response: List[Union[MCPTextContent, MCPImageContent, MCPEmbeddedResource]]
hidden_params: HiddenParams
mcp_tool_call_response: List[
Union[MCPTextContent, MCPImageContent, MCPEmbeddedResource]
]
hidden_params: HiddenParams
litellm/types/mcp_server/mcp_server_manager.py
+6 -1
View File
@@ -6,7 +6,6 @@ from typing_extensions import TypedDict
from litellm.proxy._types import MCPAuthType, MCPTransportType
from litellm.types.mcp import MCPServerCostInfo
# MCPInfo now allows arbitrary additional fields for custom metadata
MCPInfo = Dict[str, Any]
@@ -21,6 +20,12 @@ class MCPServer(BaseModel):
auth_type: Optional[MCPAuthType] = None
authentication_token: Optional[str] = None
mcp_info: Optional[MCPInfo] = None
# OAuth-specific fields
client_id: Optional[str] = None
client_secret: Optional[str] = None
scopes: Optional[List[str]] = None
authorization_url: Optional[str] = None
token_url: Optional[str] = None
# Stdio-specific fields
command: Optional[str] = None
args: Optional[List[str]] = None
tests/mcp_tests/test_mcp_server.py
+21 -21
View File
@@ -98,9 +98,9 @@ async def test_mcp_server_manager_https_server():
assert tools[0].name == f"{expected_prefix}-gmail_send_email"
# Manually set up the tool mapping for the call_tool test
mcp_server_manager.tool_name_to_mcp_server_name_mapping[
"gmail_send_email"
] = expected_prefix
mcp_server_manager.tool_name_to_mcp_server_name_mapping["gmail_send_email"] = (
expected_prefix
)
mcp_server_manager.tool_name_to_mcp_server_name_mapping[
f"{expected_prefix}-gmail_send_email"
] = expected_prefix
@@ -260,9 +260,9 @@ async def test_mcp_http_transport_call_tool_mock():
)
# Manually set up tool mapping (normally done by list_tools)
test_manager.tool_name_to_mcp_server_name_mapping[
"gmail_send_email"
] = "test_http_server"
test_manager.tool_name_to_mcp_server_name_mapping["gmail_send_email"] = (
"test_http_server"
)
# Call the tool
result = await test_manager.call_tool(
@@ -326,9 +326,9 @@ async def test_mcp_http_transport_call_tool_error_mock():
)
# Manually set up tool mapping
test_manager.tool_name_to_mcp_server_name_mapping[
"gmail_send_email"
] = "test_http_server"
test_manager.tool_name_to_mcp_server_name_mapping["gmail_send_email"] = (
"test_http_server"
)
# Call the tool with invalid data
result = await test_manager.call_tool(
@@ -792,18 +792,18 @@ async def test_get_tools_from_mcp_servers():
assert len(result) == 1, "Should only return tools from server1"
assert result[0].name == "tool1", "Should return tool from server1"
# Test Case 2: Without specific MCP servers
# Create a different mock manager for the second test case
mock_manager_2 = AsyncMock()
mock_manager_2.get_allowed_mcp_servers = AsyncMock(
return_value=["server1_id", "server2_id"]
)
mock_manager_2.get_mcp_server_by_id = mock_get_server_by_id
mock_manager_2._get_tools_from_server = AsyncMock(
side_effect=lambda server, mcp_auth_header=None: [mock_tool_1]
if server.server_id == "server1_id"
else [mock_tool_2]
)
# Test Case 2: Without specific MCP servers
# Create a different mock manager for the second test case
mock_manager_2 = AsyncMock()
mock_manager_2.get_allowed_mcp_servers = AsyncMock(
return_value=["server1_id", "server2_id"]
)
mock_manager_2.get_mcp_server_by_id = mock_get_server_by_id
mock_manager_2._get_tools_from_server = AsyncMock(
side_effect=lambda server, mcp_auth_header=None, extra_headers=None: (
[mock_tool_1] if server.server_id == "server1_id" else [mock_tool_2]
)
)
with patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
+6
View File
@@ -364,6 +364,7 @@ class TestMCPRequestHandler:
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
) = await MCPRequestHandler.process_mcp_request(scope)
# Assert the results
@@ -543,6 +544,7 @@ class TestMCPRequestHandler:
mcp_auth_header,
mcp_servers_result,
mcp_server_auth_headers,
oauth2_headers,
) = await MCPRequestHandler.process_mcp_request(scope)
assert auth_result == mock_auth_result
assert mcp_auth_header == expected_result["mcp_auth"]
@@ -716,6 +718,7 @@ class TestMCPCustomHeaderName:
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
) = await MCPRequestHandler.process_mcp_request(scope)
# Assert the results
@@ -886,6 +889,7 @@ class TestMCPAccessGroupsE2E:
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
) = await MCPRequestHandler.process_mcp_request(scope)
# Assert the results
@@ -935,6 +939,7 @@ class TestMCPAccessGroupsE2E:
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
) = await MCPRequestHandler.process_mcp_request(scope)
# Assert the results
@@ -963,6 +968,7 @@ def test_mcp_path_based_server_segregation(monkeypatch):
mcp_auth_header,
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
) = get_auth_context()
# Capture the MCP servers for testing
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
+118 -2
View File
@@ -102,7 +102,9 @@ async def test_get_tools_from_mcp_servers_continues_when_one_server_fails():
working_server if server_id == "working_server" else failing_server
)
async def mock_get_tools_from_server(server, mcp_auth_header=None):
async def mock_get_tools_from_server(
server, mcp_auth_header=None, extra_headers=None
):
if server.name == "working_server":
# Working server returns tools
tool1 = MagicMock()
@@ -184,7 +186,9 @@ async def test_get_tools_from_mcp_servers_handles_all_servers_failing():
failing_server1 if server_id == "failing_server1" else failing_server2
)
async def mock_get_tools_from_server(server, mcp_auth_header=None):
async def mock_get_tools_from_server(
server, mcp_auth_header=None, extra_headers=None
):
# All servers fail
raise Exception(f"Server {server.name} connection failed")
@@ -448,3 +452,115 @@ async def test_mcp_routing_with_conflicting_alias_and_group_name():
assert (
called_servers[0].server_id == specific_server.server_id
), "Should have contacted the specific server alias, not the group."
@pytest.mark.asyncio
async def test_oauth2_headers_passed_to_mcp_client():
"""Test that OAuth2 headers are properly passed through to the MCP client for OAuth2 servers like github_mcp"""
try:
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
from litellm.proxy._experimental.mcp_server.server import (
_get_tools_from_mcp_servers,
set_auth_context,
)
from litellm.proxy._types import MCPTransport
from litellm.types.mcp import MCPAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
except ImportError:
pytest.skip("MCP server not available")
# Clear the registry to avoid conflicts with other tests
global_mcp_server_manager.registry.clear()
# Create an OAuth2 MCP server similar to github_mcp configuration
oauth2_server = MCPServer(
server_id="github_mcp_server_id",
name="github_mcp",
alias="github_mcp",
transport=MCPTransport.http,
url="https://api.githubcopilot.com/mcp",
auth_type=MCPAuth.oauth2,
client_id="test_github_client_id",
client_secret="test_github_client_secret",
scopes=["public_repo", "user:email"],
authorization_url="https://github.com/login/oauth/authorize",
token_url="https://github.com/login/oauth/access_token",
)
global_mcp_server_manager.registry[oauth2_server.server_id] = oauth2_server
# Mock user auth
user_api_key_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user")
# Set up OAuth2 headers that would come from the client
oauth2_headers = {"Authorization": "Bearer github_oauth_token_12345"}
# Set auth context with OAuth2 headers
set_auth_context(user_api_key_auth=user_api_key_auth, oauth2_headers=oauth2_headers)
# This will capture the arguments passed to _create_mcp_client
captured_client_args = {}
def mock_create_mcp_client(server, mcp_auth_header=None, extra_headers=None):
# Capture the arguments for verification
captured_client_args.update(
{
"server": server,
"mcp_auth_header": mcp_auth_header,
"extra_headers": extra_headers,
}
)
# Return a mock client that doesn't actually connect
mock_client = MagicMock()
mock_client.disconnect = AsyncMock()
return mock_client
# Mock _fetch_tools_with_timeout to avoid actual network calls
async def mock_fetch_tools_with_timeout(client, server_name):
return [] # Return empty list of tools
with patch.object(
global_mcp_server_manager,
"_create_mcp_client",
side_effect=mock_create_mcp_client,
) as mock_create_client, patch.object(
global_mcp_server_manager,
"_fetch_tools_with_timeout",
side_effect=mock_fetch_tools_with_timeout,
), patch.object(
global_mcp_server_manager,
"get_allowed_mcp_servers",
AsyncMock(return_value=[oauth2_server.server_id]),
):
# Call _get_tools_from_mcp_servers which should eventually call _create_mcp_client
await _get_tools_from_mcp_servers(
user_api_key_auth=user_api_key_auth,
mcp_auth_header=None,
mcp_servers=None, # Will use all allowed servers
oauth2_headers=oauth2_headers,
)
# Verify that _create_mcp_client was called
assert (
mock_create_client.call_count == 1
), "Expected _create_mcp_client to be called once"
# Verify the server passed to _create_mcp_client is the OAuth2 server
assert captured_client_args["server"].server_id == oauth2_server.server_id
assert captured_client_args["server"].auth_type == MCPAuth.oauth2
# Most importantly: verify that OAuth2 headers were passed as extra_headers
assert (
captured_client_args["extra_headers"] is not None
), "Expected extra_headers to be passed for OAuth2 server"
assert (
captured_client_args["extra_headers"] == oauth2_headers
), f"Expected OAuth2 headers to be passed as extra_headers, got {captured_client_args['extra_headers']}"
# Verify the Authorization header specifically
assert "Authorization" in captured_client_args["extra_headers"]
assert (
captured_client_args["extra_headers"]["Authorization"]
== "Bearer github_oauth_token_12345"
)