diff --git a/analysis/analysis-api-standalone/src/org/jetbrains/kotlin/analysis/api/standalone/StandaloneAnalysisAPISession.kt b/analysis/analysis-api-standalone/src/org/jetbrains/kotlin/analysis/api/standalone/StandaloneAnalysisAPISession.kt
index fc23085d565..9d3f33d5b07 100644
--- a/analysis/analysis-api-standalone/src/org/jetbrains/kotlin/analysis/api/standalone/StandaloneAnalysisAPISession.kt
+++ b/analysis/analysis-api-standalone/src/org/jetbrains/kotlin/analysis/api/standalone/StandaloneAnalysisAPISession.kt
@@ -11,14 +11,14 @@ import com.intellij.openapi.application.Application
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiFile
 import com.intellij.psi.search.GlobalSearchScope
-import org.jetbrains.kotlin.analysis.project.structure.KtModule
+import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreProjectEnvironment
 import org.jetbrains.kotlin.load.kotlin.PackagePartProvider
 
 public class StandaloneAnalysisAPISession internal constructor(
     kotlinCoreProjectEnvironment: KotlinCoreProjectEnvironment,
     public val createPackagePartProvider: (GlobalSearchScope) -> PackagePartProvider,
-    modulesWithFilesProvider: () -> Map<KtModule, List<PsiFile>>
+    modulesWithFilesProvider: () -> Map<KtSourceModule, List<PsiFile>>
 ) {
     // TODO: better to limit exposure? Current usages are: addExtension, jarFileSystem
     public val coreApplicationEnvironment: CoreApplicationEnvironment = kotlinCoreProjectEnvironment.environment
@@ -33,5 +33,5 @@ public class StandaloneAnalysisAPISession internal constructor(
     )
     public val mockProject: MockProject = kotlinCoreProjectEnvironment.project
 
-    public val modulesWithFiles: Map<KtModule, List<PsiFile>> by lazy(modulesWithFilesProvider)
+    public val modulesWithFiles: Map<KtSourceModule, List<PsiFile>> by lazy(modulesWithFilesProvider)
 }
\ No newline at end of file
diff --git a/analysis/analysis-api-standalone/src/org/jetbrains/kotlin/analysis/api/standalone/StandaloneAnalysisAPISessionBuilder.kt b/analysis/analysis-api-standalone/src/org/jetbrains/kotlin/analysis/api/standalone/StandaloneAnalysisAPISessionBuilder.kt
index 4f6be1a197e..2a683fa860e 100644
--- a/analysis/analysis-api-standalone/src/org/jetbrains/kotlin/analysis/api/standalone/StandaloneAnalysisAPISessionBuilder.kt
+++ b/analysis/analysis-api-standalone/src/org/jetbrains/kotlin/analysis/api/standalone/StandaloneAnalysisAPISessionBuilder.kt
@@ -21,9 +21,11 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.services.FirSealedCla
 import org.jetbrains.kotlin.analysis.low.level.api.fir.project.structure.JvmFirDeserializedSymbolProviderFactory
 import org.jetbrains.kotlin.analysis.project.structure.KtModuleScopeProvider
 import org.jetbrains.kotlin.analysis.project.structure.KtModuleScopeProviderImpl
+import org.jetbrains.kotlin.analysis.project.structure.KtSourceModule
 import org.jetbrains.kotlin.analysis.project.structure.builder.KtModuleProviderBuilder
 import org.jetbrains.kotlin.analysis.project.structure.builder.buildProjectStructureProvider
 import org.jetbrains.kotlin.analysis.project.structure.impl.KtModuleProviderImpl
+import org.jetbrains.kotlin.analysis.project.structure.impl.KtSourceModuleImpl
 import org.jetbrains.kotlin.analysis.project.structure.impl.buildKtModuleProviderByCompilerConfiguration
 import org.jetbrains.kotlin.analysis.project.structure.impl.getPsiFilesFromPaths
 import org.jetbrains.kotlin.analysis.project.structure.impl.getSourceFilePaths
@@ -216,10 +218,11 @@ public class StandaloneAnalysisAPISessionBuilder(
             kotlinCoreProjectEnvironment,
             createPackagePartProvider,
         ) {
-            projectStructureProvider.allSourceFiles
-                .asSequence()
-                .filterIsInstance<PsiFile>()
-                .groupBy { projectStructureProvider.getModule(it.originalElement, null) }
+            projectStructureProvider.allKtModules.mapNotNull { ktModule ->
+                if (ktModule !is KtSourceModule) return@mapNotNull null
+                check(ktModule is KtSourceModuleImpl)
+                ktModule to ktModule.sourceRoots.filterIsInstance<PsiFile>()
+            }.toMap()
         }
     }
 }
@@ -231,7 +234,7 @@ public inline fun buildStandaloneAnalysisAPISession(
     unitTestMode: Boolean = false,
     withPsiDeclarationFromBinaryModuleProvider: Boolean = false,
     classLoader: ClassLoader = MockProject::class.java.classLoader,
-    init: StandaloneAnalysisAPISessionBuilder.() -> Unit
+    init: StandaloneAnalysisAPISessionBuilder.() -> Unit,
 ): StandaloneAnalysisAPISession {
     contract {
         callsInPlace(init, InvocationKind.EXACTLY_ONCE)