ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(user_api_key_auth_mcp.py): pass extra headers from clientside straight through - allow multiple clientside headers

Browse Source 
Closes LIT-952
This commit is contained in:
Krrish Dholakia 2025-09-20 20:00:25 -07:00
parent e73b631524
commit 0ff7177373
6 changed files with 89 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

38
litellm/experimental_mcp_client/client.py
View File

@ -5,7 +5,7 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers.
import asyncio
import base64
from datetime import timedelta
from typing import Dict, List, Optional
from typing import Dict, List, Optional, Union
from mcp import ClientSession, StdioServerParameters
from mcp.client.sse import sse_client
@ -44,7 +44,7 @@ class MCPClient:
server_url: str = "",
transport_type: MCPTransportType = MCPTransport.http,
auth_type: MCPAuthType = None,
auth_value: Optional[str] = None,
auth_value: Optional[Union[str, Dict[str, str]]] = None,
timeout: float = 60.0,
stdio_config: Optional[MCPStdioConfig] = None,
extra_headers: Optional[Dict[str, str]] = None,
@ -53,7 +53,7 @@ class MCPClient:
self.transport_type: MCPTransport = transport_type
self.auth_type: MCPAuthType = auth_type
self.timeout: float = timeout
self._mcp_auth_value: Optional[str] = None
self._mcp_auth_value: Optional[Union[str, Dict[str, str]]] = None
self._session: Optional[ClientSession] = None
self._context = None
self._transport_ctx = None
@ -180,28 +180,34 @@ class MCPClient:
pass
self._context = None
def update_auth_value(self, mcp_auth_value: str):
def update_auth_value(self, mcp_auth_value: Union[str, Dict[str, str]]):
"""
Set the authentication header for the MCP client.
"""
if self.auth_type == MCPAuth.basic:
# Assuming mcp_auth_value is in format "username:password", convert it when updating
mcp_auth_value = to_basic_auth(mcp_auth_value)
self._mcp_auth_value = mcp_auth_value
if isinstance(mcp_auth_value, dict):
self._mcp_auth_value = mcp_auth_value
else:
if self.auth_type == MCPAuth.basic:
# Assuming mcp_auth_value is in format "username:password", convert it when updating
mcp_auth_value = to_basic_auth(mcp_auth_value)
self._mcp_auth_value = mcp_auth_value
def _get_auth_headers(self) -> dict:
"""Generate authentication headers based on auth type."""
headers = {"MCP-Protocol-Version": "2025-06-18"}
if self._mcp_auth_value:
if self.auth_type == MCPAuth.bearer_token:
headers["Authorization"] = f"Bearer {self._mcp_auth_value}"
elif self.auth_type == MCPAuth.basic:
headers["Authorization"] = f"Basic {self._mcp_auth_value}"
elif self.auth_type == MCPAuth.api_key:
headers["X-API-Key"] = self._mcp_auth_value
elif self.auth_type == MCPAuth.authorization:
headers["Authorization"] = self._mcp_auth_value
if isinstance(self._mcp_auth_value, str):
if self.auth_type == MCPAuth.bearer_token:
headers["Authorization"] = f"Bearer {self._mcp_auth_value}"
elif self.auth_type == MCPAuth.basic:
headers["Authorization"] = f"Basic {self._mcp_auth_value}"
elif self.auth_type == MCPAuth.api_key:
headers["X-API-Key"] = self._mcp_auth_value
elif self.auth_type == MCPAuth.authorization:
headers["Authorization"] = self._mcp_auth_value
elif isinstance(self._mcp_auth_value, dict):
headers.update(self._mcp_auth_value)
# update the headers with the extra headers
if self.extra_headers:

6
litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py
View File

@ -1,4 +1,4 @@
from typing import Dict, List, Optional
from typing import Dict, List, Optional, Union
from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser
@ -22,9 +22,9 @@ class MCPAuthenticatedUser(AuthenticatedUser):
user_api_key_auth: UserAPIKeyAuth,
mcp_auth_header: Optional[str] = None,
mcp_servers: Optional[List[str]] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
mcp_protocol_version: Optional[str] = None,
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
mcp_protocol_version: Optional[str] = None,
):
self.user_api_key_auth = user_api_key_auth
self.mcp_auth_header = mcp_auth_header

25
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
View File

