fix(jetbrains): sync worktree selection with editor

.changeset/worktree-editor-selection.md

@@ -0,0 +1,5 @@
+---
+"@kilocode/kilo-jetbrains": patch
+---
+
+Keep the Agent Manager worktree selection in sync with the active editor tab.

packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/AgentManagerPanel.kt

@@ -5,6 +5,8 @@ import ai.kilocode.client.agentManager.worktree.ConfigureWorktreeDialog
 import ai.kilocode.client.agentManager.worktree.WorktreeController
 import ai.kilocode.client.agentManager.worktree.WorktreeIcons
 import ai.kilocode.client.agentManager.worktree.WorktreeNameCache
+import ai.kilocode.client.agentManager.worktree.WorktreeEditorMatchers
+import ai.kilocode.client.agentManager.worktree.WorktreeSessionEditorMatcher
 import ai.kilocode.client.agentManager.worktree.WorktreeSessionEditorKind
 import ai.kilocode.client.agentManager.worktree.ensureWorktreeSessionEditorKind
 import ai.kilocode.client.agentManager.worktree.worktreeActivityBadge
@@ -21,7 +23,6 @@ import ai.kilocode.client.ui.list.ActiveListSelection
 import ai.kilocode.client.ui.list.ActiveListSurface
 import ai.kilocode.client.ui.list.activeListToolWindowBackground
 import ai.kilocode.client.vfs.KiloVfsManager
-import ai.kilocode.client.vfs.KiloVirtualFile
 import ai.kilocode.rpc.dto.RemoveWorktreeResultDto
 import ai.kilocode.rpc.dto.WorktreeDto
 import com.intellij.icons.AllIcons
@@ -39,8 +40,12 @@ import com.intellij.openapi.actionSystem.UiDataProvider
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.components.service
 import com.intellij.openapi.fileEditor.FileEditorManager
+import com.intellij.openapi.fileEditor.FileEditorManagerEvent
+import com.intellij.openapi.fileEditor.FileEditorManagerListener
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Disposer
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.util.concurrency.annotations.RequiresEdt
 import com.intellij.util.ui.JBUI
 import com.intellij.util.ui.components.BorderLayoutPanel
 import java.awt.Color
