From 5ab0ba24b3934a04ad9dc02b0fc67cc888a34c0c Mon Sep 17 00:00:00 2001 From: kirillk Date: Mon, 3 Aug 2026 14:10:38 -0400 Subject: [PATCH] fix(jetbrains): sync worktree selection with editor --- .changeset/worktree-editor-selection.md | 5 ++ .../client/agentManager/AgentManagerPanel.kt | 47 ++++++++--- .../worktree/WorktreeEditorMatcher.kt | 35 ++++++++ .../ai/kilocode/client/ui/list/ActiveList.kt | 3 + .../kilocode/client/ui/list/ActiveListView.kt | 6 ++ .../agentManager/AgentManagerPanelTest.kt | 82 +++++++++++++++++++ 6 files changed, 167 insertions(+), 11 deletions(-) create mode 100644 .changeset/worktree-editor-selection.md create mode 100644 packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/worktree/WorktreeEditorMatcher.kt diff --git a/.changeset/worktree-editor-selection.md b/.changeset/worktree-editor-selection.md new file mode 100644 index 00000000000..9da6ad27e3c --- /dev/null +++ b/.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. diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/AgentManagerPanel.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/AgentManagerPanel.kt index e209208e27e..2e2da0bc4c6 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/AgentManagerPanel.kt +++ b/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().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().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()?.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() - .firstOrNull { it.path.kind == WorktreeSessionEditorKind.ID } - ?.path - ?.params - ?.get("path") - ?.takeIf { it.isNotBlank() } + val file = FileEditorManager.getInstance(target).selectedFiles.firstOrNull() + return target.service().match(file) } private fun item(key: String): WorktreeDto? { diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/worktree/WorktreeEditorMatcher.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/agentManager/worktree/WorktreeEditorMatcher.kt new file mode 100644 index 00000000000..d4646dc25e8 --- /dev/null +++ b/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() + + 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() } + } +} diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/list/ActiveList.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/list/ActiveList.kt index cdb3ad5018b..53f84f77a76 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/list/ActiveList.kt +++ b/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 = view.selectedItems() diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/list/ActiveListView.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/list/ActiveListView.kt index d23a6e91d25..59ff870ac49 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/ui/list/ActiveListView.kt +++ b/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 { checkEdt() diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/agentManager/AgentManagerPanelTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/agentManager/AgentManagerPanelTest.kt index eb3bbbb9661..4ede8122f4b 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/agentManager/AgentManagerPanelTest.kt +++ b/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().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().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().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