From 342e80d48d030c8d20748b7f33fd71990e5364c8 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 26 Sep 2025 15:25:42 -0700 Subject: [PATCH 01/18] feat: initial commit for v2 oauth flow --- .../mcp_server/auth/user_api_key_auth_mcp.py | 27 +-- .../mcp_server/discoverable_endpoints.py | 204 ++++++++++++++++++ .../mcp_server/mcp_server_manager.py | 58 +++-- .../mcp_server/rest_endpoints.py | 119 +++++++++- .../proxy/_experimental/mcp_server/server.py | 50 +++-- .../index.html} | 0 .../proxy/_experimental/out/onboarding.html | 1 - litellm/proxy/_new_secret_config.yaml | 13 ++ litellm/proxy/auth/user_api_key_auth.py | 47 ++-- litellm/proxy/proxy_server.py | 83 ++++++- litellm/types/mcp.py | 20 +- .../types/mcp_server/mcp_server_manager.py | 7 +- 12 files changed, 544 insertions(+), 85 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py rename litellm/proxy/_experimental/out/{model_hub_table.html => model_hub_table/index.html} (100%) delete mode 100644 litellm/proxy/_experimental/out/onboarding.html 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..4e7745392b 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 @@ -96,9 +96,12 @@ 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, @@ -359,10 +362,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 +538,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..d38edf5dae --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -0,0 +1,204 @@ +from typing import Optional +from urllib.parse import urlencode, urlparse, urlunparse + +import httpx +from fastapi import APIRouter, Form, HTTPException, Request +from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse + +STATE_MAP = {} + +router = APIRouter( + tags=["mcp"], +) + + +@router.get("/authorize") +async def authorize( + request: Request, client_id: str, redirect_uri: str, state: str = "" +): + # 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("/") + STATE_MAP[state] = base_url + params = { + "client_id": mcp_server.client_id, + "redirect_uri": f"{request_base_url}/callback", + "scope": " ".join(mcp_server.scopes), + "state": 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 with httpx.AsyncClient() as client: + resp = await 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", + }, + ) + resp.raise_for_status() + github_token = resp.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): + # Exchange code for token with GitHub + params = {"code": code, "state": state} + + # Forward token to Claude ephemeral endpoint + redirect_uri = STATE_MAP.pop(state, None) + + complete_returned_url = f"{redirect_uri}?{urlencode(params)}" + if redirect_uri: + async with httpx.AsyncClient() as client: + await client.post( + complete_returned_url, + json={}, + ) + # Return simple page so the browser doesn't hang + return HTMLResponse( + "Authentication complete. You can close this window." + ) + + # fallback if redirect_uri 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..254f036df0 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -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), @@ -607,20 +613,26 @@ class MCPServerManager: "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, + "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 @@ -834,6 +846,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 +1045,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..332d1d4e05 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1,12 +1,17 @@ import importlib from typing import Dict, List, Optional +from urllib.parse import urlencode, urlparse, urlunparse -from fastapi import APIRouter, Depends, Query, Request +import httpx +from fastapi import APIRouter, Depends, HTTPException, Query, Request +from fastapi.params import Form +from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse from litellm._logging import verbose_logger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +STATE_MAP = {} MCP_AVAILABLE: bool = True try: importlib.import_module("mcp") @@ -177,9 +182,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: @@ -333,3 +338,109 @@ if MCP_AVAILABLE: } return await _execute_with_mcp_client(request, _list_tools_operation) + + @router.get("/authorize") + async def authorize( + request: Request, client_id: str, redirect_uri: str, state: str = "" + ): + # Redirect to real GitHub OAuth + from mcp_server_manager import global_mcp_server_manager + + mcp_server = global_mcp_server_manager.get_mcp_server_by_id(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") + + request_base_url = str(request.base_url).rstrip("/") + # Parse it to remove any existing query + parsed = urlparse(redirect_uri) + base_url = urlunparse(parsed._replace(query="")) + STATE_MAP[state] = base_url + params = { + "client_id": mcp_server.client_id, + "redirect_uri": f"{request_base_url}/callback", + "scope": " ".join(mcp_server.scopes), + "state": state, + } + return RedirectResponse(f"{mcp_server.authorization_url}?{urlencode(params)}") + + @router.post("/token") + async def token_endpoint( + request: Request, + grant_type=Form(...), + code=Form(None), + client_id=Form(...), + ): + """ + Accept the authorization code from Claude and exchange it for GitHub token. + Forward the GitHub token back to Claude in standard OAuth format. + """ + from mcp_server_manager import global_mcp_server_manager + + mcp_server = global_mcp_server_manager.get_mcp_server_by_id(str(client_id)) + if mcp_server is None: + raise HTTPException(status_code=404, detail="MCP server not found") + + if mcp_server.token_url is None: + raise HTTPException( + status_code=400, detail="MCP server token url is not set" + ) + + if grant_type != "authorization_code": + raise HTTPException(status_code=400, detail="Unsupported grant_type") + + request_base_url = str(request.base_url).rstrip("/") + # Exchange code for real GitHub token + async with httpx.AsyncClient() as client: + resp = await 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"{request_base_url}/callback", + }, + ) + resp.raise_for_status() + return JSONResponse(resp.json()) + + @router.get("/callback") + async def callback(code: str, state: str): + # Exchange code for token with GitHub + params = {"code": code, "state": state} + + # Forward token to Claude ephemeral endpoint + redirect_uri = STATE_MAP.pop(state, None) + + complete_returned_url = f"{redirect_uri}?{urlencode(params)}" + if redirect_uri: + async with httpx.AsyncClient() as client: + await client.post( + complete_returned_url, + json={}, + ) + # Return simple page so the browser doesn't hang + return HTMLResponse( + "Authentication complete. You can close this window." + ) + + # fallback if redirect_uri not found + return HTMLResponse( + "Authentication incomplete. You can close this window." + ) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 2cf91c84dc..f02b079196 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -519,16 +519,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 {} @@ -638,26 +638,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] @@ -809,6 +815,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) @@ -839,14 +847,12 @@ if MCP_AVAILABLE: ) 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]], + ]: """ Get the UserAPIKeyAuth from the auth context variable. diff --git a/litellm/proxy/_experimental/out/model_hub_table.html b/litellm/proxy/_experimental/out/model_hub_table/index.html similarity index 100% rename from litellm/proxy/_experimental/out/model_hub_table.html rename to litellm/proxy/_experimental/out/model_hub_table/index.html diff --git a/litellm/proxy/_experimental/out/onboarding.html b/litellm/proxy/_experimental/out/onboarding.html deleted file mode 100644 index 1dc5fdb1f8..0000000000 --- a/litellm/proxy/_experimental/out/onboarding.html +++ /dev/null @@ -1 +0,0 @@ -LiteLLM Dashboard \ No newline at end of file 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..9edfe677ca 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -390,19 +390,18 @@ 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, - api_key=api_key, - azure_api_key_header=azure_api_key_header, - anthropic_api_key_header=anthropic_api_key_header, - google_ai_studio_api_key_header=google_ai_studio_api_key_header, - azure_apim_header=azure_apim_header, - pass_through_endpoints=pass_through_endpoints, - route=route, - request=request, - ) + # api_key, passed_in_key = get_api_key( + # custom_litellm_key_header=custom_litellm_key_header, + # api_key=api_key, + # azure_api_key_header=azure_api_key_header, + # anthropic_api_key_header=anthropic_api_key_header, + # google_ai_studio_api_key_header=google_ai_studio_api_key_header, + # azure_apim_header=azure_apim_header, + # pass_through_endpoints=pass_through_endpoints, + # route=route, + # request=request, + # ) # if user wants to pass LiteLLM_Master_Key as a custom header, example pass litellm keys as X-LiteLLM-Key: Bearer sk-1234 custom_litellm_key_header_name = general_settings.get("litellm_key_header_name") if custom_litellm_key_header_name is not None: @@ -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 472aeb140c..628af159ed 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -35,13 +35,13 @@ from litellm.constants import ( LITELLM_SETTINGS_SAFE_DB_OVERRIDES, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.utils import load_credentials_from_list from litellm.types.utils import ( ModelResponse, ModelResponseStream, TextCompletionResponse, TokenCountResponse, ) +from litellm.utils import load_credentials_from_list if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -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, ) @@ -9470,5 +9469,75 @@ 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, + ) + + 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 From 667df3c3db1423218492bd4800ad48366049c64b Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 26 Sep 2025 17:21:42 -0700 Subject: [PATCH 02/18] fix: revert commit --- litellm/proxy/auth/user_api_key_auth.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 9edfe677ca..63047f37b1 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -391,17 +391,17 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 "pass_through_endpoints", 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, - # api_key=api_key, - # azure_api_key_header=azure_api_key_header, - # anthropic_api_key_header=anthropic_api_key_header, - # google_ai_studio_api_key_header=google_ai_studio_api_key_header, - # azure_apim_header=azure_apim_header, - # pass_through_endpoints=pass_through_endpoints, - # route=route, - # request=request, - # ) + api_key, passed_in_key = get_api_key( + custom_litellm_key_header=custom_litellm_key_header, + api_key=api_key, + azure_api_key_header=azure_api_key_header, + anthropic_api_key_header=anthropic_api_key_header, + google_ai_studio_api_key_header=google_ai_studio_api_key_header, + azure_apim_header=azure_apim_header, + pass_through_endpoints=pass_through_endpoints, + route=route, + request=request, + ) # if user wants to pass LiteLLM_Master_Key as a custom header, example pass litellm keys as X-LiteLLM-Key: Bearer sk-1234 custom_litellm_key_header_name = general_settings.get("litellm_key_header_name") if custom_litellm_key_header_name is not None: From 0d36e6cbe9e5def82e1a7be5d9f293568b6d5728 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 26 Sep 2025 18:13:12 -0700 Subject: [PATCH 03/18] fix(discoverable_endpoints.py): redirect to redirect uri --- .../mcp_server/discoverable_endpoints.py | 20 +++++++++---------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index d38edf5dae..513b590b8d 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -12,9 +12,14 @@ router = APIRouter( ) +@router.get("/{mcp_server_name}/authorize") @router.get("/authorize") async def authorize( - request: Request, client_id: str, redirect_uri: str, state: str = "" + 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 ( @@ -115,17 +120,10 @@ async def callback(code: str, state: str): # Forward token to Claude ephemeral endpoint redirect_uri = STATE_MAP.pop(state, None) - complete_returned_url = f"{redirect_uri}?{urlencode(params)}" if redirect_uri: - async with httpx.AsyncClient() as client: - await client.post( - complete_returned_url, - json={}, - ) - # Return simple page so the browser doesn't hang - return HTMLResponse( - "Authentication complete. You can close this window." - ) + complete_returned_url = f"{redirect_uri}?{urlencode(params)}" + + return RedirectResponse(url=complete_returned_url, status_code=302) # fallback if redirect_uri not found return HTMLResponse( From 1c939c70e58c049203f9214a2ca419023274f114 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 26 Sep 2025 18:39:10 -0700 Subject: [PATCH 04/18] feat(discoverable_endpoints.py): use encryption + encoding to securely handle state + redirect uri without storing in db would needlessly flood the db --- .../mcp_server/discoverable_endpoints.py | 74 +++++++++++++++---- 1 file changed, 60 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 513b590b8d..05dae916af 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -1,3 +1,4 @@ +import json from typing import Optional from urllib.parse import urlencode, urlparse, urlunparse @@ -5,13 +6,54 @@ import httpx from fastapi import APIRouter, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse -STATE_MAP = {} +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( @@ -44,12 +86,15 @@ async def authorize( parsed = urlparse(redirect_uri) base_url = urlunparse(parsed._replace(query="")) request_base_url = str(request.base_url).rstrip("/") - STATE_MAP[state] = base_url + + # 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": state, + "state": encoded_state, } return RedirectResponse(f"{mcp_server.authorization_url}?{urlencode(params)}") @@ -114,21 +159,22 @@ async def token_endpoint( @router.get("/callback") async def callback(code: str, state: str): - # Exchange code for token with GitHub - params = {"code": code, "state": state} + try: + # Decode the state hash to get base_url and original state + base_url, original_state = decode_state_hash(state) - # Forward token to Claude ephemeral endpoint - redirect_uri = STATE_MAP.pop(state, None) - - if redirect_uri: - complete_returned_url = f"{redirect_uri}?{urlencode(params)}" + # 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) - # fallback if redirect_uri not found - return HTMLResponse( - "Authentication incomplete. You can close this window." - ) + except Exception: + # fallback if state hash not found + return HTMLResponse( + "Authentication incomplete. You can close this window." + ) # ------------------------------ From c73f4914de49d46256923305836f1512d05a4c87 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 26 Sep 2025 19:35:05 -0700 Subject: [PATCH 05/18] feat(mcp/): forward oauth2 headers, when applicable if server is of auth type oauth, forward along the oauth headers Allows github mcp on litellm to work via claude code --- litellm/experimental_mcp_client/client.py | 14 ++++--- .../mcp_server/auth/litellm_auth_handler.py | 16 ++++++-- .../mcp_server/auth/user_api_key_auth_mcp.py | 22 ++++++++++- .../mcp_server/mcp_server_manager.py | 14 ++++++- .../proxy/_experimental/mcp_server/server.py | 37 +++++++++++++++++-- litellm/proxy/proxy_server.py | 1 + 6 files changed, 90 insertions(+), 14 deletions(-) diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index ecb58e1822..d4db9251fc 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) @@ -186,9 +188,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 +200,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 4e7745392b..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 @@ -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 @@ -107,6 +115,7 @@ class MCPRequestHandler: mcp_auth_header, mcp_servers, mcp_server_auth_headers, + oauth2_headers, ) @staticmethod @@ -177,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: """ diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 254f036df0..bc69314bcc 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 @@ -382,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. @@ -411,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 @@ -421,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. @@ -447,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) @@ -564,6 +569,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) @@ -684,9 +690,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: diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index f02b079196..f79d40ca61 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( @@ -357,15 +361,17 @@ 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]: - """ + f""" Helper method to fetch tools from MCP servers based on server filtering criteria. Args: 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 +404,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 +416,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 +438,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 +462,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" @@ -683,6 +696,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: @@ -691,8 +705,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 @@ -705,6 +726,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}" @@ -718,6 +740,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 @@ -756,6 +779,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}" @@ -768,6 +792,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: @@ -829,6 +854,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. @@ -844,6 +870,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, ) auth_context_var.set(auth_user) @@ -852,6 +879,7 @@ if MCP_AVAILABLE: Optional[str], Optional[List[str]], Optional[Dict[str, str]], + Optional[Dict[str, str]], ]: """ Get the UserAPIKeyAuth from the auth context variable. @@ -867,8 +895,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/proxy_server.py b/litellm/proxy/proxy_server.py index 628af159ed..42bbae77f4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -9483,6 +9483,7 @@ async def dynamic_mcp_route(mcp_server_name: str, request: Request): 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: From 47ef6a7ee676064a02f21148fd2719ea1e7ae1bf Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 09:02:08 -0700 Subject: [PATCH 06/18] fix(mcp/): pass oauth2 headers on call tools --- litellm/experimental_mcp_client/client.py | 3 +++ litellm/proxy/_experimental/mcp_server/server.py | 5 +++++ 2 files changed, 8 insertions(+) diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index d4db9251fc..2e02f460b6 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -117,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), diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index f79d40ca61..7f16179390 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -270,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: @@ -505,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]]: """ @@ -552,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, ) @@ -604,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""" @@ -616,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) From 0202f1d034ea85f2342cef4d92087d65bac47e2b Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 09:22:30 -0700 Subject: [PATCH 07/18] refactor: cleanup dup code --- .../mcp_server/rest_endpoints.py | 107 ------------------ 1 file changed, 107 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 332d1d4e05..d2f16154f3 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -11,7 +11,6 @@ from litellm._logging import verbose_logger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -STATE_MAP = {} MCP_AVAILABLE: bool = True try: importlib.import_module("mcp") @@ -338,109 +337,3 @@ if MCP_AVAILABLE: } return await _execute_with_mcp_client(request, _list_tools_operation) - - @router.get("/authorize") - async def authorize( - request: Request, client_id: str, redirect_uri: str, state: str = "" - ): - # Redirect to real GitHub OAuth - from mcp_server_manager import global_mcp_server_manager - - mcp_server = global_mcp_server_manager.get_mcp_server_by_id(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") - - request_base_url = str(request.base_url).rstrip("/") - # Parse it to remove any existing query - parsed = urlparse(redirect_uri) - base_url = urlunparse(parsed._replace(query="")) - STATE_MAP[state] = base_url - params = { - "client_id": mcp_server.client_id, - "redirect_uri": f"{request_base_url}/callback", - "scope": " ".join(mcp_server.scopes), - "state": state, - } - return RedirectResponse(f"{mcp_server.authorization_url}?{urlencode(params)}") - - @router.post("/token") - async def token_endpoint( - request: Request, - grant_type=Form(...), - code=Form(None), - client_id=Form(...), - ): - """ - Accept the authorization code from Claude and exchange it for GitHub token. - Forward the GitHub token back to Claude in standard OAuth format. - """ - from mcp_server_manager import global_mcp_server_manager - - mcp_server = global_mcp_server_manager.get_mcp_server_by_id(str(client_id)) - if mcp_server is None: - raise HTTPException(status_code=404, detail="MCP server not found") - - if mcp_server.token_url is None: - raise HTTPException( - status_code=400, detail="MCP server token url is not set" - ) - - if grant_type != "authorization_code": - raise HTTPException(status_code=400, detail="Unsupported grant_type") - - request_base_url = str(request.base_url).rstrip("/") - # Exchange code for real GitHub token - async with httpx.AsyncClient() as client: - resp = await 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"{request_base_url}/callback", - }, - ) - resp.raise_for_status() - return JSONResponse(resp.json()) - - @router.get("/callback") - async def callback(code: str, state: str): - # Exchange code for token with GitHub - params = {"code": code, "state": state} - - # Forward token to Claude ephemeral endpoint - redirect_uri = STATE_MAP.pop(state, None) - - complete_returned_url = f"{redirect_uri}?{urlencode(params)}" - if redirect_uri: - async with httpx.AsyncClient() as client: - await client.post( - complete_returned_url, - json={}, - ) - # Return simple page so the browser doesn't hang - return HTMLResponse( - "Authentication complete. You can close this window." - ) - - # fallback if redirect_uri not found - return HTMLResponse( - "Authentication incomplete. You can close this window." - ) From ac09d29c0ed5748a0ee3cee94c078ac8b35cc5fa Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 09:30:07 -0700 Subject: [PATCH 08/18] test(tests/): add unit test to confirm oauth2 headers forwarded on receipt --- .../mcp_server/test_mcp_server.py | 120 +++++++++++++++++- 1 file changed, 118 insertions(+), 2 deletions(-) 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" + ) From 6451f10f2c9f685fde5cd2b463bb2bd2e79bb3ab Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 09:34:01 -0700 Subject: [PATCH 09/18] fix: fix ruff check --- .../mcp_server/mcp_server_manager.py | 140 ++++++++++-------- .../mcp_server/rest_endpoints.py | 6 +- .../proxy/_experimental/mcp_server/server.py | 2 +- 3 files changed, 79 insertions(+), 69 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index bc69314bcc..9fb39e7498 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -561,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, @@ -614,71 +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: diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index d2f16154f3..ef770a9d43 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1,11 +1,7 @@ import importlib from typing import Dict, List, Optional -from urllib.parse import urlencode, urlparse, urlunparse -import httpx -from fastapi import APIRouter, Depends, HTTPException, Query, Request -from fastapi.params import Form -from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse +from fastapi import APIRouter, Depends, Query, Request from litellm._logging import verbose_logger from litellm.proxy._types import UserAPIKeyAuth diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 7f16179390..38b0cfd99a 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -364,7 +364,7 @@ if MCP_AVAILABLE: mcp_server_auth_headers: Optional[Dict[str, str]] = None, oauth2_headers: Optional[Dict[str, str]] = None, ) -> List[MCPTool]: - f""" + """ Helper method to fetch tools from MCP servers based on server filtering criteria. Args: From 3a31e9cfa0766a5347db5aa4d2308f23249f17d3 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 09:39:13 -0700 Subject: [PATCH 10/18] fix: fix perf issue --- .../mcp_server/discoverable_endpoints.py | 31 +++++++++++-------- 1 file changed, 18 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 05dae916af..af586098be 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -6,6 +6,10 @@ import httpx 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, @@ -134,19 +138,20 @@ async def token_endpoint( proxy_base_url = str(request.base_url).rstrip("/") # Exchange code for real GitHub token - async with httpx.AsyncClient() as client: - resp = await 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", - }, - ) - resp.raise_for_status() - github_token = resp.json()["access_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 From 5787f868612ceb259b8aaf6be59646b6983f85c0 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 09:59:47 -0700 Subject: [PATCH 11/18] fix: fix ruff check --- litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index af586098be..9c4ab70f0d 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -2,7 +2,6 @@ import json from typing import Optional from urllib.parse import urlencode, urlparse, urlunparse -import httpx from fastapi import APIRouter, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse From 7278ae5305f407b6fcafc814fc7f49dcac85c142 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 10:03:54 -0700 Subject: [PATCH 12/18] test: fix tests --- .../_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py | 1 + 1 file changed, 1 insertion(+) 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..fb15dbc14e 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 @@ -543,6 +543,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"] From a569e4d7114e1f563e8a8f9dd30848677b4e8647 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 10:23:46 -0700 Subject: [PATCH 13/18] test: update test to not stop on first error (allows us to see all issues) --- .github/workflows/test-linting.yml | 2 +- .github/workflows/test-litellm.yml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index ffca305a0d..6162c53f95 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -44,7 +44,7 @@ jobs: - name: Run MyPy type checking run: | cd litellm - poetry run mypy . --ignore-missing-imports + poetry run mypy . --config-file=mypy.ini cd .. - name: Check for circular imports diff --git a/.github/workflows/test-litellm.yml b/.github/workflows/test-litellm.yml index 0d3a9f2b5d..b6b77ad326 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 -vv -n 4 From 5d6889dc928f7038fc27dcaac038b85a2b53b26e Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 10:40:06 -0700 Subject: [PATCH 14/18] test: update tests --- tests/mcp_tests/test_mcp_server.py | 42 +++++++++++++++--------------- 1 file changed, 21 insertions(+), 21 deletions(-) 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", From 82eeca78708543efef008a1dd4989a36a89e981e Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 10:49:21 -0700 Subject: [PATCH 15/18] test: update test --- .github/workflows/test-linting.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index 6162c53f95..7e169d6a27 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -44,7 +44,7 @@ jobs: - name: Run MyPy type checking run: | cd litellm - poetry run mypy . --config-file=mypy.ini + poetry run mypy . --ignore-missing-imports --disable-error-code=var-annotated cd .. - name: Check for circular imports From 09556d0a44ac1c6464d465cdbc5a48e4e73ec7ac Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 10:58:14 -0700 Subject: [PATCH 16/18] test: add debugging step to linting --- .github/workflows/test-linting.yml | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index 7e169d6a27..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 From 5cdeb63cdd6fe39ef56d0d643ed85d3d4a7fc10c Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 11:09:25 -0700 Subject: [PATCH 17/18] test: fix tests --- .github/workflows/test-litellm.yml | 2 +- .../mcp_server/auth/test_user_api_key_auth_mcp.py | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test-litellm.yml b/.github/workflows/test-litellm.yml index b6b77ad326..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 -vv -n 4 + poetry run pytest tests/test_litellm --tb=short -vv --maxfail=10 -n 4 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 fb15dbc14e..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 @@ -717,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 @@ -887,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 @@ -936,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 @@ -964,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 From 5c6a863e40d6ebd8d1920b138bfd0ced4ac2ca50 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 27 Sep 2025 11:10:44 -0700 Subject: [PATCH 18/18] docs: cleanup --- docs/my-website/sidebars.js | 1 - 1 file changed, 1 deletion(-) diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index a131e5c34e..47289ca1d4 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -562,7 +562,6 @@ const sidebars = { items: [ "set_keys", "completion/token_usage", - "sdk/headers", "sdk_custom_pricing", "embedding/async_embedding", "embedding/moderation",