diff --git a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirAnnotationArgumentsLazyResolver.kt b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirAnnotationArgumentsLazyResolver.kt
index 0dc897fd46f..f02b7cb3d30 100644
--- a/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirAnnotationArgumentsLazyResolver.kt
+++ b/analysis/low-level-api-fir/src/org/jetbrains/kotlin/analysis/low/level/api/fir/transformers/LLFirAnnotationArgumentsLazyResolver.kt
@@ -88,7 +88,7 @@ private class LLFirAnnotationArgumentsTargetResolver(
         when {
             target is FirRegularClass -> {
                 val declarationTransformer = transformer.declarationsTransformer
-                declarationTransformer.context.insideClassHeader {
+                declarationTransformer.context.withClassHeader(target) {
                     target.transformAnnotations(declarationTransformer, ResolutionMode.ContextIndependent)
                     target.transformTypeParameters(declarationTransformer, ResolutionMode.ContextIndependent)
                     target.transformSuperTypeRefs(declarationTransformer, ResolutionMode.ContextIndependent)
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt
index 4089ca7ef71..b20b8744c14 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt
@@ -98,11 +98,11 @@ class BodyResolveContext(
     var insideClassHeader: Boolean = false
 
     @OptIn(PrivateForInline::class)
-    inline fun insideClassHeader(action: () -> Unit) {
+    inline fun withClassHeader(clazz: FirRegularClass, action: () -> Unit) {
         val old = insideClassHeader
         insideClassHeader = true
         try {
-            action()
+            withContainer(clazz, action)
         } finally {
             insideClassHeader = old
         }
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt
index b75f41e3e66..7bb3fd26832 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt
@@ -664,7 +664,7 @@ open class FirDeclarationsResolveTransformer(
             }
 
             if (isLocal || !implicitTypeOnly) {
-                context.insideClassHeader {
+                context.withClassHeader(regularClass) {
                     regularClass.transformAnnotations(this, ResolutionMode.ContextIndependent)
                     regularClass.transformTypeParameters(this, ResolutionMode.ContextIndependent)
                     regularClass.transformSuperTypeRefs(this, ResolutionMode.ContextIndependent)
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirAnnotationArgumentsTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirAnnotationArgumentsTransformer.kt
index 3a365c56155..a4243727b82 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirAnnotationArgumentsTransformer.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/plugin/FirAnnotationArgumentsTransformer.kt
@@ -27,7 +27,6 @@ import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculator
 import org.jetbrains.kotlin.fir.resolve.transformers.ReturnTypeCalculatorForFullBodyResolve
 import org.jetbrains.kotlin.fir.resolve.transformers.body.resolve.*
 import org.jetbrains.kotlin.fir.symbols.impl.FirEnumEntrySymbol
-import org.jetbrains.kotlin.fir.visitors.transformSingle
 import org.jetbrains.kotlin.name.ClassId
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.StandardClassIds
@@ -185,7 +184,7 @@ private class FirDeclarationsResolveTransformerForAnnotationArguments(
     }
 
     override fun transformRegularClass(regularClass: FirRegularClass, data: ResolutionMode): FirRegularClass {
-        context.insideClassHeader {
+        context.withClassHeader(regularClass) {
             regularClass.transformAnnotations(this, ResolutionMode.ContextIndependent)
             regularClass.transformTypeParameters(this, ResolutionMode.ContextIndependent)
             regularClass.transformSuperTypeRefs(this, ResolutionMode.ContextIndependent)