diff --git a/java/kotlin-extractor2/src/main/kotlin/entities/Class.kt b/java/kotlin-extractor2/src/main/kotlin/entities/Class.kt
index 642164fd0ce..b66289a6d20 100644
--- a/java/kotlin-extractor2/src/main/kotlin/entities/Class.kt
+++ b/java/kotlin-extractor2/src/main/kotlin/entities/Class.kt
@@ -74,7 +74,7 @@ fun KotlinFileExtractor.extractClassSource(
                         if (getter == null) {
                             logger.warnElement(
                                 "Expected an annotation property to have a getter",
-                                it.psiSafe() ?: TODO()
+                                it
                             )
                         } else {
                             extractFunction(
diff --git a/java/kotlin-extractor2/src/main/kotlin/entities/Function.kt b/java/kotlin-extractor2/src/main/kotlin/entities/Function.kt
index af55a02df2a..5edae548dc7 100644
--- a/java/kotlin-extractor2/src/main/kotlin/entities/Function.kt
+++ b/java/kotlin-extractor2/src/main/kotlin/entities/Function.kt
@@ -641,6 +641,7 @@ OLD: KE1
     }
 }
 
+context(KaSession)
 fun KotlinFileExtractor.extractValueParameter(
     id: Label<out DbParam>,
     t: KaType,
diff --git a/java/kotlin-extractor2/src/main/kotlin/entities/FunctionalInterface.kt b/java/kotlin-extractor2/src/main/kotlin/entities/FunctionalInterface.kt
index c3d3f8b4975..c37068e5680 100644
--- a/java/kotlin-extractor2/src/main/kotlin/entities/FunctionalInterface.kt
+++ b/java/kotlin-extractor2/src/main/kotlin/entities/FunctionalInterface.kt
@@ -216,6 +216,7 @@ private val nullableAnyArrayType: KaType
  * Adds a function named `invoke` with the specified parameter types and return type to the
  * class identified by `parentId`.
  */
+context(KaSession)
 private fun KotlinFileExtractor.addFunctionInvoke(
     methodId: Label<DbMethod>,
     parameterTypes: List<KaType>,
@@ -237,6 +238,7 @@ private fun KotlinFileExtractor.addFunctionInvoke(
  * Extracts a function with the given name, parameter types, return type, containing type, and
  * location.
  */
+context(KaSession)
 private fun KotlinFileExtractor.addFunctionManual(
     methodId: Label<DbMethod>,
     name: String,