[kotlin] KTIJ-34957 get script configurations from workspace model

GitOrigin-RevId: a1666357b06d2e4ef26abfe72035a6591b76407f
2026-09-27 10:03:11 +07:00 · 2025-11-14 01:04:31 +00:00
parent d5018a1583
commit bdeb003fdc
24 changed files with 272 additions and 359 deletions
@@ -20,7 +20,7 @@ internal object MetadataStorageImpl : MetadataStorageBase() {
    var typeMetadata: StorageTypeMetadata

    typeMetadata = FinalClassMetadata.ObjectMetadata(
-      fqName = "org.jetbrains.kotlin.idea.core.script.k2.configurations.MainKtsScriptConfigurationProvider\$MainKtsKotlinScriptEntitySource",
+      fqName = "org.jetbrains.kotlin.idea.core.script.k2.configurations.MainKtsConfigurationProvider\$MainKtsKotlinScriptEntitySource",
      properties = listOf(OwnPropertyMetadata(isComputable = false, isKey = false, isOpen = false, name = "virtualFileUrl",
                                              valueType = ValueTypeMetadata.SimpleType.CustomType(isNullable = true,
                                                                                                  typeMetadata = FinalClassMetadata.KnownClass(
@@ -236,13 +236,13 @@ internal object MetadataStorageImpl : MetadataStorageBase() {
                    metadataHash = -1162660984)
    addMetadataHash(typeFqn = "com.intellij.platform.workspace.jps.entities.SdkId", metadataHash = 51502100)
    addMetadataHash(typeFqn = "org.jetbrains.kotlin.idea.core.script.k2.modules.ScriptDiagnosticData", metadataHash = -579968699)
-    addMetadataHash(typeFqn = "com.intellij.platform.workspace.storage.EntitySource", metadataHash = -706084288)
+    addMetadataHash(typeFqn = "com.intellij.platform.workspace.storage.EntitySource", metadataHash = 1965151115)
    addMetadataHash(
      typeFqn = "org.jetbrains.kotlin.idea.core.script.k2.configurations.DefaultScriptConfigurationHandler\$DefaultScriptEntitySource",
      metadataHash = -513260474)
    addMetadataHash(
-      typeFqn = "org.jetbrains.kotlin.idea.core.script.k2.configurations.MainKtsScriptConfigurationProvider\$MainKtsKotlinScriptEntitySource",
-      metadataHash = 2056756862)
+      typeFqn = "org.jetbrains.kotlin.idea.core.script.k2.configurations.MainKtsConfigurationProvider\$MainKtsKotlinScriptEntitySource",
+      metadataHash = 618687241)
    addMetadataHash(typeFqn = "com.intellij.platform.workspace.storage.SymbolicEntityId", metadataHash = -2087540107)
  }

@@ -5,6 +5,7 @@ package org.jetbrains.kotlin.idea.core.script.k2.modules

 import com.intellij.platform.workspace.jps.entities.SdkId
 import com.intellij.platform.workspace.storage.*
+import com.intellij.platform.workspace.storage.annotations.Default
 import com.intellij.platform.workspace.storage.impl.containers.toMutableWorkspaceList
 import com.intellij.platform.workspace.storage.url.VirtualFileUrl

@@ -3,6 +3,7 @@

 package org.jetbrains.kotlin.idea.core.script.k2.modules

+import com.intellij.openapi.util.NlsSafe
 import com.intellij.platform.workspace.storage.*
 import com.intellij.platform.workspace.storage.impl.containers.toMutableWorkspaceList
 import com.intellij.platform.workspace.storage.url.VirtualFileUrl
@@ -3,6 +3,7 @@ package org.jetbrains.kotlin.idea.core.script.k2.modules.impl

 import com.intellij.platform.workspace.jps.entities.SdkId
 import com.intellij.platform.workspace.storage.*
+import com.intellij.platform.workspace.storage.annotations.Default
 import com.intellij.platform.workspace.storage.impl.ModifiableWorkspaceEntityBase
 import com.intellij.platform.workspace.storage.impl.SoftLinkable
 import com.intellij.platform.workspace.storage.impl.WorkspaceEntityBase
@@ -1,6 +1,7 @@
 // Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package org.jetbrains.kotlin.idea.core.script.k2.modules.impl

+import com.intellij.openapi.util.NlsSafe
 import com.intellij.platform.workspace.storage.*
 import com.intellij.platform.workspace.storage.impl.ModifiableWorkspaceEntityBase
 import com.intellij.platform.workspace.storage.impl.WorkspaceEntityBase
@@ -53,8 +53,6 @@
    <projectService serviceInterface="org.jetbrains.kotlin.scripting.definitions.ScriptDefinitionProvider"
                    serviceImplementation="org.jetbrains.kotlin.idea.core.script.k2.definitions.ScriptDefinitionProviderImpl"/>

-    <projectService serviceImplementation="org.jetbrains.kotlin.idea.core.script.k2.configurations.DefaultScriptConfigurationHandler"/>
-
    <workspaceModel.fileIndexContributor implementation="org.jetbrains.kotlin.idea.core.script.k2.KotlinScriptWorkspaceFileIndexContributor"/>
    <projectView.externalLibraries.workspaceModelNodesProvider implementation="org.jetbrains.kotlin.idea.core.script.k2.KotlinScriptExternalLibrariesNodesProvider"/>

@@ -29,7 +29,7 @@ import kotlinx.coroutines.withContext
 import org.intellij.lang.annotations.Language
 import org.jetbrains.annotations.ApiStatus
 import org.jetbrains.kotlin.idea.core.script.k2.ReloadScriptConfigurationService.Companion.TOPIC
-import org.jetbrains.kotlin.idea.core.script.k2.configurations.getConfigurationResolver
+import org.jetbrains.kotlin.idea.core.script.k2.configurations.getConfigurationProviderExtension
 import org.jetbrains.kotlin.idea.core.script.k2.definitions.ScriptDefinitionsModificationTracker
 import org.jetbrains.kotlin.idea.core.script.k2.highlighting.KotlinScriptResolutionService
 import org.jetbrains.kotlin.idea.core.script.k2.modules.KotlinScriptModuleManager.Companion.removeScriptModules
@@ -112,7 +112,7 @@ class ReloadScriptConfigurationService(private val project: Project, private val
        val virtualFile = ktFile.alwaysVirtualFile

        scope.launch {
-            definition.getConfigurationResolver(project).remove(virtualFile)
+            definition.getConfigurationProviderExtension(project).remove(virtualFile)
            project.removeScriptModules(listOf(virtualFile))
            ScriptDefinitionsModificationTracker.getInstance(project).incModificationCount()
            KotlinScriptResolutionService.getInstance(project).process(ktFile)
@@ -2,65 +2,43 @@
 package org.jetbrains.kotlin.idea.core.script.k2.configurations

 import com.intellij.openapi.application.smartReadAction
+import com.intellij.openapi.components.Service
 import com.intellij.openapi.components.service
-import com.intellij.openapi.components.serviceAsync
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.roots.ProjectRootManager
 import com.intellij.openapi.vfs.VirtualFile
-import com.intellij.platform.backend.workspace.WorkspaceModel
 import com.intellij.platform.backend.workspace.toVirtualFileUrl
+import com.intellij.platform.backend.workspace.workspaceModel
 import com.intellij.platform.workspace.storage.EntitySource
 import com.intellij.platform.workspace.storage.MutableEntityStorage
+import com.intellij.platform.workspace.storage.toBuilder
+import com.intellij.platform.workspace.storage.url.VirtualFileUrl
+import com.intellij.platform.workspace.storage.url.VirtualFileUrlManager
 import kotlinx.coroutines.CoroutineScope
+import org.jetbrains.kotlin.idea.core.script.k2.asEntity
 import org.jetbrains.kotlin.idea.core.script.k2.modules.KotlinScriptEntity
 import org.jetbrains.kotlin.idea.core.script.k2.modules.KotlinScriptLibraryEntity
-import org.jetbrains.kotlin.idea.core.script.k2.modules.ScriptRefinedConfigurationResolver
-import org.jetbrains.kotlin.idea.core.script.k2.modules.ScriptWorkspaceModelManager
+import org.jetbrains.kotlin.idea.core.script.k2.modules.ScriptConfigurationProviderExtension
+import org.jetbrains.kotlin.idea.core.script.k2.modules.updateKotlinScriptEntities
 import org.jetbrains.kotlin.scripting.definitions.ScriptDefinition
-import org.jetbrains.kotlin.scripting.definitions.findScriptDefinition
 import org.jetbrains.kotlin.scripting.resolve.ScriptCompilationConfigurationResult
 import org.jetbrains.kotlin.scripting.resolve.ScriptReportSink
 import org.jetbrains.kotlin.scripting.resolve.VirtualFileScriptSource
 import org.jetbrains.kotlin.scripting.resolve.refineScriptCompilationConfiguration
 import java.io.File
-import java.util.concurrent.ConcurrentHashMap
 import kotlin.script.experimental.api.*
 import kotlin.script.experimental.api.ScriptDiagnostic.Companion.unspecifiedError
 import kotlin.script.experimental.jvm.jdkHome
 import kotlin.script.experimental.jvm.jvm

-open class DefaultScriptConfigurationHandler(
+@Service(Service.Level.PROJECT)
+class DefaultScriptConfigurationHandler(
    val project: Project, val coroutineScope: CoroutineScope
-) : ScriptWorkspaceModelManager, ScriptRefinedConfigurationResolver {
-    private val data = ConcurrentHashMap<VirtualFile, ScriptCompilationConfigurationResult>()
+) : ScriptConfigurationProviderExtension {

    override suspend fun create(virtualFile: VirtualFile, definition: ScriptDefinition): ScriptCompilationConfigurationResult {
-        val current = data[virtualFile]
-        if (current != null) return current
-
-        val configuration = resolveScriptConfiguration(virtualFile, definition)
-        data[virtualFile] = configuration
-
-        return configuration
-    }
-
-    override fun get(virtualFile: VirtualFile): ScriptCompilationConfigurationResult? = data[virtualFile]
-
-    override fun remove(virtualFile: VirtualFile) {
-        data.remove(virtualFile)
-    }
-
-    private suspend fun resolveScriptConfiguration(virtualFile: VirtualFile, definition: ScriptDefinition): ScriptCompilationConfigurationResult {
        val definitionJdk = definition.compilationConfiguration[ScriptCompilationConfiguration.jvm.jdkHome]
-        val configuration = if (definitionJdk != null) definition.compilationConfiguration
-        else {
-            val projectSdk = ProjectRootManager.getInstance(project).projectSdk?.homePath
-            definition.compilationConfiguration.with {
-                projectSdk?.let {
-                    jvm.jdkHome(File(it))
-                }
-            }
-        }
+        val configuration = getInitialConfiguration(definitionJdk, definition)

        val scriptSource = VirtualFileScriptSource(virtualFile)

@@ -79,53 +57,55 @@ open class DefaultScriptConfigurationHandler(
            }
        }

+        fun MutableEntityStorage.updatedStorage() {
+            val configuration = result.valueOrNull()?.configuration ?: return
+
+            val libraryIds = generateScriptLibraryEntities(configuration, definition, project)
+            libraryIds.filterNot {
+                this.contains(it)
+            }.forEach { (classes, sources) ->
+                this addEntity KotlinScriptLibraryEntity(classes, sources, DefaultScriptEntitySource)
+            }
+
+            this addEntity KotlinScriptEntity(
+                virtualFile.virtualFileUrl, libraryIds.toList(), DefaultScriptEntitySource
+            ) {
+                this.configuration = configuration.asEntity()
+                this.sdkId = configuration.sdkId
+            }
+        }
+
+        project.updateKotlinScriptEntities(DefaultScriptEntitySource) {
+            val builder = it.toSnapshot().toBuilder()
+            if (builder.getVirtualFileUrlIndex().findEntitiesByUrl(virtualFile.virtualFileUrl).none()) {
+                builder.updatedStorage()
+                it.applyChangesFrom(builder)
+            }
+        }
+
        project.service<ScriptReportSink>().attachReports(virtualFile, result.reports)

        return result
    }

-    override suspend fun updateWorkspaceModel(configurationPerFile: Map<VirtualFile, ScriptCompilationConfigurationResult>) {
-        val workspaceModel = project.serviceAsync<WorkspaceModel>()
-
-        workspaceModel.update("updating .kts modules") {
-            val entityStorage = getUpdatedStorage(configurationPerFile, workspaceModel)
-            it.applyChangesFrom(entityStorage)
+    private fun getInitialConfiguration(
+        definitionJdk: File?,
+        definition: ScriptDefinition
+    ): ScriptCompilationConfiguration = if (definitionJdk != null) definition.compilationConfiguration
+    else {
+        val projectSdk = ProjectRootManager.getInstance(project).projectSdk?.homePath
+        definition.compilationConfiguration.with {
+            projectSdk?.let {
+                jvm.jdkHome(File(it))
+            }
        }
    }

-    private fun getUpdatedStorage(
-        configurations: Map<VirtualFile, ScriptCompilationConfigurationResult>,
-        workspaceModel: WorkspaceModel,
-    ): MutableEntityStorage {
-        val storage = workspaceModel.currentSnapshot
-        val result = MutableEntityStorage.from(storage)
-        val index = storage.getVirtualFileUrlIndex()
+    private val urlManager: VirtualFileUrlManager
+        get() = project.workspaceModel.getVirtualFileUrlManager()

-        val fileUrlManager = workspaceModel.getVirtualFileUrlManager()
-
-        for ((scriptFile, configurationResult) in configurations) {
-            val scriptUrl = scriptFile.toVirtualFileUrl(fileUrlManager)
-            if (index.findEntitiesByUrl(scriptUrl).any()) continue
-
-            val definition = findScriptDefinition(project, VirtualFileScriptSource(scriptFile))
-            val configuration = configurationResult.valueOrNull()?.configuration ?: continue
-
-            val libraryIds = generateScriptLibraryEntities(configuration, definition, project)
-            libraryIds.filterNot {
-                result.contains(it)
-            }.forEach { (classes, sources) ->
-                result addEntity KotlinScriptLibraryEntity(classes, sources, DefaultScriptEntitySource)
-            }
-
-            result addEntity KotlinScriptEntity(
-                scriptUrl, libraryIds.toList(), DefaultScriptEntitySource
-            ) {
-                this.sdkId = configuration.sdkId
-            }
-        }
-
-        return result
-    }
+    private val VirtualFile.virtualFileUrl: VirtualFileUrl
+        get() = toVirtualFileUrl(urlManager)

    companion object {
        @JvmStatic
@@ -9,7 +9,6 @@ import com.intellij.openapi.components.service
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.roots.ProjectRootManager
 import com.intellij.openapi.vfs.VirtualFile
-import com.intellij.platform.backend.workspace.WorkspaceModel
 import com.intellij.platform.backend.workspace.toVirtualFileUrl
 import com.intellij.platform.backend.workspace.workspaceModel
 import com.intellij.platform.ide.progress.withBackgroundProgress
@@ -17,12 +16,16 @@ import com.intellij.platform.util.progress.ProgressReporter
 import com.intellij.platform.util.progress.reportProgressScope
 import com.intellij.platform.workspace.storage.EntitySource
 import com.intellij.platform.workspace.storage.MutableEntityStorage
+import com.intellij.platform.workspace.storage.toBuilder
+import com.intellij.platform.workspace.storage.url.VirtualFileUrl
+import com.intellij.platform.workspace.storage.url.VirtualFileUrlManager
 import kotlinx.coroutines.CoroutineScope
+import org.jetbrains.kotlin.idea.core.script.k2.asEntity
 import org.jetbrains.kotlin.idea.core.script.k2.configurations.DefaultScriptConfigurationHandler.DefaultScriptEntitySource
 import org.jetbrains.kotlin.idea.core.script.k2.modules.KotlinScriptEntity
 import org.jetbrains.kotlin.idea.core.script.k2.modules.KotlinScriptLibraryEntity
-import org.jetbrains.kotlin.idea.core.script.k2.modules.ScriptRefinedConfigurationResolver
-import org.jetbrains.kotlin.idea.core.script.k2.modules.ScriptWorkspaceModelManager
+import org.jetbrains.kotlin.idea.core.script.k2.modules.ScriptConfigurationProviderExtension
+import org.jetbrains.kotlin.idea.core.script.k2.modules.updateKotlinScriptEntities
 import org.jetbrains.kotlin.idea.core.script.shared.KotlinBaseScriptingBundle
 import org.jetbrains.kotlin.scripting.definitions.ScriptDefinition
 import org.jetbrains.kotlin.scripting.definitions.findScriptDefinition
@@ -31,16 +34,19 @@ import org.jetbrains.kotlin.scripting.resolve.ScriptCompilationConfigurationResu
 import org.jetbrains.kotlin.scripting.resolve.VirtualFileScriptSource
 import org.jetbrains.kotlin.scripting.resolve.refineScriptCompilationConfiguration
 import java.io.File
-import java.util.concurrent.ConcurrentHashMap
 import kotlin.script.experimental.api.valueOrNull
 import kotlin.script.experimental.api.with
 import kotlin.script.experimental.jvm.jdkHome
 import kotlin.script.experimental.jvm.jvm

@Service(Service.Level.PROJECT)
-class MainKtsScriptConfigurationProvider(val project: Project, val coroutineScope: CoroutineScope) : ScriptRefinedConfigurationResolver,
-                                                                                                     ScriptWorkspaceModelManager {
-    private val data = ConcurrentHashMap<VirtualFile, ScriptCompilationConfigurationResult>()
+class MainKtsConfigurationProvider(val project: Project, val coroutineScope: CoroutineScope) : ScriptConfigurationProviderExtension {
+    private val urlManager: VirtualFileUrlManager
+        get() = project.workspaceModel.getVirtualFileUrlManager()
+
+    private val VirtualFile.virtualFileUrl: VirtualFileUrl
+        get() = toVirtualFileUrl(urlManager)
+
    private val visitedScripts = TreeMultimap.create(COMPARATOR, COMPARATOR)
    private val visitedScriptsTraverser = Traverser.forTree<VirtualFile> { visitedScripts.get(it) }

@@ -49,11 +55,8 @@ class MainKtsScriptConfigurationProvider(val project: Project, val coroutineScop
    var reporter: ProgressReporter? = null
        private set

-    override fun get(virtualFile: VirtualFile): ScriptCompilationConfigurationResult? = data[virtualFile]
-
    override fun remove(virtualFile: VirtualFile) {
        visitedScripts.removeAll(virtualFile)
-        data.remove(virtualFile)
    }

    override suspend fun create(virtualFile: VirtualFile, definition: ScriptDefinition): ScriptCompilationConfigurationResult? {
@@ -64,22 +67,50 @@ class MainKtsScriptConfigurationProvider(val project: Project, val coroutineScop
            resolveDeeply(scriptsToResolve)
        }

-        data[virtualFile] = mainKtsConfiguration
+        fun MutableEntityStorage.updatedStorage() {
+            val configuration = mainKtsConfiguration.valueOrNull()?.configuration ?: return
+            val definition = findScriptDefinition(project, VirtualFileScriptSource(virtualFile))

-        return data[virtualFile]
+            val libraryIds = generateScriptLibraryEntities(configuration, definition, project)
+            libraryIds.filterNot {
+                this.contains(it)
+            }.forEach { (classes, sources) ->
+                this addEntity KotlinScriptLibraryEntity(classes, sources, DefaultScriptEntitySource)
+            }
+
+            this addEntity KotlinScriptEntity(
+                virtualFile.virtualFileUrl, libraryIds.toList(), MainKtsKotlinScriptEntitySource
+            ) {
+                this.configuration = configuration.asEntity()
+                this.sdkId = configuration.sdkId
+            }
+        }
+
+        project.updateKotlinScriptEntities(MainKtsKotlinScriptEntitySource) {
+            val builder = it.toSnapshot().toBuilder()
+            if (builder.getVirtualFileUrlIndex().findEntitiesByUrl(virtualFile.virtualFileUrl).none()) {
+                builder.updatedStorage()
+                it.applyChangesFrom(builder)
+            }
+        }
+
+        return mainKtsConfiguration
    }

-    suspend fun resolveDeeply(scripts: Collection<VirtualFile>) {
+    private suspend fun resolveDeeply(scripts: Collection<VirtualFile>) {
        for (script in scripts) {
            val scriptSource = VirtualFileScriptSource(script)
            val definition = findScriptDefinition(project, scriptSource)

-            val resolver = definition.getConfigurationResolver(project)
-            resolver.get(script) ?: resolver.create(script, definition)
+            val resolver = definition.getConfigurationProviderExtension(project)
+            resolver.get(project, script) ?: resolver.create(script, definition)
        }
    }

-    private suspend fun resolveMainKtsConfiguration(mainKts: VirtualFile, definition: ScriptDefinition): ScriptCompilationConfigurationResult {
+    private suspend fun resolveMainKtsConfiguration(
+        mainKts: VirtualFile,
+        definition: ScriptDefinition
+    ): ScriptCompilationConfigurationResult {
        val projectSdk = ProjectRootManager.getInstance(project).projectSdk?.homePath

        val configuration = definition.compilationConfiguration.with {
@@ -90,7 +121,7 @@ class MainKtsScriptConfigurationProvider(val project: Project, val coroutineScop

        val scriptSource = VirtualFileScriptSource(mainKts)

-        val result = withBackgroundProgress(
+        return withBackgroundProgress(
            project, title = KotlinBaseScriptingBundle.message("progress.title.dependency.resolution", mainKts.name)
        ) {
            reportProgressScope {
@@ -104,25 +135,6 @@ class MainKtsScriptConfigurationProvider(val project: Project, val coroutineScop
                }
            }
        }
-
-        return result
-    }
-
-    override suspend fun updateWorkspaceModel(configurationPerFile: Map<VirtualFile, ScriptCompilationConfigurationResult>) {
-        val mainKtsScript = configurationPerFile.entries.firstOrNull()?.key ?: return
-        for (it in getImportedScripts(mainKtsScript)) {
-            val scriptSource = VirtualFileScriptSource(it)
-            val definition = findScriptDefinition(project, scriptSource)
-            val configuration = definition.getConfigurationResolver(project).get(it) ?: continue
-            definition.getWorkspaceModelManager(project).updateWorkspaceModel(mapOf(it to configuration))
-        }
-
-        val workspaceModel = project.workspaceModel
-        workspaceModel.update("updating .main.kts modules") { targetStorage ->
-            val updatedStorage = getUpdatedStorage(project, configurationPerFile, workspaceModel)
-
-            targetStorage.applyChangesFrom(updatedStorage)
-        }
    }

    private val ScriptCompilationConfigurationResult.importedScripts: List<VirtualFile>
@@ -131,41 +143,11 @@ class MainKtsScriptConfigurationProvider(val project: Project, val coroutineScop
            return importedScripts.mapNotNull { (it as? VirtualFileScriptSource)?.virtualFile }.filterNot { it.isNonScript() }
        }

-    private fun getUpdatedStorage(
-        project: Project,
-        configurationsData: Map<VirtualFile, ScriptCompilationConfigurationResult>,
-        workspaceModel: WorkspaceModel,
-    ): MutableEntityStorage {
-        val virtualFileManager = workspaceModel.getVirtualFileUrlManager()
-        val storageToUpdate = MutableEntityStorage.from(workspaceModel.currentSnapshot)
-
-        for ((scriptFile, configurationWithSdk) in configurationsData) {
-            val configuration = configurationWithSdk.valueOrNull()?.configuration ?: continue
-            val definition = findScriptDefinition(project, VirtualFileScriptSource(scriptFile))
-            val virtualFileUrl = scriptFile.toVirtualFileUrl(virtualFileManager)
-
-            val libraryIds = generateScriptLibraryEntities(configuration, definition, project)
-            libraryIds.filterNot {
-                storageToUpdate.contains(it)
-            }.forEach { (classes, sources) ->
-                storageToUpdate addEntity KotlinScriptLibraryEntity(classes, sources, DefaultScriptEntitySource)
-            }
-
-            storageToUpdate addEntity KotlinScriptEntity(
-                virtualFileUrl, libraryIds.toList(), MainKtsKotlinScriptEntitySource
-            ) {
-                this.sdkId = configuration.sdkId
-            }
-        }
-
-        return storageToUpdate
-    }
-
    companion object {
        private val COMPARATOR = Comparator<VirtualFile> { left, right -> left.path.compareTo(right.path) }

        @JvmStatic
-        fun getInstance(project: Project): MainKtsScriptConfigurationProvider = project.service()
+        fun getInstance(project: Project): MainKtsConfigurationProvider = project.service()
    }

    object MainKtsKotlinScriptEntitySource : EntitySource
@@ -20,17 +20,13 @@ import com.intellij.workspaceModel.ide.impl.legacyBridge.sdk.customName
 import kotlinx.coroutines.CoroutineScope
 import org.jetbrains.kotlin.idea.core.script.k2.modules.KotlinScriptEntity
 import org.jetbrains.kotlin.idea.core.script.k2.modules.KotlinScriptLibraryEntity
-import org.jetbrains.kotlin.idea.core.script.k2.modules.ScriptRefinedConfigurationResolver
 import org.jetbrains.kotlin.idea.core.script.v1.ScriptDependenciesModificationTracker
 import org.jetbrains.kotlin.idea.core.script.v1.ScriptDependencyAware
 import org.jetbrains.kotlin.idea.core.script.v1.alwaysVirtualFile
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.scripting.definitions.ScriptConfigurationsProvider
-import org.jetbrains.kotlin.scripting.definitions.ScriptDefinition
 import org.jetbrains.kotlin.scripting.definitions.findScriptDefinition
 import org.jetbrains.kotlin.scripting.resolve.ScriptCompilationConfigurationResult
-import kotlin.script.experimental.api.ScriptCompilationConfiguration
-import kotlin.script.experimental.api.ide

 private class AllScriptsDependencies(
    val classes: Set<VirtualFile>,
@@ -116,12 +112,7 @@ class ScriptConfigurationsProviderImpl(project: Project, val coroutineScope: Cor

    override fun getScriptConfigurationResult(file: KtFile): ScriptCompilationConfigurationResult? {
        val definition = file.findScriptDefinition() ?: return null
-        return getConfigurationSupplier(definition).get(file.alwaysVirtualFile)
-    }
-
-    private fun getConfigurationSupplier(definition: ScriptDefinition): ScriptRefinedConfigurationResolver {
-        return definition.compilationConfiguration[ScriptCompilationConfiguration.ide.configurationResolverDelegate]?.invoke()
-            ?: DefaultScriptConfigurationHandler.getInstance(project)
+        return definition.getConfigurationProviderExtension(project).get(project, file.alwaysVirtualFile)
    }

    companion object {
@@ -6,8 +6,7 @@ import com.intellij.openapi.project.guessProjectDir
 import com.intellij.openapi.util.io.FileUtil
 import com.intellij.openapi.vfs.VfsUtilCore
 import com.intellij.openapi.vfs.VirtualFile
-import org.jetbrains.kotlin.idea.core.script.k2.modules.ScriptRefinedConfigurationResolver
-import org.jetbrains.kotlin.idea.core.script.k2.modules.ScriptWorkspaceModelManager
+import org.jetbrains.kotlin.idea.core.script.k2.modules.ScriptConfigurationProviderExtension
 import org.jetbrains.kotlin.scripting.definitions.ScriptDefinition
 import java.io.File
 import kotlin.script.experimental.api.IdeScriptCompilationConfigurationKeys
@@ -15,15 +14,10 @@ import kotlin.script.experimental.api.ScriptCompilationConfiguration
 import kotlin.script.experimental.api.ide
 import kotlin.script.experimental.util.PropertiesCollection

-val IdeScriptCompilationConfigurationKeys.configurationResolverDelegate: PropertiesCollection.Key<() -> ScriptRefinedConfigurationResolver> by PropertiesCollection.key()
-val IdeScriptCompilationConfigurationKeys.scriptWorkspaceModelManagerDelegate: PropertiesCollection.Key<() -> ScriptWorkspaceModelManager> by PropertiesCollection.key()
+val IdeScriptCompilationConfigurationKeys.configurationProviderExtension: PropertiesCollection.Key<() -> ScriptConfigurationProviderExtension> by PropertiesCollection.key()

-fun ScriptDefinition.getConfigurationResolver(project: Project): ScriptRefinedConfigurationResolver =
-    compilationConfiguration[ScriptCompilationConfiguration.ide.configurationResolverDelegate]?.invoke()
-        ?: DefaultScriptConfigurationHandler.getInstance(project)
-
-fun ScriptDefinition.getWorkspaceModelManager(project: Project): ScriptWorkspaceModelManager =
-    compilationConfiguration[ScriptCompilationConfiguration.ide.scriptWorkspaceModelManagerDelegate]?.invoke()
+fun ScriptDefinition.getConfigurationProviderExtension(project: Project): ScriptConfigurationProviderExtension =
+    compilationConfiguration[ScriptCompilationConfiguration.ide.configurationProviderExtension]?.invoke()
        ?: DefaultScriptConfigurationHandler.getInstance(project)

 fun VirtualFile.relativeLocation(project: Project): String {
@@ -3,9 +3,8 @@ package org.jetbrains.kotlin.idea.core.script.k2.definitions

 import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.idea.base.plugin.artifacts.KotlinArtifacts
-import org.jetbrains.kotlin.idea.core.script.k2.configurations.MainKtsScriptConfigurationProvider
-import org.jetbrains.kotlin.idea.core.script.k2.configurations.configurationResolverDelegate
-import org.jetbrains.kotlin.idea.core.script.k2.configurations.scriptWorkspaceModelManagerDelegate
+import org.jetbrains.kotlin.idea.core.script.k2.configurations.MainKtsConfigurationProvider
+import org.jetbrains.kotlin.idea.core.script.k2.configurations.configurationProviderExtension
 import org.jetbrains.kotlin.idea.core.script.v1.NewScriptFileInfo
 import org.jetbrains.kotlin.idea.core.script.v1.kotlinScriptTemplateInfo
 import org.jetbrains.kotlin.idea.core.script.v1.loggingReporter
@@ -27,11 +26,11 @@ class MainKtsScriptDefinitionSource(val project: Project) : ScriptDefinitionsSou
                ::loggingReporter
            ).definitions

-            val mainKtsScriptConfigurationProvider = MainKtsScriptConfigurationProvider.getInstance(project)
+            val mainKtsConfigurationProvider = MainKtsConfigurationProvider.getInstance(project)

            return discoveredDefinitions.map { definition ->
                val compilationConfiguration = definition.compilationConfiguration.withTransformedResolvers {
-                    ReportingExternalDependenciesResolver(it, mainKtsScriptConfigurationProvider)
+                    ReportingExternalDependenciesResolver(it, mainKtsConfigurationProvider)
                }.with {
                    ide.dependenciesSources(JvmDependency(KotlinArtifacts.kotlinStdlibSources))
                    ide {
@@ -40,11 +39,8 @@ class MainKtsScriptDefinitionSource(val project: Project) : ScriptDefinitionsSou
                            title = ".main.kts"
                            templateName = "Kotlin Script MainKts"
                        })
-                        configurationResolverDelegate {
-                            mainKtsScriptConfigurationProvider
-                        }
-                        scriptWorkspaceModelManagerDelegate {
-                            mainKtsScriptConfigurationProvider
+                        configurationProviderExtension {
+                            mainKtsConfigurationProvider
                        }
                    }
                }
@@ -1,8 +1,8 @@
 // Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package org.jetbrains.kotlin.idea.core.script.k2.definitions

+import org.jetbrains.kotlin.idea.core.script.k2.configurations.MainKtsConfigurationProvider
 import org.jetbrains.kotlin.idea.core.script.shared.KotlinBaseScriptingBundle
-import org.jetbrains.kotlin.idea.core.script.k2.configurations.MainKtsScriptConfigurationProvider
 import java.io.File
 import kotlin.script.experimental.api.ResultWithDiagnostics
 import kotlin.script.experimental.api.SourceCode
@@ -11,7 +11,7 @@ import kotlin.script.experimental.dependencies.RepositoryCoordinates

 internal class ReportingExternalDependenciesResolver(
    private val delegate: ExternalDependenciesResolver,
-    private val configurationProvider: MainKtsScriptConfigurationProvider,
+    private val configurationProvider: MainKtsConfigurationProvider,
 ) : ExternalDependenciesResolver {

    override fun acceptsArtifact(artifactCoordinates: String): Boolean =
@@ -8,15 +8,14 @@ import com.intellij.openapi.components.service
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.io.relativizeToClosestAncestor
 import com.intellij.openapi.util.io.toNioPathOrNull
+import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.platform.backend.observation.launchTracked
 import com.intellij.util.application
 import com.intellij.util.concurrency.ThreadingAssertions
 import kotlinx.coroutines.CoroutineScope
 import org.jetbrains.kotlin.analysis.api.platform.modification.publishGlobalModuleStateModificationEvent
 import org.jetbrains.kotlin.analysis.api.platform.modification.publishGlobalScriptModuleStateModificationEvent
-import org.jetbrains.kotlin.idea.core.script.k2.configurations.getConfigurationResolver
-import org.jetbrains.kotlin.idea.core.script.k2.configurations.getWorkspaceModelManager
-import org.jetbrains.kotlin.idea.core.script.k2.highlighting.KotlinScriptResolutionService.Companion.dropKotlinScriptCachesUnderWriteAction
+import org.jetbrains.kotlin.idea.core.script.k2.configurations.getConfigurationProviderExtension
 import org.jetbrains.kotlin.idea.core.script.shared.KotlinScriptProcessingFilter
 import org.jetbrains.kotlin.idea.core.script.v1.ScriptDependenciesModificationTracker
 import org.jetbrains.kotlin.idea.core.script.v1.alwaysVirtualFile
@@ -26,6 +25,7 @@ import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.scripting.definitions.ScriptDefinition
 import org.jetbrains.kotlin.scripting.definitions.findScriptDefinition
 import org.jetbrains.kotlin.scripting.resolve.KtFileScriptSource
+import org.jetbrains.kotlin.scripting.resolve.ScriptCompilationConfigurationResult


 /**
@@ -115,37 +115,24 @@ class KotlinScriptResolutionService(
     * - Tries to get an existing configuration for the file; if absent, creates it.
     * - Accumulates configurations and applies them to the workspace model in a single batch update.
     *
-     * After updating the model, Kotlin script-related caches are dropped via [dropKotlinScriptCachesUnderWriteAction].
-     *
     * @param definitionByFile mapping of script files to their resolved [ScriptDefinition].
     */
    private suspend fun prepareHighlighting(definitionByFile: Map<KtFile, ScriptDefinition>) {
-        val configurationsSupplier = definitionByFile.firstNotNullOf { it.value.getConfigurationResolver(project) }
-        val projectModelUpdater = definitionByFile.firstNotNullOf { it.value.getWorkspaceModelManager(project) }
-
-        val configurationPerVirtualFile = definitionByFile.entries.associate { (file, definition) ->
-            val virtualFile = file.virtualFile
-            val configuration = configurationsSupplier.get(virtualFile) ?: configurationsSupplier.create(virtualFile, definition) ?: return
-
-            virtualFile to configuration
-        }
+        val configurationProviderExtension = definitionByFile.firstNotNullOf { it.value.getConfigurationProviderExtension(project) }

        assert(!application.isWriteAccessAllowed)
-
-        dropKotlinScriptCachesUnderWriteAction(project)
-        projectModelUpdater.updateWorkspaceModel(configurationPerVirtualFile)
+        val configurationByFile = mutableMapOf<VirtualFile, ScriptCompilationConfigurationResult>()
+        for ((file, definition) in definitionByFile) {
+            val virtualFile = file.virtualFile
+            configurationByFile[virtualFile] = configurationProviderExtension.get(project, virtualFile)
+                ?: configurationProviderExtension.create(virtualFile, definition) ?: continue
+        }
    }

    companion object {
        @JvmStatic
        fun getInstance(project: Project): KotlinScriptResolutionService = project.service()

-        suspend fun dropKotlinScriptCachesUnderWriteAction(project: Project) {
-            edtWriteAction {
-                dropKotlinScriptCaches(project)
-            }
-        }
-
        fun dropKotlinScriptCaches(project: Project) {
            ThreadingAssertions.assertWriteAccess()

@@ -20,8 +20,7 @@ internal class ScriptTrafficLightRendererContributor : TrafficLightRendererContr
        return ScriptTrafficLightRenderer(ktFile.project, editor, ktFile)
    }

-    class ScriptTrafficLightRenderer(project: Project, editor: Editor, private val file: KtFile) :
-      TrafficLightRenderer(project, editor) {
+    class ScriptTrafficLightRenderer(project: Project, editor: Editor, private val file: KtFile) : TrafficLightRenderer(project, editor) {
        override fun getDaemonCodeAnalyzerStatus(severityRegistrar: SeverityRegistrar): DaemonCodeAnalyzerStatus {
            val status = super.getDaemonCodeAnalyzerStatus(severityRegistrar)

@@ -5,13 +5,13 @@ import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.platform.backend.workspace.toVirtualFileUrl
 import com.intellij.platform.backend.workspace.workspaceModel
-import org.jetbrains.kotlin.idea.core.script.k2.configurations.MainKtsScriptConfigurationProvider
+import org.jetbrains.kotlin.idea.core.script.k2.configurations.MainKtsConfigurationProvider

 class MainKtsScriptDependenciesProvider : K2IdeScriptAdditionalIdeaDependenciesProvider {
    override fun getRelatedModules(
        file: VirtualFile, project: Project
    ): List<VirtualFile> {
-        return MainKtsScriptConfigurationProvider.getInstance(project).getImportedScripts(file)
+        return MainKtsConfigurationProvider.getInstance(project).getImportedScripts(file)
    }

    override fun getRelatedLibraries(
@@ -1,12 +1,39 @@
 // Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package org.jetbrains.kotlin.idea.core.script.k2.modules

+import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.platform.backend.workspace.toVirtualFileUrl
+import com.intellij.platform.backend.workspace.workspaceModel
+import com.intellij.platform.workspace.storage.EntitySource
+import com.intellij.platform.workspace.storage.MutableEntityStorage
+import org.jetbrains.kotlin.idea.core.script.k2.highlighting.KotlinScriptResolutionService.Companion.dropKotlinScriptCaches
+import org.jetbrains.kotlin.idea.core.script.k2.toConfigurationResult
 import org.jetbrains.kotlin.scripting.definitions.ScriptDefinition
 import org.jetbrains.kotlin.scripting.resolve.ScriptCompilationConfigurationResult

-interface ScriptRefinedConfigurationResolver {
-    suspend fun create(virtualFile: VirtualFile, definition: ScriptDefinition): ScriptCompilationConfigurationResult?
-    fun get(virtualFile: VirtualFile): ScriptCompilationConfigurationResult?
+interface ScriptConfigurationProviderExtension {
+    fun get(project: Project, virtualFile: VirtualFile): ScriptCompilationConfigurationResult? =
+        getScriptConfigurationFromWorkspaceModel(project, virtualFile)
+
+    suspend fun create(virtualFile: VirtualFile, definition: ScriptDefinition): ScriptCompilationConfigurationResult? = null
+
    fun remove(virtualFile: VirtualFile): Unit = Unit
+
+    companion object {
+        fun getScriptConfigurationFromWorkspaceModel(project: Project, virtualFile: VirtualFile): ScriptCompilationConfigurationResult? {
+            val virtualFileUrl = virtualFile.toVirtualFileUrl(project.workspaceModel.getVirtualFileUrlManager())
+            val entity = project.workspaceModel.currentSnapshot.getVirtualFileUrlIndex().findEntitiesByUrl(virtualFileUrl)
+                .singleOrNull { it is KotlinScriptEntity } as? KotlinScriptEntity
+
+            return entity?.toConfigurationResult()
+        }
+    }
+}
+
+suspend fun Project.updateKotlinScriptEntities(entitySource: EntitySource, updater: (MutableEntityStorage) -> Unit) {
+    workspaceModel.update("updating kotlin script entities [$entitySource]") {
+        updater(it)
+        dropKotlinScriptCaches(this)
+    }
 }
@@ -15,7 +15,7 @@ import kotlin.script.experimental.api.ScriptDiagnostic.Severity
 import kotlin.script.experimental.api.SourceCode

 interface ScriptWorkspaceModelManager {
-    suspend fun updateWorkspaceModel(configurationPerFile: Map<VirtualFile, ScriptCompilationConfigurationResult>)
+    suspend fun updateWorkspaceModel(configurationPerFile: Map<VirtualFile, ScriptCompilationConfigurationResult>) {}
 }

@Service(Service.Level.PROJECT)
@@ -1,19 +1,56 @@
 // Copyright 2000-2025 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
-package org.jetbrains.kotlin.gradle.scripting.k2.workspaceModel
+package org.jetbrains.kotlin.idea.core.script.k2

+import com.intellij.platform.backend.workspace.virtualFile
+import org.jetbrains.kotlin.idea.core.script.k2.modules.KotlinScriptEntity
 import org.jetbrains.kotlin.idea.core.script.k2.modules.ScriptCompilationConfigurationEntity
 import org.jetbrains.kotlin.idea.core.script.k2.modules.ScriptEvaluationConfigurationEntity
 import org.jetbrains.kotlin.idea.core.script.k2.modules.ScriptingHostConfigurationEntity
+import org.jetbrains.kotlin.scripting.resolve.ScriptCompilationConfigurationResult
+import org.jetbrains.kotlin.scripting.resolve.ScriptCompilationConfigurationWrapper
+import org.jetbrains.kotlin.scripting.resolve.VirtualFileScriptSource
 import java.io.ByteArrayInputStream
 import java.io.ByteArrayOutputStream
 import java.io.ObjectInputStream
 import java.io.ObjectOutputStream
+import kotlin.script.experimental.api.ResultWithDiagnostics
 import kotlin.script.experimental.api.ScriptCompilationConfiguration
 import kotlin.script.experimental.api.ScriptEvaluationConfiguration
 import kotlin.script.experimental.host.ScriptingHostConfiguration
 import kotlin.script.experimental.util.PropertiesCollection
 import kotlin.script.experimental.util.PropertiesCollection.Key

+fun KotlinScriptEntity.toConfigurationResult(): ScriptCompilationConfigurationResult? {
+    val virtualFile = virtualFileUrl.virtualFile ?: return null
+
+    val result = if (configuration == null) {
+        ResultWithDiagnostics.Failure(listOf())
+    } else {
+        ResultWithDiagnostics.Success<ScriptCompilationConfigurationWrapper>(
+            ScriptCompilationConfigurationWrapper.FromCompilationConfiguration(
+                VirtualFileScriptSource(virtualFile), configuration?.deserialize()
+            ), reports.map { report -> report.toScriptDiagnostic() })
+    }
+
+    return result
+}
+
+fun ScriptCompilationConfiguration.asEntity(): ScriptCompilationConfigurationEntity = ScriptCompilationConfigurationEntity(this.serialize())
+
+fun ScriptCompilationConfigurationEntity.deserialize(): ScriptCompilationConfiguration? {
+    val params = ByteArrayInputStream(data).use { bis ->
+        ObjectInputStream(bis).use { ois ->
+            ois.readObject() as? LinkedHashMap<Key<Any>, Any?>
+        }
+    } ?: return null
+
+    return ScriptCompilationConfiguration {
+        params.forEach { (key, any) ->
+            key.putIfNotNull(any)
+        }
+    }
+}
+
 private fun PropertiesCollection.serialize(): ByteArray = ByteArrayOutputStream().use { bos ->
    ObjectOutputStream(bos).use { oos ->
        oos.writeObject(notTransientData)
@@ -38,23 +75,6 @@ fun ScriptEvaluationConfigurationEntity.deserialize(): ScriptEvaluationConfigura
    }
 }

-fun ScriptCompilationConfiguration.asEntity(): ScriptCompilationConfigurationEntity = ScriptCompilationConfigurationEntity(this.serialize())
-
-fun ScriptCompilationConfigurationEntity.deserialize(): ScriptCompilationConfiguration? {
-    val params = ByteArrayInputStream(data).use { bis ->
-        ObjectInputStream(bis).use { ois ->
-            ois.readObject() as? LinkedHashMap<Key<Any>, Any?>
-        }
-    } ?: return null
-
-    return ScriptCompilationConfiguration {
-        params.forEach { (key, any) ->
-            key.putIfNotNull(any)
-        }
-    }
-}
-
-
 fun ScriptingHostConfiguration.asEntity() = ScriptingHostConfigurationEntity(this.serialize())

 fun ScriptingHostConfigurationEntity.deserialize(): ScriptingHostConfiguration? {
@@ -11,9 +11,7 @@ import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.platform.backend.observation.launchTracked
 import com.intellij.platform.backend.workspace.WorkspaceModelChangeListener
 import com.intellij.platform.backend.workspace.toVirtualFileUrl
-import com.intellij.platform.backend.workspace.virtualFile
 import com.intellij.platform.backend.workspace.workspaceModel
-import com.intellij.platform.workspace.storage.EntityChange
 import com.intellij.platform.workspace.storage.MutableEntityStorage
 import com.intellij.platform.workspace.storage.VersionedStorageChange
 import com.intellij.platform.workspace.storage.toBuilder
@@ -25,41 +23,32 @@ import org.jetbrains.kotlin.gradle.scripting.k2.importing.GradleScriptData
 import org.jetbrains.kotlin.gradle.scripting.k2.importing.GradleScriptModel
 import org.jetbrains.kotlin.gradle.scripting.k2.workspaceModel.GradleScriptDefinitionEntity
 import org.jetbrains.kotlin.gradle.scripting.k2.workspaceModel.KotlinGradleScriptEntitySource
-import org.jetbrains.kotlin.gradle.scripting.k2.workspaceModel.asEntity
-import org.jetbrains.kotlin.gradle.scripting.k2.workspaceModel.deserialize
 import org.jetbrains.kotlin.gradle.scripting.shared.definition.GradleScriptDefinition
 import org.jetbrains.kotlin.gradle.scripting.shared.definition.loadGradleDefinitions
+import org.jetbrains.kotlin.idea.core.script.k2.asEntity
 import org.jetbrains.kotlin.idea.core.script.k2.configurations.sdkId
 import org.jetbrains.kotlin.idea.core.script.k2.configurations.toVirtualFileUrl
 import org.jetbrains.kotlin.idea.core.script.k2.definitions.ScriptDefinitionsModificationTracker
-import org.jetbrains.kotlin.idea.core.script.k2.highlighting.KotlinScriptResolutionService
+import org.jetbrains.kotlin.idea.core.script.k2.highlighting.KotlinScriptResolutionService.Companion.dropKotlinScriptCaches
 import org.jetbrains.kotlin.idea.core.script.k2.modules.*
 import org.jetbrains.kotlin.idea.core.script.v1.indexSourceRootsEagerly
 import org.jetbrains.kotlin.idea.core.script.v1.scriptingDebugLog
 import org.jetbrains.kotlin.idea.core.script.v1.scriptingWarnLog
 import org.jetbrains.kotlin.scripting.definitions.ScriptDefinition
-import org.jetbrains.kotlin.scripting.resolve.*
+import org.jetbrains.kotlin.scripting.resolve.ScriptCompilationConfigurationResult
+import org.jetbrains.kotlin.scripting.resolve.VirtualFileScriptSource
+import org.jetbrains.kotlin.scripting.resolve.adjustByDefinition
+import org.jetbrains.kotlin.scripting.resolve.refineScriptCompilationConfiguration
 import org.jetbrains.kotlin.utils.addIfNotNull
 import java.io.File
-import java.util.concurrent.ConcurrentHashMap
 import java.util.function.Predicate
 import kotlin.script.experimental.api.*
 import kotlin.script.experimental.jvm.JvmDependency
 import kotlin.script.experimental.jvm.jdkHome
 import kotlin.script.experimental.jvm.jvm

-private sealed interface CachedCompilationConfigurationResult {
-    data class Some(val value: ScriptCompilationConfigurationResult) : CachedCompilationConfigurationResult
-    data object None : CachedCompilationConfigurationResult
-}
-
-private fun ScriptCompilationConfigurationResult?.cached(): CachedCompilationConfigurationResult =
-    if (this == null) CachedCompilationConfigurationResult.None else CachedCompilationConfigurationResult.Some(this)
-
@Service(Service.Level.PROJECT)
-class GradleKotlinScriptService(val project: Project) : ScriptRefinedConfigurationResolver, ScriptWorkspaceModelManager {
-    private val data = ConcurrentHashMap<VirtualFileUrl, CachedCompilationConfigurationResult>()
-
+class GradleKotlinScriptService(val project: Project) : ScriptConfigurationProviderExtension {
    private val urlManager: VirtualFileUrlManager
        get() = project.workspaceModel.getVirtualFileUrlManager()

@@ -67,21 +56,15 @@ class GradleKotlinScriptService(val project: Project) : ScriptRefinedConfigurati
        get() = toVirtualFileUrl(urlManager)

    override suspend fun create(virtualFile: VirtualFile, definition: ScriptDefinition): ScriptCompilationConfigurationResult {
-        val updatedConfiguration = refineScriptCompilationConfiguration(VirtualFileScriptSource(virtualFile), definition, project)
-        data[virtualFile.virtualFileUrl] = updatedConfiguration.cached()
+        val configuration = refineScriptCompilationConfiguration(VirtualFileScriptSource(virtualFile), definition, project)

-        return updatedConfiguration
-    }
-
-    override fun get(virtualFile: VirtualFile): ScriptCompilationConfigurationResult? {
-        val wrapped = data.computeIfAbsent(virtualFile.virtualFileUrl) { key ->
-            val entity = project.workspaceModel.currentSnapshot.getVirtualFileUrlIndex().findEntitiesByUrl(key)
-                .singleOrNull { it.entitySource is KotlinGradleScriptEntitySource && it is KotlinScriptEntity } as? KotlinScriptEntity
-
-            entity?.toConfigurationResult().cached()
+        val currentStorage = project.workspaceModel.currentSnapshot.toBuilder()
+        project.updateKotlinScriptEntities(KotlinGradleScriptEntitySource) { storage ->
+            currentStorage.updateStorage(virtualFile, configuration)
+            storage.applyChangesFrom(currentStorage)
        }

-        return (wrapped as? CachedCompilationConfigurationResult.Some)?.value
+        return configuration
    }

    fun updateStorage(
@@ -132,7 +115,7 @@ class GradleKotlinScriptService(val project: Project) : ScriptRefinedConfigurati
            configurations[model.virtualFile] = refineScriptCompilationConfiguration(sourceCode, definition, project, configuration)
        }

-        storage.enrichStorage(configurations)
+        configurations.forEach { (virtualFile, configurationResult) -> storage.updateStorage(virtualFile, configurationResult) }
    }

    private fun String?.resolveSdk(): Sdk? {
@@ -145,65 +128,53 @@ class GradleKotlinScriptService(val project: Project) : ScriptRefinedConfigurati
        }
    }

-    override suspend fun updateWorkspaceModel(configurationPerFile: Map<VirtualFile, ScriptCompilationConfigurationResult>) {
-        val currentStorage = project.workspaceModel.currentSnapshot.toBuilder()
-        project.workspaceModel.update("updating .gradle.kts scripts") { storage ->
-            currentStorage.enrichStorage(configurationPerFile) // under writeAction from workspaceModel.update
-            storage.applyChangesFrom(currentStorage)
+    private fun MutableEntityStorage.updateStorage(virtualFile: VirtualFile, configurationResult: ScriptCompilationConfigurationResult) {
+        val configurationWrapper = configurationResult.valueOrNull() ?: return
+        if (getVirtualFileUrlIndex().findEntitiesByUrl(virtualFile.virtualFileUrl).any()) return
+
+        val classes =
+            configurationWrapper.dependenciesClassPath.sorted().map { it.path.toVirtualFileUrl(urlManager) }.toMutableSet()
+        val sources = configurationWrapper.dependenciesSources.sorted().map { it.path.toVirtualFileUrl(urlManager) }.toMutableSet()
+
+        val dependencies = buildList {
+            addIfNotNull(
+                extractRootsByPredicate(classes, sources) {
+                    it.url.contains("kotlin-stdlib")
+                })
+
+            addIfNotNull(
+                extractRootsByPredicate(classes, sources) {
+                    it.url.contains("accessors")
+                })
+
+            addIfNotNull(
+                extractRootsByPredicate(classes, sources) {
+                    it.url.contains("kotlin-gradle-plugin")
+                })
+
+            if (indexSourceRootsEagerly() || AdvancedSettings.getBoolean("gradle.attach.scripts.dependencies.sources")) {
+                addAll(extractDependenciesWithSources(classes, sources))
+
+                groupSourcesByParent(sources)
+
+                addAll(
+                    classes.map {
+                        getOrCreateScriptLibrary(it, sources)
+                    })
+            } else {
+                addAll(
+                    classes.map {
+                        getOrCreateScriptLibrary(it)
+                    })
+            }
        }
-    }

-    fun MutableEntityStorage.enrichStorage(
-        configurations: Map<VirtualFile, ScriptCompilationConfigurationResult>,
-    ) {
-        for ((scriptFile, configurationResult) in configurations) {
-            val configurationWrapper = configurationResult.valueOrNull() ?: continue
-            if (getVirtualFileUrlIndex().findEntitiesByUrl(scriptFile.virtualFileUrl).any()) continue
-
-            val classes =
-                configurationWrapper.dependenciesClassPath.sorted().map { it.path.toVirtualFileUrl(urlManager) }.toMutableSet()
-            val sources = configurationWrapper.dependenciesSources.sorted().map { it.path.toVirtualFileUrl(urlManager) }.toMutableSet()
-
-            val dependencies = buildList {
-                addIfNotNull(
-                    extractRootsByPredicate(classes, sources) {
-                        it.url.contains("kotlin-stdlib")
-                    })
-
-                addIfNotNull(
-                    extractRootsByPredicate(classes, sources) {
-                        it.url.contains("accessors")
-                    })
-
-                addIfNotNull(
-                    extractRootsByPredicate(classes, sources) {
-                        it.url.contains("kotlin-gradle-plugin")
-                    })
-
-                if (indexSourceRootsEagerly() || AdvancedSettings.getBoolean("gradle.attach.scripts.dependencies.sources")) {
-                    addAll(extractDependenciesWithSources(classes, sources))
-
-                    groupSourcesByParent(sources)
-
-                    addAll(
-                        classes.map {
-                            getOrCreateScriptLibrary(it, sources)
-                        })
-                } else {
-                    addAll(
-                        classes.map {
-                            getOrCreateScriptLibrary(it)
-                        })
-                }
-            }
-
-            this addEntity KotlinScriptEntity(
-                scriptFile.virtualFileUrl, dependencies, KotlinGradleScriptEntitySource
-            ) {
-                this.configuration = configurationWrapper.configuration?.asEntity()
-                this.reports = configurationResult.reports.map { it.toData() }.toMutableList()
-                this.sdkId = configurationWrapper.configuration?.sdkId
-            }
+        this addEntity KotlinScriptEntity(
+            virtualFile.virtualFileUrl, dependencies, KotlinGradleScriptEntitySource
+        ) {
+            this.configuration = configurationWrapper.configuration?.asEntity()
+            this.reports = configurationResult.reports.map { it.toData() }.toMutableList()
+            this.sdkId = configurationWrapper.configuration?.sdkId
        }
    }

@@ -320,7 +291,7 @@ class GradleKotlinScriptService(val project: Project) : ScriptRefinedConfigurati
            val configurationChanges = event.getChanges(KotlinScriptEntity::class.java)

            if (definitionChanges.any() || configurationChanges.any()) {
-                KotlinScriptResolutionService.dropKotlinScriptCaches(project)
+                dropKotlinScriptCaches(project)
            }
        }

@@ -332,22 +303,6 @@ class GradleKotlinScriptService(val project: Project) : ScriptRefinedConfigurati
            if (definitionChanges.any()) {
                ScriptDefinitionsModificationTracker.getInstance(project).incModificationCount()
            }
-
-            val updatedData = mutableMapOf<VirtualFileUrl, ScriptCompilationConfigurationResult>()
-            for (entityChange in configurationChanges) {
-                if (entityChange is EntityChange.Removed) continue
-                val entity = entityChange.newEntity ?: continue
-                if (entity.entitySource !is KotlinGradleScriptEntitySource) continue
-
-                val result = entity.toConfigurationResult()
-                if (result != null) {
-                    updatedData[entity.virtualFileUrl] = result
-                }
-            }
-
-            if (updatedData.any()) {
-                getInstance(project).data.clear()
-            }
        }
    }

@@ -356,19 +311,3 @@ class GradleKotlinScriptService(val project: Project) : ScriptRefinedConfigurati
        fun getInstance(project: Project): GradleKotlinScriptService = project.service()
    }
 }
-
-
-fun KotlinScriptEntity.toConfigurationResult(): ScriptCompilationConfigurationResult? {
-    val virtualFile = virtualFileUrl.virtualFile ?: return null
-
-    val result = if (configuration == null) {
-        ResultWithDiagnostics.Failure(listOf())
-    } else {
-        ResultWithDiagnostics.Success<ScriptCompilationConfigurationWrapper>(
-            ScriptCompilationConfigurationWrapper.FromCompilationConfiguration(
-                VirtualFileScriptSource(virtualFile), configuration?.deserialize()
-            ), reports.map { report -> report.toScriptDiagnostic() })
-    }
-
-    return result
-}
@@ -158,7 +158,7 @@ internal class GradleScriptNotificationProvider : EditorNotificationProvider {
    }

    private fun isImported(virtualFile: VirtualFile, project: Project): Boolean =
-        GradleKotlinScriptService.getInstance(project).get(virtualFile) != null
+        GradleKotlinScriptService.getInstance(project).get(project, virtualFile) != null

    private fun linkProject(
        project: Project,
@@ -4,9 +4,9 @@ package org.jetbrains.kotlin.gradle.scripting.k2.definition
 import com.intellij.openapi.project.Project
 import com.intellij.platform.backend.workspace.workspaceModel
 import org.jetbrains.kotlin.gradle.scripting.k2.workspaceModel.GradleScriptDefinitionEntity
-import org.jetbrains.kotlin.gradle.scripting.k2.workspaceModel.deserialize
 import org.jetbrains.kotlin.gradle.scripting.shared.definition.ErrorGradleScriptDefinition
 import org.jetbrains.kotlin.gradle.scripting.shared.definition.GradleScriptDefinition
+import org.jetbrains.kotlin.idea.core.script.k2.deserialize
 import org.jetbrains.kotlin.scripting.definitions.ScriptDefinition
 import org.jetbrains.kotlin.scripting.definitions.ScriptDefinitionsSource

@@ -4,8 +4,7 @@ package org.jetbrains.kotlin.gradle.scripting.k2.definition
 import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.gradle.scripting.k2.GradleKotlinScriptService
 import org.jetbrains.kotlin.gradle.scripting.shared.definition.GradleScriptDefinition
-import org.jetbrains.kotlin.idea.core.script.k2.configurations.configurationResolverDelegate
-import org.jetbrains.kotlin.idea.core.script.k2.configurations.scriptWorkspaceModelManagerDelegate
+import org.jetbrains.kotlin.idea.core.script.k2.configurations.configurationProviderExtension
 import org.jetbrains.kotlin.idea.core.script.v1.NewScriptFileInfo
 import org.jetbrains.kotlin.idea.core.script.v1.kotlinScriptTemplateInfo
 import kotlin.script.experimental.api.ide
@@ -18,10 +17,7 @@ fun GradleScriptDefinition.withIdeKeys(project: Project): GradleScriptDefinition
            title = ".gradle.kts"
            templateName = "Kotlin Script Gradle"
        })
-        configurationResolverDelegate {
-            GradleKotlinScriptService.getInstance(project)
-        }
-        scriptWorkspaceModelManagerDelegate {
+        configurationProviderExtension {
            GradleKotlinScriptService.getInstance(project)
        }
    }
@@ -7,7 +7,7 @@ import com.intellij.openapi.project.Project
 import com.intellij.openapi.vfs.VirtualFile
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.launch
-import org.jetbrains.kotlin.idea.core.script.k2.configurations.getConfigurationResolver
+import org.jetbrains.kotlin.idea.core.script.k2.configurations.getConfigurationProviderExtension
 import org.jetbrains.kotlin.idea.core.script.k2.highlighting.KotlinScriptResolutionService
 import org.jetbrains.kotlin.idea.core.script.k2.modules.KotlinScriptModuleManager.Companion.removeScriptModules
 import org.jetbrains.kotlin.idea.core.script.v1.ScriptRelatedModuleNameFile
@@ -23,7 +23,7 @@ class K2KotlinScratchFile(project: Project, virtualFile: VirtualFile, val corout
        val psiFile = ktFile ?: return

        coroutineScope.launch {
-            psiFile.findScriptDefinition()?.getConfigurationResolver(project)?.remove(virtualFile)
+            psiFile.findScriptDefinition()?.getConfigurationProviderExtension(project)?.remove(virtualFile)
            project.removeScriptModules(listOf(virtualFile))
            KotlinScriptResolutionService.getInstance(project).process(psiFile)
        }