scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

[MCP Gateway] LiteLLM Fix MCP gateway key auth (#13630)

Browse source 
* Fix - add safe divide by 0 for most places to prevent crash

* Enhance MCPRequestHandler to support permission inheritance and intersection logic for access groups. Added integration tests to verify behavior when keys have no permissions and when both keys and teams have overlapping permissions.

* Remove redundant assertions for permission checks in test_user_api_key_auth_mcp.py to streamline test logic.

* Refactor integration tests for MCPRequestHandler to simplify mocking. Replace complex database mocks with direct function mocks for permission inheritance and intersection scenarios, improving test clarity and maintainability.

* Revert "Fix - add safe divide by 0 for most places to prevent crash"

This reverts commit 265d40e390.
This commit is contained in:
Jugal D. Bhatt 2025-08-14 16:32:18 -07:00 • committed by GitHub
parent aea0605eed
commit b8fe5f7b17
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 195 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
View file

@ -236,12 +236,17 @@ class MCPRequestHandler:
)
#########################################################
# If team has mcp_servers, then key must have a subset of the team's mcp_servers
# If team has mcp_servers, handle inheritance and intersection logic
#########################################################
if len(allowed_mcp_servers_for_team) > 0:
for _mcp_server in allowed_mcp_servers_for_key:
if _mcp_server in allowed_mcp_servers_for_team:
allowed_mcp_servers.append(_mcp_server)
if len(allowed_mcp_servers_for_key) > 0:
# Key has its own MCP permissions - use intersection with team permissions
for _mcp_server in allowed_mcp_servers_for_key:
if _mcp_server in allowed_mcp_servers_for_team:
allowed_mcp_servers.append(_mcp_server)
else:
# Key has no MCP permissions - inherit from team
allowed_mcp_servers = allowed_mcp_servers_for_team
else:
allowed_mcp_servers = allowed_mcp_servers_for_key

106
tests/mcp_tests/test_mcp_server.py
View file

@ -1773,3 +1773,109 @@ async def test_list_tool_rest_api_all_servers_with_auth():
assert calls[1][0][2] == "2025-06-18" # mcp_protocol_version
@pytest.mark.asyncio
async def test_mcp_access_group_permission_inheritance_integration():
"""Integration test for MCP access group permission inheritance"""
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler
from litellm.proxy._types import UserAPIKeyAuth
# Test scenario: team has access groups, key has no permissions -> should inherit
# Use direct mocking of the helper functions instead of complex database mocking
with patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_key") as mock_key:
with patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_team") as mock_team:
# Key has no permissions, team has servers
mock_key.return_value = [] # Key inherits nothing directly
mock_team.return_value = ["staff-server-1", "staff-server-2", "ops-server-1"] # Team has servers
# Create user auth object
user_auth = UserAPIKeyAuth(
api_key="test-key",
user_id="test-user",
team_id="team-staff",
object_permission_id=None # Key has no explicit permissions
)
# Test the inheritance logic
allowed_servers = await MCPRequestHandler.get_allowed_mcp_servers(user_auth)
# Should inherit all team servers since key has no permissions
expected_servers = ["staff-server-1", "staff-server-2", "ops-server-1"]
assert sorted(allowed_servers) == sorted(expected_servers)
@pytest.mark.asyncio
async def test_mcp_access_group_permission_intersection_integration():
"""Integration test for MCP access group permission intersection"""
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler
from litellm.proxy._types import UserAPIKeyAuth
# Test scenario: both team and key have access groups -> should intersect
# Use direct mocking of the helper functions instead of complex database mocking
with patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_key") as mock_key:
with patch.object(MCPRequestHandler, "_get_allowed_mcp_servers_for_team") as mock_team:
# Both key and team have permissions - should intersect
mock_key.return_value = ["ops-server", "external-server"] # Key has these servers
mock_team.return_value = ["staff-server", "ops-server", "admin-server"] # Team has these servers
# Create user auth object
user_auth = UserAPIKeyAuth(
api_key="test-key",
user_id="test-user",
team_id="team-staff",
object_permission_id="key-permission-id" # Key has explicit permissions
)
# Test the intersection logic
allowed_servers = await MCPRequestHandler.get_allowed_mcp_servers(user_auth)
# Should only get intersection (ops-server is common)
expected_servers = ["ops-server"]
assert sorted(allowed_servers) == sorted(expected_servers)
@pytest.mark.asyncio
async def test_mcp_server_manager_with_access_groups_integration():
"""Integration test for MCPServerManager with access group filtering"""
from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler
from litellm.proxy._types import UserAPIKeyAuth
# Create a test manager
test_manager = MCPServerManager()
# Load servers with access groups
test_manager.load_servers_from_config({
"staff_server": {
"url": "https://staff-server.com/mcp",
"access_groups": ["staff"],
"transport": MCPTransport.http,
},
"ops_server": {
"url": "https://ops-server.com/mcp",
"access_groups": ["ops"],
"transport": MCPTransport.http,
},
"admin_server": {
"url": "https://admin-server.com/mcp",
"access_groups": ["admin"],
"transport": MCPTransport.http,
}
})
# Mock user with specific access groups
user_auth = UserAPIKeyAuth(
api_key="test-key",
user_id="test-user",
team_id="team-staff"
)
# Mock the permission lookup to return staff access group
with patch.object(MCPRequestHandler, "get_allowed_mcp_servers") as mock_get_allowed:
mock_get_allowed.return_value = ["staff-server-id", "ops-server-id"] # User has access to staff and ops
allowed_servers = await test_manager.get_allowed_mcp_servers(user_auth)
# Should only get servers user has access to
assert len(allowed_servers) >= 0 # At least verify no errors
mock_get_allowed.assert_called_once_with(user_auth)

