diff --git a/analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/KtFe10AnalysisSession.kt b/analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/KtFe10AnalysisSession.kt
index 3468a49c1ed..8e887a9b321 100644
--- a/analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/KtFe10AnalysisSession.kt
+++ b/analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/KtFe10AnalysisSession.kt
@@ -6,7 +6,6 @@
 package org.jetbrains.kotlin.analysis.api.descriptors
 
 import com.intellij.openapi.components.ServiceManager
-import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
 import org.jetbrains.kotlin.analysis.api.components.*
 import org.jetbrains.kotlin.analysis.api.descriptors.components.*
@@ -17,7 +16,7 @@ import org.jetbrains.kotlin.psi.KtFile
 
 @Suppress("LeakingThis")
 class KtFe10AnalysisSession(
-    private val contextElement: PsiElement, token: ValidityToken
+    private val contextElement: KtElement, token: ValidityToken
 ) : KtAnalysisSession(token), Fe10AnalysisFacade by ServiceManager.getService(contextElement.project, Fe10AnalysisFacade::class.java) {
     override val smartCastProviderImpl: KtSmartCastProvider = KtFe10SmartCastProvider(this)
     override val diagnosticProviderImpl: KtDiagnosticProvider = KtFe10DiagnosticProvider(this)
diff --git a/analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/KtFe10CliAnalysisSessionProvider.kt b/analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/KtFe10CliAnalysisSessionProvider.kt
index b236cab5c27..113a972be44 100644
--- a/analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/KtFe10CliAnalysisSessionProvider.kt
+++ b/analysis/analysis-api-fe10/src/org/jetbrains/kotlin/analysis/api/descriptors/KtFe10CliAnalysisSessionProvider.kt
@@ -26,7 +26,7 @@ class KtFe10AnalysisSessionProvider : KtAnalysisSessionProvider() {
             return contextSymbol.analysisSession
         } else {
             val contextElement = contextSymbol.psi
-            if (contextElement != null) {
+            if (contextElement is KtElement) {
                 return KtFe10AnalysisSession(contextElement, contextSymbol.token)
             }
         }