diff --git a/docs/my-website/docs/mcp.md b/docs/my-website/docs/mcp.md index 3a4de87bec..3b687cccb6 100644 --- a/docs/my-website/docs/mcp.md +++ b/docs/my-website/docs/mcp.md @@ -83,7 +83,6 @@ mcp_servers: - ## Using your MCP @@ -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: ``` + - +## 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 -/mcp + + + +```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": "/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. - + -#### 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 '/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": "/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 -``` + -**Usage:** + -```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 = "/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": "/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. + - ## 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. 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 a9782db994..a16928bf3d 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 @@ -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}") diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 8dd4a6be6b..de488e6ca7 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -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 ############# diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 7737bf73a2..e36e1267a4 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -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)}") + + 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 564326a28c..8944933b6c 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_connect.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_connect.tsx @@ -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 = ({ + icon, + title, + description, + children, + serverName, +}) => { + const [useServerHeader, setUseServerHeader] = useState(false); + + const getHeadersConfig = () => { + const headers: Record = { + "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 ( + +
+ {icon} +
+ + {title} + + {description} +
+
+ {serverName && (title === "Implementation Example" || title === "Configuration") && ( + +
+ + Segregate tools to just use {serverName} tools +
+ + )} + {React.Children.map(children, child => { + if (React.isValidElement(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; + })} + + ); +}; const MCPConnect: React.FC = () => { const proxyBaseUrl = getProxyBaseUrl(); const [copiedStates, setCopiedStates] = useState>({}); + const [serverHeaders, setServerHeaders] = useState>({ + 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 = { + "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 = () => { ); - const FeatureCard: React.FC<{ - icon: React.ReactNode; - title: string; - description: string; - children: React.ReactNode; - }> = ({ icon, title, description, children }) => ( - -
-
- {icon} -
-
- {title} - {description} -
-
- {children} - - ); - const StepCard: React.FC<{ step: number; title: string; @@ -167,6 +255,7 @@ const MCPConnect: React.FC = () => { icon={} title="Implementation Example" description="Complete cURL example for using the LiteLLM Proxy Responses API" + serverName={currentServer} > { icon={} title="Implementation Example" description="Complete cURL example for using the Responses API" + serverName="Zapier Gmail" > { title="Add Configuration" > Copy the JSON configuration below and paste it into Cursor, then save with Cmd+S or Ctrl+S - } + title="Configuration" + description="Cursor MCP configuration" + serverName="Zapier Gmail" + > + + copyKey="cursor-config" + /> + @@ -360,6 +457,13 @@ copyKey="cursor-config" code={`${proxyBaseUrl}/mcp`} copyKey="http-server-url" /> +