feat(mcp/): initial commit raising correct oauth error

This commit is contained in:
Krrish Dholakia
2025-10-08 17:33:11 -07:00
parent ecb6c0e814
commit 0b27c361fe
6 changed files with 108 additions and 79 deletions
@@ -5,7 +5,12 @@ from starlette.requests import Request
from starlette.types import Scope
from litellm._logging import verbose_logger
from litellm.proxy._types import LiteLLM_TeamTable, SpecialHeaders, UserAPIKeyAuth
from litellm.proxy._types import (
LiteLLM_TeamTable,
LitellmUserRoles,
SpecialHeaders,
UserAPIKeyAuth,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
@@ -109,10 +114,21 @@ class MCPRequestHandler:
request.body = mock_body # type: ignore
if ".well-known" in str(request.url): # public routes
validated_user_api_key_auth = UserAPIKeyAuth()
elif litellm_api_key == "":
from fastapi import HTTPException
raise HTTPException(
status_code=401,
detail="LiteLLM API key is missing. Please add it or use OAuth authentication.",
headers={
"WWW-Authenticate": f'Bearer resource_metadata=f"{request.base_url}/.well-known/oauth-protected-resource"',
},
)
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,
@@ -344,14 +360,14 @@ class MCPRequestHandler:
proxy_logging_obj,
user_api_key_cache,
)
if not user_api_key_auth:
return None
# Already loaded
if user_api_key_auth.object_permission:
return user_api_key_auth.object_permission
# Need to fetch from DB
if user_api_key_auth.object_permission_id and prisma_client:
return await get_object_permission(
@@ -361,7 +377,7 @@ class MCPRequestHandler:
parent_otel_span=user_api_key_auth.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
return None
@staticmethod
@@ -369,16 +385,19 @@ class MCPRequestHandler:
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
):
"""Helper to get team object_permission from cache or DB."""
from litellm.proxy.auth.auth_checks import get_object_permission, get_team_object
from litellm.proxy.auth.auth_checks import (
get_object_permission,
get_team_object,
)
from litellm.proxy.proxy_server import (
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
if not user_api_key_auth or not user_api_key_auth.team_id or not prisma_client:
return None
# First get the team object (which may have object_permission already loaded)
team_obj: Optional[LiteLLM_TeamTable] = await get_team_object(
team_id=user_api_key_auth.team_id,
@@ -387,14 +406,14 @@ class MCPRequestHandler:
parent_otel_span=user_api_key_auth.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
if not team_obj:
return None
# Already loaded
if team_obj.object_permission:
return team_obj.object_permission
# Need to fetch from DB using object_permission_id
if team_obj.object_permission_id:
return await get_object_permission(
@@ -404,7 +423,7 @@ class MCPRequestHandler:
parent_otel_span=user_api_key_auth.parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
return None
@staticmethod
@@ -415,26 +434,38 @@ class MCPRequestHandler:
"""
Get list of allowed tool names for a specific server based on key/team permissions.
Follows same inheritance logic as get_allowed_mcp_servers.
Args:
server_id: Server ID to check permissions for
user_api_key_auth: User auth
Returns:
List[str] if restrictions exist, None if no restrictions (allow all)
"""
if not user_api_key_auth:
return None
try:
# Get key and team object permissions
key_obj_perm = await MCPRequestHandler._get_key_object_permission(user_api_key_auth)
team_obj_perm = await MCPRequestHandler._get_team_object_permission(user_api_key_auth)
key_obj_perm = await MCPRequestHandler._get_key_object_permission(
user_api_key_auth
)
team_obj_perm = await MCPRequestHandler._get_team_object_permission(
user_api_key_auth
)
# Extract tool permissions for this server
key_tools = key_obj_perm.mcp_tool_permissions.get(server_id) if key_obj_perm and key_obj_perm.mcp_tool_permissions else None
team_tools = team_obj_perm.mcp_tool_permissions.get(server_id) if team_obj_perm and team_obj_perm.mcp_tool_permissions else None
key_tools = (
key_obj_perm.mcp_tool_permissions.get(server_id)
if key_obj_perm and key_obj_perm.mcp_tool_permissions
else None
)
team_tools = (
team_obj_perm.mcp_tool_permissions.get(server_id)
if team_obj_perm and team_obj_perm.mcp_tool_permissions
else None
)
# Apply same inheritance logic as get_allowed_mcp_servers
if team_tools:
if key_tools:
@@ -446,7 +477,7 @@ class MCPRequestHandler:
else:
# No team restrictions → use key restrictions
return key_tools
except Exception as e:
verbose_logger.warning(f"Failed to get allowed tools for server: {str(e)}")
return None
@@ -459,12 +490,12 @@ class MCPRequestHandler:
) -> bool:
"""
Check if a specific tool is allowed for a server based on key/team permissions.
Args:
tool_name: Name of the tool to check
server_id: Server ID
user_api_key_auth: User auth
Returns:
True if allowed, False if blocked
"""
@@ -472,15 +503,15 @@ class MCPRequestHandler:
server_id=server_id,
user_api_key_auth=user_api_key_auth,
)
# None means no restrictions (allow all)
if allowed_tools is None:
return True
# Empty list means no tools allowed
if not allowed_tools:
return False
# Check if tool is in allowed list
return tool_name in allowed_tools
@@ -555,7 +586,7 @@ class MCPRequestHandler:
) -> List[str]:
"""
Get allowed MCP servers for a team.
Uses the helper _get_team_object_permission which:
1. First checks if object_permission is already loaded on the team
2. If not, fetches from DB using object_permission_id if it exists
@@ -571,7 +602,7 @@ class MCPRequestHandler:
object_permissions = await MCPRequestHandler._get_team_object_permission(
user_api_key_auth
)
if object_permissions is None:
return []
@@ -364,25 +364,25 @@ if MCP_AVAILABLE:
def _tool_name_matches(tool_name: str, filter_list: List[str]) -> bool:
"""
Check if a tool name matches any name in the filter list.
Checks both the full tool name and unprefixed version (without server prefix).
This allows users to configure simple tool names regardless of prefixing.
Args:
tool_name: The tool name to check (may be prefixed like "server-tool_name")
filter_list: List of tool names to match against
Returns:
True if the tool name (prefixed or unprefixed) is in the filter list
"""
from litellm.proxy._experimental.mcp_server.utils import (
get_server_name_prefix_tool_mcp,
)
# Check if the full name is in the list
if tool_name in filter_list:
return True
# Check if the unprefixed name is in the list
unprefixed_name, _ = get_server_name_prefix_tool_mcp(tool_name)
return unprefixed_name in filter_list
@@ -393,34 +393,36 @@ if MCP_AVAILABLE:
) -> List[MCPTool]:
"""
Filter tools by allowed/disallowed tools configuration.
If allowed_tools is set, only tools in that list are returned.
If disallowed_tools is set, tools in that list are excluded.
Tool names are matched with and without server prefixes for flexibility.
Args:
tools: List of tools to filter
mcp_server: Server configuration with allowed_tools/disallowed_tools
Returns:
Filtered list of tools
"""
tools_to_return = tools
# Filter by allowed_tools (whitelist)
if mcp_server.allowed_tools:
tools_to_return = [
tool for tool in tools
tool
for tool in tools
if _tool_name_matches(tool.name, mcp_server.allowed_tools)
]
# Filter by disallowed_tools (blacklist)
if mcp_server.disallowed_tools:
tools_to_return = [
tool for tool in tools_to_return
tool
for tool in tools_to_return
if not _tool_name_matches(tool.name, mcp_server.disallowed_tools)
]
return tools_to_return
async def _get_tools_from_mcp_servers(
@@ -497,17 +499,17 @@ if MCP_AVAILABLE:
extra_headers=extra_headers,
add_prefix=add_prefix,
)
filtered_tools = filter_tools_by_allowed_tools(tools, server)
filtered_tools = await filter_tools_by_key_team_permissions(
tools=filtered_tools,
server_id=server_id,
user_api_key_auth=user_api_key_auth,
)
all_tools.extend(filtered_tools)
verbose_logger.debug(
f"Successfully fetched {len(tools)} tools from server {server.name}, {len(filtered_tools)} after filtering"
)
@@ -529,7 +531,7 @@ if MCP_AVAILABLE:
) -> List[MCPTool]:
"""
Filter tools based on key/team mcp_tool_permissions.
Note: Tool names in the DB are stored without server prefixes,
but tool names from MCP servers are prefixed. We need to strip
the prefix before comparing.
@@ -551,7 +553,7 @@ if MCP_AVAILABLE:
else:
# No restrictions, return all tools
filtered_tools = tools
return filtered_tools
async def _list_mcp_tools(
@@ -906,6 +908,7 @@ if MCP_AVAILABLE:
await session_manager.handle_request(scope, receive, send)
except Exception as e:
raise e
verbose_logger.exception(f"Error handling MCP request: {e}")
# Instead of re-raising, try to send a graceful error response
try:
File diff suppressed because one or more lines are too long
+9 -17
View File
@@ -17,23 +17,15 @@ model_list:
api_key: dummy
mcp_servers:
my_api_mcp:
url: "http://0.0.0.0:8090"
spec_path: "/Users/krrishdholakia/Documents/temp_py_folder/example_openapi.json"
auth_type: none
allowed_tools: ["getpetbyid", "my_api_mcp-findpetsbystatus"]
# Configure allowed parameters per tool
# Key: tool name (with or without prefix)
# Value: list of allowed parameter names
allowed_params:
# Using unprefixed tool name
"getpetbyid": ["petId"]
# Using prefixed tool name (both formats work)
"my_api_mcp-findpetsbystatus": ["status", "limit"]
# Example: allow only specific params for another tool
# "another_tool": ["param1", "param2"]
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"]
litellm_settings:
callbacks: ["prometheus"]
custom_prometheus_metadata_labels: ["metadata.initiative", "metadata.business-unit"]
+17 -13
View File
@@ -1883,7 +1883,9 @@ class ProxyConfig:
elif key == "priority_reservation_settings":
from litellm.types.utils import PriorityReservationSettings
litellm.priority_reservation_settings = PriorityReservationSettings(**value)
litellm.priority_reservation_settings = PriorityReservationSettings(
**value
)
elif key == "callbacks":
initialize_callbacks_on_proxy(
value=value,
@@ -2966,32 +2968,32 @@ class ProxyConfig:
) -> bool:
"""
Check if an object type should be loaded from the database based on general_settings.supported_db_objects.
Args:
object_type: Type of object to check (e.g., SupportedDBObjectType.MODELS, "models", etc.)
Returns:
True if the object should be loaded, False otherwise
"""
global general_settings
# Get the supported_db_objects configuration
supported_db_objects = general_settings.get("supported_db_objects", None)
# If supported_db_objects is not set, load all objects (default behavior)
if supported_db_objects is None:
return True
# If supported_db_objects is set, only load specified objects
if not isinstance(supported_db_objects, list):
verbose_proxy_logger.warning(
f"supported_db_objects is not a list, got {type(supported_db_objects)}. Loading all objects."
)
return True
# Convert object_type to string for comparison (handles both str and enum)
object_type_str = str(object_type)
# Check if the object type is in the list (supports both str and enum values)
return any(str(obj) == object_type_str for obj in supported_db_objects)
@@ -3063,19 +3065,19 @@ class ProxyConfig:
"""
if self._should_load_db_object(object_type="guardrails"):
await self._init_guardrails_in_db(prisma_client=prisma_client)
if self._should_load_db_object(object_type="vector_stores"):
await self._init_vector_stores_in_db(prisma_client=prisma_client)
if self._should_load_db_object(object_type="mcp"):
await self._init_mcp_servers_in_db()
if self._should_load_db_object(object_type="pass_through_endpoints"):
await self._init_pass_through_endpoints_in_db()
if self._should_load_db_object(object_type="prompts"):
await self._init_prompts_in_db(prisma_client=prisma_client)
if self._should_load_db_object(object_type="model_cost_map"):
await self._check_and_reload_model_cost_map(prisma_client=prisma_client)
@@ -9708,6 +9710,8 @@ async def dynamic_mcp_route(mcp_server_name: str, request: Request):
media_type=headers_dict.get("content-type", "application/json"),
)
except HTTPException as e:
raise e
except Exception as e:
verbose_proxy_logger.error(
f"Error handling dynamic MCP route for {mcp_server_name}: {str(e)}"