From e944217dab60b889bb98cb432265073d6f4145c1 Mon Sep 17 00:00:00 2001 From: Arvind Kumar Date: Fri, 23 Oct 2026 11:12:40 +0530 Subject: [PATCH] test(mcp): regression tests for signer vs server credential (#47988) --- .../test_mcp_jwt_signer_server_credential.py | 165 ++++++++++++++++++ 1 file changed, 165 insertions(+) create mode 100644 tests/unit/proxy/_experimental/mcp_server/test_mcp_jwt_signer_server_credential.py diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_jwt_signer_server_credential.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_jwt_signer_server_credential.py new file mode 100644 index 00000000000..b27b28dd51a --- /dev/null +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_jwt_signer_server_credential.py @@ -0,0 +1,165 @@ +"""MCPJWTSigner must not replace an MCP server's own Authorization credential. + +Regression tests for BerriAI/litellm#47988: a server configured with bearer_token / token / +authorization / basic plus authentication_token lost its key to the signer's JWT on both +tools/list and tools/call. +""" + +from typing import Any, Dict, Optional +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + MCPServerManager, + _server_credential_occupies_authorization, +) +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.mcp import MCPAuth, MCPTransport +from litellm.types.mcp_server.mcp_server_manager import MCPServer + +SIGNER_SINGLETON = "litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer.get_mcp_jwt_signer" +SIGNER_INJECT = "litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer.inject_mcp_jwt_headers_for_upstream" + +AUTHORIZATION_AUTH_TYPES = [MCPAuth.bearer_token, MCPAuth.token, MCPAuth.authorization, MCPAuth.basic] + + +def _server( + auth_type: Optional[MCPAuth], + authentication_token: Optional[str] = "sk-upstream-key", + upstream_token_header: Optional[str] = None, +) -> MCPServer: + return MCPServer( + server_id="srv", + name="srv", + server_name="srv", + url="https://upstream.example/mcp", + transport=MCPTransport.http, + auth_type=auth_type, + authentication_token=authentication_token, + upstream_token_header=upstream_token_header, + ) + + +@pytest.mark.parametrize( + ("auth_type", "token", "header", "expected"), + [ + pytest.param(MCPAuth.bearer_token, "k", None, True, id="bearer-default-slot"), + pytest.param(MCPAuth.token, "k", None, True, id="token-default-slot"), + pytest.param(MCPAuth.authorization, "k", None, True, id="authorization-default-slot"), + pytest.param(MCPAuth.basic, "user:pass", None, True, id="basic-default-slot"), + pytest.param(MCPAuth.api_key, "k", None, False, id="api-key-uses-x-api-key"), + pytest.param(MCPAuth.api_key, "k", "Authorization", True, id="api-key-moved-to-authorization"), + pytest.param(MCPAuth.bearer_token, "k", "X-Upstream-Key", False, id="bearer-moved-off-authorization"), + pytest.param(MCPAuth.bearer_token, "k", "authorization", True, id="header-name-case-insensitive"), + pytest.param(MCPAuth.bearer_token, None, None, False, id="no-credential-configured"), + pytest.param(MCPAuth.none, "k", None, False, id="auth-none"), + pytest.param(None, "k", None, False, id="auth-unset"), + ], +) +def test_server_credential_occupies_authorization(auth_type, token, header, expected): + assert _server_credential_occupies_authorization(_server(auth_type, token, header)) is expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("auth_type", AUTHORIZATION_AUTH_TYPES) +async def test_tools_list_skips_signer_when_server_key_rides_authorization(auth_type): + manager = MCPServerManager() + manager.create_mcp_client = AsyncMock(return_value=object()) + manager._fetch_tools_with_timeout = AsyncMock(return_value=[]) + user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice") + + with ( + patch( # test-quality-ok: the signer is a process-wide singleton the manager reads, no injection seam + SIGNER_SINGLETON, return_value=MagicMock() + ), + patch( # test-quality-ok: same singleton's header injection, asserted on by call + SIGNER_INJECT, new=AsyncMock(return_value={"Authorization": "Bearer signed-jwt"}) + ) as mock_inject, + ): + await manager._get_tools_from_server(_server(auth_type), user_api_key_auth=user_auth) + + mock_inject.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "server", + [ + pytest.param(_server(MCPAuth.api_key), id="api-key-rides-x-api-key"), + pytest.param(_server(MCPAuth.bearer_token, authentication_token=None), id="no-server-credential"), + pytest.param( + _server(MCPAuth.bearer_token, upstream_token_header="X-Upstream-Key"), + id="credential-moved-off-authorization", + ), + ], +) +async def test_tools_list_still_runs_signer_when_authorization_is_free(server): + manager = MCPServerManager() + manager.create_mcp_client = AsyncMock(return_value=object()) + manager._fetch_tools_with_timeout = AsyncMock(return_value=[]) + user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice") + + with ( + patch( # test-quality-ok: the signer is a process-wide singleton the manager reads, no injection seam + SIGNER_SINGLETON, return_value=MagicMock() + ), + patch( # test-quality-ok: same singleton's header injection, asserted on by call + SIGNER_INJECT, new=AsyncMock(return_value={"Authorization": "Bearer signed-jwt"}) + ) as mock_inject, + ): + await manager._get_tools_from_server(server, user_api_key_auth=user_auth) + + mock_inject.assert_awaited_once() + + +async def _call_tool_and_capture_extra_headers(server: MCPServer, hook_headers: Dict[str, str]) -> Dict[str, str]: + manager = MCPServerManager() + captured: Dict[str, Any] = {} + + async def fake_create_mcp_client(server, mcp_auth_header=None, extra_headers=None, stdio_env=None, **kwargs): + captured["extra_headers"] = 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=None, + raw_headers=None, + proxy_logging_obj=None, + hook_extra_headers=hook_headers, + ) + except Exception: + pass # the mocked client's result is not what is under test + + assert "extra_headers" in captured, "create_mcp_client was never reached" + return captured["extra_headers"] or {} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("auth_type", AUTHORIZATION_AUTH_TYPES) +async def test_tools_call_drops_hook_authorization_when_server_key_rides_authorization(auth_type): + headers = await _call_tool_and_capture_extra_headers( + _server(auth_type), {"Authorization": "Bearer hook-jwt", "X-Trace-Id": "trace-123"} + ) + + assert "Authorization" not in headers + assert headers.get("X-Trace-Id") == "trace-123" + + +@pytest.mark.asyncio +async def test_tools_call_keeps_hook_authorization_when_server_key_rides_x_api_key(): + headers = await _call_tool_and_capture_extra_headers( + _server(MCPAuth.api_key), {"Authorization": "Bearer hook-jwt"} + ) + + assert headers.get("Authorization") == "Bearer hook-jwt" \ No newline at end of file