KTIJ-35157 Scripts: importScripts doesn't import files - fix

plugins/kotlin/base/scripting/scripting.k2/resources/intellij.kotlin.base.scripting.xml

@@ -82,6 +82,9 @@
             order="last"
             implementation="org.jetbrains.kotlin.idea.core.script.k2.definitions.BundledScriptDefinitionSource"/>
 
+    <k2IdeScriptAdditionalIdeaDependenciesProvider
+            implementation="org.jetbrains.kotlin.idea.core.script.k2.modules.DefaultKotlinScriptDependenciesProvider"/>
+
     <k2IdeScriptAdditionalIdeaDependenciesProvider
             implementation="org.jetbrains.kotlin.idea.core.script.k2.modules.MainKtsScriptDependenciesProvider"/>
 

plugins/kotlin/base/scripting/scripting.k2/src/org/jetbrains/kotlin/idea/core/script/k2/configurations/DefaultKotlinScriptEntityProvider.kt

@@ -1,6 +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.configurations
 
+import com.google.common.collect.TreeMultimap
+import com.google.common.graph.Traverser
 import com.intellij.openapi.application.smartReadAction
 import com.intellij.openapi.components.Service
 import com.intellij.openapi.components.service
@@ -12,6 +14,7 @@ import com.intellij.platform.workspace.storage.EntitySource
 import com.intellij.platform.workspace.storage.MutableEntityStorage
 import kotlinx.coroutines.CoroutineScope
 import org.jetbrains.kotlin.idea.core.script.k2.getOrCreateScriptConfigurationId
+import org.jetbrains.kotlin.idea.core.script.k2.highlighting.KotlinScriptResolutionService
 import org.jetbrains.kotlin.idea.core.script.k2.modules.KotlinScriptEntity
 import org.jetbrains.kotlin.idea.core.script.k2.modules.KotlinScriptEntityProvider
 import org.jetbrains.kotlin.idea.core.script.k2.modules.KotlinScriptLibraryEntity
