diff --git a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/diagnostics/KtAbstractFirDiagnostic.kt b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/diagnostics/KtAbstractFirDiagnostic.kt
index 3844111f8ac..f1708a79f02 100644
--- a/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/diagnostics/KtAbstractFirDiagnostic.kt
+++ b/analysis/analysis-api-fir/src/org/jetbrains/kotlin/analysis/api/fir/diagnostics/KtAbstractFirDiagnostic.kt
@@ -7,20 +7,22 @@ package org.jetbrains.kotlin.analysis.api.fir.diagnostics
 
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.analysis.api.ValidityTokenOwner
 import org.jetbrains.kotlin.diagnostics.Severity
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDefaultErrorMessages
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirDiagnostic
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirPsiDiagnostic
 import org.jetbrains.kotlin.analysis.api.diagnostics.KtDiagnosticWithPsi
+import org.jetbrains.kotlin.analysis.api.withValidityAssertion
 
-internal interface KtAbstractFirDiagnostic<PSI : PsiElement> : KtDiagnosticWithPsi<PSI> {
+internal interface KtAbstractFirDiagnostic<PSI : PsiElement> : KtDiagnosticWithPsi<PSI>, ValidityTokenOwner {
     val firDiagnostic: FirPsiDiagnostic
 
     override val factoryName: String
-        get() = firDiagnostic.factory.name
+        get() = withValidityAssertion { firDiagnostic.factory.name }
 
     override val defaultMessage: String
-        get() {
+        get() = withValidityAssertion {
             val diagnostic = firDiagnostic as FirDiagnostic
 
             val firDiagnosticRenderer = FirDefaultErrorMessages.getRendererForDiagnostic(diagnostic)
@@ -28,12 +30,12 @@ internal interface KtAbstractFirDiagnostic<PSI : PsiElement> : KtDiagnosticWithP
         }
 
     override val textRanges: Collection<TextRange>
-        get() = firDiagnostic.textRanges
+        get() = withValidityAssertion { firDiagnostic.textRanges }
 
     @Suppress("UNCHECKED_CAST")
     override val psi: PSI
-        get() = firDiagnostic.psiElement as PSI
+        get() = withValidityAssertion { firDiagnostic.psiElement as PSI }
 
     override val severity: Severity
-        get() = firDiagnostic.severity
+        get() = withValidityAssertion { firDiagnostic.severity }
 }