diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/FirModuleResolveStateDepended.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/FirModuleResolveStateDepended.kt
index c9ca8ccf9b9..fa565514c27 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/FirModuleResolveStateDepended.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/FirModuleResolveStateDepended.kt
@@ -28,7 +28,7 @@ import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtLambdaExpression
 
 internal class FirModuleResolveStateDepended(
-    private val originalState: FirModuleResolveStateImpl,
+    val originalState: FirModuleResolveStateImpl,
     val towerProviderBuiltUponElement: FirTowerContextProvider,
     private val ktToFirMapping: KtToFirMapping?,
 ) : FirModuleResolveState() {
diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/LowLevelFirApiFacadeForResolveOnAir.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/LowLevelFirApiFacadeForResolveOnAir.kt
index 9637f6fe3fb..65d4d50b127 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/LowLevelFirApiFacadeForResolveOnAir.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/api/LowLevelFirApiFacadeForResolveOnAir.kt
@@ -102,7 +102,6 @@ object LowLevelFirApiFacadeForResolveOnAir {
         place: KtElement,
     ): FirTowerContextProvider {
         require(state is FirModuleResolveStateImpl)
-        require(place.isPhysical)
 
         return if (place is KtFile) {
             FileTowerProvider(place, onAirGetTowerContextForFile(state, place))
@@ -328,13 +327,11 @@ object LowLevelFirApiFacadeForResolveOnAir {
 
     private class TowerProviderForElementForState(private val state: FirModuleResolveState) : FirTowerContextProvider {
         override fun getClosestAvailableParentContext(ktElement: KtElement): FirTowerDataContext? {
-            return if (ktElement.isPhysical) {
-                onAirGetTowerContextProvider(state, ktElement).getClosestAvailableParentContext(ktElement)
-            } else {
-                require(state is FirModuleResolveStateDepended) {
-                    "Invalid resolve state ${this::class.simpleName} but have to be ${FirModuleResolveStateDepended::class.simpleName}"
-                }
+            return if (state is FirModuleResolveStateDepended) {
                 state.towerProviderBuiltUponElement.getClosestAvailableParentContext(ktElement)
+                    ?: onAirGetTowerContextProvider(state.originalState, ktElement).getClosestAvailableParentContext(ktElement)
+            } else {
+                onAirGetTowerContextProvider(state, ktElement).getClosestAvailableParentContext(ktElement)
             }
         }
     }