80
tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
View file

@ -118,6 +118,86 @@ class TestMCPRequestHandler:
if not user_api_key_auth or not user_api_key_auth.object_permission_id:
mock_find_unique.assert_not_called()
@pytest.mark.parametrize(
"team_servers,key_servers,expected_servers,scenario",
[
# Test case 1: Key has no permissions, should inherit from team
(["server1", "server2"], [], ["server1", "server2"], "inherit_from_team"),
# Test case 2: Key has permissions, should use intersection with team
(["server1", "server2", "server3"], ["server2", "server4"], ["server2"], "intersection_logic"),
# Test case 3: Key has permissions but no overlap with team
(["server1", "server2"], ["server3", "server4"], [], "no_overlap"),
# Test case 4: Team has no permissions, use key permissions
([], ["server1", "server2"], ["server1", "server2"], "no_team_permissions"),
# Test case 5: Both team and key have no permissions
([], [], [], "no_permissions"),
# Test case 6: Team has permissions, key has subset
(["server1", "server2", "server3"], ["server1", "server3"], ["server1", "server3"], "key_subset"),
# Test case 7: Team has permissions, key has superset (intersection should limit)
(["server1", "server2"], ["server1", "server2", "server3"], ["server1", "server2"], "key_superset"),
],
)
async def test_get_allowed_mcp_servers_inheritance_logic(
self, team_servers, key_servers, expected_servers, scenario
):
"""Test the inheritance and intersection logic in get_allowed_mcp_servers"""
# Create mock user_api_key_auth
user_api_key_auth = UserAPIKeyAuth(
api_key="test-key",
user_id="test-user",
team_id="test-team" if team_servers else None,
object_permission_id="test-permission" if key_servers else None
)
# Mock the helper functions
with patch.object(
MCPRequestHandler, "_get_allowed_mcp_servers_for_key"
) as mock_key_servers:
with patch.object(
MCPRequestHandler, "_get_allowed_mcp_servers_for_team"
) as mock_team_servers:
# Configure mocks to return the test data
mock_key_servers.return_value = key_servers
mock_team_servers.return_value = team_servers
# Call the method
result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth)
# Assert the result (order-independent comparison)
assert sorted(result) == sorted(expected_servers)
# Verify the mock functions were called correctly
mock_key_servers.assert_called_once_with(user_api_key_auth)
mock_team_servers.assert_called_once_with(user_api_key_auth)
async def test_permission_inheritance_edge_cases(self):
"""Test edge cases in permission inheritance"""
# Test case: None values in database
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_objectpermissiontable.find_unique.return_value = None
mock_prisma_client.db.litellm_teamtable.find_unique.return_value = None
user_api_key_auth = UserAPIKeyAuth(
api_key="test-key",
user_id="test-user",
team_id="test-team",
object_permission_id="test-permission"
)
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client):
result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth)
assert result == []
# Test case: Exception handling
mock_prisma_client.db.litellm_objectpermissiontable.find_unique.side_effect = Exception("DB Error")
with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client):
result = await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth)
assert result == [] # Should handle exception gracefully
@pytest.mark.parametrize(
"headers,expected_api_key,expected_mcp_auth_header,expected_server_auth_headers",
[