diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 7469848e2f2..a075de13fb1 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -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 diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 43130e62d0d..8a390412495 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -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) + + diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 2249b9a6fa9..a9f1f8b12d2 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -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", [