feat(toolwindow): save state of the branches across files

src/main/kotlin/io/kotest/plugin/intellij/toolwindow/TestExplorerWindow.kt

@@ -9,6 +9,7 @@ import com.intellij.openapi.fileEditor.FileEditorManagerEvent
 import com.intellij.openapi.fileEditor.FileEditorManagerListener
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.ui.SimpleToolWindowPanel
+import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.openapi.vfs.VirtualFileManager
 import com.intellij.openapi.vfs.newvfs.BulkFileListener
 import com.intellij.openapi.vfs.newvfs.events.VFileEvent
@@ -104,6 +105,11 @@ class TestExplorerWindow(private val project: Project) : SimpleToolWindowPanel(t
       project.messageBus.connect().subscribe(
          FileEditorManagerListener.FILE_EDITOR_MANAGER,
          object : FileEditorManagerListener {
+            override fun fileClosed(source: FileEditorManager, file: VirtualFile) {
+               // when a file is closed, reset the one-time expanded state so reopening expands all again
+               tree.markFileClosed(file)
+            }
+
             override fun selectionChanged(event: FileEditorManagerEvent) {
                val file = fileEditorManager.selectedEditor?.file
                if (file != null) {

src/main/kotlin/io/kotest/plugin/intellij/toolwindow/TestFileTree.kt

@@ -3,11 +3,18 @@ package io.kotest.plugin.intellij.toolwindow
 import com.intellij.ide.util.treeView.NodeRenderer
 import com.intellij.ide.util.treeView.PresentableNodeDescriptor
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.ui.TreeUIHelper
 import javax.swing.tree.DefaultMutableTreeNode
 import javax.swing.tree.TreeModel
 import javax.swing.tree.TreeSelectionModel
 
+private data class FileTreeState(
+   val allKeys: Set<String>,
+   val expandedKeys: Set<String>,
+   var initiallyExpanded: Boolean,
+)
+
 class TestFileTree(
    project: Project,
 ) : com.intellij.ui.treeStructure.Tree(),
@@ -17,6 +24,8 @@ class TestFileTree(
    private val kotestTestExplorerService: KotestTestExplorerService =
       project.getService(KotestTestExplorerService::class.java)
    private var initialized = false
+   private var lastFileKey: String? = null
+   private val stateByFileKey = mutableMapOf<String, FileTreeState>()
 
    init {
       selectionModel.selectionMode = TreeSelectionModel.SINGLE_TREE_SELECTION
@@ -40,13 +49,64 @@ class TestFileTree(
          super.setModel(treeModel)
          return
       }
-      val expanded = isExpanded(0)
+      val newFileKey = currentFileKey()
+
+      // If switching away from a file, save its state first
+      if (lastFileKey != null && lastFileKey != newFileKey) {
+         val prevAll = collectAllPathKeys()
+         val prevExpanded = collectExpandedPathKeys()
+         val prevInit = stateByFileKey[lastFileKey!!]?.initiallyExpanded ?: false
+         stateByFileKey[lastFileKey!!] = FileTreeState(prevAll, prevExpanded, prevInit)
+      }
+
+      val sameFile = newFileKey == lastFileKey
+      val prevStateForNew = if (newFileKey != null) stateByFileKey[newFileKey] else null
+      val firstOpenForFile = newFileKey != null && prevStateForNew == null
+
+      // Baselines (use live tree for same file; fallback to stored state when switching)
+      val prevAllKeysForThisFile: Set<String> = when {
+         firstOpenForFile -> emptySet()
+         sameFile -> collectAllPathKeys()
+         newFileKey != null -> prevStateForNew?.allKeys ?: emptySet()
+         else -> emptySet()
+      }
+      val expandedKeysToRestore: Set<String> = when {
+         firstOpenForFile -> emptySet()
+         sameFile -> collectExpandedPathKeys()
+         newFileKey != null -> prevStateForNew?.expandedKeys ?: emptySet()
+         else -> emptySet()
+      }
+
       super.setModel(treeModel)
-      expandAllNodes()
-      setModuleGroupNodeExpandedState(expanded)
+
+      // Compute added nodes relative to the previous snapshot of this file (if any)
+      val newAllKeys = collectAllPathKeys()
+      if (!firstOpenForFile) {
+         val addedKeys = newAllKeys - prevAllKeysForThisFile
+         if (addedKeys.isNotEmpty()) expandAncestorPrefixesFor(addedKeys)
+      }
+
+      if (firstOpenForFile) {
+         // First time this file is shown in the tool window: expand everything
+         expandAllNodes()
+         stateByFileKey[newFileKey] = FileTreeState(newAllKeys, collectExpandedPathKeys(), initiallyExpanded = true)
+      } else {
+         // Restore previous expansion state for this file
+         if (expandedKeysToRestore.isNotEmpty()) expandPathsByKeys(expandedKeysToRestore)
+         if (newFileKey != null) {
+            val init = prevStateForNew?.initiallyExpanded ?: true
+            stateByFileKey[newFileKey] = FileTreeState(newAllKeys, collectExpandedPathKeys(), init)
+         }
+      }
+
+      lastFileKey = newFileKey
    }
 
-   private fun setModuleGroupNodeExpandedState(expanded: Boolean) {
-      if (expanded) expandRow(0) else collapseRow(0)
+   fun markFileClosed(file: VirtualFile) {
+      stateByFileKey.remove(file.path)
+      if (lastFileKey == file.path) lastFileKey = null
    }
+
+   private fun currentFileKey(): String? = kotestTestExplorerService.currentFile?.path
+
 }

src/main/kotlin/io/kotest/plugin/intellij/toolwindow/treeutils.kt

@@ -34,3 +34,104 @@ fun TreePath.nodeDescriptor(): PresentableNodeDescriptor<*>? {
       else -> null
    }
 }
+
+// --- Expansion state helpers ---
+
+/**
+ * Builds a logical key for a node to preserve expansion state across model rebuilds.
+ * Only nodes that can have children are keyed (root, modules, module, tags, file, spec, container test).
+ */
+private fun DefaultMutableTreeNode.expansionKeyOrNull(): String? {
+   return when (val descriptor = userObject) {
+      is KotestRootNodeDescriptor -> "root"
+      is ModulesNodeDescriptor -> "modules"
+      is ModuleNodeDescriptor -> "module:${descriptor.module.name}"
+      is TagsNodeDescriptor -> "tags"
+      is TestFileNodeDescriptor -> "file"
+      is SpecNodeDescriptor -> "spec:${descriptor.fqn.asString()}"
+      is TestNodeDescriptor -> {
+         // Only container tests can have children; still safe to key all tests
+         "test:${descriptor.test.test.descriptorPath()}"
+      }
+      else -> null
+   }
+}
+
+/**
+ * Returns a set of expansion path keys for the currently expanded nodes.
+ */
+fun JTree.collectExpandedPathKeys(): Set<String> {
+   val root = model.root as? DefaultMutableTreeNode ?: return emptySet()
+   val expanded = mutableSetOf<String>()
+   // Enumerate all expanded descendants starting from root
+   val enumeration = getExpandedDescendants(TreePath(root.path)) ?: return emptySet()
+   while (enumeration.hasMoreElements()) {
+      val path = enumeration.nextElement()
+      val key = pathToExpansionKey(path)
+      if (key != null) expanded.add(key)
+   }
+   return expanded
+}
+
+/**
+ * Expands nodes in the current model whose logical expansion keys appear in [keys].
+ */
+fun JTree.expandPathsByKeys(keys: Set<String>) {
+   val root = model.root as? DefaultMutableTreeNode ?: return
+   fun recurse(node: DefaultMutableTreeNode, prefix: String?) {
+      val key = node.expansionKeyOrNull()
+      val pathKey = if (key == null) prefix else listOfNotNull(prefix, key).joinToString("/")
+      if (pathKey != null && keys.contains(pathKey)) {
+         expandPath(TreePath(node.path))
+      }
+      val children = node.children()
+      while (children.hasMoreElements()) {
+         val child = children.nextElement() as DefaultMutableTreeNode
+         recurse(child, pathKey)
+      }
+   }
+   recurse(root, null)
+}
+
+private fun pathToExpansionKey(path: TreePath): String? {
+   val parts = path.path
+      .mapNotNull { it as? DefaultMutableTreeNode }
+      .mapNotNull { it.expansionKeyOrNull() }
+   return if (parts.isEmpty()) null else parts.joinToString("/")
+}
+
+/**
+ * Returns a set of path keys for all nodes in the current model.
+ */
+fun JTree.collectAllPathKeys(): Set<String> {
+   val root = model.root as? DefaultMutableTreeNode ?: return emptySet()
+   val keys = mutableSetOf<String>()
+   fun recurse(node: DefaultMutableTreeNode) {
+      val path = TreePath(node.path)
+      val key = pathToExpansionKey(path)
+      if (key != null) keys.add(key)
+      val children = node.children()
+      while (children.hasMoreElements()) {
+         recurse(children.nextElement() as DefaultMutableTreeNode)
+      }
+   }
+   recurse(root)
+   return keys
+}
+
+/**
+ * Expands all ancestor prefixes for the given full path keys.
+ */
+fun JTree.expandAncestorPrefixesFor(keys: Set<String>) {
+   if (keys.isEmpty()) return
+   val toExpand = mutableSetOf<String>()
+   keys.forEach { key ->
+      val parts = key.split('/')
+      val acc = mutableListOf<String>()
+      parts.forEach { part ->
+         acc.add(part)
+         toExpand.add(acc.joinToString("/"))
+      }
+   }
+   expandPathsByKeys(toExpand)
+}