diff --git a/utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/language/JavaLanguage.kt b/utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/language/JavaLanguage.kt
index f2fa701e19..4895f6ca52 100644
--- a/utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/language/JavaLanguage.kt
+++ b/utbot-intellij/src/main/kotlin/org/utbot/intellij/plugin/language/JavaLanguage.kt
@@ -109,7 +109,7 @@ object JvmLanguageAssistant : LanguageAssistant() {
                     }
                 }
             } else {
-                val someSelection = e.getData(PlatformDataKeys.SELECTED_ITEMS)?: return null
+                val someSelection = e.getData(PlatformDataKeys.PSI_ELEMENT_ARRAY)?: return null
                 someSelection.forEach {
                     when(it) {
                         is PsiFileSystemItem  -> srcClasses += getAllClasses(project, arrayOf(it.virtualFile))
diff --git a/utbot-ui-commons/src/main/kotlin/org/utbot/intellij/plugin/language/agnostic/LanguageAssistant.kt b/utbot-ui-commons/src/main/kotlin/org/utbot/intellij/plugin/language/agnostic/LanguageAssistant.kt
index 98b6c8ab9a..87d7314ee4 100644
--- a/utbot-ui-commons/src/main/kotlin/org/utbot/intellij/plugin/language/agnostic/LanguageAssistant.kt
+++ b/utbot-ui-commons/src/main/kotlin/org/utbot/intellij/plugin/language/agnostic/LanguageAssistant.kt
@@ -45,7 +45,7 @@ abstract class LanguageAssistant {
                         element.containingFile?.let { getLanguageFromFile(it) }
                     }
                     else -> {
-                        val someSelection = e.getData(PlatformDataKeys.SELECTED_ITEMS)?: return null
+                        val someSelection = e.getData(PlatformDataKeys.PSI_ELEMENT_ARRAY)?: return null
                         someSelection.firstNotNullOfOrNull {
                             when(it) {
                                 is PsiFileSystemItem -> findLanguageRecursively(project, arrayOf(it.virtualFile))