scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

Merge pull request #15002 from BerriAI/litellm_dev_09_27_2025_p3

Browse source 
MCP - specify forwardable headers, specify allowed/disallowed tools for MCP servers
This commit is contained in:
Krish Dholakia 2025-09-28 09:21:52 -07:00 • committed by GitHub
commit c328b4a036
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 1032 additions and 65 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

271
docs/my-website/docs/mcp.md
View file

@ -137,6 +137,7 @@ mcp_servers:
| `basic` | `Authorization: Basic <auth_value>` |
| `authorization` | `Authorization: <auth_value>` |
- **Extra Headers**: Optional list of additional header names that should be forwarded from client to the MCP server
- **Spec Version**: Optional MCP specification version (defaults to `2025-06-18`)
Examples for each auth type:
@ -162,6 +163,13 @@ mcp_servers:
url: "https://my-mcp-server.com/mcp"
auth_type: "authorization"
auth_value: "Token example123" # headers={"Authorization": "Token example123"}
# Example with extra headers forwarding
github_mcp:
url: "https://api.githubcopilot.com/mcp"
auth_type: "bearer_token"
auth_value: "ghp_example_token"
extra_headers: ["custom_key", "x-custom-header"] # These headers will be forwarded from client
```
@ -191,6 +199,65 @@ litellm_settings:
</TabItem>
</Tabs>
## MCP Tool Filtering
Control which tools are available from your MCP servers. You can either allow only specific tools or block dangerous ones.
<Tabs>
<TabItem value="allowed" label="Only Allow Specific Tools">
Use `allowed_tools` to specify exactly which tools users can access. All other tools will be blocked.
```yaml title="config.yaml" showLineNumbers
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"]
# only list_tools will be available
```
**Use this when:**
- You want strict control over which tools are available
- You're in a high-security environment
- You're testing a new MCP server with limited tools
</TabItem>
<TabItem value="blocked" label="Block Specific Tools">
Use `disallowed_tools` to block specific tools. All other tools will be available.
```yaml title="config.yaml" showLineNumbers
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"]
disallowed_tools: ["repo_delete"]
# only repo_delete will be blocked
```
**Use this when:**
- Most tools are safe, but you want to block a few dangerous ones
- You want to prevent expensive API calls
- You're gradually adding restrictions to an existing server
</TabItem>
</Tabs>
### Important Notes
- If you specify both `allowed_tools` and `disallowed_tools`, the allowed list takes priority
- Tool names are case-sensitive
## Using your MCP
@ -771,6 +838,203 @@ When creating API keys, you can assign them to specific access groups for permis
/>
## Forwarding Custom Headers to MCP Servers
LiteLLM supports forwarding additional custom headers from MCP clients to backend MCP servers using the `extra_headers` configuration parameter. This allows you to pass custom authentication tokens, API keys, or other headers that your MCP server requires.
### Configuration
<Tabs>
<TabItem value="config" label="config.yaml">
Configure `extra_headers` in your MCP server configuration to specify which header names should be forwarded:
```yaml title="config.yaml with extra_headers" showLineNumbers
mcp_servers:
github_mcp:
url: "https://api.githubcopilot.com/mcp"
auth_type: "bearer_token"
auth_value: "ghp_default_token"
extra_headers: ["custom_key", "x-custom-header", "Authorization"]
description: "GitHub MCP server with custom header forwarding"
```
</TabItem>
<TabItem value="clientside" label="Dynamically on Client Side">
Use this when giving users access to a [group of MCP servers](#grouping-mcps-access-groups).
**Format:** `x-mcp-{server_alias}-{header_name}: value`
This allows you to use different authentication for different MCP servers.
**Examples:**
- `x-mcp-github-authorization: Bearer ghp_xxxxxxxxx` - GitHub MCP server with Bearer token
- `x-mcp-zapier-x-api-key: sk-xxxxxxxxx` - Zapier MCP server with API key
- `x-mcp-deepwiki-authorization: Basic base64_encoded_creds` - DeepWiki MCP server with Basic auth
```python title="Python Client with Server-Specific Auth" showLineNumbers
from fastmcp import Client
import asyncio
# Standard MCP configuration with multiple servers
config = {
"mcpServers": {
"mcp_group": {
"url": "http://localhost:4000/mcp",
"headers": {
"x-mcp-servers": "dev_group", # assume this gives access to github, zapier and deepwiki
"x-litellm-api-key": "Bearer sk-1234",
"x-mcp-github-authorization": "Bearer gho_token",
"x-mcp-zapier-x-api-key": "sk-xxxxxxxxx",
"x-mcp-deepwiki-authorization": "Basic base64_encoded_creds",
"custom_key": "value"
}
}
}
}
# Create a client that connects to all servers
client = Client(config)
async def main():
async with client:
tools = await client.list_tools()
print(f"Available tools: {tools}")
# call mcp
await client.call_tool(
name="github_mcp-search_issues",
arguments={'query': 'created:>2024-01-01', 'sort': 'created', 'order': 'desc', 'perPage': 30}
)
if __name__ == "__main__":
asyncio.run(main())
```
**Benefits:**
- **Server-specific authentication**: Each MCP server can use different auth methods
- **Better security**: No need to share the same auth token across all servers
- **Flexible header names**: Support for different auth header types (authorization, x-api-key, etc.)
- **Clean separation**: Each server's auth is clearly identified
</TabItem>
</Tabs>
### Client Usage
When connecting from MCP clients, include the custom headers that match the `extra_headers` configuration:
<Tabs>
<TabItem value="fastmcp" label="Python FastMCP">
```python title="FastMCP Client with Custom Headers" showLineNumbers
from fastmcp import Client
import asyncio
# MCP client configuration with custom headers
config = {
"mcpServers": {
"github": {
"url": "http://localhost:4000/github_mcp/mcp",
"headers": {
"x-litellm-api-key": "Bearer sk-1234",
"Authorization": "Bearer gho_token",
"custom_key": "custom_value",
"x-custom-header": "additional_data"
}
}
}
}
# Create a client that connects to the server
client = Client(config)
async def main():
async with client:
# List available tools
tools = await client.list_tools()
print(f"Available tools: {tools}")
# Call a tool if available
if tools:
result = await client.call_tool(tools[0].name, {})
print(f"Tool result: {result}")
# Run the client
asyncio.run(main())
```
</TabItem>
<TabItem value="cursor" label="Cursor IDE">
```json title="Cursor MCP Configuration with Custom Headers" showLineNumbers
{
"mcpServers": {
"GitHub": {
"url": "http://localhost:4000/github_mcp/mcp",
"headers": {
"x-litellm-api-key": "Bearer $LITELLM_API_KEY",
"Authorization": "Bearer $GITHUB_TOKEN",
"custom_key": "custom_value",
"x-custom-header": "additional_data"
}
}
}
}
```
</TabItem>
<TabItem value="http" label="HTTP Client">
```bash title="cURL with Custom Headers" showLineNumbers
curl --location 'http://localhost:4000/github_mcp/mcp' \
--header 'Content-Type: application/json' \
--header 'x-litellm-api-key: Bearer sk-1234' \
--header 'Authorization: Bearer gho_token' \
--header 'custom_key: custom_value' \
--header 'x-custom-header: additional_data' \
--data '{
"jsonrpc": "2.0",
"id": 1,
"method": "tools/list"
}'
```
</TabItem>
</Tabs>
### How It Works
1. **Configuration**: Define `extra_headers` in your MCP server config with the header names you want to forward
2. **Client Headers**: Include the corresponding headers in your MCP client requests
3. **Header Forwarding**: LiteLLM automatically forwards matching headers to the backend MCP server
4. **Authentication**: The backend MCP server receives both the configured auth headers and the custom headers
### Use Cases
- **Custom Authentication**: Forward custom API keys or tokens required by specific MCP servers
- **Request Context**: Pass user identification, session data, or request tracking headers
- **Third-party Integration**: Include headers required by external services that your MCP server integrates with
- **Multi-tenant Systems**: Forward tenant-specific headers for proper request routing
### Security Considerations
- Only headers listed in `extra_headers` are forwarded to maintain security
- Sensitive headers should be passed through environment variables when possible
- Consider using server-specific auth headers for better security isolation
---
## Using your MCP with client side credentials
Use this if you want to pass a client side authentication token to LiteLLM to then pass to your MCP to auth to your MCP.
@ -780,13 +1044,6 @@ Use this if you want to pass a client side authentication token to LiteLLM to th
You can specify MCP auth tokens using server-specific headers in the format `x-mcp-{server_alias}-{header_name}`. This allows you to use different authentication for different MCP servers.
**Format:** `x-mcp-{server_alias}-{header_name}: value`
**Examples:**
- `x-mcp-github-authorization: Bearer ghp_xxxxxxxxx` - GitHub MCP server with Bearer token
- `x-mcp-zapier-x-api-key: sk-xxxxxxxxx` - Zapier MCP server with API key
- `x-mcp-deepwiki-authorization: Basic base64_encoded_creds` - DeepWiki MCP server with Basic auth
**Benefits:**
- **Server-specific authentication**: Each MCP server can use different auth methods
- **Better security**: No need to share the same auth token across all servers

38
litellm/experimental_mcp_client/client.py
View file

@ -5,7 +5,7 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers.
import asyncio
import base64
from datetime import timedelta
from typing import Dict, List, Optional
from typing import Dict, List, Optional, Union
from mcp import ClientSession, StdioServerParameters
from mcp.client.sse import sse_client
@ -44,7 +44,7 @@ class MCPClient:
server_url: str = "",
transport_type: MCPTransportType = MCPTransport.http,
auth_type: MCPAuthType = None,
auth_value: Optional[str] = None,
auth_value: Optional[Union[str, Dict[str, str]]] = None,
timeout: float = 60.0,
stdio_config: Optional[MCPStdioConfig] = None,
extra_headers: Optional[Dict[str, str]] = None,
@ -53,7 +53,7 @@ class MCPClient:
self.transport_type: MCPTransport = transport_type
self.auth_type: MCPAuthType = auth_type
self.timeout: float = timeout
self._mcp_auth_value: Optional[str] = None
self._mcp_auth_value: Optional[Union[str, Dict[str, str]]] = None
self._session: Optional[ClientSession] = None
self._context = None
self._transport_ctx = None
@ -180,28 +180,34 @@ class MCPClient:
pass
self._context = None
def update_auth_value(self, mcp_auth_value: str):
def update_auth_value(self, mcp_auth_value: Union[str, Dict[str, str]]):
"""
Set the authentication header for the MCP client.
"""
if self.auth_type == MCPAuth.basic:
# Assuming mcp_auth_value is in format "username:password", convert it when updating
mcp_auth_value = to_basic_auth(mcp_auth_value)
self._mcp_auth_value = mcp_auth_value
if isinstance(mcp_auth_value, dict):
self._mcp_auth_value = mcp_auth_value
else:
if self.auth_type == MCPAuth.basic:
# Assuming mcp_auth_value is in format "username:password", convert it when updating
mcp_auth_value = to_basic_auth(mcp_auth_value)
self._mcp_auth_value = mcp_auth_value
def _get_auth_headers(self) -> dict:
"""Generate authentication headers based on auth type."""
headers = {"MCP-Protocol-Version": "2025-06-18"}
if self._mcp_auth_value:
if self.auth_type == MCPAuth.bearer_token:
headers["Authorization"] = f"Bearer {self._mcp_auth_value}"
elif self.auth_type == MCPAuth.basic:
headers["Authorization"] = f"Basic {self._mcp_auth_value}"
elif self.auth_type == MCPAuth.api_key:
headers["X-API-Key"] = self._mcp_auth_value
elif self.auth_type == MCPAuth.authorization:
headers["Authorization"] = self._mcp_auth_value
if isinstance(self._mcp_auth_value, str):
if self.auth_type == MCPAuth.bearer_token:
headers["Authorization"] = f"Bearer {self._mcp_auth_value}"
elif self.auth_type == MCPAuth.basic:
headers["Authorization"] = f"Basic {self._mcp_auth_value}"
elif self.auth_type == MCPAuth.api_key:
headers["X-API-Key"] = self._mcp_auth_value
elif self.auth_type == MCPAuth.authorization:
headers["Authorization"] = self._mcp_auth_value
elif isinstance(self._mcp_auth_value, dict):
headers.update(self._mcp_auth_value)
# update the headers with the extra headers
if self.extra_headers:

7
litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py
View file

@ -15,6 +15,7 @@ class MCPAuthenticatedUser(AuthenticatedUser):
3. MCP server configuration (can include access groups)
4. Server-specific authentication headers
5. OAuth2 headers
6. Raw headers - allows forwarding specific headers to the MCP server, specified by the admin.
"""
def __init__(
@ -22,9 +23,10 @@ class MCPAuthenticatedUser(AuthenticatedUser):
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,
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
mcp_protocol_version: Optional[str] = None,
raw_headers: Optional[Dict[str, str]] = None,
):
self.user_api_key_auth = user_api_key_auth
self.mcp_auth_header = mcp_auth_header
@ -32,3 +34,4 @@ class MCPAuthenticatedUser(AuthenticatedUser):
self.mcp_server_auth_headers = mcp_server_auth_headers or {}
self.mcp_protocol_version = mcp_protocol_version
self.oauth2_headers = oauth2_headers
self.raw_headers = raw_headers

29
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
View file

@ -39,6 +39,7 @@ class MCPRequestHandler:
UserAPIKeyAuth,
Optional[str],
Optional[List[str]],
Optional[Dict[str, Dict[str, str]]],
Optional[Dict[str, str]],
Optional[Dict[str, str]],
]:
@ -49,6 +50,7 @@ class MCPRequestHandler:
2. Processing MCP server configuration
3. Handling MCP-specific headers
4. Handling oauth2 headers
5. Raw headers - allows forwarding specific headers to the MCP server, specified by the admin.
Args:
scope: ASGI scope containing request information
@ -58,7 +60,8 @@ class MCPRequestHandler:
mcp_auth_header: Optional[str] MCP auth header to be passed to the MCP server (deprecated)
mcp_servers: Optional[List[str]] List of MCP servers and access groups to use
mcp_server_auth_headers: Optional[Dict[str, str]] Server-specific auth headers in format {server_alias: auth_value}
oauth2_headers: Optional[Dict[str, str]] OAuth2 headers
raw_headers: Optional[Dict[str, str]] Raw headers to be forwarded to the MCP server
Raises:
HTTPException: If headers are invalid or missing required headers
"""
@ -116,6 +119,7 @@ class MCPRequestHandler:
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
dict(headers),
)
@staticmethod
@ -145,7 +149,9 @@ class MCPRequestHandler:
return auth_header
@staticmethod
def _get_mcp_server_auth_headers_from_headers(headers: Headers) -> Dict[str, str]:
def _get_mcp_server_auth_headers_from_headers(
headers: Headers,
) -> Dict[str, Dict[str, str]]:
"""
Parse server-specific MCP auth headers from the request headers.
@ -156,7 +162,7 @@ class MCPRequestHandler:
- x-mcp-deepwiki-authorization: Basic base64_encoded_creds
Returns:
Dict[str, str]: Mapping of server alias to auth value
Dict[str, Dict[str, str]]: Mapping of server alias to header dict
"""
server_auth_headers = {}
prefix = "x-mcp-"
@ -175,11 +181,22 @@ class MCPRequestHandler:
# Extract server_alias and header_name from x-mcp-{server_alias}-{header_name}
remaining = header_name[len(prefix) :].lower()
if "-" in remaining:
# Split on the last dash to separate server_alias from header_name
parts = remaining.rsplit("-", 1)
# Split on the first dash to separate server_alias from header_name
parts = remaining.split("-", 1)
if len(parts) == 2:
server_alias, auth_header_name = parts
server_auth_headers[server_alias] = header_value
# Convert common header names to proper case
if auth_header_name == "authorization":
auth_header_name = "Authorization"
# Initialize server dict if not exists
if server_alias not in server_auth_headers:
server_auth_headers[server_alias] = {}
server_auth_headers[server_alias][
auth_header_name
] = header_value
verbose_logger.debug(
f"Found server auth header: {server_alias} -> {auth_header_name}: {header_value[:10]}..."
)

52
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View file

@ -10,7 +10,7 @@ import asyncio
import datetime
import hashlib
import json
from typing import Any, Dict, List, Optional, cast
from typing import Any, Dict, List, Optional, Union, cast
from fastapi import HTTPException
from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
@ -212,6 +212,9 @@ class MCPServerManager:
"authentication_token", server_config.get("auth_value", None)
),
mcp_info=mcp_info,
extra_headers=server_config.get("extra_headers", None),
allowed_tools=server_config.get("allowed_tools", None),
disallowed_tools=server_config.get("disallowed_tools", None),
access_groups=server_config.get("access_groups", None),
)
self.config_mcp_servers[server_id] = new_server
@ -264,11 +267,20 @@ class MCPServerManager:
transport=cast(MCPTransportType, mcp_server.transport),
auth_type=cast(MCPAuthType, mcp_server.auth_type),
mcp_info=mcp_info,
extra_headers=getattr(mcp_server, "extra_headers", None),
# oauth specific fields
client_id=getattr(mcp_server, "client_id", None),
client_secret=getattr(mcp_server, "client_secret", None),
scopes=getattr(mcp_server, "scopes", None),
authorization_url=getattr(mcp_server, "authorization_url", None),
token_url=getattr(mcp_server, "token_url", None),
# Stdio-specific fields
command=getattr(mcp_server, "command", None),
args=getattr(mcp_server, "args", None) or [],
env=env_dict,
access_groups=getattr(mcp_server, "mcp_access_groups", None),
allowed_tools=getattr(mcp_server, "allowed_tools", None),
disallowed_tools=getattr(mcp_server, "disallowed_tools", None),
)
self.registry[mcp_server.server_id] = new_server
verbose_logger.debug(f"Added MCP Server: {name_for_prefix}")
@ -319,7 +331,7 @@ class MCPServerManager:
self,
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
mcp_auth_header: Optional[str] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
mcp_server_auth_headers: Optional[Dict[str, Union[str, Dict[str, str]]]] = None,
) -> List[MCPTool]:
"""
List all tools available across all MCP Servers.
@ -381,7 +393,7 @@ class MCPServerManager:
def _create_mcp_client(
self,
server: MCPServer,
mcp_auth_header: Optional[str] = None,
mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None,
extra_headers: Optional[Dict[str, str]] = None,
) -> MCPClient:
"""
@ -429,7 +441,7 @@ class MCPServerManager:
async def _get_tools_from_server(
self,
server: MCPServer,
mcp_auth_header: Optional[str] = None,
mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None,
extra_headers: Optional[Dict[str, str]] = None,
) -> List[MCPTool]:
"""
@ -561,6 +573,16 @@ class MCPServerManager:
)
return prefixed_tools
def check_allowed_or_banned_tools(self, tool_name: str, server: MCPServer) -> bool:
"""
Check if the tool is allowed or banned for the given server
"""
if server.allowed_tools:
return tool_name in server.allowed_tools
if server.disallowed_tools:
return tool_name not in server.disallowed_tools
return True
async def pre_call_tool_check(
self,
name: str,
@ -568,7 +590,18 @@ class MCPServerManager:
server_name_from_prefix: str,
user_api_key_auth: Optional[UserAPIKeyAuth],
proxy_logging_obj: ProxyLogging,
server: MCPServer,
):
## check if the tool is allowed or banned for the given server
if not self.check_allowed_or_banned_tools(name, server):
raise HTTPException(
status_code=403,
detail={
"error": f"Tool {name} is not allowed for server {server.name}. Contact proxy admin to allow this tool."
},
)
pre_hook_kwargs = {
"name": name,
"arguments": arguments,
@ -638,9 +671,10 @@ class MCPServerManager:
arguments: Dict[str, Any],
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
mcp_auth_header: Optional[str] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None,
proxy_logging_obj: Optional[ProxyLogging] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
raw_headers: Optional[Dict[str, str]] = None,
) -> CallToolResult:
"""
Call a tool with the given name and arguments (handles prefixed tool names)
@ -691,6 +725,7 @@ class MCPServerManager:
server_name_from_prefix=server_name_from_prefix,
user_api_key_auth=user_api_key_auth,
proxy_logging_obj=proxy_logging_obj,
server=mcp_server,
)
# Get server-specific auth header if available
@ -709,6 +744,13 @@ class MCPServerManager:
if mcp_server.auth_type == MCPAuth.oauth2:
extra_headers = oauth2_headers
if mcp_server.extra_headers and raw_headers:
if extra_headers is None:
extra_headers = {}
for header in mcp_server.extra_headers:
if header in raw_headers:
extra_headers[header] = raw_headers[header]
client = self._create_mcp_client(
server=mcp_server,
mcp_auth_header=server_auth_header,

65
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -180,6 +180,7 @@ if MCP_AVAILABLE:
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
raw_headers,
) = get_auth_context()
verbose_logger.debug(
f"MCP list_tools - User API Key Auth from context: {user_api_key_auth}"
@ -198,6 +199,7 @@ if MCP_AVAILABLE:
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
)
verbose_logger.info(
f"MCP list_tools - Successfully returned {len(tools)} tools"
@ -239,6 +241,7 @@ if MCP_AVAILABLE:
_,
mcp_server_auth_headers,
oauth2_headers,
raw_headers,
) = get_auth_context()
verbose_logger.debug(
@ -271,6 +274,7 @@ if MCP_AVAILABLE:
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
**data, # for logging
)
except BlockedPiiEntityError as e:
@ -357,12 +361,31 @@ if MCP_AVAILABLE:
return allowed_mcp_servers
def filter_tools_by_allowed_tools(
tools: List[MCPTool],
mcp_server: MCPServer,
) -> List[MCPTool]:
"""
Filter tools by allowed tools
"""
tools_to_return = tools
if mcp_server.allowed_tools:
tools_to_return = [
tool for tool in tools if tool.name in mcp_server.allowed_tools
]
if mcp_server.disallowed_tools:
tools_to_return = [
tool for tool in tools if tool.name not in mcp_server.disallowed_tools
]
return tools_to_return
async def _get_tools_from_mcp_servers(
user_api_key_auth: Optional[UserAPIKeyAuth],
mcp_auth_header: Optional[str],
mcp_servers: Optional[List[str]],
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
raw_headers: Optional[Dict[str, str]] = None,
) -> List[MCPTool]:
"""
Helper method to fetch tools from MCP servers based on server filtering criteria.
@ -371,8 +394,8 @@ if MCP_AVAILABLE:
user_api_key_auth: User authentication info for access control
mcp_auth_header: Optional auth header for MCP server (deprecated)
mcp_servers: Optional list of server names/aliases to filter by
mcp_server_auth_headers: Optional dict of server-specific auth headers
oauth2_headers: Optional dict of oauth2 headers
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
@ -409,6 +432,13 @@ if MCP_AVAILABLE:
if server.auth_type == MCPAuth.oauth2:
extra_headers = oauth2_headers
if server.extra_headers and raw_headers:
if extra_headers is None:
extra_headers = {}
for header in server.extra_headers:
if header in raw_headers:
extra_headers[header] = raw_headers[header]
# 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
@ -419,7 +449,7 @@ if MCP_AVAILABLE:
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
)
all_tools.extend(tools)
all_tools.extend(filter_tools_by_allowed_tools(tools, server))
verbose_logger.debug(
f"Successfully fetched {len(tools)} tools from server {server.name}"
)
@ -438,8 +468,9 @@ if MCP_AVAILABLE:
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
mcp_auth_header: Optional[str] = None,
mcp_servers: Optional[List[str]] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
raw_headers: Optional[Dict[str, str]] = None,
) -> List[MCPTool]:
"""
List all available MCP tools.
@ -464,6 +495,7 @@ if MCP_AVAILABLE:
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
)
verbose_logger.debug(
f"Successfully fetched {len(managed_tools)} tools from managed MCP servers"
@ -505,8 +537,9 @@ if MCP_AVAILABLE:
arguments: Optional[Dict[str, Any]] = None,
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
mcp_auth_header: Optional[str] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
raw_headers: Optional[Dict[str, str]] = None,
**kwargs: Any,
) -> List[Union[TextContent, ImageContent, EmbeddedResource]]:
"""
@ -555,6 +588,7 @@ if MCP_AVAILABLE:
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
litellm_logging_obj=litellm_logging_obj,
)
@ -606,8 +640,9 @@ if MCP_AVAILABLE:
arguments: Dict[str, Any],
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
mcp_auth_header: Optional[str] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
raw_headers: Optional[Dict[str, str]] = None,
litellm_logging_obj: Optional[Any] = None,
) -> List[Union[TextContent, ImageContent, EmbeddedResource]]:
"""Handle tool execution for managed server tools"""
@ -621,6 +656,7 @@ if MCP_AVAILABLE:
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
proxy_logging_obj=proxy_logging_obj,
)
verbose_logger.debug("CALL TOOL RESULT: %s", call_tool_result)
@ -702,6 +738,7 @@ if MCP_AVAILABLE:
_,
mcp_server_auth_headers,
oauth2_headers,
raw_headers,
) = await MCPRequestHandler.process_mcp_request(scope)
mcp_servers = mcp_servers_from_path
else:
@ -711,6 +748,7 @@ if MCP_AVAILABLE:
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
raw_headers,
) = await MCPRequestHandler.process_mcp_request(scope)
return (
user_api_key_auth,
@ -718,6 +756,7 @@ if MCP_AVAILABLE:
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
raw_headers,
)
async def handle_streamable_http_mcp(
@ -732,6 +771,7 @@ if MCP_AVAILABLE:
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
raw_headers,
) = await extract_mcp_auth_context(scope, path)
verbose_logger.debug(
f"MCP request mcp_servers (header/path): {mcp_servers}"
@ -746,6 +786,7 @@ if MCP_AVAILABLE:
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
)
# Ensure session managers are initialized
@ -785,6 +826,7 @@ if MCP_AVAILABLE:
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
raw_headers,
) = await extract_mcp_auth_context(scope, path)
verbose_logger.debug(
f"MCP request mcp_servers (header/path): {mcp_servers}"
@ -798,6 +840,7 @@ if MCP_AVAILABLE:
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
)
if not _SESSION_MANAGERS_INITIALIZED:
@ -858,8 +901,9 @@ if MCP_AVAILABLE:
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_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
raw_headers: Optional[Dict[str, str]] = None,
) -> None:
"""
Set the UserAPIKeyAuth in the auth context variable.
@ -876,6 +920,7 @@ if MCP_AVAILABLE:
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
)
auth_context_var.set(auth_user)
@ -883,6 +928,7 @@ if MCP_AVAILABLE:
Optional[UserAPIKeyAuth],
Optional[str],
Optional[List[str]],
Optional[Dict[str, Dict[str, str]]],
Optional[Dict[str, str]],
Optional[Dict[str, str]],
]:
@ -901,8 +947,9 @@ if MCP_AVAILABLE:
auth_user.mcp_servers,
auth_user.mcp_server_auth_headers,
auth_user.oauth2_headers,
auth_user.raw_headers,
)
return None, None, None, None, None
return None, None, None, None, None, None
########################################################
############ End of Auth Context Functions #############

2
litellm/proxy/_new_secret_config.yaml
View file

@ -25,6 +25,6 @@ mcp_servers:
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"]
allowed_tools: ["list_tools"]
# disallowed_tools: ["repo_delete"]

5
litellm/types/mcp_server/mcp_server_manager.py
View file

@ -20,6 +20,11 @@ class MCPServer(BaseModel):
auth_type: Optional[MCPAuthType] = None
authentication_token: Optional[str] = None
mcp_info: Optional[MCPInfo] = None
extra_headers: Optional[List[str]] = (
None # allow admin to specify which headers to forward to the MCP server
)
allowed_tools: Optional[List[str]] = None
disallowed_tools: Optional[List[str]] = None
# OAuth-specific fields
client_id: Optional[str] = None
client_secret: Optional[str] = None

311
tests/mcp_tests/test_mcp_server.py
View file

@ -1349,6 +1349,11 @@ def test_add_update_server_with_alias():
mock_mcp_server.command = None
mock_mcp_server.args = []
mock_mcp_server.env = None
# OAuth fields - set explicitly to None to avoid MagicMock objects
mock_mcp_server.client_id = None
mock_mcp_server.client_secret = None
mock_mcp_server.authorization_url = None
mock_mcp_server.token_url = None
# Add server to manager
test_manager.add_update_server(mock_mcp_server)
@ -1380,6 +1385,11 @@ def test_add_update_server_without_alias():
mock_mcp_server.command = None
mock_mcp_server.args = []
mock_mcp_server.env = None
# OAuth fields - set explicitly to None to avoid MagicMock objects
mock_mcp_server.client_id = None
mock_mcp_server.client_secret = None
mock_mcp_server.authorization_url = None
mock_mcp_server.token_url = None
# Add server to manager
test_manager.add_update_server(mock_mcp_server)
@ -1411,6 +1421,11 @@ def test_add_update_server_fallback_to_server_id():
mock_mcp_server.command = None
mock_mcp_server.args = []
mock_mcp_server.env = None
# OAuth fields - set explicitly to None to avoid MagicMock objects
mock_mcp_server.client_id = None
mock_mcp_server.client_secret = None
mock_mcp_server.authorization_url = None
mock_mcp_server.token_url = None
# Add server to manager
test_manager.add_update_server(mock_mcp_server)
@ -1982,6 +1997,302 @@ async def test_list_tool_rest_api_all_servers_with_auth():
assert calls[1][0][1] == "Bearer slack_token" # server_auth_header
@pytest.mark.asyncio
async def test_filter_tools_by_allowed_tools_integration():
"""Test that filter_tools_by_allowed_tools works correctly via _get_tools_from_mcp_servers"""
from litellm.proxy._experimental.mcp_server.server import (
_get_tools_from_mcp_servers,
)
from litellm.proxy._types import UserAPIKeyAuth
from mcp.types import Tool as MCPTool
# Create a mock user auth
mock_user_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user")
# Create mock tools that will be returned by the server
mock_tools = [
MCPTool(
name="allowed_tool_1",
description="This tool should be allowed",
inputSchema={"type": "object"},
),
MCPTool(
name="allowed_tool_2",
description="This tool should also be allowed",
inputSchema={"type": "object"},
),
MCPTool(
name="blocked_tool_1",
description="This tool should be blocked",
inputSchema={"type": "object"},
),
MCPTool(
name="blocked_tool_2",
description="This tool should also be blocked",
inputSchema={"type": "object"},
),
]
# Create a mock server with allowed_tools restriction
mock_server = MCPServer(
server_id="test-server-123",
name="test_server_with_allowed_tools",
url="https://test-server.com/mcp",
transport=MCPTransport.http,
allowed_tools=[
"allowed_tool_1",
"allowed_tool_2",
], # Only these tools should be returned
disallowed_tools=None,
)
# Create a mock MCPClient that returns all tools
mock_client = AsyncMock()
mock_client.list_tools = AsyncMock(return_value=mock_tools)
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=None)
def mock_client_constructor(*args, **kwargs):
return mock_client
# Mock the global MCP server manager
with patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager"
) as mock_manager:
# Mock manager methods
mock_manager.get_allowed_mcp_servers = AsyncMock(
return_value=["test-server-123"]
)
mock_manager.get_mcp_server_by_id = MagicMock(return_value=mock_server)
# Mock the _get_tools_from_server method to return all tools
mock_manager._get_tools_from_server = AsyncMock(return_value=mock_tools)
# Mock the MCPClient constructor
with patch(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPClient",
mock_client_constructor,
):
# Call _get_tools_from_mcp_servers which should apply the filtering
filtered_tools = await _get_tools_from_mcp_servers(
user_api_key_auth=mock_user_auth,
mcp_auth_header="Bearer test_token",
mcp_servers=None, # Get from all servers
)
# Verify that only allowed tools are returned
assert (
len(filtered_tools) == 2
), f"Expected 2 tools, got {len(filtered_tools)}"
tool_names = [tool.name for tool in filtered_tools]
assert (
"allowed_tool_1" in tool_names
), "allowed_tool_1 should be in filtered results"
assert (
"allowed_tool_2" in tool_names
), "allowed_tool_2 should be in filtered results"
assert (
"blocked_tool_1" not in tool_names
), "blocked_tool_1 should be filtered out"
assert (
"blocked_tool_2" not in tool_names
), "blocked_tool_2 should be filtered out"
# Verify the manager methods were called correctly
mock_manager.get_allowed_mcp_servers.assert_called_once_with(mock_user_auth)
mock_manager.get_mcp_server_by_id.assert_called_once_with("test-server-123")
mock_manager._get_tools_from_server.assert_called_once()
@pytest.mark.asyncio
async def test_filter_tools_by_disallowed_tools_integration():
"""Test that filter_tools_by_allowed_tools works correctly with disallowed_tools via _get_tools_from_mcp_servers"""
from litellm.proxy._experimental.mcp_server.server import (
_get_tools_from_mcp_servers,
)
from litellm.proxy._types import UserAPIKeyAuth
from mcp.types import Tool as MCPTool
# Create a mock user auth
mock_user_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user")
# Create mock tools that will be returned by the server
mock_tools = [
MCPTool(
name="safe_tool_1",
description="This tool should be allowed",
inputSchema={"type": "object"},
),
MCPTool(
name="safe_tool_2",
description="This tool should also be allowed",
inputSchema={"type": "object"},
),
MCPTool(
name="dangerous_tool_1",
description="This tool should be blocked",
inputSchema={"type": "object"},
),
MCPTool(
name="dangerous_tool_2",
description="This tool should also be blocked",
inputSchema={"type": "object"},
),
]
# Create a mock server with disallowed_tools restriction
mock_server = MCPServer(
server_id="test-server-456",
name="test_server_with_disallowed_tools",
url="https://test-server.com/mcp",
transport=MCPTransport.http,
allowed_tools=None,
disallowed_tools=[
"dangerous_tool_1",
"dangerous_tool_2",
], # These tools should be filtered out
)
# Create a mock MCPClient that returns all tools
mock_client = AsyncMock()
mock_client.list_tools = AsyncMock(return_value=mock_tools)
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=None)
def mock_client_constructor(*args, **kwargs):
return mock_client
# Mock the global MCP server manager
with patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager"
) as mock_manager:
# Mock manager methods
mock_manager.get_allowed_mcp_servers = AsyncMock(
return_value=["test-server-456"]
)
mock_manager.get_mcp_server_by_id = MagicMock(return_value=mock_server)
# Mock the _get_tools_from_server method to return all tools
mock_manager._get_tools_from_server = AsyncMock(return_value=mock_tools)
# Mock the MCPClient constructor
with patch(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPClient",
mock_client_constructor,
):
# Call _get_tools_from_mcp_servers which should apply the filtering
filtered_tools = await _get_tools_from_mcp_servers(
user_api_key_auth=mock_user_auth,
mcp_auth_header="Bearer test_token",
mcp_servers=None, # Get from all servers
)
# Verify that only safe tools are returned (dangerous tools filtered out)
assert (
len(filtered_tools) == 2
), f"Expected 2 tools, got {len(filtered_tools)}"
tool_names = [tool.name for tool in filtered_tools]
assert (
"safe_tool_1" in tool_names
), "safe_tool_1 should be in filtered results"
assert (
"safe_tool_2" in tool_names
), "safe_tool_2 should be in filtered results"
assert (
"dangerous_tool_1" not in tool_names
), "dangerous_tool_1 should be filtered out"
assert (
"dangerous_tool_2" not in tool_names
), "dangerous_tool_2 should be filtered out"
# Verify the manager methods were called correctly
mock_manager.get_allowed_mcp_servers.assert_called_once_with(mock_user_auth)
mock_manager.get_mcp_server_by_id.assert_called_once_with("test-server-456")
mock_manager._get_tools_from_server.assert_called_once()
@pytest.mark.asyncio
async def test_filter_tools_no_restrictions_integration():
"""Test that filter_tools_by_allowed_tools returns all tools when no restrictions are set"""
from litellm.proxy._experimental.mcp_server.server import (
_get_tools_from_mcp_servers,
)
from litellm.proxy._types import UserAPIKeyAuth
from mcp.types import Tool as MCPTool
# Create a mock user auth
mock_user_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user")
# Create mock tools that will be returned by the server
mock_tools = [
MCPTool(
name="tool_1",
description="Tool 1",
inputSchema={"type": "object"},
),
MCPTool(
name="tool_2",
description="Tool 2",
inputSchema={"type": "object"},
),
]
# Create a mock server with no tool restrictions
mock_server = MCPServer(
server_id="test-server-000",
name="test_server_no_restrictions",
url="https://test-server.com/mcp",
transport=MCPTransport.http,
allowed_tools=None, # No restrictions
disallowed_tools=None, # No restrictions
)
# Create a mock MCPClient that returns all tools
mock_client = AsyncMock()
mock_client.list_tools = AsyncMock(return_value=mock_tools)
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=None)
def mock_client_constructor(*args, **kwargs):
return mock_client
# Mock the global MCP server manager
with patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager"
) as mock_manager:
# Mock manager methods
mock_manager.get_allowed_mcp_servers = AsyncMock(
return_value=["test-server-000"]
)
mock_manager.get_mcp_server_by_id = MagicMock(return_value=mock_server)
# Mock the _get_tools_from_server method to return all tools
mock_manager._get_tools_from_server = AsyncMock(return_value=mock_tools)
# Mock the MCPClient constructor
with patch(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPClient",
mock_client_constructor,
):
# Call _get_tools_from_mcp_servers which should apply the filtering
filtered_tools = await _get_tools_from_mcp_servers(
user_api_key_auth=mock_user_auth,
mcp_auth_header="Bearer test_token",
mcp_servers=None, # Get from all servers
)
# Should return all tools when no restrictions
assert (
len(filtered_tools) == 2
), f"Expected 2 tools, got {len(filtered_tools)}"
tool_names = [tool.name for tool in filtered_tools]
assert "tool_1" in tool_names, "tool_1 should be in filtered results"
assert "tool_2" in tool_names, "tool_2 should be in filtered results"
@pytest.mark.asyncio
async def test_mcp_access_group_permission_inheritance_integration():
"""Integration test for MCP access group permission inheritance"""

80
tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
View file

@ -286,7 +286,10 @@ class TestMCPRequestHandler:
],
"test-api-key-123",
None,
{"github": "Bearer github-token", "zapier_x_api": "zapier-api-key"},
{
"github": {"Authorization": "Bearer github-token"},
"zapier_x_api": {"key": "zapier-api-key"},
},
),
# Test case 10: Both legacy and server-specific auth headers
(
@ -297,7 +300,7 @@ class TestMCPRequestHandler:
],
"test-api-key-123",
"legacy-token",
{"github": "Bearer github-token"},
{"github": {"Authorization": "Bearer github-token"}},
),
# Test case 11: Server-specific auth headers with different header types
(
@ -308,7 +311,10 @@ class TestMCPRequestHandler:
],
"test-api-key-123",
None,
{"deepwiki": "Basic base64-encoded", "custom_x_custom": "custom-value"},
{
"deepwiki": {"Authorization": "Basic base64-encoded"},
"custom_x_custom": {"header": "custom-value"},
},
),
# Test case 12: Case insensitive server-specific headers
(
@ -318,7 +324,7 @@ class TestMCPRequestHandler:
],
"test-api-key-123",
None,
{"github": "Bearer github-token"},
{"github": {"Authorization": "Bearer github-token"}},
),
],
)
@ -365,6 +371,7 @@ class TestMCPRequestHandler:
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
raw_headers,
) = await MCPRequestHandler.process_mcp_request(scope)
# Assert the results
@ -751,7 +758,7 @@ class TestMCPCustomHeaderName:
}
)
result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers)
assert result == {"github": "Bearer github-token"}
assert result == {"github": {"Authorization": "Bearer github-token"}}
# Test case 3: Multiple server-specific headers
headers = Headers(
@ -764,9 +771,9 @@ class TestMCPCustomHeaderName:
)
result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers)
expected = {
"github": "Bearer github-token",
"zapier_x_api": "zapier-api-key",
"deepwiki": "Basic base64-encoded",
"github": {"Authorization": "Bearer github-token"},
"zapier_x_api": {"key": "zapier-api-key"},
"deepwiki": {"Authorization": "Basic base64-encoded"},
}
assert result == expected
@ -775,11 +782,14 @@ class TestMCPCustomHeaderName:
{
"x-litellm-api-key": "test-key",
"X-MCP-GITHUB-AUTHORIZATION": "Bearer github-token",
"x-mcp-ZAPIER_x_api-key": "zapier-api-key",
"x-mcp-ZAPIER-x-api-key": "zapier-api-key",
}
)
result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers)
expected = {"github": "Bearer github-token", "zapier_x_api": "zapier-api-key"}
expected = {
"github": {"Authorization": "Bearer github-token"},
"zapier": {"x-api-key": "zapier-api-key"},
}
assert result == expected
# Test case 5: Invalid format headers (should be ignored)
@ -792,7 +802,7 @@ class TestMCPCustomHeaderName:
}
)
result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers)
assert result == {"github": "Bearer github-token"}
assert result == {"github": {"Authorization": "Bearer github-token"}}
# Test case 6: Edge case - header with multiple hyphens in server alias
headers = Headers(
@ -804,8 +814,8 @@ class TestMCPCustomHeaderName:
)
result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers)
expected = {
"github_mcp": "Bearer github-mcp-token",
"gh_mcp2": "Bearer gh-mcp2-token",
"github_mcp": {"Authorization": "Bearer github-mcp-token"},
"gh_mcp2": {"Authorization": "Bearer gh-mcp2-token"},
}
assert result == expected
@ -817,14 +827,14 @@ class TestMCPCustomHeaderName:
}
)
result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers)
assert result == {"github_mcp": "Bearer github-mcp-token"}
assert result == {"github_mcp": {"Authorization": "Bearer github-mcp-token"}}
# Test case 8: Edge case - empty header value
headers = Headers(
{"x-litellm-api-key": "test-key", "x-mcp-github-authorization": ""}
)
result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers)
assert result == {"github": ""}
assert result == {"github": {"Authorization": ""}}
# Test case 9: Edge case - very long header value
long_token = "Bearer " + "x" * 1000
@ -832,7 +842,7 @@ class TestMCPCustomHeaderName:
{"x-litellm-api-key": "test-key", "x-mcp-github-authorization": long_token}
)
result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers)
assert result == {"github": long_token}
assert result == {"github": {"Authorization": long_token}}
# Test case 10: Edge case - special characters in server alias
headers = Headers(
@ -844,8 +854,8 @@ class TestMCPCustomHeaderName:
)
result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers)
expected = {
"github-123": "Bearer github-123-token",
"github_test": "Bearer github-test-token",
"github": {"123-authorization": "Bearer github-123-token"},
"github_test": {"Authorization": "Bearer github-test-token"},
}
assert result == expected
@ -890,6 +900,7 @@ class TestMCPAccessGroupsE2E:
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
raw_headers,
) = await MCPRequestHandler.process_mcp_request(scope)
# Assert the results
@ -940,6 +951,7 @@ class TestMCPAccessGroupsE2E:
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
raw_headers,
) = await MCPRequestHandler.process_mcp_request(scope)
# Assert the results
@ -1014,3 +1026,35 @@ def test_mcp_path_based_server_segregation(monkeypatch):
# The context should have mcp_servers set to ["zapier", "group1"]
assert list(captured_mcp_servers.values())[0] == ["zapier", "group1"]
@pytest.mark.parametrize(
"headers,expected_result",
[
(
Headers(
{
"x-litellm-api-key": "test-key",
"x-mcp-github-authorization": "Bearer github-token",
}
),
{"github": {"Authorization": "Bearer github-token"}},
),
(
Headers(
{
"x-litellm-api-key": "test-key",
"x-mcp-github-x-api-key": "Basic base64-encoded-creds",
}
),
{"github": {"x-api-key": "Basic base64-encoded-creds"}},
),
],
)
def test_get_mcp_server_auth_headers_from_headers(headers, expected_result):
"""Test _get_mcp_server_auth_headers_from_headers method"""
from starlette.datastructures import Headers
headers = Headers(headers)
result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers)
assert result == expected_result

237
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
View file

@ -1,8 +1,9 @@
import sys
from datetime import datetime
from unittest.mock import MagicMock, AsyncMock
from unittest.mock import AsyncMock, MagicMock
import pytest
from fastapi import HTTPException
# Add the parent directory to the path so we can import litellm
sys.path.insert(0, "../../../../../")
@ -420,6 +421,240 @@ class TestMCPServerManager:
assert result["status"] == "healthy"
assert result["tools_count"] == 1
@pytest.mark.asyncio
async def test_pre_call_tool_check_allowed_tools_list_allows_tool(self):
"""Test pre_call_tool_check allows tool when it's in allowed_tools list"""
manager = MCPServerManager()
# Create server with allowed_tools list
server = MCPServer(
server_id="test-server",
name="test-server",
transport=MCPTransport.stdio,
allowed_tools=["allowed_tool", "another_allowed_tool"],
disallowed_tools=None,
)
# Mock dependencies
user_api_key_auth = MagicMock()
proxy_logging_obj = MagicMock()
# Mock the async methods that pre_call_tool_check calls
proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(
return_value={}
)
proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={})
proxy_logging_obj.pre_call_hook = AsyncMock(return_value={})
# This should not raise an exception
await manager.pre_call_tool_check(
name="allowed_tool",
arguments={"param": "value"},
server_name_from_prefix="test-server",
user_api_key_auth=user_api_key_auth,
proxy_logging_obj=proxy_logging_obj,
server=server,
)
@pytest.mark.asyncio
async def test_pre_call_tool_check_allowed_tools_list_blocks_tool(self):
"""Test pre_call_tool_check blocks tool when it's not in allowed_tools list"""
manager = MCPServerManager()
# Create server with allowed_tools list
server = MCPServer(
server_id="test-server",
name="test-server",
transport=MCPTransport.stdio,
allowed_tools=["allowed_tool", "another_allowed_tool"],
disallowed_tools=None,
)
# Mock dependencies
user_api_key_auth = MagicMock()
proxy_logging_obj = MagicMock()
# This should raise an HTTPException
with pytest.raises(HTTPException) as exc_info:
await manager.pre_call_tool_check(
name="blocked_tool",
arguments={"param": "value"},
server_name_from_prefix="test-server",
user_api_key_auth=user_api_key_auth,
proxy_logging_obj=proxy_logging_obj,
server=server,
)
assert exc_info.value.status_code == 403
assert (
"Tool blocked_tool is not allowed for server test-server"
in exc_info.value.detail["error"]
)
assert (
"Contact proxy admin to allow this tool" in exc_info.value.detail["error"]
)
@pytest.mark.asyncio
async def test_pre_call_tool_check_disallowed_tools_list_allows_tool(self):
"""Test pre_call_tool_check allows tool when it's not in disallowed_tools list"""
manager = MCPServerManager()
# Create server with disallowed_tools list
server = MCPServer(
server_id="test-server",
name="test-server",
transport=MCPTransport.stdio,
allowed_tools=None,
disallowed_tools=["banned_tool", "another_banned_tool"],
)
# Mock dependencies
user_api_key_auth = MagicMock()
proxy_logging_obj = MagicMock()
# Mock the async methods that pre_call_tool_check calls
proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(
return_value={}
)
proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={})
proxy_logging_obj.pre_call_hook = AsyncMock(return_value={})
# This should not raise an exception
await manager.pre_call_tool_check(
name="allowed_tool",
arguments={"param": "value"},
server_name_from_prefix="test-server",
user_api_key_auth=user_api_key_auth,
proxy_logging_obj=proxy_logging_obj,
server=server,
)
@pytest.mark.asyncio
async def test_pre_call_tool_check_disallowed_tools_list_blocks_tool(self):
"""Test pre_call_tool_check blocks tool when it's in disallowed_tools list"""
manager = MCPServerManager()
# Create server with disallowed_tools list
server = MCPServer(
server_id="test-server",
name="test-server",
transport=MCPTransport.stdio,
allowed_tools=None,
disallowed_tools=["banned_tool", "another_banned_tool"],
)
# Mock dependencies
user_api_key_auth = MagicMock()
proxy_logging_obj = MagicMock()
# This should raise an HTTPException
with pytest.raises(HTTPException) as exc_info:
await manager.pre_call_tool_check(
name="banned_tool",
arguments={"param": "value"},
server_name_from_prefix="test-server",
user_api_key_auth=user_api_key_auth,
proxy_logging_obj=proxy_logging_obj,
server=server,
)
assert exc_info.value.status_code == 403
assert (
"Tool banned_tool is not allowed for server test-server"
in exc_info.value.detail["error"]
)
assert (
"Contact proxy admin to allow this tool" in exc_info.value.detail["error"]
)
@pytest.mark.asyncio
async def test_pre_call_tool_check_no_restrictions_allows_any_tool(self):
"""Test pre_call_tool_check allows any tool when no restrictions are set"""
manager = MCPServerManager()
# Create server with no tool restrictions
server = MCPServer(
server_id="test-server",
name="test-server",
transport=MCPTransport.stdio,
allowed_tools=None,
disallowed_tools=None,
)
# Mock dependencies
user_api_key_auth = MagicMock()
proxy_logging_obj = MagicMock()
# Mock the async methods that pre_call_tool_check calls
proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(
return_value={}
)
proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={})
proxy_logging_obj.pre_call_hook = AsyncMock(return_value={})
# This should not raise an exception
await manager.pre_call_tool_check(
name="any_tool",
arguments={"param": "value"},
server_name_from_prefix="test-server",
user_api_key_auth=user_api_key_auth,
proxy_logging_obj=proxy_logging_obj,
server=server,
)
@pytest.mark.asyncio
async def test_pre_call_tool_check_allowed_tools_takes_precedence(self):
"""Test that allowed_tools list takes precedence over disallowed_tools list"""
manager = MCPServerManager()
# Create server with both allowed_tools and disallowed_tools
# Note: The logic in check_allowed_or_banned_tools prioritizes allowed_tools
server = MCPServer(
server_id="test-server",
name="test-server",
transport=MCPTransport.stdio,
allowed_tools=["tool1", "tool2"],
disallowed_tools=["tool2", "tool3"], # tool2 is in both lists
)
# Mock dependencies
user_api_key_auth = MagicMock()
proxy_logging_obj = MagicMock()
# Mock the async methods that pre_call_tool_check calls
proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(
return_value={}
)
proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={})
proxy_logging_obj.pre_call_hook = AsyncMock(return_value={})
# tool2 should be allowed since it's in allowed_tools (takes precedence)
await manager.pre_call_tool_check(
name="tool2",
arguments={"param": "value"},
server_name_from_prefix="test-server",
user_api_key_auth=user_api_key_auth,
proxy_logging_obj=proxy_logging_obj,
server=server,
)
# tool3 should be blocked since it's not in allowed_tools
with pytest.raises(HTTPException) as exc_info:
await manager.pre_call_tool_check(
name="tool3",
arguments={"param": "value"},
server_name_from_prefix="test-server",
user_api_key_auth=user_api_key_auth,
proxy_logging_obj=proxy_logging_obj,
server=server,
)
assert exc_info.value.status_code == 403
assert (
"Tool tool3 is not allowed for server test-server"
in exc_info.value.detail["error"]
)
if __name__ == "__main__":
pytest.main([__file__])