tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-06 06:24:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(mcp/): allow specifying forwardable headers

Browse Source 
allows admin to specify which clientside headers to forward to the backend mcp server

easier than requiring client to specify `x-mcp-{server_alias}-key`
This commit is contained in:
Krrish Dholakia
2025-09-27 17:30:46 -07:00
parent 0ff7177373
commit 1c71d4cdbb
5 changed files with 50 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/_experimental/mcp_server/auth/litellm_auth_handler.py
+3
View File
@@ -15,6 +15,7 @@ class MCPAuthenticatedUser(AuthenticatedUser):
3. MCP server configuration (can include access groups)
4. Server-specific authentication headers
5. OAuth2 headers
6. Raw headers - allows forwarding specific headers to the MCP server, specified by the admin.
"""
def __init__(
@@ -25,6 +26,7 @@ class MCPAuthenticatedUser(AuthenticatedUser):
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
mcp_protocol_version: Optional[str] = None,
raw_headers: Optional[Dict[str, str]] = None,
):
self.user_api_key_auth = user_api_key_auth
self.mcp_auth_header = mcp_auth_header
@@ -32,3 +34,4 @@ class MCPAuthenticatedUser(AuthenticatedUser):
self.mcp_server_auth_headers = mcp_server_auth_headers or {}
self.mcp_protocol_version = mcp_protocol_version
self.oauth2_headers = oauth2_headers
self.raw_headers = raw_headers
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
+5 -1
View File
@@ -41,6 +41,7 @@ class MCPRequestHandler:
Optional[List[str]],
Optional[Dict[str, Dict[str, str]]],
Optional[Dict[str, str]],
Optional[Dict[str, str]],
]:
"""
Process and validate MCP request headers from the ASGI scope.
@@ -49,6 +50,7 @@ class MCPRequestHandler:
2. Processing MCP server configuration
3. Handling MCP-specific headers
4. Handling oauth2 headers
5. Raw headers - allows forwarding specific headers to the MCP server, specified by the admin.
Args:
scope: ASGI scope containing request information
@@ -58,7 +60,8 @@ class MCPRequestHandler:
mcp_auth_header: Optional[str] MCP auth header to be passed to the MCP server (deprecated)
mcp_servers: Optional[List[str]] List of MCP servers and access groups to use
mcp_server_auth_headers: Optional[Dict[str, str]] Server-specific auth headers in format {server_alias: auth_value}
oauth2_headers: Optional[Dict[str, str]] OAuth2 headers
raw_headers: Optional[Dict[str, str]] Raw headers to be forwarded to the MCP server
Raises:
HTTPException: If headers are invalid or missing required headers
"""
@@ -116,6 +119,7 @@ class MCPRequestHandler:
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
dict(headers),
)
@staticmethod
litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
+16
View File
@@ -212,6 +212,7 @@ class MCPServerManager:
"authentication_token", server_config.get("auth_value", None)
),
mcp_info=mcp_info,
forwardable_headers=server_config.get("forwardable_headers", None),
access_groups=server_config.get("access_groups", None),
)
self.config_mcp_servers[server_id] = new_server
@@ -264,6 +265,13 @@ class MCPServerManager:
transport=cast(MCPTransportType, mcp_server.transport),
auth_type=cast(MCPAuthType, mcp_server.auth_type),
mcp_info=mcp_info,
forwardable_headers=getattr(mcp_server, "forwardable_headers", None),
# oauth specific fields
client_id=getattr(mcp_server, "client_id", None),
client_secret=getattr(mcp_server, "client_secret", None),
scopes=getattr(mcp_server, "scopes", None),
authorization_url=getattr(mcp_server, "authorization_url", None),
token_url=getattr(mcp_server, "token_url", None),
# Stdio-specific fields
command=getattr(mcp_server, "command", None),
args=getattr(mcp_server, "args", None) or [],
@@ -641,6 +649,7 @@ class MCPServerManager:
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None,
proxy_logging_obj: Optional[ProxyLogging] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
raw_headers: Optional[Dict[str, str]] = None,
) -> CallToolResult:
"""
Call a tool with the given name and arguments (handles prefixed tool names)
@@ -709,6 +718,13 @@ class MCPServerManager:
if mcp_server.auth_type == MCPAuth.oauth2:
extra_headers = oauth2_headers
if mcp_server.forwardable_headers and raw_headers:
if extra_headers is None:
extra_headers = {}
for header in mcp_server.forwardable_headers:
if header in raw_headers:
extra_headers[header] = raw_headers[header]
client = self._create_mcp_client(
server=mcp_server,
mcp_auth_header=server_auth_header,
litellm/proxy/_experimental/mcp_server/server.py
+23 -1
View File
@@ -180,6 +180,7 @@ if MCP_AVAILABLE:
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
raw_headers,
) = get_auth_context()
verbose_logger.debug(
f"MCP list_tools - User API Key Auth from context: {user_api_key_auth}"
@@ -198,6 +199,7 @@ if MCP_AVAILABLE:
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
)
verbose_logger.info(
f"MCP list_tools - Successfully returned {len(tools)} tools"
@@ -239,6 +241,7 @@ if MCP_AVAILABLE:
_,
mcp_server_auth_headers,
oauth2_headers,
raw_headers,
) = get_auth_context()
verbose_logger.debug(
@@ -271,6 +274,7 @@ if MCP_AVAILABLE:
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
**data, # for logging
)
except BlockedPiiEntityError as e:
@@ -363,6 +367,7 @@ if MCP_AVAILABLE:
mcp_servers: Optional[List[str]],
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
raw_headers: Optional[Dict[str, str]] = None,
) -> List[MCPTool]:
"""
Helper method to fetch tools from MCP servers based on server filtering criteria.
@@ -440,6 +445,7 @@ if MCP_AVAILABLE:
mcp_servers: Optional[List[str]] = None,
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
raw_headers: Optional[Dict[str, str]] = None,
) -> List[MCPTool]:
"""
List all available MCP tools.
@@ -464,6 +470,7 @@ if MCP_AVAILABLE:
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
)
verbose_logger.debug(
f"Successfully fetched {len(managed_tools)} tools from managed MCP servers"
@@ -507,6 +514,7 @@ if MCP_AVAILABLE:
mcp_auth_header: Optional[str] = None,
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
raw_headers: Optional[Dict[str, str]] = None,
**kwargs: Any,
) -> List[Union[TextContent, ImageContent, EmbeddedResource]]:
"""
@@ -555,6 +563,7 @@ if MCP_AVAILABLE:
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
litellm_logging_obj=litellm_logging_obj,
)
@@ -608,6 +617,7 @@ if MCP_AVAILABLE:
mcp_auth_header: Optional[str] = None,
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
raw_headers: Optional[Dict[str, str]] = None,
litellm_logging_obj: Optional[Any] = None,
) -> List[Union[TextContent, ImageContent, EmbeddedResource]]:
"""Handle tool execution for managed server tools"""
@@ -621,6 +631,7 @@ if MCP_AVAILABLE:
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
proxy_logging_obj=proxy_logging_obj,
)
verbose_logger.debug("CALL TOOL RESULT: %s", call_tool_result)
@@ -702,6 +713,7 @@ if MCP_AVAILABLE:
_,
mcp_server_auth_headers,
oauth2_headers,
raw_headers,
) = await MCPRequestHandler.process_mcp_request(scope)
mcp_servers = mcp_servers_from_path
else:
@@ -711,6 +723,7 @@ if MCP_AVAILABLE:
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
raw_headers,
) = await MCPRequestHandler.process_mcp_request(scope)
return (
user_api_key_auth,
@@ -718,6 +731,7 @@ if MCP_AVAILABLE:
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
raw_headers,
)
async def handle_streamable_http_mcp(
@@ -732,6 +746,7 @@ if MCP_AVAILABLE:
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
raw_headers,
) = await extract_mcp_auth_context(scope, path)
verbose_logger.debug(
f"MCP request mcp_servers (header/path): {mcp_servers}"
@@ -746,6 +761,7 @@ if MCP_AVAILABLE:
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
)
# Ensure session managers are initialized
@@ -785,6 +801,7 @@ if MCP_AVAILABLE:
mcp_servers,
mcp_server_auth_headers,
oauth2_headers,
raw_headers,
) = await extract_mcp_auth_context(scope, path)
verbose_logger.debug(
f"MCP request mcp_servers (header/path): {mcp_servers}"
@@ -798,6 +815,7 @@ if MCP_AVAILABLE:
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
)
if not _SESSION_MANAGERS_INITIALIZED:
@@ -860,6 +878,7 @@ if MCP_AVAILABLE:
mcp_servers: Optional[List[str]] = None,
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None,
oauth2_headers: Optional[Dict[str, str]] = None,
raw_headers: Optional[Dict[str, str]] = None,
) -> None:
"""
Set the UserAPIKeyAuth in the auth context variable.
@@ -876,6 +895,7 @@ if MCP_AVAILABLE:
mcp_servers=mcp_servers,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
)
auth_context_var.set(auth_user)
@@ -885,6 +905,7 @@ if MCP_AVAILABLE:
Optional[List[str]],
Optional[Dict[str, Dict[str, str]]],
Optional[Dict[str, str]],
Optional[Dict[str, str]],
]:
"""
Get the UserAPIKeyAuth from the auth context variable.
@@ -901,8 +922,9 @@ if MCP_AVAILABLE:
auth_user.mcp_servers,
auth_user.mcp_server_auth_headers,
auth_user.oauth2_headers,
auth_user.raw_headers,
)
return None, None, None, None, None
return None, None, None, None, None, None
########################################################
############ End of Auth Context Functions #############
litellm/types/mcp_server/mcp_server_manager.py
+3
View File
@@ -20,6 +20,9 @@ class MCPServer(BaseModel):
auth_type: Optional[MCPAuthType] = None
authentication_token: Optional[str] = None
mcp_info: Optional[MCPInfo] = None
forwardable_headers: Optional[List[str]] = (
None # allow admin to specify which headers to forward to the MCP server
)
# OAuth-specific fields
client_id: Optional[str] = None
client_secret: Optional[str] = None