fix(jetbrains): persist worktree order

This commit is contained in:
kirillk
2026-08-03 15:58:51 -04:00
parent aafde0bf7e
commit 916548a8cc
3 changed files with 165 additions and 13 deletions
+5
View File
@@ -0,0 +1,5 @@
---
"@kilocode/kilo-jetbrains": patch
---
Keep JetBrains Agent Manager worktrees in a stable creation order after switching panels or reloading.
@@ -13,9 +13,12 @@ import com.intellij.execution.configurations.GeneralCommandLine
import com.intellij.execution.process.CapturingProcessHandler
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import kotlinx.serialization.Serializable
import kotlinx.serialization.builtins.MapSerializer
import kotlinx.serialization.builtins.serializer
import kotlinx.serialization.json.Json
import kotlinx.serialization.json.JsonObject
import kotlinx.serialization.json.decodeFromJsonElement
import java.nio.file.Files
import java.nio.file.Path
import java.nio.file.StandardCopyOption
@@ -32,8 +35,9 @@ class KiloWorktreeRpcApiImpl : KiloWorktreeRpcApi {
if (!res.ok) return@withContext WorktreeListDto()
val items = managedWorktrees(parseWorktreeList(res.stdout))
val store = worktreeNameStore(items)
val names = store?.let(::readWorktreeNames).orEmpty()
WorktreeListDto(overlayWorktreeNames(items, names))
val state = store?.let { syncWorktreeState(it, worktreePaths(items)) } ?: WorktreeState()
val named = overlayWorktreeNames(items, state.names)
WorktreeListDto(orderWorktrees(named, state.worktreeOrder))
}
override suspend fun listBranches(directory: String): WorktreeBranchesDto = withContext(Dispatchers.IO) {
@@ -63,6 +67,11 @@ class KiloWorktreeRpcApiImpl : KiloWorktreeRpcApi {
} else {
LOG.info("worktree created: branch=$branch dir=$dir")
val path = dir.toRealPath().toString()
val list = runGit(base, "worktree", "list", "--porcelain")
val items = if (list.ok) managedWorktrees(parseWorktreeList(list.stdout)) else emptyList()
val store = worktreeNameStore(items) ?: base.resolve(".kilo").resolve(WORKTREE_NAMES_FILE)
val paths = worktreePaths(items).ifEmpty { listOf(path) }
appendWorktreeOrder(store, path, paths)
CreateWorktreeResultDto(
worktree = WorktreeDto(path, dir.fileName.toString(), branch, path),
)
@@ -73,6 +82,9 @@ class KiloWorktreeRpcApiImpl : KiloWorktreeRpcApi {
withContext(Dispatchers.IO) {
val base = Path.of(directory).normalize()
LOG.info("worktree remove requested: path=$path branch=${branch ?: "(none)"} force=$force base=$base")
val list = runGit(base, "worktree", "list", "--porcelain")
val store = (if (list.ok) worktreeNameStore(managedWorktrees(parseWorktreeList(list.stdout))) else null)
?: base.resolve(".kilo").resolve(WORKTREE_NAMES_FILE)
// Force means the user accepted removing a locked worktree; unlock first so the plain
// remove succeeds. Unlock fails harmlessly when the tree isn't actually locked.
if (force) {
@@ -94,6 +106,7 @@ class KiloWorktreeRpcApiImpl : KiloWorktreeRpcApi {
if (!del.ok) LOG.warn("worktree branch delete failed: branch=$it exit=${del.exit} stderr=${del.stderr.trim()}")
}
LOG.info("worktree removed: path=$path branch=${branch ?: "(none)"}")
removeWorktreeState(store, path)
RemoveWorktreeResultDto(ok = true)
}
@@ -110,9 +123,10 @@ class KiloWorktreeRpcApiImpl : KiloWorktreeRpcApi {
val target = items.firstOrNull { samePath(it.path, path) && !it.main }
?: return@withContext RenameWorktreeResultDto(error = "Worktree not found")
return@withContext try {
val names = readWorktreeNames(store).toMutableMap()
val state = readWorktreeState(store).reconcile(worktreePaths(items))
val names = state.names.toMutableMap()
names[target.path] = title
writeWorktreeNames(store, names)
writeWorktreeState(store, state.copy(names = names))
RenameWorktreeResultDto(worktree = target.copy(name = title))
} catch (e: Exception) {
LOG.warn("worktree rename failed: path=$path message=${e.message}", e)
@@ -133,12 +147,13 @@ class KiloWorktreeRpcApiImpl : KiloWorktreeRpcApi {
val target = items.firstOrNull { samePath(it.path, path) && !it.main }
?: return@withContext RenameWorktreeResultDto(error = "Worktree not found")
return@withContext try {
val names = readWorktreeNames(store).toMutableMap()
val state = readWorktreeState(store).reconcile(worktreePaths(items))
val names = state.names.toMutableMap()
// Only adopt while the worktree is still default. A recorded name means the user (or a
// prior adoption) already titled it, so leave it untouched and report a no-op.
if (!names[target.path].isNullOrBlank()) return@withContext RenameWorktreeResultDto()
names[target.path] = title
writeWorktreeNames(store, names)
writeWorktreeState(store, state.copy(names = names))
LOG.info("worktree name adopted: path=$path name=$title")
RenameWorktreeResultDto(worktree = target.copy(name = title))
} catch (e: Exception) {
@@ -162,10 +177,28 @@ class KiloWorktreeRpcApiImpl : KiloWorktreeRpcApi {
}
}
private val json = Json { prettyPrint = true }
private val json = Json { prettyPrint = true; ignoreUnknownKeys = true }
private val codec = MapSerializer(String.serializer(), String.serializer())
private const val WORKTREE_NAMES_FILE = "worktree-names.json"
@Serializable
private data class WorktreeNamesFile(
val names: Map<String, String> = emptyMap(),
val worktreeOrder: List<String> = emptyList(),
)
internal data class WorktreeState(
val names: Map<String, String> = emptyMap(),
val worktreeOrder: List<String> = emptyList(),
) {
fun reconcile(paths: List<String>): WorktreeState {
val set = paths.toSet()
val order = (worktreeOrder.filter { it in set } + paths.filter { it !in worktreeOrder }).distinct()
val next = names.filterKeys { it in set }
return WorktreeState(next, order)
}
}
/** Parse `git worktree list --porcelain`. First entry is the main working tree. */
internal fun parseWorktreeList(raw: String): List<WorktreeDto> {
val out = mutableListOf<WorktreeDto>()
@@ -218,24 +251,50 @@ internal fun overlayWorktreeNames(items: List<WorktreeDto>, names: Map<String, S
}
}
internal fun orderWorktrees(items: List<WorktreeDto>, order: List<String>): List<WorktreeDto> {
if (order.isEmpty()) return items
val rank = order.withIndex().associate { it.value to it.index }
val main = items.filter { it.main }
val extra = items.filter { !it.main }
.sortedWith(compareBy<WorktreeDto> { rank[it.path] ?: Int.MAX_VALUE }.thenBy { it.path })
return main + extra
}
internal fun readWorktreeNames(file: Path): Map<String, String> {
if (!Files.exists(file)) return emptyMap()
return readWorktreeState(file).names
}
internal fun readWorktreeState(file: Path): WorktreeState {
if (!Files.exists(file)) return WorktreeState()
return try {
val raw = Files.readString(file)
json.decodeFromString(codec, raw)
.filterValues { it.isNotBlank() }
val element = json.parseToJsonElement(raw)
if (element is JsonObject && ("names" in element || "worktreeOrder" in element)) {
val data = json.decodeFromJsonElement<WorktreeNamesFile>(element)
return WorktreeState(data.names.filterValues { it.isNotBlank() }, data.worktreeOrder.filter { it.isNotBlank() })
}
val names = json.decodeFromJsonElement(codec, element).filterValues { it.isNotBlank() }
WorktreeState(names, names.keys.toList())
} catch (e: Exception) {
KiloWorktreeRpcApiImpl.LOG.warn("worktree names read failed: file=$file message=${e.message}", e)
emptyMap()
WorktreeState()
}
}
internal fun writeWorktreeNames(file: Path, names: Map<String, String>) {
val order = readWorktreeState(file).worktreeOrder
writeWorktreeState(file, WorktreeState(names, order))
}
internal fun writeWorktreeState(file: Path, state: WorktreeState) {
Files.createDirectories(file.parent)
val data = names.filterValues { it.isNotBlank() }
val data = WorktreeNamesFile(
names = state.names.filterValues { it.isNotBlank() },
worktreeOrder = state.worktreeOrder.filter { it.isNotBlank() }.distinct(),
)
val tmp = Files.createTempFile(file.parent, ".worktree-names", ".tmp")
try {
Files.writeString(tmp, json.encodeToString(codec, data))
Files.writeString(tmp, json.encodeToString(WorktreeNamesFile.serializer(), data))
try {
Files.move(tmp, file, StandardCopyOption.ATOMIC_MOVE, StandardCopyOption.REPLACE_EXISTING)
} catch (_: Exception) {
@@ -246,6 +305,39 @@ internal fun writeWorktreeNames(file: Path, names: Map<String, String>) {
}
}
private fun syncWorktreeState(file: Path, paths: List<String>): WorktreeState {
val state = readWorktreeState(file)
val next = state.reconcile(paths)
if (next == state) return next
try {
writeWorktreeState(file, next)
} catch (e: Exception) {
KiloWorktreeRpcApiImpl.LOG.warn("worktree state sync failed: file=$file message=${e.message}", e)
}
return next
}
private fun appendWorktreeOrder(file: Path, path: String, paths: List<String>) {
val state = readWorktreeState(file)
val set = paths.toSet()
val order = state.worktreeOrder.filter { it in set && !samePath(it, path) } +
paths.filter { it !in state.worktreeOrder && !samePath(it, path) } +
path
writeWorktreeState(file, state.copy(worktreeOrder = order.distinct()))
}
private fun removeWorktreeState(file: Path, path: String) {
val state = readWorktreeState(file)
val names = state.names.filterKeys { !samePath(it, path) }
val order = state.worktreeOrder.filter { !samePath(it, path) }
if (names == state.names && order == state.worktreeOrder) return
writeWorktreeState(file, state.copy(names = names, worktreeOrder = order))
}
private fun worktreePaths(items: List<WorktreeDto>): List<String> {
return items.filter { !it.main }.map { it.path }
}
private fun worktreeNameStore(items: List<WorktreeDto>): Path? {
val main = items.firstOrNull { it.main } ?: return null
return Path.of(main.path).normalize().resolve(".kilo").resolve(WORKTREE_NAMES_FILE)
@@ -133,11 +133,38 @@ class KiloWorktreeRpcApiImplTest {
writeWorktreeNames(file, mapOf("/repo/.kilo/worktrees/feature-x" to "Feature Label", "/blank" to ""))
assertEquals(mapOf("/repo/.kilo/worktrees/feature-x" to "Feature Label"), readWorktreeNames(file))
assertEquals(emptyList(), readWorktreeState(file).worktreeOrder)
Files.writeString(file, "not json")
assertTrue(readWorktreeNames(file).isEmpty())
}
@Test
fun `worktree state round trips and migrates legacy names`() {
val file = repo.resolve(".kilo").resolve("worktree-names.json")
val first = "/repo/.kilo/worktrees/zebra"
val second = "/repo/.kilo/worktrees/alpha"
writeWorktreeState(file, WorktreeState(mapOf(first to "Zebra", second to "Alpha"), listOf(first, second)))
assertEquals(WorktreeState(mapOf(first to "Zebra", second to "Alpha"), listOf(first, second)), readWorktreeState(file))
Files.writeString(file, """{"$second":"Alpha","$first":"Zebra","/blank":""}""")
assertEquals(WorktreeState(mapOf(second to "Alpha", first to "Zebra"), listOf(second, first)), readWorktreeState(file))
}
@Test
fun `orderWorktrees keeps main first and sorts worktrees by persisted order`() {
val main = WorktreeDto("/repo", "repo", "main", "/repo", main = true)
val first = WorktreeDto("/repo/.kilo/worktrees/zebra", "zebra", "zebra", "/repo/.kilo/worktrees/zebra")
val second = WorktreeDto("/repo/.kilo/worktrees/alpha", "alpha", "alpha", "/repo/.kilo/worktrees/alpha")
val third = WorktreeDto("/repo/.kilo/worktrees/beta", "beta", "beta", "/repo/.kilo/worktrees/beta")
val out = orderWorktrees(listOf(main, second, third, first), listOf(first.path, second.path))
assertEquals(listOf(main.path, first.path, second.path, third.path), out.map { it.path })
}
@Test
fun `remove reports locked and force removes a locked worktree`() = runBlocking {
initRepo()
@@ -186,6 +213,34 @@ class KiloWorktreeRpcApiImplTest {
assertFalse(after.any { it.branch == "feature/x" }, "removed worktree should be gone")
}
@Test
fun `create records order so reload keeps creation order`() = runBlocking {
initRepo()
val first = assertNotNull(api.create(repo.toString(), CreateWorktreeRequestDto("zebra")).worktree)
val second = assertNotNull(api.create(repo.toString(), CreateWorktreeRequestDto("alpha")).worktree)
val listed = api.list(repo.toString()).worktrees.filter { !it.main }
assertEquals(listOf(first.path, second.path), listed.map { it.path })
assertEquals(listOf(first.path, second.path), readWorktreeState(repo.resolve(".kilo").resolve("worktree-names.json")).worktreeOrder)
}
@Test
fun `remove prunes names and order from worktree state`() = runBlocking {
initRepo()
val first = assertNotNull(api.create(repo.toString(), CreateWorktreeRequestDto("zebra")).worktree)
val second = assertNotNull(api.create(repo.toString(), CreateWorktreeRequestDto("alpha")).worktree)
assertNotNull(api.rename(repo.toString(), first.path, "First").worktree)
assertNotNull(api.rename(repo.toString(), second.path, "Second").worktree)
val removed = api.remove(repo.toString(), first.path, first.branch)
assertTrue(removed.ok, "remove should report success: ${removed.error}")
val state = readWorktreeState(repo.resolve(".kilo").resolve("worktree-names.json"))
assertEquals(mapOf(second.path to "Second"), state.names)
assertEquals(listOf(second.path), state.worktreeOrder)
}
@Test
fun `rename persists a custom worktree name and list overlays it`() = runBlocking {
initRepo()