mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
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:
parent
a155ea1e8a
commit
8300657af9
4 changed files with 600 additions and 11 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue