From 529f67f5448127cd7cfe9ff1cc5cdb5e62467bc9 Mon Sep 17 00:00:00 2001 From: "Jugal D. Bhatt" <55304795+jugaldb@users.noreply.github.com> Date: Sat, 5 Jul 2025 12:38:43 +0530 Subject: [PATCH] Add mcp server segregation comma separated support (#12326) * add mcp comma separated support * dont support legacy json array * lint --- docs/my-website/docs/mcp.md | 16 ++--- .../mcp_server/auth/user_api_key_auth_mcp.py | 17 +++--- .../auth/test_user_api_key_auth_mcp.py | 60 ++++++++++++------- .../src/components/mcp_tools/mcp_connect.tsx | 27 ++++++++- 4 files changed, 83 insertions(+), 37 deletions(-) diff --git a/docs/my-website/docs/mcp.md b/docs/my-website/docs/mcp.md index 3b687cccb6..efa51be866 100644 --- a/docs/my-website/docs/mcp.md +++ b/docs/my-website/docs/mcp.md @@ -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 @@ -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 '/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 '/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. @@ -253,7 +255,7 @@ This configuration restricts the request to only use tools from the specified MC "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): ``` - + \ No newline at end of file diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index a16928bf3d..7fface60ed 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -1,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 [] \ No newline at end of file diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 2712815e3e..3d3f7f336d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -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"] \ No newline at end of file diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_connect.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_connect.tsx index 8944933b6c..344f1a90b8 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_connect.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_connect.tsx @@ -72,7 +72,8 @@ const FeatureCard: React.FC = ({ 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 = ({ /> Segregate tools to just use {serverName} tools + {useServerHeader && ( + +

Specify one or more MCP servers using a comma-separated list:

+
    +
  • Single server: "Server1"
    • +
  • Multiple servers: "Server1,Server2,Server3"
    • +
+

Note: Server names with spaces will be automatically converted to use underscores.

+ + } + /> + )} )} {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;