diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index ffca305a0d..e2c6500b98 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -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 diff --git a/.github/workflows/test-litellm.yml b/.github/workflows/test-litellm.yml index 0d3a9f2b5d..b7f4a25d59 100644 --- a/.github/workflows/test-litellm.yml +++ b/.github/workflows/test-litellm.yml @@ -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 diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 07110b938d..6716cfe664 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -584,6 +584,7 @@ const sidebars = { "budget_manager", "caching/all_caches", "completion/token_usage", + "sdk_custom_pricing", "embedding/async_embedding", "embedding/moderation", "migration", diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index ecb58e1822..2e02f460b6 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -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]: 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 058f45d712..56a22040f0 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py +++ b/litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py @@ -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 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 6afed97fe9..281edf38ad 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,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") diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py new file mode 100644 index 0000000000..9c4ab70f0d --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -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( + "Authentication incomplete. You can close this window." + ) + + +# ------------------------------ +# 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"], + } diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index ab5d1b10bf..9fb39e7498 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -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 diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 2a9174717d..ef770a9d43 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -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: diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 2cf91c84dc..38b0cfd99a 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -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 ############# diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 8ffe5e8aef..c50d391e05 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -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"] + diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 6dea634a80..63047f37b1 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -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( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f51b65c915..6d2e56c444 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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) diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index 4cd32915ce..bbb9e7d7b4 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -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 \ No newline at end of file + + mcp_tool_call_response: List[ + Union[MCPTextContent, MCPImageContent, MCPEmbeddedResource] + ] + hidden_params: HiddenParams diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index bd09ef9c19..0dec1b23c6 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -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 diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index e55fdc8516..d0489bac58 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -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", 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 2a5cc7bcb7..afafa510a4 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 @@ -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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 8fa9964b1f..10a6ab8cb0 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -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" + )