tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-06 18:23:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add mcp server segregation comma separated support (#12326)

Browse Source 
* add mcp comma separated support

* dont support legacy json array

* lint
This commit is contained in:
Jugal D. Bhatt
2025-07-05 00:08:43 -07:00
committed by GitHub
parent 4b3b4f8dff
commit 529f67f544
4 changed files with 83 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/my-website/docs/mcp.md
+9 -7
View File
@@ -180,9 +180,11 @@ You can choose to access specific MCP servers and only list their tools using th
- Limit tool access to one or more specific MCP servers
- Control which tools are available in different environments or use cases
The header accepts a JSON array of server names, where:
The header accepts either:
1. A comma-separated list of server names: `"Zapier_Gmail,Server2,Server3"`
Notes:
- 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
<Tabs>
@@ -202,7 +204,7 @@ curl --location 'https://api.openai.com/v1/responses' \
"require_approval": "never",
"headers": {
"x-litellm-api-key": "Bearer YOUR_LITELLM_API_KEY",
"x-mcp-servers": "[\"Zapier_Gmail\"]"
"x-mcp-servers": "Zapier_Gmail"
}
}
],
@@ -231,7 +233,7 @@ curl --location '<your-litellm-proxy-base-url>/v1/responses' \
"require_approval": "never",
"headers": {
"x-litellm-api-key": "Bearer YOUR_LITELLM_API_KEY",
"x-mcp-servers": "[\"Zapier_Gmail\"]"
"x-mcp-servers": "Zapier_Gmail,Server2"
}
}
],
@@ -240,7 +242,7 @@ curl --location '<your-litellm-proxy-base-url>/v1/responses' \
}'
```
This configuration restricts the request to only use tools from the specified MCP server.
This configuration restricts the request to only use tools from the specified MCP servers.
</TabItem>
@@ -253,7 +255,7 @@ This configuration restricts the request to only use tools from the specified MC
"url": "<your-litellm-proxy-base-url>/mcp",
"headers": {
"x-litellm-api-key": "Bearer $LITELLM_API_KEY",
"x-mcp-servers": "[\"Zapier_Gmail\"]"
"x-mcp-servers": "Zapier_Gmail,Server2"
}
}
}
@@ -715,4 +717,4 @@ async with stdio_client(server_params) as (read, write):
```
</TabItem>
</Tabs>
</Tabs>
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
+9 -8
View File
@@ -1,5 +1,4 @@
from typing import List, Optional, Tuple
import json
from starlette.datastructures import Headers
from starlette.requests import Request
@@ -55,17 +54,19 @@ class MCPRequestHandler:
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:
if mcp_servers_header is not None: # Changed from 'if mcp_servers_header:' to handle empty strings
try:
mcp_servers = json.loads(mcp_servers_header)
# Parse as comma-separated list
mcp_servers = [s.strip() for s in mcp_servers_header.split(",") if s.strip()]
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:
except Exception as e:
verbose_logger.debug(f"Error parsing mcp_servers header: {e}")
mcp_servers = None
# If we got an empty string or parsing resulted in no servers, return empty list
if mcp_servers_header == "" or (mcp_servers is not None and len(mcp_servers) == 0):
mcp_servers = []
# Create a proper Request object with mock body method to avoid ASGI receive channel issues
request = Request(scope=scope)
@@ -219,4 +220,4 @@ class MCPRequestHandler:
if object_permissions is None:
return []
return object_permissions.mcp_servers or []
return object_permissions.mcp_servers or []
tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
+40 -20
View File
@@ -228,7 +228,7 @@ class TestMCPRequestHandler:
[
(b"x-litellm-api-key", b"test-api-key"),
(b"x-mcp-auth", b"test-mcp-auth"),
(b"x-mcp-servers", b'["server1", "server2"]'),
(b"x-mcp-servers", b"server1,server2"),
],
{
"api_key": "test-api-key",
@@ -245,35 +245,35 @@ class TestMCPRequestHandler:
"mcp_servers": None,
}
),
# Test case 3: Invalid JSON in mcp_servers
# Test case 3: Invalid format in mcp_servers
(
[
(b"x-litellm-api-key", b"test-api-key"),
(b"x-mcp-servers", b'invalid-json'),
(b"x-mcp-servers", b"[invalid,format]"),
],
{
"api_key": "test-api-key",
"mcp_auth": None,
"mcp_servers": None,
"mcp_servers": ["[invalid", "format]"],
}
),
# Test case 4: mcp_servers not a list
# Test case 4: Single server
(
[
(b"x-litellm-api-key", b"test-api-key"),
(b"x-mcp-servers", b'{"key": "value"}'),
(b"x-mcp-servers", b"server1"),
],
{
"api_key": "test-api-key",
"mcp_auth": None,
"mcp_servers": None,
"mcp_servers": ["server1"],
}
),
# Test case 5: Empty mcp_servers list
# Test case 5: Empty server string
(
[
(b"x-litellm-api-key", b"test-api-key"),
(b"x-mcp-servers", b'[]'),
(b"x-mcp-servers", b""),
],
{
"api_key": "test-api-key",
@@ -285,7 +285,7 @@ class TestMCPRequestHandler:
(
[
(b"authorization", b"Bearer test-api-key"),
(b"x-mcp-servers", b'["server1"]'),
(b"x-mcp-servers", b"server1"),
],
{
"api_key": "Bearer test-api-key",
@@ -298,7 +298,7 @@ class TestMCPRequestHandler:
[
(b"X-LITELLM-API-KEY", b"test-api-key"),
(b"X-MCP-AUTH", b"test-mcp-auth"),
(b"X-MCP-SERVERS", b'["server1"]'),
(b"X-MCP-SERVERS", b"server1"),
],
{
"api_key": "test-api-key",
@@ -306,6 +306,18 @@ class TestMCPRequestHandler:
"mcp_servers": ["server1"],
}
),
# Test case 8: Multiple servers with spaces
(
[
(b"x-litellm-api-key", b"test-api-key"),
(b"x-mcp-servers", b"server1, server2, server3"),
],
{
"api_key": "test-api-key",
"mcp_auth": None,
"mcp_servers": ["server1", "server2", "server3"],
}
),
]
)
async def test_header_extraction(self, headers, expected_result):
@@ -332,16 +344,24 @@ class TestMCPRequestHandler:
# Verify MCP servers
mcp_servers_header = extracted_headers.get(SpecialHeaders.mcp_servers.value)
if mcp_servers_header:
mcp_servers = None
if mcp_servers_header is not None: # Changed from 'if mcp_servers_header:' to handle empty strings
try:
mcp_servers = json.loads(mcp_servers_header)
if not isinstance(mcp_servers, list):
mcp_servers = None
except:
# First try to parse as JSON array for backward compatibility
try:
mcp_servers = json.loads(mcp_servers_header)
if not isinstance(mcp_servers, list):
mcp_servers = None
except (json.JSONDecodeError, TypeError, ValueError):
# If JSON parsing fails, treat as comma-separated list
mcp_servers = [s.strip() for s in mcp_servers_header.split(",") if s.strip()]
except Exception:
mcp_servers = None
else:
mcp_servers = None
# If we got an empty string or parsing resulted in no servers, return empty list
if mcp_servers_header == "" or (mcp_servers is not None and len(mcp_servers) == 0):
mcp_servers = []
assert mcp_servers == expected_result["mcp_servers"]
# Test the full process_mcp_request method
@@ -362,4 +382,4 @@ class TestMCPRequestHandler:
# Assert the results
assert auth_result == mock_auth_result
assert mcp_auth_header == expected_result["mcp_auth"]
assert mcp_servers_result == expected_result["mcp_servers"]
assert mcp_servers_result == expected_result["mcp_servers"]
ui/litellm-dashboard/src/components/mcp_tools/mcp_connect.tsx
+25 -2
View File
@@ -72,7 +72,8 @@ const FeatureCard: React.FC<FeatureCardProps> = ({
if (useServerHeader && serverName) {
// Replace spaces with underscores in server name
const formattedServerName = serverName.replace(/\s+/g, '_');
headers["x-mcp-servers"] = `["${formattedServerName}"]`;
// Use simple comma-separated format
headers["x-mcp-servers"] = formattedServerName;
}
return headers;
@@ -101,6 +102,24 @@ const FeatureCard: React.FC<FeatureCardProps> = ({
/>
<Text className="text-sm">Segregate tools to just use {serverName} tools</Text>
</div>
{useServerHeader && (
<Alert
className="mt-2"
type="info"
showIcon
message="MCP Server Header Format"
description={
<div>
<p>Specify one or more MCP servers using a comma-separated list:</p>
<ul>
<li><strong>Single server:</strong> "Server1"</li>
<li><strong>Multiple servers:</strong> "Server1,Server2,Server3"</li>
</ul>
<p>Note: Server names with spaces will be automatically converted to use underscores.</p>
</div>
}
/>
)}
</Form.Item>
)}
{React.Children.map(children, child => {
@@ -153,7 +172,11 @@ const MCPConnect: React.FC = () => {
};
if (serverHeaders[type]?.length > 0) {
headers["x-mcp-servers"] = JSON.stringify(serverHeaders[type]);
// Format server names (replace spaces with underscores)
const formattedServers = serverHeaders[type].map(s => s.replace(/\s+/g, '_'));
// Use comma-separated format
headers["x-mcp-servers"] = formattedServers.join(',');
}
return headers;