diff --git a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/KtFirAnalysisSessionProvider.kt b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/KtFirAnalysisSessionProvider.kt
index 3a250d9a8f7..aa85626fb12 100644
--- a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/KtFirAnalysisSessionProvider.kt
+++ b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/KtFirAnalysisSessionProvider.kt
@@ -21,10 +21,6 @@ import org.jetbrains.kotlin.analysis.utils.errors.buildErrorWithAttachment
 
 @OptIn(KtAnalysisApiInternals::class)
 class KtFirAnalysisSessionProvider(project: Project) : CachingKtAnalysisSessionProvider<LLFirResolveSession>(project) {
-    override fun getFirResolveSession(contextElement: KtElement): LLFirResolveSession {
-        return contextElement.getFirResolveSession()
-    }
-
     override fun getFirResolveSession(contextSymbol: KtSymbol): LLFirResolveSession {
         return when (contextSymbol) {
             is KtFirSymbol<*> -> contextSymbol.firResolveSession
diff --git a/analysis/analysis-api-impl-base/src/org/jetbrains/kotlin/analysis/api/impl/base/CachingKtAnalysisSessionProvider.kt b/analysis/analysis-api-impl-base/src/org/jetbrains/kotlin/analysis/api/impl/base/CachingKtAnalysisSessionProvider.kt
index aeeb2f6d76f..19ff66c6aa1 100644
--- a/analysis/analysis-api-impl-base/src/org/jetbrains/kotlin/analysis/api/impl/base/CachingKtAnalysisSessionProvider.kt
+++ b/analysis/analysis-api-impl-base/src/org/jetbrains/kotlin/analysis/api/impl/base/CachingKtAnalysisSessionProvider.kt
@@ -17,6 +17,7 @@ import org.jetbrains.kotlin.analysis.api.KtAnalysisSession
 import org.jetbrains.kotlin.analysis.api.session.KtAnalysisSessionProvider
 import org.jetbrains.kotlin.analysis.api.symbols.KtSymbol
 import org.jetbrains.kotlin.analysis.project.structure.KtModule
+import org.jetbrains.kotlin.analysis.project.structure.getKtModule
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeToken
 import org.jetbrains.kotlin.analysis.api.lifetime.KtLifetimeTokenFactory
 import org.jetbrains.kotlin.psi.KtElement
@@ -25,9 +26,8 @@ import kotlin.reflect.KClass
 
 @KtAnalysisApiInternals
 abstract class CachingKtAnalysisSessionProvider<State : Any>(project: Project) : KtAnalysisSessionProvider(project) {
-    private val cache = KtAnalysisSessionCache<Pair<State, KClass<out KtLifetimeToken>>>(project)
+    private val cache = KtAnalysisSessionCache(project)
 
-    protected abstract fun getFirResolveSession(contextElement: KtElement): State
     protected abstract fun getFirResolveSession(contextSymbol: KtSymbol): State
 
     protected abstract fun getFirResolveSession(contextModule: KtModule): State
@@ -38,16 +38,12 @@ abstract class CachingKtAnalysisSessionProvider<State : Any>(project: Project) :
     ): KtAnalysisSession
 
     final override fun getAnalysisSession(useSiteKtElement: KtElement, factory: KtLifetimeTokenFactory): KtAnalysisSession {
-        val firResolveSession = getFirResolveSession(useSiteKtElement)
-        return cache.getAnalysisSession(firResolveSession to factory.identifier) {
-            val token = factory.create(project)
-            createAnalysisSession(firResolveSession, token)
-        }
+        return getAnalysisSessionByUseSiteKtModule(useSiteKtElement.getKtModule(project), factory)
     }
 
     final override fun getAnalysisSessionByUseSiteKtModule(useSiteKtModule: KtModule, factory: KtLifetimeTokenFactory): KtAnalysisSession {
-        val firResolveSession = getFirResolveSession(useSiteKtModule)
-        return cache.getAnalysisSession(firResolveSession to factory.identifier) {
+        return cache.getAnalysisSession(useSiteKtModule to factory.identifier) {
+            val firResolveSession = getFirResolveSession(useSiteKtModule)
             val validityToken = factory.create(project)
             createAnalysisSession(firResolveSession, validityToken)
         }
@@ -63,10 +59,10 @@ abstract class CachingKtAnalysisSessionProvider<State : Any>(project: Project) :
     }
 }
 
-private class KtAnalysisSessionCache<KEY : Any>(project: Project) {
+private class KtAnalysisSessionCache(project: Project) {
     private val cache = CachedValuesManager.getManager(project).createCachedValue {
         CachedValueProvider.Result(
-            ConcurrentHashMap<KEY, KtAnalysisSession>(),
+            ConcurrentHashMap<Pair<KtModule, KClass<out KtLifetimeToken>>, KtAnalysisSession>(),
             PsiModificationTracker.MODIFICATION_COUNT,
             ProjectRootModificationTracker.getInstance(project),
             project.createProjectWideOutOfBlockModificationTracker()
@@ -78,7 +74,7 @@ private class KtAnalysisSessionCache<KEY : Any>(project: Project) {
         cache.value.clear()
     }
 
-    inline fun getAnalysisSession(key: KEY, create: () -> KtAnalysisSession): KtAnalysisSession =
+    inline fun getAnalysisSession(key: Pair<KtModule, KClass<out KtLifetimeToken>>, create: () -> KtAnalysisSession): KtAnalysisSession =
         cache.value.getOrPut(key) { create() }
 
     fun getCachedAnalysisSession(key: KEY): KtAnalysisSession? =