From 4c9220bdec68d5f6f551077716351db193c627c9 Mon Sep 17 00:00:00 2001 From: Milan Date: Tue, 10 Mar 2026 14:48:48 +0200 Subject: [PATCH] feat(mcp): add token authentication support for MCP servers - Add 'token' to MCPAuth enum for custom token auth format - Implement token auth in MCP client (_get_auth_headers) - Add token auth support for OpenAPI-based MCP tools - Add comprehensive unit tests to existing test_mcp_client.py - Fixes issue where MCP servers expecting 'Authorization: token ' header could not connect --- litellm/experimental_mcp_client/client.py | 2 + .../mcp_server/mcp_server_manager.py | 2 + litellm/types/mcp.py | 2 + .../test_mcp_client.py | 68 ++++++++++++++++++- .../mcp_tools/create_mcp_server.tsx | 4 +- .../components/mcp_tools/mcp_server_edit.tsx | 3 +- .../src/components/mcp_tools/types.tsx | 1 + 7 files changed, 79 insertions(+), 3 deletions(-) diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 849ce023109..e4f241880d8 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -212,6 +212,8 @@ class MCPClient: headers["Authorization"] = self._mcp_auth_value elif self.auth_type == MCPAuth.oauth2: headers["Authorization"] = f"Bearer {self._mcp_auth_value}" + elif self.auth_type == MCPAuth.token: + headers["Authorization"] = f"token {self._mcp_auth_value}" elif isinstance(self._mcp_auth_value, dict): headers.update(self._mcp_auth_value) diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 0b58009fcf6..1946e69fd68 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -418,6 +418,8 @@ class MCPServerManager: headers["Authorization"] = f"ApiKey {server.authentication_token}" elif server.auth_type == MCPAuth.basic: headers["Authorization"] = f"Basic {server.authentication_token}" + elif server.auth_type == MCPAuth.token: + headers["Authorization"] = f"token {server.authentication_token}" # Add any static headers from server config. # diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index 884dfefb42c..7b2ea820acc 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -35,6 +35,7 @@ class MCPAuth(str, enum.Enum): basic = "basic" authorization = "authorization" oauth2 = "oauth2" + token = "token" # MCP Literals @@ -50,6 +51,7 @@ MCPAuthType = Optional[ MCPAuth.basic, MCPAuth.authorization, MCPAuth.oauth2, + MCPAuth.token, ] ] diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index febc7c454bd..13a09f54e68 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -11,7 +11,7 @@ sys.path.insert(0, "../../../") import litellm.experimental_mcp_client.client as mcp_client_module from litellm.experimental_mcp_client.client import MCPClient -from litellm.types.mcp import MCPStdioConfig, MCPTransport +from litellm.types.mcp import MCPAuth, MCPStdioConfig, MCPTransport class TestMCPClient: @@ -245,6 +245,72 @@ class TestMCPClient: assert test_client.headers is not None await test_client.aclose() + def test_token_auth_header_generation(self): + """Test that token auth generates correct Authorization header""" + client = MCPClient( + server_url="http://example.com/sse", + transport_type="sse", + auth_type=MCPAuth.token, + auth_value="my-secret-token" + ) + + headers = client._get_auth_headers() + + assert "Authorization" in headers + assert headers["Authorization"] == "token my-secret-token" + + def test_token_auth_compatibility_with_existing_auth_types(self): + """Verify existing auth types are not affected by token auth addition""" + # Test bearer token + client = MCPClient( + server_url="http://example.com/sse", + transport_type="sse", + auth_type=MCPAuth.bearer_token, + auth_value="bearer-token" + ) + headers = client._get_auth_headers() + assert headers["Authorization"] == "Bearer bearer-token" + + # Test API key + client = MCPClient( + server_url="http://example.com/sse", + transport_type="sse", + auth_type=MCPAuth.api_key, + auth_value="api-key" + ) + headers = client._get_auth_headers() + assert headers["X-API-Key"] == "api-key" + + # Test basic auth (gets base64 encoded) + client = MCPClient( + server_url="http://example.com/sse", + transport_type="sse", + auth_type=MCPAuth.basic, + auth_value="user:pass" + ) + headers = client._get_auth_headers() + assert headers["Authorization"].startswith("Basic ") + + def test_token_auth_with_extra_headers(self): + """Test that token auth works alongside extra headers""" + client = MCPClient( + server_url="http://example.com/sse", + transport_type="sse", + auth_type=MCPAuth.token, + auth_value="my-token", + extra_headers={"X-Custom-Header": "custom-value"} + ) + + headers = client._get_auth_headers() + + assert headers["Authorization"] == "token my-token" + assert headers["X-Custom-Header"] == "custom-value" + + def test_token_auth_enum_value(self): + """Test that MCPAuth.token enum exists and has correct value""" + assert hasattr(MCPAuth, "token") + assert MCPAuth.token.value == "token" + if __name__ == "__main__": pytest.main([__file__]) diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx index 6ca58ffae24..6f1e103bc6e 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx @@ -29,7 +29,7 @@ interface CreateMCPServerProps { onBackToDiscovery?: () => void; } -const AUTH_TYPES_REQUIRING_AUTH_VALUE = [AUTH_TYPE.API_KEY, AUTH_TYPE.BEARER_TOKEN, AUTH_TYPE.BASIC]; +const AUTH_TYPES_REQUIRING_AUTH_VALUE = [AUTH_TYPE.API_KEY, AUTH_TYPE.BEARER_TOKEN, AUTH_TYPE.TOKEN, AUTH_TYPE.BASIC]; const AUTH_TYPES_REQUIRING_CREDENTIALS = [...AUTH_TYPES_REQUIRING_AUTH_VALUE, AUTH_TYPE.OAUTH2]; const CREATE_OAUTH_UI_STATE_KEY = "litellm-mcp-oauth-create-state"; @@ -654,6 +654,7 @@ const CreateMCPServer: React.FC = ({ User keys will be sent as:{" "} {getFieldValue("auth_type") === "bearer_token" && "Authorization: Bearer {key}"} + {getFieldValue("auth_type") === "token" && "Authorization: token {key}"} {getFieldValue("auth_type") === "api_key" && "x-api-key: {key}"} {getFieldValue("auth_type") === "basic" && "Authorization: Basic {key}"} {getFieldValue("auth_type") === "authorization" && "Authorization: {key}"} @@ -718,6 +719,7 @@ const CreateMCPServer: React.FC = ({ None API Key Bearer Token + Token Basic Auth OAuth diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx index 00060658ea1..fc55542a0c9 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx @@ -20,7 +20,7 @@ interface MCPServerEditProps { availableAccessGroups: string[]; } -const AUTH_TYPES_REQUIRING_AUTH_VALUE = [AUTH_TYPE.API_KEY, AUTH_TYPE.BEARER_TOKEN, AUTH_TYPE.BASIC]; +const AUTH_TYPES_REQUIRING_AUTH_VALUE = [AUTH_TYPE.API_KEY, AUTH_TYPE.BEARER_TOKEN, AUTH_TYPE.TOKEN, AUTH_TYPE.BASIC]; const AUTH_TYPES_REQUIRING_CREDENTIALS = [...AUTH_TYPES_REQUIRING_AUTH_VALUE, AUTH_TYPE.OAUTH2]; const EDIT_OAUTH_UI_STATE_KEY = "litellm-mcp-oauth-edit-state"; @@ -658,6 +658,7 @@ const MCPServerEdit: React.FC = ({ None API Key Bearer Token + Token Basic Auth OAuth diff --git a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx index 6ba25012197..27b96fd5635 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/types.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/types.tsx @@ -9,6 +9,7 @@ export const AUTH_TYPE = { NONE: "none", API_KEY: "api_key", BEARER_TOKEN: "bearer_token", + TOKEN: "token", BASIC: "basic", OAUTH2: "oauth2", };