diff --git a/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/KtAnalysisSessionProvider.kt b/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/KtAnalysisSessionProvider.kt
index c94600539ad..a2da38d8c29 100644
--- a/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/KtAnalysisSessionProvider.kt
+++ b/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/KtAnalysisSessionProvider.kt
@@ -35,33 +35,24 @@ abstract class KtAnalysisSessionProvider {
     inline fun <R> analyseInFakeAnalysisSession(originalFile: KtFile, fakeExpresion: KtElement, action: KtAnalysisSession.() -> R): R {
         val fakeAnalysisSession = getAnalysisSession(originalFile, ReadActionConfinementValidityTokenFactory)
             .createContextDependentCopy(originalFile, fakeExpresion)
-        return analyse(fakeAnalysisSession, action)
+        return analyse(fakeAnalysisSession, ReadActionConfinementValidityTokenFactory, action)
     }
 
     @InvalidWayOfUsingAnalysisSession
     inline fun <R> analyse(contextElement: KtElement, tokenFactory: ValidityTokenFactory, action: KtAnalysisSession.() -> R): R =
-        analyse(getAnalysisSession(contextElement, tokenFactory), action)
+        analyse(getAnalysisSession(contextElement, tokenFactory), tokenFactory, action)
 
     @OptIn(KtAnalysisSessionProviderInternals::class)
     @InvalidWayOfUsingAnalysisSession
-    inline fun <R> analyse(analysisSession: KtAnalysisSession, action: KtAnalysisSession.() -> R): R {
-        currentAnalysisContextEnteringCount.set(currentAnalysisContextEnteringCount.get() + 1)
+    inline fun <R> analyse(analysisSession: KtAnalysisSession, factory: ValidityTokenFactory, action: KtAnalysisSession.() -> R): R {
+        factory.beforeEnteringAnalysisContext()
         return try {
             analysisSession.action()
         } finally {
-            currentAnalysisContextEnteringCount.set(currentAnalysisContextEnteringCount.get() - 1)
+            factory.afterLeavingAnalysisContext()
         }
     }
 
-    companion object {
-        @KtAnalysisSessionProviderInternals
-        val currentAnalysisContextEnteringCount = object : ThreadLocal<Int>() {
-            override fun initialValue() = 0
-        }
-
-        @OptIn(KtAnalysisSessionProviderInternals::class)
-        fun isInsideAnalysisContext() = currentAnalysisContextEnteringCount.get() > 0
-    }
 }
 
 /**
diff --git a/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/tokens/ReadActionConfinementValidityToken.kt b/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/tokens/ReadActionConfinementValidityToken.kt
index 2919bceeeeb..9e0e414b8f3 100644
--- a/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/tokens/ReadActionConfinementValidityToken.kt
+++ b/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/tokens/ReadActionConfinementValidityToken.kt
@@ -31,7 +31,7 @@ class ReadActionConfinementValidityToken(project: Project) : ValidityToken() {
         if (application.isDispatchThread && !allowOnEdt.get()) return false
         if (ForbidKtResolve.resovleIsForbidenInActionWithName.get() != null) return false
         if (!application.isReadAccessAllowed) return false
-        if (!KtAnalysisSessionProvider.isInsideAnalysisContext()) return false
+        if (!ReadActionConfinementValidityTokenFactory.isInsideAnalysisContext()) return false
         return true
     }
 
@@ -43,7 +43,7 @@ class ReadActionConfinementValidityToken(project: Project) : ValidityToken() {
         ForbidKtResolve.resovleIsForbidenInActionWithName.get()?.let { actionName ->
             return "Resolve is forbidden in $actionName"
         }
-        if (!KtAnalysisSessionProvider.isInsideAnalysisContext()) return "Called outside analyse method"
+        if (!ReadActionConfinementValidityTokenFactory.isInsideAnalysisContext()) return "Called outside analyse method"
         error("Getting inaccessibility reason for validity token when it is accessible")
     }
 
@@ -58,6 +58,18 @@ object ReadActionConfinementValidityTokenFactory : ValidityTokenFactory() {
     override val identifier: KClass<out ValidityToken> = ReadActionConfinementValidityToken::class
 
     override fun create(project: Project): ValidityToken = ReadActionConfinementValidityToken(project)
+
+    override fun beforeEnteringAnalysisContext() {
+        currentAnalysisContextEnteringCount.set(currentAnalysisContextEnteringCount.get() + 1)
+    }
+
+    override fun afterLeavingAnalysisContext() {
+        currentAnalysisContextEnteringCount.set(currentAnalysisContextEnteringCount.get() - 1)
+    }
+
+    private val currentAnalysisContextEnteringCount = ThreadLocal.withInitial { 0 }
+
+    internal fun isInsideAnalysisContext() = currentAnalysisContextEnteringCount.get() > 0
 }
 
 @RequiresOptIn("All frontend related work should not be allowed to be ran from EDT thread. Only use it as a temporary solution")
diff --git a/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/tokens/ValidityToken.kt b/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/tokens/ValidityToken.kt
index 256d6ebfe5e..282b58c0301 100644
--- a/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/tokens/ValidityToken.kt
+++ b/idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/tokens/ValidityToken.kt
@@ -19,6 +19,9 @@ abstract class ValidityToken {
 abstract class ValidityTokenFactory {
     abstract val identifier: KClass<out ValidityToken>
     abstract fun create(project: Project): ValidityToken
+
+    open fun beforeEnteringAnalysisContext() {}
+    open fun afterLeavingAnalysisContext() {}
 }
 
 
@@ -32,8 +35,8 @@ inline fun ValidityToken.assertIsValidAndAccessible() {
     }
 }
 
-abstract class BadEntityAccessException(): IllegalStateException()
+abstract class BadEntityAccessException() : IllegalStateException()
 
 class InvalidEntityAccessException(override val message: String) : BadEntityAccessException()
-class InaccessibleEntityAccessException(override val message: String): BadEntityAccessException()
+class InaccessibleEntityAccessException(override val message: String) : BadEntityAccessException()