@ -39,7 +39,7 @@ class MCPRequestHandler:
UserAPIKeyAuth,
Optional[str],
Optional[List[str]],
Optional[Dict[str, str]],
Optional[Dict[str, Dict[str, str]]],
Optional[Dict[str, str]],
]:
"""
@ -145,7 +145,9 @@ class MCPRequestHandler:
return auth_header
@staticmethod
def _get_mcp_server_auth_headers_from_headers(headers: Headers) -> Dict[str, str]:
def _get_mcp_server_auth_headers_from_headers(
headers: Headers,
) -> Dict[str, Dict[str, str]]:
"""
Parse server-specific MCP auth headers from the request headers.
@ -156,7 +158,7 @@ class MCPRequestHandler:
- x-mcp-deepwiki-authorization: Basic base64_encoded_creds
Returns:
Dict[str, str]: Mapping of server alias to auth value
Dict[str, Dict[str, str]]: Mapping of server alias to header dict
"""
server_auth_headers = {}
prefix = "x-mcp-"
@ -175,11 +177,22 @@ class MCPRequestHandler:
# Extract server_alias and header_name from x-mcp-{server_alias}-{header_name}
remaining = header_name[len(prefix) :].lower()
if "-" in remaining:
# Split on the last dash to separate server_alias from header_name
parts = remaining.rsplit("-", 1)
# Split on the first dash to separate server_alias from header_name
parts = remaining.split("-", 1)
if len(parts) == 2:
server_alias, auth_header_name = parts
server_auth_headers[server_alias] = header_value
# Convert header name to proper case (e.g., "authorization" -> "Authorization")
if auth_header_name == "authorization":
auth_header_name = "Authorization"
# Initialize server dict if not exists
if server_alias not in server_auth_headers:
server_auth_headers[server_alias] = {}
server_auth_headers[server_alias][
auth_header_name
] = header_value
verbose_logger.debug(
f"Found server auth header: {server_alias} -> {auth_header_name}: {header_value[:10]}..."
)

10
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
View File

@ -10,7 +10,7 @@ import asyncio
import datetime
import hashlib
import json
from typing import Any, Dict, List, Optional, cast
from typing import Any, Dict, List, Optional, Union, cast
from fastapi import HTTPException
from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
@ -319,7 +319,7 @@ class MCPServerManager:
self,
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
mcp_auth_header: Optional[str] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
mcp_server_auth_headers: Optional[Dict[str, Union[str, Dict[str, str]]]] = None,
) -> List[MCPTool]:
"""
List all tools available across all MCP Servers.
@ -381,7 +381,7 @@ class MCPServerManager:
def _create_mcp_client(
self,
server: MCPServer,
mcp_auth_header: Optional[str] = None,
mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None,
extra_headers: Optional[Dict[str, str]] = None,
) -> MCPClient:
"""
@ -429,7 +429,7 @@ class MCPServerManager:
async def _get_tools_from_server(
self,
server: MCPServer,
mcp_auth_header: Optional[str] = None,
mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None,
extra_headers: Optional[Dict[str, str]] = None,
) -> List[MCPTool]:
"""
@ -638,7 +638,7 @@ class MCPServerManager:
arguments: Dict[str, Any],
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
mcp_auth_header: Optional[str] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None,
proxy_logging_obj: Optional[ProxyLogging] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
) -> CallToolResult:

16
litellm/proxy/_experimental/mcp_server/server.py
View File

@ -361,7 +361,7 @@ if MCP_AVAILABLE:
user_api_key_auth: Optional[UserAPIKeyAuth],
mcp_auth_header: Optional[str],
mcp_servers: Optional[List[str]],
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
) -> List[MCPTool]:
"""
@ -371,8 +371,8 @@ if MCP_AVAILABLE:
user_api_key_auth: User authentication info for access control
mcp_auth_header: Optional auth header for MCP server (deprecated)
mcp_servers: Optional list of server names/aliases to filter by
mcp_server_auth_headers: Optional dict of server-specific auth headers
oauth2_headers: Optional dict of oauth2 headers
mcp_server_auth_headers: Optional dict of server-specific auth headers
oauth2_headers: Optional dict of oauth2 headers
Returns:
List[MCPTool]: Combined list of tools from filtered servers
@ -438,7 +438,7 @@ if MCP_AVAILABLE:
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
mcp_auth_header: Optional[str] = None,
mcp_servers: Optional[List[str]] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
) -> List[MCPTool]:
"""
@ -505,7 +505,7 @@ if MCP_AVAILABLE:
arguments: Optional[Dict[str, Any]] = None,
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
mcp_auth_header: Optional[str] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
**kwargs: Any,
) -> List[Union[TextContent, ImageContent, EmbeddedResource]]:
@ -606,7 +606,7 @@ if MCP_AVAILABLE:
arguments: Dict[str, Any],
user_api_key_auth: Optional[UserAPIKeyAuth] = None,
mcp_auth_header: Optional[str] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
litellm_logging_obj: Optional[Any] = None,
) -> List[Union[TextContent, ImageContent, EmbeddedResource]]:
@ -858,7 +858,7 @@ if MCP_AVAILABLE:
user_api_key_auth: UserAPIKeyAuth,
mcp_auth_header: Optional[str] = None,
mcp_servers: Optional[List[str]] = None,
mcp_server_auth_headers: Optional[Dict[str, str]] = None,
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
) -> None:
"""
@ -883,7 +883,7 @@ if MCP_AVAILABLE:
Optional[UserAPIKeyAuth],
Optional[str],
Optional[List[str]],
Optional[Dict[str, str]],
Optional[Dict[str, Dict[str, str]]],
Optional[Dict[str, str]],
]:
"""

32
tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
View File

@ -1014,3 +1014,35 @@ def test_mcp_path_based_server_segregation(monkeypatch):
# The context should have mcp_servers set to ["zapier", "group1"]
assert list(captured_mcp_servers.values())[0] == ["zapier", "group1"]
@pytest.mark.parametrize(
"headers,expected_result",
[
(
Headers(
{
"x-litellm-api-key": "test-key",
"x-mcp-github-authorization": "Bearer github-token",
}
),
{"github": {"Authorization": "Bearer github-token"}},
),
(
Headers(
{
"x-litellm-api-key": "test-key",
"x-mcp-github-x-api-key": "Basic base64-encoded-creds",
}
),
{"github": {"x-api-key": "Basic base64-encoded-creds"}},
),
],
)
def test_get_mcp_server_auth_headers_from_headers(headers, expected_result):
"""Test _get_mcp_server_auth_headers_from_headers method"""
from starlette.datastructures import Headers
headers = Headers(headers)
result = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(headers)
assert result == expected_result