diff --git a/java/kotlin-extractor2/src/main/kotlin/KotlinFileExtractor.kt b/java/kotlin-extractor2/src/main/kotlin/KotlinFileExtractor.kt
index ae10f57c2ec..14ea36b079b 100644
--- a/java/kotlin-extractor2/src/main/kotlin/KotlinFileExtractor.kt
+++ b/java/kotlin-extractor2/src/main/kotlin/KotlinFileExtractor.kt
@@ -1,5 +1,6 @@
 package com.github.codeql
 
+import org.jetbrains.kotlin.analysis.api.KaSession
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
 
@@ -73,6 +74,7 @@ import org.jetbrains.kotlin.util.OperatorNameConventions
 import org.jetbrains.kotlin.utils.addToStdlib.firstIsInstanceOrNull
 */
 
+context (KaSession)
 open class KotlinFileExtractor(
     override val logger: FileLogger,
     override val tw: FileTrapWriter,
diff --git a/java/kotlin-extractor2/src/main/kotlin/KotlinUsesExtractor.kt b/java/kotlin-extractor2/src/main/kotlin/KotlinUsesExtractor.kt
index 0019725ad10..9bf98efa5c7 100644
--- a/java/kotlin-extractor2/src/main/kotlin/KotlinUsesExtractor.kt
+++ b/java/kotlin-extractor2/src/main/kotlin/KotlinUsesExtractor.kt
@@ -1,6 +1,7 @@
 package com.github.codeql
 
 import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.analysis.api.KaSession
 import org.jetbrains.kotlin.psi.*
 
 /*
@@ -33,6 +34,7 @@ import org.jetbrains.kotlin.types.Variance
 import org.jetbrains.kotlin.util.OperatorNameConventions
 */
 
+context (KaSession)
 open class KotlinUsesExtractor(
     open val logger: Logger,
     open val tw: TrapWriter,
@@ -1822,6 +1824,9 @@ OLD: KE1
         classTypeArgsIncludingOuterClasses: List<IrTypeArgument>?
 */
     ): Label<out T> {
+        println("=== useFunction")
+        println(f)
+        println(f.returnType)
         val label = getFunctionLabel(javaFun, parentId /* TODO , classTypeArgsIncludingOuterClasses */)
         val id: Label<T> =
             tw.getLabelFor(label) {