@@ -35,6 +38,17 @@ import kotlin.script.experimental.jvm.jvm
 class DefaultKotlinScriptEntityProvider(
     override val project: Project, val coroutineScope: CoroutineScope
 ) : KotlinScriptEntityProvider(project) {
+
+    private val visitedScripts = TreeMultimap.create(COMPARATOR, COMPARATOR)
+    private val visitedScriptsTraverser = Traverser.forGraph<VirtualFile> { visitedScripts.get(it) }
+
+    fun getImportedScripts(kts: VirtualFile): List<VirtualFile> = visitedScriptsTraverser.breadthFirst(kts) - kts
+
+    override suspend fun removeKotlinScriptEntity(virtualFile: VirtualFile) {
+        visitedScripts.removeAll(virtualFile)
+        super.removeKotlinScriptEntity(virtualFile)
+    }
+
     override suspend fun updateWorkspaceModel(
         virtualFile: VirtualFile, definition: ScriptDefinition
     ) {
@@ -55,6 +69,13 @@ class DefaultKotlinScriptEntityProvider(
             }
         }
 
+        result.importedScripts.let { scriptsToResolve ->
+            if (scriptsToResolve.isNotEmpty()) {
+                visitedScripts.putAll(virtualFile, scriptsToResolve)
+                KotlinScriptResolutionService.getInstance(project).process(scriptsToResolve - visitedScripts.keys())
+            }
+        }
+
         fun updateStorage(storage: MutableEntityStorage) {
             val configuration = result.valueOrNull()?.configuration ?: return
             val definition = findScriptDefinition(project, VirtualFileScriptSource(virtualFile))

plugins/kotlin/base/scripting/scripting.k2/src/org/jetbrains/kotlin/idea/core/script/k2/configurations/MainKtsEntityProvider.kt

@@ -42,7 +42,7 @@ class MainKtsEntityProvider(
     val coroutineScope: CoroutineScope
 ) : KotlinScriptEntityProvider(project) {
     private val visitedScripts = TreeMultimap.create(COMPARATOR, COMPARATOR)
-    private val visitedScriptsTraverser = Traverser.forTree<VirtualFile> { visitedScripts.get(it) }
+    private val visitedScriptsTraverser = Traverser.forGraph<VirtualFile> { visitedScripts.get(it) }
 
     fun getImportedScripts(mainKts: VirtualFile): List<VirtualFile> = visitedScriptsTraverser.breadthFirst(mainKts) - mainKts
 
@@ -62,10 +62,11 @@ class MainKtsEntityProvider(
         if (project.workspaceModel.currentSnapshot.containsScriptEntity(scriptUrl)) return
 
         val mainKtsConfiguration = resolveMainKtsConfiguration(virtualFile, definition)
-        val scriptsToResolve = mainKtsConfiguration.importedScripts - visitedScripts.keys()
-        if (scriptsToResolve.isNotEmpty()) {
-            visitedScripts.putAll(virtualFile, scriptsToResolve)
-            KotlinScriptResolutionService.getInstance(project).process(scriptsToResolve)
+        mainKtsConfiguration.importedScripts.let { scriptsToResolve ->
+            if (scriptsToResolve.isNotEmpty()) {
+                visitedScripts.putAll(virtualFile, scriptsToResolve)
+                KotlinScriptResolutionService.getInstance(project).process(scriptsToResolve - visitedScripts.keys())
+            }
         }
 
         fun updateStorage(storage: MutableEntityStorage) {
@@ -138,14 +139,7 @@ class MainKtsEntityProvider(
         }
     }
 
-    private val ScriptCompilationConfigurationResult.importedScripts: List<VirtualFile>
-        get() {
-            val importedScripts = this.valueOrNull()?.importedScripts ?: return emptyList()
-            return importedScripts.mapNotNull { (it as? VirtualFileScriptSource)?.virtualFile }.filterNot { it.isNonScript() }
-        }
-
     companion object {
-        private val COMPARATOR = Comparator<VirtualFile> { left, right -> left.path.compareTo(right.path) }
 
         @JvmStatic
         fun getInstance(project: Project): MainKtsEntityProvider = project.service()

plugins/kotlin/base/scripting/scripting.k2/src/org/jetbrains/kotlin/idea/core/script/k2/modules/DefaultKotlinScriptDependenciesProvider.kt

@@ -0,0 +1,28 @@
+// 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 org.jetbrains.kotlin.idea.core.script.k2.configurations.DefaultKotlinScriptEntityProvider
+
+class DefaultKotlinScriptDependenciesProvider : K2IdeScriptAdditionalIdeaDependenciesProvider {
+    override fun getRelatedModules(
+        file: VirtualFile, project: Project
+    ): List<VirtualFile> {
+        return DefaultKotlinScriptEntityProvider.getInstance(project).getImportedScripts(file)
+    }
+
+    override fun getRelatedLibraries(
+        file: VirtualFile, project: Project
+    ): List<KotlinScriptLibraryEntity> {
+        val currentSnapshot = project.workspaceModel.currentSnapshot
+        val index = currentSnapshot.getVirtualFileUrlIndex()
+        val manager = project.workspaceModel.getVirtualFileUrlManager()
+
+        return getRelatedModules(file, project).flatMap {
+            index.findEntitiesByUrl(it.toVirtualFileUrl(manager))
+        }.filterIsInstance<KotlinScriptEntity>().flatMap { it.dependencies.mapNotNull { currentSnapshot.resolve(it) } }
+    }
+}

plugins/kotlin/base/scripting/scripting.k2/src/org/jetbrains/kotlin/idea/core/script/k2/modules/KotlinScriptEntityProvider.kt

@@ -11,6 +11,10 @@ import com.intellij.platform.workspace.storage.ImmutableEntityStorage
 import com.intellij.platform.workspace.storage.MutableEntityStorage
 import com.intellij.platform.workspace.storage.url.VirtualFileUrl
 import org.jetbrains.kotlin.scripting.definitions.ScriptDefinition
+import org.jetbrains.kotlin.scripting.definitions.isNonScript
+import org.jetbrains.kotlin.scripting.resolve.ScriptCompilationConfigurationResult
+import org.jetbrains.kotlin.scripting.resolve.VirtualFileScriptSource
+import kotlin.script.experimental.api.valueOrNull
 
 /**
  * Base contract for providing and maintaining [KotlinScriptEntity] instances in the Workspace Model.
@@ -39,6 +43,12 @@ abstract class KotlinScriptEntityProvider(
     protected val VirtualFile.virtualFileUrl: VirtualFileUrl
         get() = toVirtualFileUrl(project.workspaceModel.getVirtualFileUrlManager())
 
+    protected val ScriptCompilationConfigurationResult.importedScripts: List<VirtualFile>
+        get() {
+            val importedScripts = this.valueOrNull()?.importedScripts ?: return emptyList()
+            return importedScripts.mapNotNull { (it as? VirtualFileScriptSource)?.virtualFile }.filterNot { it.isNonScript() }
+        }
+
     /**
      * Finds a [KotlinScriptEntity] associated with the given [virtualFile] in the current snapshot.
      *
@@ -81,4 +91,10 @@ abstract class KotlinScriptEntityProvider(
             updater(it)
         }
     }
+
+    protected companion object {
+
+        @JvmStatic
+        protected val COMPARATOR = Comparator<VirtualFile> { left, right -> left.path.compareTo(right.path) }
+    }
 }