mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
test(mcp): cover JWT signer + tool-call resolution branches
Adds unit tests for the new MCPServerManager helpers (_resolve_mcp_server_for_tool_call, _resolve_oauth2_headers_for_tool_call) and the new MCPJWTSigner paths (_build_scope call_type branches and inject_mcp_jwt_headers_for_upstream). Brings patch coverage above the auto target without changing behavior. Co-authored-by: Claude <claude@anthropic.com>
This commit is contained in:
Claude
parent
d429e845c9
commit
c0cbf6a440
2 changed files with 300 additions and 0 deletions
|
|
@ -1852,6 +1852,193 @@ class TestMCPServerManager:
|
|||
)
|
||||
mock_inject.assert_awaited_once()
|
||||
|
||||
def test_resolve_mcp_server_for_tool_call_via_prefixed_name(self):
|
||||
"""Resolution succeeds when the prefixed tool name is in the mapping."""
|
||||
manager = MCPServerManager()
|
||||
server = MCPServer(
|
||||
server_id="jira",
|
||||
name="jira",
|
||||
transport=MCPTransport.http,
|
||||
)
|
||||
manager.registry = {"jira": server}
|
||||
manager.tool_name_to_mcp_server_name_mapping["jira-search_issues"] = "jira"
|
||||
manager.tool_name_to_mcp_server_name_mapping["search_issues"] = "jira"
|
||||
|
||||
resolved = manager._resolve_mcp_server_for_tool_call("jira", "search_issues")
|
||||
assert resolved is server
|
||||
|
||||
def test_resolve_mcp_server_for_tool_call_via_alias(self):
|
||||
"""Resolution falls back to alias/server_name match in the registry."""
|
||||
manager = MCPServerManager()
|
||||
server = MCPServer(
|
||||
server_id="srv-uuid-123",
|
||||
name="zapier",
|
||||
alias="zapier-alias",
|
||||
transport=MCPTransport.http,
|
||||
)
|
||||
manager.registry = {"srv-uuid-123": server}
|
||||
|
||||
resolved = manager._resolve_mcp_server_for_tool_call(
|
||||
"zapier-alias", "create_zap"
|
||||
)
|
||||
assert resolved is server
|
||||
|
||||
def test_resolve_mcp_server_for_tool_call_fallback_to_unprefixed_lookup(self):
|
||||
"""Fallback to unprefixed _get_mcp_server_from_tool_name when other paths fail."""
|
||||
manager = MCPServerManager()
|
||||
server = MCPServer(
|
||||
server_id="linear",
|
||||
name="linear",
|
||||
transport=MCPTransport.http,
|
||||
)
|
||||
manager.registry = {"linear": server}
|
||||
manager.tool_name_to_mcp_server_name_mapping["create_issue"] = "linear"
|
||||
|
||||
# server_name is empty so the fallback unprefixed lookup runs and matches.
|
||||
resolved = manager._resolve_mcp_server_for_tool_call("", "create_issue")
|
||||
assert resolved is server
|
||||
|
||||
def test_resolve_mcp_server_for_tool_call_raises_when_not_found(self):
|
||||
"""ValueError is raised when no resolution path finds the tool."""
|
||||
manager = MCPServerManager()
|
||||
with pytest.raises(ValueError, match="Tool .* not found"):
|
||||
manager._resolve_mcp_server_for_tool_call("nonexistent", "ghost_tool")
|
||||
|
||||
def test_resolve_mcp_server_for_tool_call_unknown_tool_with_known_server(self):
|
||||
"""Server-name match alone must not let unknown tools slip through.
|
||||
|
||||
If the registry has tools for this server but neither the prefixed nor
|
||||
unprefixed tool name is in the mapping, raise rather than returning the
|
||||
server (would otherwise allow tool enumeration via name spoofing).
|
||||
"""
|
||||
manager = MCPServerManager()
|
||||
server = MCPServer(
|
||||
server_id="github",
|
||||
name="github",
|
||||
transport=MCPTransport.http,
|
||||
)
|
||||
manager.registry = {"github": server}
|
||||
# Mapping has *some* tools for github but not "missing_tool".
|
||||
manager.tool_name_to_mcp_server_name_mapping["github-list_repos"] = "github"
|
||||
manager.tool_name_to_mcp_server_name_mapping["list_repos"] = "github"
|
||||
|
||||
with pytest.raises(ValueError, match="Tool missing_tool not found"):
|
||||
manager._resolve_mcp_server_for_tool_call("github", "missing_tool")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_oauth2_headers_skipped_when_not_user_oauth(self):
|
||||
"""Returns input headers unchanged when server does not need user OAuth."""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
manager = MCPServerManager()
|
||||
server = MCPServer(
|
||||
server_id="plain",
|
||||
name="plain",
|
||||
transport=MCPTransport.http,
|
||||
)
|
||||
# needs_user_oauth_token defaults to False.
|
||||
user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="bob")
|
||||
|
||||
result = await manager._resolve_oauth2_headers_for_tool_call(
|
||||
server, oauth2_headers=None, user_api_key_auth=user_auth
|
||||
)
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_oauth2_headers_returns_client_supplied_token(self):
|
||||
"""Returns the client's oauth2_headers as-is when already set."""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
manager = MCPServerManager()
|
||||
server = MCPServer(
|
||||
server_id="oauth-srv",
|
||||
name="oauth-srv",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
)
|
||||
assert server.needs_user_oauth_token is True
|
||||
user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice")
|
||||
supplied = {"Authorization": "Bearer client-supplied"}
|
||||
|
||||
result = await manager._resolve_oauth2_headers_for_tool_call(
|
||||
server, oauth2_headers=supplied, user_api_key_auth=user_auth
|
||||
)
|
||||
assert result is supplied
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_oauth2_headers_looks_up_stored_token(self):
|
||||
"""Falls back to stored per-user OAuth headers when no token is supplied."""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
manager = MCPServerManager()
|
||||
server = MCPServer(
|
||||
server_id="oauth-srv",
|
||||
name="oauth-srv",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
)
|
||||
user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice")
|
||||
stored = {"Authorization": "Bearer stored-user-token"}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db",
|
||||
new=AsyncMock(return_value=stored),
|
||||
) as mock_lookup:
|
||||
result = await manager._resolve_oauth2_headers_for_tool_call(
|
||||
server, oauth2_headers=None, user_api_key_auth=user_auth
|
||||
)
|
||||
|
||||
assert result == stored
|
||||
mock_lookup.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_oauth2_headers_swallows_lookup_exception(self):
|
||||
"""Returns supplied headers (None) when the stored-token lookup raises."""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
manager = MCPServerManager()
|
||||
server = MCPServer(
|
||||
server_id="oauth-srv",
|
||||
name="oauth-srv",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
)
|
||||
user_auth = UserAPIKeyAuth(api_key="sk-test", user_id="alice")
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db",
|
||||
new=AsyncMock(side_effect=RuntimeError("redis down")),
|
||||
):
|
||||
result = await manager._resolve_oauth2_headers_for_tool_call(
|
||||
server, oauth2_headers=None, user_api_key_auth=user_auth
|
||||
)
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_oauth2_headers_no_user_id(self):
|
||||
"""Skip lookup entirely when user_api_key_auth has no user_id."""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
manager = MCPServerManager()
|
||||
server = MCPServer(
|
||||
server_id="oauth-srv",
|
||||
name="oauth-srv",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
)
|
||||
# user_id is None -> lookup must not happen
|
||||
user_auth = UserAPIKeyAuth(api_key="sk-test")
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.server._get_user_oauth_extra_headers_from_db",
|
||||
new=AsyncMock(return_value={"Authorization": "Bearer x"}),
|
||||
) as mock_lookup:
|
||||
result = await manager._resolve_oauth2_headers_for_tool_call(
|
||||
server, oauth2_headers=None, user_api_key_auth=user_auth
|
||||
)
|
||||
assert result is None
|
||||
mock_lookup.assert_not_called()
|
||||
|
||||
def test_create_prefixed_tools_updates_mapping_for_both_forms(self):
|
||||
"""_create_prefixed_tools should populate mapping for prefixed and original names even when not adding prefix in output."""
|
||||
manager = MCPServerManager()
|
||||
|
|
|
|||
|
|
@ -1156,3 +1156,116 @@ async def test_hook_raises_401_when_jwt_verification_fails():
|
|||
)
|
||||
|
||||
assert exc_info.value.status_code == 401
|
||||
|
||||
|
||||
# --- _build_scope branches: call_mcp_tool with empty tool name, list_mcp_tools ---
|
||||
|
||||
|
||||
def test_build_scope_call_type_call_mcp_tool_without_tool_name():
|
||||
"""call_mcp_tool with empty tool name emits a generic mcp:tools/call only."""
|
||||
signer = _make_signer()
|
||||
scope = signer._build_scope("", call_type="call_mcp_tool")
|
||||
scopes = set(scope.split())
|
||||
assert scopes == {"mcp:tools/call"}
|
||||
|
||||
|
||||
def test_build_scope_call_type_list_mcp_tools_only_list():
|
||||
"""list_mcp_tools (no tool) emits only mcp:tools/list, never tools/call."""
|
||||
signer = _make_signer()
|
||||
scope = signer._build_scope("", call_type="list_mcp_tools")
|
||||
scopes = set(scope.split())
|
||||
assert scopes == {"mcp:tools/list"}
|
||||
|
||||
|
||||
def test_build_scope_default_is_list_only_when_no_call_type():
|
||||
"""No call_type and no tool falls through to tools/list (least-privilege default)."""
|
||||
signer = _make_signer()
|
||||
scope = signer._build_scope("")
|
||||
scopes = set(scope.split())
|
||||
assert "mcp:tools/list" in scopes
|
||||
assert "mcp:tools/call" not in scopes
|
||||
|
||||
|
||||
# --- inject_mcp_jwt_headers_for_upstream ---
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inject_mcp_jwt_returns_unchanged_when_signer_not_configured():
|
||||
"""No signer configured -> return a fresh copy of extra_headers untouched."""
|
||||
import litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer as mod
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
mod._mcp_jwt_signer_instance = None
|
||||
headers = {"X-Trace-Id": "abc"}
|
||||
user_dict = UserAPIKeyAuth(api_key="sk-test", user_id="alice")
|
||||
|
||||
result = await mod.inject_mcp_jwt_headers_for_upstream(
|
||||
user_api_key_dict=user_dict,
|
||||
extra_headers=headers,
|
||||
)
|
||||
assert result == headers
|
||||
assert result is not headers # must be a copy
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inject_mcp_jwt_returns_unchanged_when_user_dict_none():
|
||||
"""No user_api_key_dict -> short-circuit without invoking the signer."""
|
||||
from litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer import (
|
||||
inject_mcp_jwt_headers_for_upstream,
|
||||
)
|
||||
|
||||
_make_signer() # ensure instance is created
|
||||
result = await inject_mcp_jwt_headers_for_upstream(
|
||||
user_api_key_dict=None,
|
||||
extra_headers={"X-Trace-Id": "abc"},
|
||||
)
|
||||
assert result == {"X-Trace-Id": "abc"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inject_mcp_jwt_signs_for_list_tools_path():
|
||||
"""When for_list_tools=True, signer is invoked with list_mcp_tools call_type."""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer import (
|
||||
inject_mcp_jwt_headers_for_upstream,
|
||||
)
|
||||
|
||||
_make_signer(issuer="https://litellm.example.com", audience="mcp", ttl_seconds=300)
|
||||
user_dict = UserAPIKeyAuth(api_key="sk-test", user_id="alice")
|
||||
|
||||
result = await inject_mcp_jwt_headers_for_upstream(
|
||||
user_api_key_dict=user_dict,
|
||||
extra_headers={"X-Trace": "1"},
|
||||
raw_headers={"Authorization": "Bearer incoming.opaque.token"},
|
||||
for_list_tools=True,
|
||||
)
|
||||
assert result["X-Trace"] == "1"
|
||||
assert result["Authorization"].startswith("Bearer ")
|
||||
token = result["Authorization"].removeprefix("Bearer ")
|
||||
decoded = _decode_unverified(token)
|
||||
scopes = set(decoded["scope"].split())
|
||||
assert scopes == {"mcp:tools/list"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inject_mcp_jwt_signs_for_tool_call_path():
|
||||
"""for_list_tools=False with a tool name signs a call_mcp_tool JWT."""
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.guardrails.guardrail_hooks.mcp_jwt_signer.mcp_jwt_signer import (
|
||||
inject_mcp_jwt_headers_for_upstream,
|
||||
)
|
||||
|
||||
_make_signer(issuer="https://litellm.example.com", audience="mcp", ttl_seconds=300)
|
||||
user_dict = UserAPIKeyAuth(api_key="sk-test", user_id="alice")
|
||||
|
||||
result = await inject_mcp_jwt_headers_for_upstream(
|
||||
user_api_key_dict=user_dict,
|
||||
for_list_tools=False,
|
||||
mcp_tool_name="search_web",
|
||||
)
|
||||
assert result["Authorization"].startswith("Bearer ")
|
||||
token = result["Authorization"].removeprefix("Bearer ")
|
||||
decoded = _decode_unverified(token)
|
||||
scopes = set(decoded["scope"].split())
|
||||
assert "mcp:tools/call" in scopes
|
||||
assert "mcp:tools/search_web:call" in scopes
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue