From afd382d09f154b8127efcf855707ca5889b4acde Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 9 Jul 2025 12:30:20 -0700 Subject: [PATCH] [Feat] MCP Gateway - Allow customizing what client side header to use (#12460) * add _get_mcp_auth_header_from_headers * test_process_mcp_request_with_custom_auth_header * Using a different Authentication Header * fix customize MCP Auth header name --- docs/my-website/docs/mcp.md | 52 ++++++ .../mcp_server/auth/user_api_key_auth_mcp.py | 38 +++- .../auth/test_user_api_key_auth_mcp.py | 175 +++++++++++++++++- 3 files changed, 263 insertions(+), 2 deletions(-) diff --git a/docs/my-website/docs/mcp.md b/docs/my-website/docs/mcp.md index 6601766796..3ef2238753 100644 --- a/docs/my-website/docs/mcp.md +++ b/docs/my-website/docs/mcp.md @@ -442,6 +442,58 @@ if __name__ == "__main__": +### Customize the MCP Auth Header Name + +By default, LiteLLM uses `x-mcp-auth` to pass your credentials to MCP servers. You can change this header name in one of the following ways: +1. Set the `LITELLM_MCP_CLIENT_SIDE_AUTH_HEADER_NAME` environment variable + +```bash title="Environment Variable" showLineNumbers +export LITELLM_MCP_CLIENT_SIDE_AUTH_HEADER_NAME="authorization" +``` + + +2. Set the `mcp_client_side_auth_header_name` in the general settings on the config.yaml file + +```yaml title="config.yaml" showLineNumbers +model_list: + - model_name: gpt-4o + litellm_params: + model: openai/gpt-4o + api_key: sk-xxxxxxx + +general_settings: + mcp_client_side_auth_header_name: "authorization" +``` + +#### Using the authorization header + +In this example the `authorization` header will be passed to the MCP server for authentication. + +```bash title="cURL with authorization header" 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", + "authorization": "Bearer sk-zapier-token-123" + } + } + ], + "input": "Run available tools", + "tool_choice": "required" +}' +``` + + + ## ✨ MCP Permission Management LiteLLM supports managing permissions for MCP Servers by Keys, Teams, Organizations (entities) on LiteLLM. When a MCP client attempts to list tools, LiteLLM will only return the tools the entity has permissions to access. 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 fba7928e7a..cfe357709c 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 @@ -51,7 +51,7 @@ class MCPRequestHandler: litellm_api_key = ( MCPRequestHandler.get_litellm_api_key_from_headers(headers) or "" ) - mcp_auth_header = headers.get(MCPRequestHandler.LITELLM_MCP_AUTH_HEADER_NAME) + mcp_auth_header = MCPRequestHandler._get_mcp_auth_header_from_headers(headers) 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 @@ -83,6 +83,42 @@ class MCPRequestHandler: ) return validated_user_api_key_auth, mcp_auth_header, mcp_servers + + + @staticmethod + def _get_mcp_auth_header_from_headers(headers: Headers) -> Optional[str]: + """ + Get the header passed to LiteLLM to pass to downstream MCP servers + + By default litellm will check for the header `x-mcp-auth` by setting one of the following: + 1. `LITELLM_MCP_CLIENT_SIDE_AUTH_HEADER_NAME` as an environment variable + 2. `mcp_client_side_auth_header_name` in the general settings on the config.yaml file + + Support this auth: https://docs.litellm.ai/docs/mcp#using-your-mcp-with-client-side-credentials + + If you want to use a different header name, you can set the `LITELLM_MCP_CLIENT_SIDE_AUTH_HEADER_NAME` in the secret manager or `mcp_client_side_auth_header_name` in the general settings. + """ + mcp_client_side_auth_header_name: str = MCPRequestHandler._get_mcp_client_side_auth_header_name() + return headers.get(mcp_client_side_auth_header_name) + + @staticmethod + def _get_mcp_client_side_auth_header_name() -> str: + """ + Get the header name used to pass the MCP auth header to the MCP server + + By default litellm will check for the header `x-mcp-auth` by setting one of the following: + 1. `LITELLM_MCP_CLIENT_SIDE_AUTH_HEADER_NAME` as an environment variable + 2. `mcp_client_side_auth_header_name` in the general settings on the config.yaml file + """ + from litellm.proxy.proxy_server import general_settings + from litellm.secret_managers.main import get_secret_str + MCP_CLIENT_SIDE_AUTH_HEADER_NAME: str = MCPRequestHandler.LITELLM_MCP_AUTH_HEADER_NAME + if get_secret_str("LITELLM_MCP_CLIENT_SIDE_AUTH_HEADER_NAME") is not None: + MCP_CLIENT_SIDE_AUTH_HEADER_NAME = get_secret_str("LITELLM_MCP_CLIENT_SIDE_AUTH_HEADER_NAME") or MCP_CLIENT_SIDE_AUTH_HEADER_NAME + elif general_settings.get("mcp_client_side_auth_header_name") is not None: + MCP_CLIENT_SIDE_AUTH_HEADER_NAME = general_settings.get("mcp_client_side_auth_header_name") or MCP_CLIENT_SIDE_AUTH_HEADER_NAME + return MCP_CLIENT_SIDE_AUTH_HEADER_NAME + @staticmethod def get_litellm_api_key_from_headers(headers: Headers) -> Optional[str]: 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 67c7a9c122..c50183fa5f 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 @@ -382,4 +382,177 @@ 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"] \ No newline at end of file + assert mcp_servers_result == expected_result["mcp_servers"] + + +class TestMCPCustomHeaderName: + """Test suite for custom MCP authentication header name functionality""" + + @pytest.mark.parametrize( + "env_var,general_setting,expected_header_name", + [ + # Test case 1: Default behavior (no custom settings) + (None, None, "x-mcp-auth"), + # Test case 2: Environment variable set + ("custom-mcp-header", None, "custom-mcp-header"), + # Test case 3: General setting set (env var takes precedence) + (None, "settings-mcp-header", "settings-mcp-header"), + # Test case 4: Both set (env var takes precedence) + ("env-mcp-header", "settings-mcp-header", "env-mcp-header"), + # Test case 5: Empty env var (should fallback to default due to 'or' logic) + ("", "settings-mcp-header", "x-mcp-auth"), + # Test case 6: Empty general setting (should fallback to default) + (None, "", "x-mcp-auth"), + ], + ) + def test_get_mcp_client_side_auth_header_name( + self, env_var, general_setting, expected_header_name + ): + """Test that custom header name configuration works correctly""" + + # Mock the secret manager and general settings + with patch("litellm.secret_managers.main.get_secret_str") as mock_get_secret: + with patch("litellm.proxy.proxy_server.general_settings") as mock_general_settings: + + # Configure mocks + mock_get_secret.return_value = env_var + mock_general_settings.get.return_value = general_setting + + # Call the method + result = MCPRequestHandler._get_mcp_client_side_auth_header_name() + + # Assert the result + assert result == expected_header_name + + # Verify secret manager was called (the function calls it twice) + expected_secret_calls = 2 if env_var is not None else 1 + assert mock_get_secret.call_count == expected_secret_calls + + # Verify all calls were with the correct parameter + for call in mock_get_secret.call_args_list: + assert call.args == ("LITELLM_MCP_CLIENT_SIDE_AUTH_HEADER_NAME",) + + # Verify general settings was called based on env var value + if env_var is None: + # When env var is None, general settings should be checked (twice if not None) + expected_general_calls = 2 if general_setting is not None else 1 + assert mock_general_settings.get.call_count == expected_general_calls + for call in mock_general_settings.get.call_args_list: + assert call.args == ("mcp_client_side_auth_header_name",) + else: + # If env var is set (even empty string), general settings shouldn't be checked + mock_general_settings.get.assert_not_called() + + @pytest.mark.parametrize( + "custom_header_name,headers,expected_auth_header", + [ + # Test case 1: Default header name + ( + "x-mcp-auth", + [(b"x-mcp-auth", b"default-auth-token")], + "default-auth-token" + ), + # Test case 2: Custom header name + ( + "custom-auth-header", + [(b"custom-auth-header", b"custom-auth-token")], + "custom-auth-token" + ), + # Test case 3: Custom header name with case insensitive + ( + "Custom-Auth-Header", + [(b"custom-auth-header", b"case-insensitive-token")], + "case-insensitive-token" + ), + # Test case 4: Header not present + ( + "missing-header", + [(b"x-mcp-auth", b"wrong-header-token")], + None + ), + # Test case 5: Multiple headers, only custom one should be used + ( + "my-custom-auth", + [ + (b"x-mcp-auth", b"default-token"), + (b"my-custom-auth", b"custom-token") + ], + "custom-token" + ), + ], + ) + def test_get_mcp_auth_header_from_headers_with_custom_name( + self, custom_header_name, headers, expected_auth_header + ): + """Test that MCP auth header extraction uses custom header name""" + + # Mock the header name method + with patch.object( + MCPRequestHandler, + '_get_mcp_client_side_auth_header_name', + return_value=custom_header_name + ): + # Create headers from the test data + scope = { + "type": "http", + "method": "POST", + "path": "/test", + "headers": headers, + } + extracted_headers = MCPRequestHandler._safe_get_headers_from_scope(scope) + + # Call the method + result = MCPRequestHandler._get_mcp_auth_header_from_headers(extracted_headers) + + # Assert the result + assert result == expected_auth_header + + @pytest.mark.asyncio + async def test_process_mcp_request_with_custom_auth_header(self): + """Test that process_mcp_request works with custom authentication header""" + + custom_header_name = "x-custom-mcp-auth" + custom_auth_token = "custom-auth-token-123" + api_key = "test-api-key" + + # Create ASGI scope with custom header + scope = { + "type": "http", + "method": "POST", + "path": "/test", + "headers": [ + (b"x-litellm-api-key", api_key.encode()), + (custom_header_name.encode(), custom_auth_token.encode()), + ], + } + + # Mock the custom header name method + with patch.object( + MCPRequestHandler, + '_get_mcp_client_side_auth_header_name', + return_value=custom_header_name + ): + # Mock user_api_key_auth + mock_auth_result = UserAPIKeyAuth( + api_key=api_key, + user_id="test-user-id", + team_id="test-team-id", + ) + + with patch( + "litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp.user_api_key_auth" + ) as mock_user_api_key_auth: + mock_user_api_key_auth.return_value = mock_auth_result + + # Call the method + auth_result, mcp_auth_header, mcp_servers = await MCPRequestHandler.process_mcp_request(scope) + + # Assert the results + assert auth_result == mock_auth_result + assert mcp_auth_header == custom_auth_token + assert mcp_servers is None + + # Verify user_api_key_auth was called with correct API key + mock_user_api_key_auth.assert_called_once() + call_args = mock_user_api_key_auth.call_args + assert call_args.kwargs["api_key"] == api_key \ No newline at end of file