hanzo/chat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Segregate MCP tools on connections using headers (#12296)

Browse Source 
* Add get tools segregation

* add ui changes (#12302)

* resolve comments

* add mapped tests

* remove advanced settings (#12323)
This commit is contained in:
Jugal D. Bhatt
2025-07-04 11:45:51 -07:00
committed by GitHub
parent a94c91d3d6
commit c75b416621
5 changed files with 338 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/my-website/docs/mcp.md
+73 -73
View File
@@ -83,7 +83,6 @@ mcp_servers:
</Tabs>
## Using your MCP
<Tabs>
@@ -159,7 +158,7 @@ Use tools directly from Cursor IDE with LiteLLM MCP:
2. **Navigate to MCP Tools**: Go to the "MCP Tools" tab and click "New MCP Server"
3. **Add Configuration**: Copy and paste the JSON configuration below, then save with `Cmd+S` or `Ctrl+S`
```json title="Cursor MCP Configuration" showLineNumbers
```json title="Basic Cursor MCP Configuration" showLineNumbers
{
"mcpServers": {
"LiteLLM": {
@@ -173,98 +172,99 @@ Use tools directly from Cursor IDE with LiteLLM MCP:
```
</TabItem>
</Tabs>
<TabItem value="http" label="Streamable HTTP">
## Segregating MCP Server Access
#### Connect via Streamable HTTP Transport
You can choose to access specific MCP servers and only list their tools using the `x-mcp-servers` header. This header allows you to:
- Limit tool access to one or more specific MCP servers
- Control which tools are available in different environments or use cases
Connect to LiteLLM MCP using HTTP transport. Compatible with any MCP client that supports HTTP streaming:
The header accepts a JSON array of server names, where:
- Server names with spaces should be replaced with underscores
- Multiple servers can be specified: `["Server1", "Server2", "Server3"]`
- If the header is not provided, tools from all available MCP servers will be accessible
**Server URL:**
```text showLineNumbers
<your-litellm-proxy-base-url>/mcp
<Tabs>
<TabItem value="openai" label="OpenAI API">
```bash title="cURL Example with Server Segregation" showLineNumbers
curl --location 'https://api.openai.com/v1/responses' \
--header 'Content-Type: application/json' \
--header "Authorization: Bearer $OPENAI_API_KEY" \
--data '{
"model": "gpt-4o",
"tools": [
{
"type": "mcp",
"server_label": "litellm",
"server_url": "<your-litellm-proxy-base-url>/mcp",
"require_approval": "never",
"headers": {
"x-litellm-api-key": "Bearer YOUR_LITELLM_API_KEY",
"x-mcp-servers": "[\"Zapier_Gmail\"]"
}
}
],
"input": "Run available tools",
"tool_choice": "required"
}'
```
**Headers:**
```text showLineNumbers
x-litellm-api-key: Bearer YOUR_LITELLM_API_KEY
```
This URL can be used with any MCP client that supports HTTP transport. Refer to your client documentation to determine the appropriate transport method.
In this example, the request will only have access to tools from the "Zapier_Gmail" MCP server.
</TabItem>
<TabItem value="fastmcp" label="Python FastMCP">
<TabItem value="litellm" label="LiteLLM Proxy">
#### Connect via Python FastMCP Client
Use the Python FastMCP client to connect to your LiteLLM MCP server:
**Installation:**
```bash title="Install FastMCP" showLineNumbers
pip install fastmcp
```bash title="cURL Example with Server Segregation" showLineNumbers
curl --location '<your-litellm-proxy-base-url>/v1/responses' \
--header 'Content-Type: application/json' \
--header "Authorization: Bearer $LITELLM_API_KEY" \
--data '{
"model": "gpt-4o",
"tools": [
{
"type": "mcp",
"server_label": "litellm",
"server_url": "<your-litellm-proxy-base-url>/mcp",
"require_approval": "never",
"headers": {
"x-litellm-api-key": "Bearer YOUR_LITELLM_API_KEY",
"x-mcp-servers": "[\"Zapier_Gmail\"]"
}
}
],
"input": "Run available tools",
"tool_choice": "required"
}'
```
or with uv:
This configuration restricts the request to only use tools from the specified MCP server.
```bash title="Install with uv" showLineNumbers
uv pip install fastmcp
```
</TabItem>
**Usage:**
<TabItem value="cursor" label="Cursor IDE">
```python title="Python FastMCP Example" showLineNumbers
import asyncio
import json
from fastmcp import Client
from fastmcp.client.transports import StreamableHttpTransport
# Create the transport with your LiteLLM MCP server URL
server_url = "<your-litellm-proxy-base-url>/mcp"
transport = StreamableHttpTransport(
server_url,
headers={
"x-litellm-api-key": "Bearer YOUR_LITELLM_API_KEY"
```json title="Cursor MCP Configuration with Server Segregation" showLineNumbers
{
"mcpServers": {
"LiteLLM": {
"url": "<your-litellm-proxy-base-url>/mcp",
"headers": {
"x-litellm-api-key": "Bearer $LITELLM_API_KEY",
"x-mcp-servers": "[\"Zapier_Gmail\"]"
}
}
)
# Initialize the client with the transport
client = Client(transport=transport)
async def main():
# Connection is established here
print("Connecting to LiteLLM MCP server...")
async with client:
print(f"Client connected: {client.is_connected()}")
# Make MCP calls within the context
print("Fetching available tools...")
tools = await client.list_tools()
print(f"Available tools: {json.dumps([t.name for t in tools], indent=2)}")
# Example: Call a tool (replace 'tool_name' with an actual tool name)
if tools:
tool_name = tools[0].name
print(f"Calling tool: {tool_name}")
# Call the tool with appropriate arguments
result = await client.call_tool(tool_name, arguments={})
print(f"Tool result: {result}")
# Run the example
if __name__ == "__main__":
asyncio.run(main())
}
}
```
This configuration in Cursor IDE settings will limit tool access to only the specified MCP server.
</TabItem>
</Tabs>
## 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.
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
+3
View File
@@ -53,11 +53,14 @@ class MCPRequestHandler:
)
mcp_auth_header = headers.get(MCPRequestHandler.LITELLM_MCP_AUTH_HEADER_NAME)
mcp_servers_header = headers.get(MCPRequestHandler.LITELLM_MCP_SERVERS_HEADER_NAME)
verbose_logger.debug(f"Raw MCP servers header: {mcp_servers_header}")
mcp_servers = None
if mcp_servers_header:
try:
mcp_servers = json.loads(mcp_servers_header)
verbose_logger.debug(f"Parsed MCP servers: {mcp_servers}")
if not isinstance(mcp_servers, list):
verbose_logger.debug("MCP servers header is not a list, setting to None")
mcp_servers = None
except (json.JSONDecodeError, TypeError, ValueError) as e:
verbose_logger.debug(f"Error parsing mcp_servers header: {e}")
litellm/proxy/_experimental/mcp_server/server.py
+62 -12
View File
@@ -19,7 +19,7 @@ from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
from litellm.proxy._experimental.mcp_server.utils import (
LITELLM_MCP_SERVER_NAME,
LITELLM_MCP_SERVER_VERSION,
LITELLM_MCP_SERVER_DESCRIPTION,
LITELLM_MCP_SERVER_DESCRIPTION, normalize_server_name,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.mcp_server.mcp_server_manager import MCPInfo
@@ -169,13 +169,15 @@ if MCP_AVAILABLE:
List all available tools
"""
# Get user authentication from context variable
user_api_key_auth, mcp_auth_header = get_auth_context()
user_api_key_auth, mcp_auth_header, mcp_servers = get_auth_context()
verbose_logger.debug(
f"MCP list_tools - User API Key Auth from context: {user_api_key_auth}"
)
# Get mcp_servers from context variable
return await _list_mcp_tools(
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
)
@server.call_tool()
@@ -196,7 +198,7 @@ if MCP_AVAILABLE:
HTTPException: If tool not found or arguments missing
"""
# Validate arguments
user_api_key_auth, mcp_auth_header = get_auth_context()
user_api_key_auth, mcp_auth_header, _ = get_auth_context()
verbose_logger.debug(
f"MCP mcp_server_tool_call - User API Key Auth from context: {user_api_key_auth}"
)
@@ -216,15 +218,53 @@ if MCP_AVAILABLE:
############ Helper Functions ##########################
########################################################
async def _get_tools_from_mcp_servers(
user_api_key_auth: Optional[UserAPIKeyAuth],
mcp_auth_header: Optional[str],
mcp_servers: Optional[List[str]]
) -> List[MCPTool]:
"""
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
mcp_servers: Optional list of server names to filter by
Returns:
List[MCPTool]: List of tools from the specified or all allowed MCP servers
"""
if mcp_servers:
# If mcp_servers header is present, only get tools from specified servers
tools = []
for server_id in await global_mcp_server_manager.get_allowed_mcp_servers(user_api_key_auth):
server = global_mcp_server_manager.get_mcp_server_by_id(server_id)
if server and any(normalize_server_name(server.name) == normalize_server_name(s) for s in mcp_servers):
server_tools = await global_mcp_server_manager._get_tools_from_server(
server=server,
mcp_auth_header=mcp_auth_header,
)
tools.extend(server_tools)
return tools
else:
# If no mcp_servers header, get tools from all allowed servers
return await global_mcp_server_manager.list_tools(
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
)
async def _list_mcp_tools(
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
mcp_auth_header: Optional[str] = None,
mcp_servers: Optional[List[str]] = None,
) -> List[MCPTool]:
"""
List all available tools
Args:
user_api_key_auth: User authentication info for access control
mcp_auth_header: Optional auth header for MCP server
mcp_servers: Optional list of server names to filter by
"""
tools = []
for tool in global_mcp_tool_registry.list_tools():
@@ -239,12 +279,13 @@ if MCP_AVAILABLE:
"GLOBAL MCP TOOLS: %s", global_mcp_tool_registry.list_tools()
)
tools_from_mcp_servers: List[MCPTool] = (
await global_mcp_server_manager.list_tools(
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
)
# Get tools from MCP servers
tools_from_mcp_servers = await _get_tools_from_mcp_servers(
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers
)
verbose_logger.debug("TOOLS FROM MCP SERVERS: %s", tools_from_mcp_servers)
if tools_from_mcp_servers is not None:
tools.extend(tools_from_mcp_servers)
@@ -364,10 +405,12 @@ if MCP_AVAILABLE:
user_api_key_auth, mcp_auth_header, mcp_servers = (
await MCPRequestHandler.process_mcp_request(scope)
)
verbose_logger.debug(f"MCP request headers - mcp_servers: {mcp_servers}")
# Set the auth context variable for easy access in MCP functions
set_auth_context(
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
)
# Ensure session managers are initialized
@@ -392,6 +435,7 @@ if MCP_AVAILABLE:
set_auth_context(
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
)
# Ensure session managers are initialized
@@ -432,21 +476,27 @@ if MCP_AVAILABLE:
############ Auth Context Functions ####################
########################################################
def set_auth_context(user_api_key_auth: UserAPIKeyAuth, mcp_auth_header: Optional[str] = None) -> None:
def set_auth_context(
user_api_key_auth: UserAPIKeyAuth,
mcp_auth_header: Optional[str] = None,
mcp_servers: Optional[List[str]] = None,
) -> None:
"""
Set the UserAPIKeyAuth in the auth context variable.
Args:
user_api_key_auth: UserAPIKeyAuth object
mcp_auth_header: MCP auth header to be passed to the MCP server
mcp_servers: Optional list of server names to filter by
"""
auth_user = MCPAuthenticatedUser(
user_api_key_auth=user_api_key_auth,
mcp_auth_header=mcp_auth_header,
mcp_servers=mcp_servers,
)
auth_context_var.set(auth_user)
def get_auth_context() -> Tuple[Optional[UserAPIKeyAuth], Optional[str]]:
def get_auth_context() -> Tuple[Optional[UserAPIKeyAuth], Optional[str], Optional[List[str]]]:
"""
Get the UserAPIKeyAuth from the auth context variable.
@@ -455,8 +505,8 @@ if MCP_AVAILABLE:
"""
auth_user = auth_context_var.get()
if auth_user and isinstance(auth_user, MCPAuthenticatedUser):
return auth_user.user_api_key_auth, auth_user.mcp_auth_header
return None, None
return auth_user.user_api_key_auth, auth_user.mcp_auth_header, auth_user.mcp_servers
return None, None, None
########################################################
############ End of Auth Context Functions #############
tests/mcp_tests/test_mcp_server.py
+71
View File
@@ -648,3 +648,74 @@ async def test_list_tools_rest_api_success():
global_mcp_server_manager.tool_name_to_mcp_server_name_mapping = original_tool_mapping
@pytest.mark.asyncio
async def test_get_tools_from_mcp_servers():
"""Test _get_tools_from_mcp_servers function with both specific and no server filters"""
from litellm.proxy._experimental.mcp_server.server import _get_tools_from_mcp_servers
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServer, MCPTransport, MCPSpecVersion
# Mock data
mock_user_auth = UserAPIKeyAuth(api_key="test_key", user_id="test_user")
mock_auth_header = "Bearer test_token"
mock_server_1 = MCPServer(
server_id="server1_id",
name="server1",
url="http://test1.com",
transport=MCPTransport.http,
spec_version=MCPSpecVersion.nov_2024
)
mock_server_2 = MCPServer(
server_id="server2_id",
name="server2",
url="http://test2.com",
transport=MCPTransport.http,
spec_version=MCPSpecVersion.nov_2024
)
mock_tool_1 = MCPTool(name="tool1", description="test tool 1", inputSchema={})
mock_tool_2 = MCPTool(name="tool2", description="test tool 2", inputSchema={})
# Test Case 1: With specific MCP servers
try:
# Mock the necessary methods
def mock_get_server_by_id(server_id):
if server_id == "server1_id":
return mock_server_1
elif server_id == "server2_id":
return mock_server_2
return None
with patch('litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPServerManager.get_allowed_mcp_servers',
new_callable=AsyncMock, return_value=["server1_id", "server2_id"]), \
patch('litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPServerManager.get_mcp_server_by_id',
side_effect=mock_get_server_by_id), \
patch('litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPServerManager._get_tools_from_server',
new_callable=AsyncMock, return_value=[mock_tool_1]):
# Test with specific servers
result = await _get_tools_from_mcp_servers(
user_api_key_auth=mock_user_auth,
mcp_auth_header=mock_auth_header,
mcp_servers=["server1"]
)
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
with patch('litellm.proxy._experimental.mcp_server.mcp_server_manager.MCPServerManager.list_tools',
new_callable=AsyncMock, return_value=[mock_tool_1, mock_tool_2]):
result = await _get_tools_from_mcp_servers(
user_api_key_auth=mock_user_auth,
mcp_auth_header=mock_auth_header,
mcp_servers=None
)
assert len(result) == 2, "Should return tools from all servers"
assert result[0].name == "tool1" and result[1].name == "tool2", "Should return tools from all servers"
except AssertionError as e:
pytest.fail(f"Test failed: {str(e)}")
except Exception as e:
pytest.fail(f"Unexpected error in tests: {str(e)}")
ui/litellm-dashboard/src/components/mcp_tools/mcp_connect.tsx
+129 -25
View File
@@ -8,6 +8,10 @@ import {
Alert,
Button,
message,
Switch,
Input,
Form,
Collapse,
} from "antd";
import {
TabPanel,
@@ -33,10 +37,102 @@ import {
import { getProxyBaseUrl } from "../networking";
const { Title, Text } = Typography;
const { Panel } = Collapse;
interface CodeBlockProps {
code: string;
title?: string;
copyKey: string;
className?: string;
}
interface FeatureCardProps {
icon: React.ReactNode;
title: string;
description: string;
children: React.ReactNode;
serverName?: string;
}
const FeatureCard: React.FC<FeatureCardProps> = ({
icon,
title,
description,
children,
serverName,
}) => {
const [useServerHeader, setUseServerHeader] = useState(false);
const getHeadersConfig = () => {
const headers: Record<string, any> = {
"x-litellm-api-key": "Bearer YOUR_LITELLM_API_KEY"
};
if (useServerHeader && serverName) {
// Replace spaces with underscores in server name
const formattedServerName = serverName.replace(/\s+/g, '_');
headers["x-mcp-servers"] = `["${formattedServerName}"]`;
}
return headers;
};
return (
<Card className="border border-gray-200">
<div className="flex items-center gap-3 mb-3">
<span className="p-2 rounded-lg bg-gray-50">{icon}</span>
<div>
<Title level={5} className="mb-0">
{title}
</Title>
<Text className="text-gray-600">{description}</Text>
</div>
</div>
{serverName && (title === "Implementation Example" || title === "Configuration") && (
<Form.Item
className="mb-4"
>
<div className="flex items-center gap-2 mb-2">
<Switch
size="small"
checked={useServerHeader}
onChange={setUseServerHeader}
/>
<Text className="text-sm">Segregate tools to just use {serverName} tools</Text>
</div>
</Form.Item>
)}
{React.Children.map(children, child => {
if (React.isValidElement<CodeBlockProps>(child) &&
child.props.hasOwnProperty('code') &&
child.props.hasOwnProperty('copyKey')) {
const code = child.props.code;
if (code && code.includes('"headers":')) {
return React.cloneElement(child, {
code: code.replace(
/"headers":\s*{[^}]*}/,
`"headers": ${JSON.stringify(getHeadersConfig(), null, 8)}`
)
});
}
}
return child;
})}
</Card>
);
};
const MCPConnect: React.FC = () => {
const proxyBaseUrl = getProxyBaseUrl();
const [copiedStates, setCopiedStates] = useState<Record<string, boolean>>({});
const [serverHeaders, setServerHeaders] = useState<Record<string, string[]>>({
openai: [],
litellm: [],
cursor: [],
http: []
});
const [currentServer] = useState("Zapier Gmail"); // This should match the current server being viewed
const copyToClipboard = async (text: string, key: string) => {
try {
@@ -51,6 +147,18 @@ const MCPConnect: React.FC = () => {
}
};
const getHeadersConfig = (type: string) => {
const headers: Record<string, any> = {
"x-litellm-api-key": "Bearer YOUR_LITELLM_API_KEY"
};
if (serverHeaders[type]?.length > 0) {
headers["x-mcp-servers"] = JSON.stringify(serverHeaders[type]);
}
return headers;
};
const CodeBlock: React.FC<{
code: string;
copyKey: string;
@@ -83,26 +191,6 @@ const MCPConnect: React.FC = () => {
</div>
);
const FeatureCard: React.FC<{
icon: React.ReactNode;
title: string;
description: string;
children: React.ReactNode;
}> = ({ icon, title, description, children }) => (
<Card className="border border-gray-200 shadow-sm hover:shadow-md transition-shadow duration-200">
<div className="flex items-start gap-3 mb-4">
<div className="flex-shrink-0 w-8 h-8 bg-blue-50 rounded-lg flex items-center justify-center">
{icon}
</div>
<div className="flex-1">
<Title level={5} className="mb-1 text-gray-800">{title}</Title>
<Text className="text-gray-600">{description}</Text>
</div>
</div>
{children}
</Card>
);
const StepCard: React.FC<{
step: number;
title: string;
@@ -167,6 +255,7 @@ const MCPConnect: React.FC = () => {
icon={<Code className="text-emerald-600" size={16} />}
title="Implementation Example"
description="Complete cURL example for using the LiteLLM Proxy Responses API"
serverName={currentServer}
>
<CodeBlock
code={`curl --location '${proxyBaseUrl}/v1/responses' \\
@@ -252,6 +341,7 @@ const MCPConnect: React.FC = () => {
icon={<Code className="text-blue-600" size={16} />}
title="Implementation Example"
description="Complete cURL example for using the Responses API"
serverName="Zapier Gmail"
>
<CodeBlock
code={`curl --location 'https://api.openai.com/v1/responses' \\
@@ -315,19 +405,26 @@ const MCPConnect: React.FC = () => {
title="Add Configuration"
>
<Text className="text-gray-600 mb-3">Copy the JSON configuration below and paste it into Cursor, then save with <code className="bg-gray-100 px-2 py-1 rounded">Cmd+S</code> or <code className="bg-gray-100 px-2 py-1 rounded">Ctrl+S</code></Text>
<CodeBlock
code={`{
<FeatureCard
icon={<Code className="text-purple-600" size={16} />}
title="Configuration"
description="Cursor MCP configuration"
serverName="Zapier Gmail"
>
<CodeBlock
code={`{
"mcpServers": {
"LiteLLM": {
"url": "${proxyBaseUrl}/mcp",
"headers": {
"x-litellm-api-key": "Bearer $LITELLM_API_KEY"
"x-litellm-api-key": "Bearer YOUR_LITELLM_API_KEY"
}
}
}
}`}
copyKey="cursor-config"
/>
copyKey="cursor-config"
/>
</FeatureCard>
</StepCard>
</Space>
</Card>
@@ -360,6 +457,13 @@ copyKey="cursor-config"
code={`${proxyBaseUrl}/mcp`}
copyKey="http-server-url"
/>
<CodeBlock
title="Headers Configuration"
code={JSON.stringify({
"x-litellm-api-key": "Bearer YOUR_LITELLM_API_KEY"
}, null, 2)}
copyKey="http-headers"
/>
<div className="mt-4">
<Button
type="link"