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

Merge branch 'BerriAI:main' into main

Browse source
This commit is contained in:
Tasmay Pankaj Tibrewal 2025-08-15 05:08:32 +05:30 • committed by GitHub
commit a642bf936c
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",
[