Merge pull request #17963 from BerriAI/litellm_feat_rest-mcp-list-tools-auth-header

add MCP auth header propagation

Author: uc4w6c

litellm/proxy/_experimental/mcp_server/rest_endpoints.py

@@ -7,6 +7,7 @@
 from litellm._logging import verbose_logger
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+from litellm.types.mcp import MCPAuth
 
 MCP_AVAILABLE: bool = True
 try:
@@ -297,6 +298,7 @@ async def call_tool_rest_api(
     async def _execute_with_mcp_client(
         request: NewMCPServerRequest,
         operation,
+        mcp_auth_header: Optional[Union[str, Dict[str, str]]] = None,
         oauth2_headers: Optional[Dict[str, str]] = None,
     ):
         """
@@ -319,7 +321,7 @@ async def _execute_with_mcp_client(
                     auth_type=request.auth_type,
                     mcp_info=request.mcp_info,
                 ),
-                mcp_auth_header=None,
+                mcp_auth_header=mcp_auth_header,
                 extra_headers=oauth2_headers,
             )
 
@@ -365,7 +367,21 @@ async def test_tools_list(
         )
 
         headers = request.headers
-        oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers)
+
+        mcp_auth_header: Optional[str] = None
+        if new_mcp_server_request.auth_type in {
+            MCPAuth.api_key,
+            MCPAuth.bearer_token,
+            MCPAuth.basic,
+            MCPAuth.authorization,
+        }:
+            credentials = getattr(new_mcp_server_request, "credentials", None)
+            if isinstance(credentials, dict):
+                mcp_auth_header = credentials.get("auth_value")
+
+        oauth2_headers: Optional[Dict[str, str]] = None
+        if new_mcp_server_request.auth_type == MCPAuth.oauth2:
+            oauth2_headers = MCPRequestHandler._get_oauth2_headers_from_headers(headers)
 
         async def _list_tools_operation(client):
             async def _list_tools_session_operation(session):
@@ -385,5 +401,8 @@ async def _list_tools_session_operation(session):
             }
 
         return await _execute_with_mcp_client(
-            new_mcp_server_request, _list_tools_operation, oauth2_headers
+            new_mcp_server_request,
+            _list_tools_operation,
+            mcp_auth_header=mcp_auth_header,
+            oauth2_headers=oauth2_headers,
         )

tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py

@@ -0,0 +1,131 @@
+from typing import Dict, Optional
+
+import pytest
+from starlette.requests import Request
+
+from litellm.proxy._experimental.mcp_server import rest_endpoints
+from litellm.proxy._experimental.mcp_server.auth import (
+    user_api_key_auth_mcp as auth_mcp,
+)
+from litellm.proxy._types import NewMCPServerRequest, UserAPIKeyAuth
+from litellm.types.mcp import MCPAuth
+
+
+def _build_request(headers: Optional[Dict[str, str]] = None) -> Request:
+    headers = headers or {}
+    raw_headers = [
+        (key.lower().encode("latin-1"), value.encode("latin-1"))
+        for key, value in headers.items()
+    ]
+    scope = {
+        "type": "http",
+        "http_version": "1.1",
+        "method": "POST",
+        "path": "/mcp-rest/test/tools/list",
+        "headers": raw_headers,
+    }
+
+    async def receive():
+        return {"type": "http.request", "body": b"", "more_body": False}
+
+    return Request(scope, receive=receive)
+
+
+@pytest.mark.asyncio
+async def test_test_tools_list_forwards_mcp_auth_header(monkeypatch):
+    """Ensure credential-based auth forwards the auth_value to the MCP client."""
+
+    captured: dict = {}
+
+    async def fake_execute(request, operation, mcp_auth_header=None, oauth2_headers=None):
+        captured["mcp_auth_header"] = mcp_auth_header
+        captured["oauth2_headers"] = oauth2_headers
+        return {
+            "tools": [],
+            "error": None,
+            "message": "Successfully retrieved tools",
+        }
+
+    monkeypatch.setattr(
+        rest_endpoints, "_execute_with_mcp_client", fake_execute, raising=False
+    )
+
+    oauth_call_counter = {"count": 0}
+
+    def fake_oauth(headers):
+        oauth_call_counter["count"] += 1
+        return {"Authorization": "Bearer oauth"}
+
+    monkeypatch.setattr(
+        auth_mcp.MCPRequestHandler,
+        "_get_oauth2_headers_from_headers",
+        staticmethod(fake_oauth),
+        raising=False,
+    )
+
+    request = _build_request()
+    payload = NewMCPServerRequest(
+        server_name="example",
+        url="https://example.com",
+        auth_type=MCPAuth.api_key,
+        credentials={"auth_value": "secret-key"},
+    )
+
+    result = await rest_endpoints.test_tools_list(
+        request, payload, user_api_key_dict=UserAPIKeyAuth()
+    )
+
+    assert result["message"] == "Successfully retrieved tools"
+    assert captured["mcp_auth_header"] == "secret-key"
+    assert captured["oauth2_headers"] is None
+    assert oauth_call_counter["count"] == 0
+
+
+@pytest.mark.asyncio
+async def test_test_tools_list_extracts_oauth2_headers(monkeypatch):
+    """Ensure oauth2 auth type pulls oauth headers and omits MCP auth header."""
+
+    captured: dict = {}
+
+    async def fake_execute(request, operation, mcp_auth_header=None, oauth2_headers=None):
+        captured["mcp_auth_header"] = mcp_auth_header
+        captured["oauth2_headers"] = oauth2_headers
+        return {
+            "tools": [],
+            "error": None,
+            "message": "Successfully retrieved tools",
+        }
+
+    monkeypatch.setattr(
+        rest_endpoints, "_execute_with_mcp_client", fake_execute, raising=False
+    )
+
+    oauth_headers = {"Authorization": "Bearer oauth"}
+    oauth_call_counter = {"count": 0}
+
+    def fake_oauth(headers):
+        oauth_call_counter["count"] += 1
+        return oauth_headers
+
+    monkeypatch.setattr(
+        auth_mcp.MCPRequestHandler,
+        "_get_oauth2_headers_from_headers",
+        staticmethod(fake_oauth),
+        raising=False,
+    )
+
+    request = _build_request({"authorization": "Bearer incoming"})
+    payload = NewMCPServerRequest(
+        server_name="example",
+        url="https://example.com",
+        auth_type=MCPAuth.oauth2,
+    )
+
+    result = await rest_endpoints.test_tools_list(
+        request, payload, user_api_key_dict=UserAPIKeyAuth()
+    )
+
+    assert result["message"] == "Successfully retrieved tools"
+    assert captured["mcp_auth_header"] is None
+    assert captured["oauth2_headers"] == oauth_headers
+    assert oauth_call_counter["count"] == 1