@@ -94,6 +99,7 @@ class AgentManagerPanel(
         controller.onCreateFailure = { err -> notifyCreateFailed(err) }
         controller.onRemoveSuccess = { item -> close(item) }
         controller.onActivityChanged = { sync() }
+        bindEditorSelection()
         // Reflect names adopted or renamed in a worktree session editor tab in the list live.
         service<WorktreeNameCache>().addListener(this) { path, name -> controller.applyName(path, name) }
         ActionManager.getInstance().getAction("RenameElement")?.shortcutSet?.let { set ->
@@ -106,7 +112,7 @@ class AgentManagerPanel(
     override fun getBackground(): Color = activeListToolWindowBackground()
 
     fun refresh() {
-        selected = activeWorktreeKey() ?: selected
+        selected = currentEditorWorktree()
         controller.reload()
     }
 
@@ -226,6 +232,16 @@ class AgentManagerPanel(
         })
     }
 
+    private fun bindEditorSelection() {
+        val target = project ?: return
+        target.service<WorktreeEditorMatchers>().register(WorktreeSessionEditorMatcher)
+        val bus = target.messageBus.connect(this)
+        bus.subscribe(FileEditorManagerListener.FILE_EDITOR_MANAGER, object : FileEditorManagerListener {
+            override fun selectionChanged(event: FileEditorManagerEvent) = track(event.newFile)
+        })
+        track(FileEditorManager.getInstance(target).selectedFiles.firstOrNull())
+    }
+
     private fun bindModel() {
         val listener = object : ListDataListener {
             override fun intervalAdded(e: ListDataEvent) = sync()
@@ -247,18 +263,27 @@ class AgentManagerPanel(
             },
             ActiveListSelection.PreserveNoScroll,
         )
-        if (key != null) list.select(key, scroll = false)
+        if (key != null) {
+            if (!list.select(key, scroll = false)) list.clearSelection()
+            return
+        }
+        list.clearSelection()
+        selected = null
     }
 
-    private fun activeWorktreeKey(): String? {
+    @RequiresEdt
+    private fun track(file: VirtualFile?) {
+        val key = project?.service<WorktreeEditorMatchers>()?.match(file)
+        selected = key
+        if (key != null && list.select(key, scroll = false)) return
+        list.clearSelection()
+    }
+
+    @RequiresEdt
+    private fun currentEditorWorktree(): String? {
         val target = project ?: return null
-        return FileEditorManager.getInstance(target).selectedFiles
-            .filterIsInstance<KiloVirtualFile>()
-            .firstOrNull { it.path.kind == WorktreeSessionEditorKind.ID }
-            ?.path
-            ?.params
-            ?.get("path")
-            ?.takeIf { it.isNotBlank() }
+        val file = FileEditorManager.getInstance(target).selectedFiles.firstOrNull()
+        return target.service<WorktreeEditorMatchers>().match(file)
     }
 
     private fun item(key: String): WorktreeDto? {

packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/worktree/WorktreeEditorMatcher.kt

@@ -0,0 +1,35 @@
+package ai.kilocode.client.agentManager.worktree
+
+import ai.kilocode.client.vfs.KiloVirtualFile
+import com.intellij.openapi.components.Service
+import com.intellij.openapi.vfs.VirtualFile
+import com.intellij.util.concurrency.annotations.RequiresEdt
+import java.util.concurrent.CopyOnWriteArrayList
+
+fun interface WorktreeEditorMatcher {
+    @RequiresEdt
+    fun match(file: VirtualFile): String?
+}
+
+@Service(Service.Level.PROJECT)
+class WorktreeEditorMatchers {
+    private val matchers = CopyOnWriteArrayList<WorktreeEditorMatcher>()
+
+    fun register(matcher: WorktreeEditorMatcher) {
+        matchers.addIfAbsent(matcher)
+    }
+
+    @RequiresEdt
+    fun match(file: VirtualFile?): String? {
+        if (file == null) return null
+        return matchers.firstNotNullOfOrNull { it.match(file) }
+    }
+}
+
+object WorktreeSessionEditorMatcher : WorktreeEditorMatcher {
+    override fun match(file: VirtualFile): String? {
+        val kilo = file as? KiloVirtualFile ?: return null
+        if (kilo.path.kind != WorktreeSessionEditorKind.ID) return null
+        return kilo.path.params["path"]?.takeIf { it.isNotBlank() }
+    }
+}

packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/list/ActiveList.kt

@@ -94,6 +94,9 @@ internal class ActiveList(
     @RequiresEdt
     fun selected(): ActiveListItem? = view.selected()
 
+    @RequiresEdt
+    fun clearSelection() = view.clearSelection()
+
     @RequiresEdt
     fun selectedItems(): List<ActiveListItem> = view.selectedItems()
 

packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/list/ActiveListView.kt

@@ -169,6 +169,12 @@ internal class ActiveListView(
         return list.selectedValue
     }
 
+    @RequiresEdt
+    fun clearSelection() {
+        checkEdt()
+        list.clearSelection()
+    }
+
     @RequiresEdt
     fun selectedItems(): List<ActiveListItem> {
         checkEdt()

packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/agentManager/AgentManagerPanelTest.kt

@@ -2,6 +2,8 @@ package ai.kilocode.client.agentManager
 
 import ai.kilocode.client.agentManager.worktree.KiloWorktreeService
 import ai.kilocode.client.agentManager.worktree.WorktreeController
+import ai.kilocode.client.agentManager.worktree.WorktreeEditorMatcher
+import ai.kilocode.client.agentManager.worktree.WorktreeEditorMatchers
 import ai.kilocode.client.agentManager.worktree.WorktreeSessionEditorKind
 import ai.kilocode.client.agentManager.worktree.ensureWorktreeSessionEditorKind
 import ai.kilocode.client.agentManager.worktree.worktreeSessionParams
@@ -22,6 +24,7 @@ import ai.kilocode.rpc.dto.SessionActivityKindDto
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.components.service
 import com.intellij.openapi.fileEditor.FileEditorManager
+import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.ui.SearchTextField
 import com.intellij.testFramework.fixtures.BasePlatformTestCase
 import com.intellij.ui.components.JBList
@@ -188,6 +191,85 @@ class AgentManagerPanelTest : BasePlatformTestCase() {
         assertEquals(second.id, edt { (list.selectedValue as ActiveListItem).key })
     }
 
+    fun `test selecting worktree editor tab selects its worktree row`() {
+        val first = WorktreeDto("/repo/.kilo/worktrees/feature-x", "feature-x", "feature/x", "/repo/.kilo/worktrees/feature-x")
+        val second = WorktreeDto("/repo/.kilo/worktrees/feature-y", "feature-y", "feature/y", "/repo/.kilo/worktrees/feature-y")
+        rpc.listed += first
+        rpc.listed += second
+        val controller = WorktreeController(service, "/test", coroutines.scope)
+        val panel = edt { AgentManagerPanel(testRootDisposable, controller, project) }
+        edt { controller.reload() }
+        flush()
+
+        val list = edt { UIUtil.findComponentOfType(panel, JBList::class.java)!! }
+        edt {
+            ensureWorktreeSessionEditorKind()
+            project.service<KiloVfsManager>().open(WorktreeSessionEditorKind.ID, worktreeSessionParams(second), focus = true)
+        }
+        pump()
+
+        assertEquals(second.id, edt { (list.selectedValue as ActiveListItem).key })
+    }
+
+    fun `test selecting non worktree editor tab clears worktree row selection`() {
+        val item = WorktreeDto("/repo/.kilo/worktrees/feature-x", "feature-x", "feature/x", "/repo/.kilo/worktrees/feature-x")
+        rpc.listed += item
+        val controller = WorktreeController(service, "/test", coroutines.scope)
+        val panel = edt { AgentManagerPanel(testRootDisposable, controller, project) }
+        edt { controller.reload() }
+        flush()
+
+        val list = edt { UIUtil.findComponentOfType(panel, JBList::class.java)!! }
+        edt {
+            ensureWorktreeSessionEditorKind()
+            project.service<KiloVfsManager>().open(WorktreeSessionEditorKind.ID, worktreeSessionParams(item), focus = true)
+        }
+        pump()
+        assertEquals(item.id, edt { (list.selectedValue as ActiveListItem).key })
+
+        val file = myFixture.addFileToProject("src/Main.kt", "fun main() = Unit").virtualFile
+        edt { FileEditorManager.getInstance(project).openFile(file, true) }
+        pump()
+
+        assertEquals(-1, edt { list.selectedIndex })
+    }
+
+    fun `test panel starts with no selection when active editor tab is not worktree`() {
+        val item = WorktreeDto("/repo/.kilo/worktrees/feature-x", "feature-x", "feature/x", "/repo/.kilo/worktrees/feature-x")
+        rpc.listed += item
+        val file = myFixture.addFileToProject("src/Current.kt", "fun current() = Unit").virtualFile
+        edt { FileEditorManager.getInstance(project).openFile(file, true) }
+        pump()
+        val controller = WorktreeController(service, "/test", coroutines.scope)
+        val panel = edt { AgentManagerPanel(testRootDisposable, controller, project) }
+        edt { controller.reload() }
+        flush()
+
+        val list = edt { UIUtil.findComponentOfType(panel, JBList::class.java)!! }
+        assertEquals(-1, edt { list.selectedIndex })
+    }
+
+    fun `test custom worktree editor matcher can select a row for another editor kind`() {
+        val item = WorktreeDto("/repo/.kilo/worktrees/feature-x", "feature-x", "feature/x", "/repo/.kilo/worktrees/feature-x")
+        rpc.listed += item
+        val controller = WorktreeController(service, "/test", coroutines.scope)
+        val panel = edt { AgentManagerPanel(testRootDisposable, controller, project) }
+        val file = myFixture.addFileToProject("src/Diff.kt", "fun diff() = Unit").virtualFile
+        edt {
+            project.service<WorktreeEditorMatchers>().register(WorktreeEditorMatcher { current: VirtualFile ->
+                if (current == file) item.path else null
+            })
+            controller.reload()
+        }
+        flush()
+
+        val list = edt { UIUtil.findComponentOfType(panel, JBList::class.java)!! }
+        edt { FileEditorManager.getInstance(project).openFile(file, true) }
+        pump()
+
+        assertEquals(item.id, edt { (list.selectedValue as ActiveListItem).key })
+    }
+
     fun `test deleting a worktree closes and releases its worktree session editor`() {
         val item = WorktreeDto("/repo/.kilo/worktrees/feature-x", "feature-x", "feature/x", "/repo/.kilo/worktrees/feature-x")
         rpc.listed += item