From 73e00c138260f27cae51b93d999c6e91957fbded Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Mon, 15 Dec 2025 07:45:35 +0900 Subject: [PATCH] fix: add MCP auth header propagation --- .../mcp_server/rest_endpoints.py | 25 +++- .../mcp_server/test_rest_endpoints.py | 131 ++++++++++++++++++ 2 files changed, 153 insertions(+), 3 deletions(-) create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 6f293a298c3..032331ece02 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -7,6 +7,7 @@ from fastapi import APIRouter, Depends, Query, Request 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 @@ if MCP_AVAILABLE: 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 @@ if MCP_AVAILABLE: 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 @@ if MCP_AVAILABLE: ) 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 @@ if MCP_AVAILABLE: } 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, ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py new file mode 100644 index 00000000000..a0c09663a88 --- /dev/null +++ b/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