fix(mcp): preserve oauth2 m2m auth for tools routes (#26871)

* Fix tool/list M2M creds issue

* Fix tool call creds issue

* Fix greptile review

* Fix lint

* Fix lint

* Fix lint

* Fix lint
This commit is contained in:
Sameer Kankute 2026-05-01 22:56:10 +05:30 committed by GitHub
parent a155ea1e8a
commit 8300657af9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 600 additions and 11 deletions

View file

@ -169,6 +169,37 @@ def _deserialize_json_dict(data: Any) -> Optional[Dict[str, str]]:
class MCPServerManager:
_STDIO_ENV_TEMPLATE_PATTERN = re.compile(r"^\$\{(X-[^}]+)\}$")
@staticmethod
def _resolve_oauth2_flow(
*,
auth_type: Optional[MCPAuthType],
oauth2_flow: Optional[str],
token_url: Optional[str],
authorization_url: Optional[str],
client_id: Optional[str],
client_secret: Optional[str],
) -> Optional[Literal["client_credentials", "authorization_code"]]:
"""Infer oauth2_flow for legacy records that omit the field.
DB rows created before oauth2_flow support may have OAuth2 client
credentials + token_url but a null oauth2_flow. Treat these as M2M,
unless authorization_url is present (interactive OAuth).
"""
if oauth2_flow in ("client_credentials", "authorization_code"):
return cast(
Literal["client_credentials", "authorization_code"], oauth2_flow
)
if oauth2_flow:
# Ignore unknown/untyped values and continue legacy inference.
return None
if auth_type != MCPAuth.oauth2:
return None
if authorization_url:
return None
if token_url and client_id and client_secret:
return "client_credentials"
return None
def __init__(self):
self.registry: Dict[str, MCPServer] = {}
self.config_mcp_servers: Dict[str, MCPServer] = {}
@ -342,7 +373,14 @@ class MCPServerManager:
# oauth specific fields
client_id=server_config.get("client_id", None),
client_secret=server_config.get("client_secret", None),
oauth2_flow=server_config.get("oauth2_flow", None),
oauth2_flow=self._resolve_oauth2_flow(
auth_type=auth_type,
oauth2_flow=server_config.get("oauth2_flow", None),
token_url=resolved_token_url,
authorization_url=resolved_authorization_url,
client_id=server_config.get("client_id", None),
client_secret=server_config.get("client_secret", None),
),
scopes=resolved_scopes,
authorization_url=resolved_authorization_url,
token_url=resolved_token_url,
@ -679,7 +717,17 @@ class MCPServerManager:
client_id=client_id_value or getattr(mcp_server, "client_id", None),
client_secret=client_secret_value
or getattr(mcp_server, "client_secret", None),
oauth2_flow=getattr(mcp_server, "oauth2_flow", None),
oauth2_flow=self._resolve_oauth2_flow(
auth_type=auth_type,
oauth2_flow=getattr(mcp_server, "oauth2_flow", None),
token_url=mcp_server.token_url
or getattr(mcp_oauth_metadata, "token_url", None),
authorization_url=mcp_server.authorization_url
or getattr(mcp_oauth_metadata, "authorization_url", None),
client_id=client_id_value or getattr(mcp_server, "client_id", None),
client_secret=client_secret_value
or getattr(mcp_server, "client_secret", None),
),
scopes=resolved_scopes,
authorization_url=mcp_server.authorization_url
or getattr(mcp_oauth_metadata, "authorization_url", None),
@ -2426,7 +2474,7 @@ class MCPServerManager:
)
)
async def _call_regular_mcp_tool(
async def _call_regular_mcp_tool( # noqa: PLR0915
self,
mcp_server: MCPServer,
original_tool_name: str,
@ -2489,7 +2537,11 @@ class MCPServerManager:
# oauth2 headers
extra_headers: Optional[Dict[str, str]] = None
if mcp_server.auth_type == MCPAuth.oauth2:
extra_headers = oauth2_headers
if mcp_server.has_client_credentials:
# For M2M OAuth servers, Authorization must come from token fetch.
extra_headers = None
else:
extra_headers = oauth2_headers
if mcp_server.extra_headers and raw_headers:
if extra_headers is None:
@ -2501,6 +2553,11 @@ class MCPServerManager:
for header in mcp_server.extra_headers:
if not isinstance(header, str):
continue
if (
mcp_server.has_client_credentials
and header.lower() == "authorization"
):
continue
header_value = normalized_raw_headers.get(header.lower())
if header_value is None:
continue
@ -2536,6 +2593,10 @@ class MCPServerManager:
)
extra_headers.update(hook_extra_headers)
# Reset to None if no headers were actually added
if extra_headers is not None and len(extra_headers) == 0:
extra_headers = None
stdio_env = self._build_stdio_env(mcp_server, raw_headers)
client = await self._create_mcp_client(

View file

@ -153,6 +153,7 @@ if MCP_AVAILABLE:
MCPAuthenticatedUser,
)
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
MCPServerManager,
global_mcp_server_manager,
)
from litellm.proxy._experimental.mcp_server.openapi_to_mcp_generator import (
@ -900,6 +901,20 @@ if MCP_AVAILABLE:
allowed_mcp_server_id
)
if mcp_server is not None:
# Apply oauth2_flow resolution for legacy DB rows where it may be NULL
resolved_flow = MCPServerManager._resolve_oauth2_flow(
auth_type=mcp_server.auth_type,
oauth2_flow=mcp_server.oauth2_flow,
token_url=mcp_server.token_url,
authorization_url=mcp_server.authorization_url,
client_id=mcp_server.client_id,
client_secret=mcp_server.client_secret,
)
if resolved_flow and resolved_flow != mcp_server.oauth2_flow:
# Create a new instance with the resolved flow for this request
mcp_server = mcp_server.model_copy(
update={"oauth2_flow": resolved_flow}
)
allowed_mcp_servers.append(mcp_server)
if mcp_servers is not None:
@ -1100,8 +1115,13 @@ if MCP_AVAILABLE:
extra_headers: Optional[Dict[str, str]] = None
if server.auth_type == MCPAuth.oauth2:
# Copy to avoid mutating the original dict (important for parallel fetching)
extra_headers = oauth2_headers.copy() if oauth2_headers else None
# For OAuth2 M2M servers, upstream Authorization must come from
# client_credentials token fetch, never from caller headers.
if server.has_client_credentials:
extra_headers = None
else:
# Copy to avoid mutating the original dict (important for parallel fetching)
extra_headers = oauth2_headers.copy() if oauth2_headers else None
if server.extra_headers and raw_headers:
if extra_headers is None:
@ -1114,11 +1134,17 @@ if MCP_AVAILABLE:
for header in server.extra_headers:
if not isinstance(header, str):
continue
if server.has_client_credentials and header.lower() == "authorization":
continue
header_value = normalized_raw_headers.get(header.lower())
if header_value is None:
continue
extra_headers[header] = header_value
# Reset to None if no headers were actually added
if extra_headers is not None and len(extra_headers) == 0:
extra_headers = None
if server_auth_header is None:
server_auth_header = mcp_auth_header
@ -1377,11 +1403,19 @@ if MCP_AVAILABLE:
spend_meta["per_server_tool_counts"] = per_server_tool_counts
end_time = datetime.now()
await litellm_logging_obj.async_success_handler(
result=all_tools,
start_time=list_tools_start_time,
end_time=end_time,
)
try:
await litellm_logging_obj.async_success_handler(
result=all_tools,
start_time=list_tools_start_time,
end_time=end_time,
)
except Exception as log_exc:
# list_tools responses must not be dropped due to non-blocking
# observability/serialization failures.
verbose_logger.warning(
"MCP list_tools success logging failed (continuing): %s",
log_exc,
)
verbose_logger.info(
f"Successfully fetched {len(all_tools)} tools total from all MCP servers"

View file

@ -673,6 +673,106 @@ class TestHookHeaderMergePriority:
assert headers["X-OAuth"] == "yes"
assert headers["X-Trace-Id"] == "trace-123"
@pytest.mark.asyncio
async def test_m2m_oauth2_does_not_forward_litellm_caller_authorization(self):
"""M2M must not put caller Bearer (LiteLLM API key) into extra_headers (#23652)."""
manager = MCPServerManager()
server = MCPServer(
server_id="test-id",
name="Test Server",
server_name="test_server",
url="https://example.com",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
oauth2_flow="client_credentials",
token_url="https://auth.example.com/token",
)
captured_extra_headers: Dict[str, Any] = {}
async def fake_create_mcp_client(
server, mcp_auth_header=None, extra_headers=None, stdio_env=None
):
captured_extra_headers["value"] = extra_headers
mock_client = MagicMock()
mock_client.call_tool = AsyncMock(return_value=MagicMock())
return mock_client
with patch.object(
manager, "_create_mcp_client", side_effect=fake_create_mcp_client
):
with patch.object(manager, "_build_stdio_env", return_value=None):
try:
await manager._call_regular_mcp_tool(
mcp_server=server,
original_tool_name="test_tool",
arguments={"key": "val"},
tasks=[],
mcp_auth_header=None,
mcp_server_auth_headers=None,
oauth2_headers={"Authorization": "Bearer sk-1234"},
raw_headers={"authorization": "Bearer sk-1234"},
proxy_logging_obj=None,
hook_extra_headers=None,
)
except Exception:
pass
assert captured_extra_headers.get("value") is None
@pytest.mark.asyncio
async def test_m2m_oauth2_skips_authorization_in_configured_extra_headers(self):
"""M2M must not take Authorization from raw_headers even if extra_headers lists it."""
manager = MCPServerManager()
server = MCPServer(
server_id="test-id",
name="Test Server",
server_name="test_server",
url="https://example.com",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
oauth2_flow="client_credentials",
token_url="https://auth.example.com/token",
extra_headers=["Authorization", "X-Custom"],
)
captured_extra_headers: Dict[str, Any] = {}
async def fake_create_mcp_client(
server, mcp_auth_header=None, extra_headers=None, stdio_env=None
):
captured_extra_headers["value"] = extra_headers
mock_client = MagicMock()
mock_client.call_tool = AsyncMock(return_value=MagicMock())
return mock_client
with patch.object(
manager, "_create_mcp_client", side_effect=fake_create_mcp_client
):
with patch.object(manager, "_build_stdio_env", return_value=None):
try:
await manager._call_regular_mcp_tool(
mcp_server=server,
original_tool_name="test_tool",
arguments={"key": "val"},
tasks=[],
mcp_auth_header=None,
mcp_server_auth_headers=None,
oauth2_headers={"Authorization": "Bearer sk-1234"},
raw_headers={
"authorization": "Bearer sk-1234",
"x-custom": "from-client",
},
proxy_logging_obj=None,
hook_extra_headers=None,
)
except Exception:
pass
headers = captured_extra_headers.get("value") or {}
assert "Authorization" not in headers
assert headers.get("X-Custom") == "from-client"
class TestUserAPIKeyAuthJwtClaims:
"""Tests that UserAPIKeyAuth correctly carries jwt_claims."""

View file

@ -17,6 +17,7 @@ from litellm.proxy._types import (
MCPTransport,
UserAPIKeyAuth,
)
from litellm.types.mcp import MCPAuth
from litellm.types.mcp_server.mcp_server_manager import MCPServer
@ -135,6 +136,152 @@ def test_prepare_mcp_server_headers_case_insensitive_extra_headers():
assert extra_headers == {"Authorization": "Bearer token"}
def test_prepare_mcp_server_headers_oauth2_m2m_omits_litellm_caller_authorization():
"""M2M OAuth must not put caller Bearer (LiteLLM API key) into extra_headers (#23652)."""
try:
from litellm.proxy._experimental.mcp_server.server import (
_prepare_mcp_server_headers,
)
except ImportError:
pytest.skip("MCP server not available")
server = MCPServer(
server_id="m2m-server",
name="m2m",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
oauth2_flow="client_credentials",
token_url="https://auth.example.com/token",
)
caller_key = {"Authorization": "Bearer sk-litellm-caller"}
server_auth_header, extra_headers = _prepare_mcp_server_headers(
server=server,
mcp_server_auth_headers=None,
mcp_auth_header=None,
oauth2_headers=caller_key,
raw_headers=None,
)
assert server_auth_header is None
assert extra_headers is None
def test_prepare_mcp_server_headers_oauth2_interactive_copies_oauth2_headers():
"""Interactive OAuth still forwards the user's OAuth token in extra_headers."""
try:
from litellm.proxy._experimental.mcp_server.server import (
_prepare_mcp_server_headers,
)
except ImportError:
pytest.skip("MCP server not available")
user_oauth = {"Authorization": "Bearer upstream-user-token"}
server = MCPServer(
server_id="3lo-server",
name="3lo",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
oauth2_flow=None,
)
server_auth_header, extra_headers = _prepare_mcp_server_headers(
server=server,
mcp_server_auth_headers=None,
mcp_auth_header=None,
oauth2_headers=user_oauth,
raw_headers=None,
)
assert server_auth_header is None
assert extra_headers == user_oauth
def test_prepare_mcp_server_headers_m2m_skips_authorization_from_raw_extra_headers():
"""M2M must not merge caller Authorization from raw_headers when extra_headers lists it."""
try:
from litellm.proxy._experimental.mcp_server.server import (
_prepare_mcp_server_headers,
)
except ImportError:
pytest.skip("MCP server not available")
server = MCPServer(
server_id="m2m-raw",
name="m2m",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
oauth2_flow="client_credentials",
token_url="https://auth.example.com/token",
extra_headers=["Authorization", "X-Custom"],
)
server_auth_header, extra_headers = _prepare_mcp_server_headers(
server=server,
mcp_server_auth_headers=None,
mcp_auth_header=None,
oauth2_headers={"Authorization": "Bearer sk-1234"},
raw_headers={
"authorization": "Bearer sk-1234",
"x-custom": "trace",
},
)
assert server_auth_header is None
assert extra_headers is not None
assert "Authorization" not in extra_headers
assert extra_headers.get("X-Custom") == "trace"
@pytest.mark.asyncio
async def test_call_tool_m2m_skips_authorization_headers():
"""M2M call_tool must not forward caller Authorization in oauth2/raw headers."""
try:
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
MCPServerManager,
)
except ImportError:
pytest.skip("MCP server not available")
manager = MCPServerManager()
server = MCPServer(
server_id="m2m-call-tool",
name="m2m-call-tool",
server_name="m2m-call-tool",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
oauth2_flow="client_credentials",
token_url="https://auth.example.com/token",
client_id="cid",
client_secret="csecret",
extra_headers=["Authorization", "X-Custom"],
)
mock_client = MagicMock()
mock_client.call_tool = AsyncMock(return_value=MagicMock())
with patch.object(
manager, "_create_mcp_client", new=AsyncMock(return_value=mock_client)
) as create_client_mock:
await manager._call_regular_mcp_tool(
mcp_server=server,
original_tool_name="echo",
arguments={"message": "hello"},
tasks=[],
mcp_auth_header=None,
mcp_server_auth_headers=None,
oauth2_headers={"Authorization": "Bearer sk-1234"},
raw_headers={"authorization": "Bearer sk-1234", "x-custom": "trace"},
proxy_logging_obj=None,
)
create_kwargs = create_client_mock.await_args.kwargs
extra_headers = create_kwargs["extra_headers"] or {}
assert "Authorization" not in extra_headers
assert extra_headers.get("X-Custom") == "trace"
@pytest.mark.asyncio
async def test_get_prompts_from_mcp_servers_success():
try:
@ -2288,6 +2435,79 @@ async def test_get_tools_from_mcp_servers_logs_list_tools_to_spendlogs_when_enab
assert spend_meta["per_server_tool_counts"]["server_a"] == 1
@pytest.mark.asyncio
async def test_get_tools_from_mcp_servers_returns_tools_when_success_logging_fails():
"""
Regression test: list_tools should still return fetched tools even if
async_success_handler raises (e.g. serialization errors in logging path).
"""
try:
from litellm.proxy._experimental.mcp_server.server import (
_get_tools_from_mcp_servers,
)
from litellm.proxy._types import UserAPIKeyAuth
except ImportError:
pytest.skip("MCP server not available")
user_auth = UserAPIKeyAuth(api_key="test-key", user_id="test-user")
server_a = MagicMock(name="server_a_obj")
server_a.name = "server_a"
server_a.alias = "server_a"
server_a.server_name = "server_a"
server_a.server_id = "a"
server_a.auth_type = None
server_a.extra_headers = None
tool_1 = MagicMock()
tool_1.name = "server_a-tool_1"
dummy_logging_obj = MagicMock()
dummy_logging_obj.model_call_details = {"metadata": {"spend_logs_metadata": {}}}
dummy_logging_obj.async_success_handler = AsyncMock(
side_effect=TypeError("Object of type Tool is not JSON serializable")
)
with (
patch(
"litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
new=AsyncMock(return_value=[server_a]),
),
patch(
"litellm.proxy._experimental.mcp_server.server._prepare_mcp_server_headers",
return_value=(None, None),
),
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
) as mock_manager,
patch(
"litellm.proxy._experimental.mcp_server.server.filter_tools_by_allowed_tools",
side_effect=lambda tools, _server: tools,
),
patch(
"litellm.proxy._experimental.mcp_server.server.filter_tools_by_key_team_permissions",
new=AsyncMock(side_effect=lambda tools, **_: tools),
),
patch(
"litellm.proxy._experimental.mcp_server.server.function_setup",
return_value=(dummy_logging_obj, None),
),
):
mock_manager._get_tools_from_server = AsyncMock(return_value=[tool_1])
tools = await _get_tools_from_mcp_servers(
user_api_key_auth=user_auth,
mcp_auth_header=None,
mcp_servers=["server_a"],
mcp_server_auth_headers=None,
log_list_tools_to_spendlogs=True,
list_tools_log_source="mcp_protocol",
)
assert tools == [tool_1]
dummy_logging_obj.async_success_handler.assert_awaited_once()
def test_tool_name_matches_case_insensitive():
"""Test that _tool_name_matches performs case-insensitive comparison.
@ -2719,3 +2939,177 @@ class TestGatewayCreateInitializationOptions:
_mcp_gateway_initialize_instructions.reset(tok)
opts = server.create_initialization_options()
assert getattr(opts, "instructions", None) is None
@pytest.mark.asyncio
async def test_list_tools_with_legacy_db_m2m_server_resolves_oauth2_flow():
"""
P1 Regression: list_tools path must apply _resolve_oauth2_flow to legacy DB
rows where oauth2_flow is NULL but M2M credentials are present.
Without this fix, has_client_credentials returns False and the caller's
Authorization header is forwarded upstream instead of being blocked.
"""
try:
from litellm.proxy._experimental.mcp_server.server import (
_get_tools_from_mcp_servers,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.mcp import MCPAuth
except ImportError:
pytest.skip("MCP server not available")
user_auth = UserAPIKeyAuth(api_key="sk-1234", user_id="test-user")
# Simulate a legacy DB row: OAuth2 with M2M credentials but oauth2_flow=None
legacy_server = MagicMock(name="legacy_m2m_server")
legacy_server.name = "legacy_m2m"
legacy_server.alias = "legacy_m2m"
legacy_server.server_name = "legacy_m2m"
legacy_server.server_id = "legacy-m2m-id"
legacy_server.auth_type = MCPAuth.oauth2
legacy_server.oauth2_flow = None # Legacy: field not set in DB
legacy_server.token_url = "https://oauth.example.com/token"
legacy_server.authorization_url = None
legacy_server.client_id = "client-id"
legacy_server.client_secret = "client-secret"
legacy_server.extra_headers = None
legacy_server.has_client_credentials = False # This is the bug: should be True
legacy_server.model_copy = MagicMock(
side_effect=lambda update: MCPServer(
server_id=legacy_server.server_id,
name=legacy_server.name,
transport=MCPTransport.http,
auth_type=legacy_server.auth_type,
oauth2_flow=update.get("oauth2_flow", legacy_server.oauth2_flow),
token_url=legacy_server.token_url,
authorization_url=legacy_server.authorization_url,
client_id=legacy_server.client_id,
client_secret=legacy_server.client_secret,
)
)
tool_1 = MagicMock()
tool_1.name = "legacy_m2m-tool"
captured_extra_headers = None
async def capture_extra_headers(*args, **kwargs):
nonlocal captured_extra_headers
captured_extra_headers = kwargs.get("extra_headers")
return [tool_1]
with (
patch(
"litellm.proxy._experimental.mcp_server.server.global_mcp_server_manager",
) as mock_manager,
patch(
"litellm.proxy._experimental.mcp_server.server.filter_tools_by_allowed_tools",
side_effect=lambda tools, _server: tools,
),
patch(
"litellm.proxy._experimental.mcp_server.server.filter_tools_by_key_team_permissions",
new=AsyncMock(side_effect=lambda tools, **_: tools),
),
):
mock_manager.get_allowed_mcp_servers = AsyncMock(return_value=["legacy-m2m-id"])
mock_manager.get_mcp_server_by_id = MagicMock(return_value=legacy_server)
mock_manager.filter_server_ids_by_ip_with_info = MagicMock(
return_value=(["legacy-m2m-id"], 0)
)
mock_manager._get_tools_from_server = AsyncMock(
side_effect=capture_extra_headers
)
tools = await _get_tools_from_mcp_servers(
user_api_key_auth=user_auth,
mcp_auth_header=None,
mcp_servers=["legacy_m2m"],
mcp_server_auth_headers=None,
oauth2_headers={"Authorization": "Bearer sk-1234"}, # Caller's token
)
# With P1 fix: _get_allowed_mcp_servers applies _resolve_oauth2_flow,
# so has_client_credentials becomes True and extra_headers should be None
# (caller's Authorization blocked)
assert captured_extra_headers is None, (
"P1 security issue: caller's Authorization header was forwarded to M2M server. "
"Expected None, got: " + str(captured_extra_headers)
)
assert tools == [tool_1]
@pytest.mark.asyncio
async def test_call_tool_empty_extra_headers_returns_none():
"""
P2 Regression: When all configured extra_headers are filtered out (e.g.
Authorization for M2M), the resulting extra_headers should be None, not {}.
Downstream code that checks `if extra_headers is None` will behave
differently if an empty dict is passed instead.
"""
try:
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
MCPServerManager,
)
from litellm.types.mcp import MCPAuth
except ImportError:
pytest.skip("MCP server not available")
manager = MCPServerManager()
# M2M server with only Authorization in extra_headers
m2m_server = MCPServer(
server_id="m2m-srv",
name="m2m_test",
transport=MCPTransport.http,
auth_type=MCPAuth.oauth2,
oauth2_flow="client_credentials",
token_url="https://oauth.example.com/token",
client_id="client-id",
client_secret="client-secret",
extra_headers=["Authorization"], # Will be filtered out for M2M
)
raw_headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
captured_extra_headers = None
async def capture_create_mcp_client(*args, **kwargs):
nonlocal captured_extra_headers
captured_extra_headers = kwargs.get("extra_headers")
# Return a mock client
mock_client = AsyncMock()
mock_client.call_tool = AsyncMock(return_value=MagicMock(content=[]))
return mock_client
with (
patch.object(
manager,
"_create_mcp_client",
side_effect=capture_create_mcp_client,
),
patch.object(
manager,
"get_mcp_server_by_id",
return_value=m2m_server,
),
):
try:
await manager._call_regular_mcp_tool(
mcp_server=m2m_server,
original_tool_name="test_tool",
arguments={},
mcp_auth_header=None,
oauth2_headers=None,
raw_headers=raw_headers,
)
except Exception:
pass # We only care about the captured headers
# With P2 fix: extra_headers should be None (not {}) when all headers filtered
assert captured_extra_headers is None, (
"P2 API consistency issue: expected None for empty extra_headers, got: "
+ str(captured_extra_headers)
)