feat(jetbrains): move sessions to a worktree from the worktree editor

.changeset/jetbrains-worktree-editor-move.md

@@ -0,0 +1,5 @@
+---
+"@kilocode/kilo-jetbrains": minor
+---
+
+Move a session into a new worktree directly from the base checkout's worktree editor tab: a "Move to Worktree" entry now leads its session list's row menu, and sessions there get the same New Worktree / Move to Worktree toolbar above the prompt that the tool window shows. The tab's header now also reports the base checkout's uncommitted changes, the ones a move would carry.

packages/kilo-jetbrains/backend/src/main/kotlin/ai/kilocode/backend/rpc/KiloWorktreeRpcApiImpl.kt

@@ -165,10 +165,16 @@ class KiloWorktreeRpcApiImpl : KiloWorktreeRpcApi {
         WorktreeStatsListDto(parallel(items.filter { !it.main }) { item -> stats(item, fallback) })
     }
 
+    /**
+     * Uncommitted counts for every working tree of [directory]'s repo, the main checkout included: its
+     * own session editor tab shows them in its header, and they are what a move to a worktree carries.
+     * Unlike [stats], which compares a worktree against the base branch and so has nothing to say about
+     * the checkout that branch lives on, this comparison is local to each working tree.
+     */
     override suspend fun dirty(directory: String): WorktreeDirtyListDto = withContext(Dispatchers.IO) {
         val root = Path.of(directory).normalize()
         val items = sync(root) ?: return@withContext WorktreeDirtyListDto()
-        WorktreeDirtyListDto(parallel(items.filter { !it.main }) { item -> dirty(item) })
+        WorktreeDirtyListDto(parallel(items) { item -> dirty(item) })
     }
 
     /**

packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/rpc/BranchLocalDiffTest.kt

@@ -80,11 +80,11 @@ class BranchLocalDiffTest {
         git(dir, "push", "-u", "origin", "feature")
 
         assertEquals(listOf("feature.txt"), parity(dir).first.map { it.file })
-        assertEquals(0, trees.dirty(repo.toString()).items.single().unpushed)
+        assertEquals(0, unpushed(dir))
         Files.writeString(dir.resolve("feature.txt"), "one\ntwo\n")
         commit(dir)
         assertEquals(2, parity(dir).first.single().additions)
-        assertEquals(1, trees.dirty(repo.toString()).items.single().unpushed)
+        assertEquals(1, unpushed(dir))
     }
 
     @Test
@@ -295,7 +295,7 @@ class BranchLocalDiffTest {
     }
 
     @Test
-    fun `bulk stats and dirty retain managed-only scope excluding primary`() = runBlocking {
+    fun `bulk stats and dirty retain managed-only scope`() = runBlocking {
         init()
         val dir = worktree()
         val outside = root.resolve("outside")
@@ -304,8 +304,12 @@ class BranchLocalDiffTest {
         Files.writeString(outside.resolve("outside.txt"), "outside\n")
         Files.writeString(dir.resolve("managed.txt"), "managed\n")
 
+        // The primary checkout holds the branch the worktrees are compared against, so it has no base
+        // stats of its own but does have uncommitted counts, which its session editor tab reports. A
+        // worktree nobody manages stays out of both.
         assertEquals(listOf(dir.toString()), trees.stats(outside.toString()).items.map { it.path })
-        assertEquals(listOf(dir.toString()), trees.dirty(outside.toString()).items.map { it.path })
+        assertEquals(listOf(repo.toString(), dir.toString()), trees.dirty(outside.toString()).items.map { it.path })
+        assertEquals(1, trees.dirty(outside.toString()).items.single { it.path == repo.toString() }.untracked)
         assertEquals(listOf("managed.txt"), parity(dir).second.map { it.file })
         assertEquals(listOf("outside.txt"), api.localDiff(outside.toString(), false).map { it.file })
     }
@@ -371,6 +375,9 @@ class BranchLocalDiffTest {
         commit(repo)
     }
 
+    private suspend fun unpushed(dir: Path): Int =
+        trees.dirty(repo.toString()).items.single { it.path == dir.toString() }.unpushed
+
     private fun worktree(): Path {
         val dir = repo.resolve(".kilo/worktrees/feature")
         git(repo, "worktree", "add", "-b", "feature", dir.toString())

packages/kilo-jetbrains/backend/src/test/kotlin/ai/kilocode/backend/rpc/KiloWorktreeRpcApiImplTest.kt

@@ -872,6 +872,35 @@ class KiloWorktreeRpcApiImplTest {
         assertEquals(0, item.unpushed, "no upstream means no unpushed count")
     }
 
+    @Test
+    fun `dirty reports the main checkout too`() = runBlocking {
+        initRepo()
+        api.create(repo.toString(), CreateWorktreeRequestDto("feature/x"))
+        val root = repo.toRealPath()
+        // Creating a worktree leaves its own traces in the main checkout, so the edits below are
+        // measured as a delta rather than against an assumed-clean starting point.
+        val items = api.dirty(repo.toString()).items
+        val before = assertNotNull(items.singleOrNull { Path.of(it.path) == root }, "main checkout missing from $items")
+
+        Files.writeString(repo.resolve("README.md"), "hello there\n")
+        Files.writeString(repo.resolve("untracked.txt"), "u\n")
+        val after = assertNotNull(api.dirty(repo.toString()).items.singleOrNull { Path.of(it.path) == root })
+
+        assertEquals(before.files + 2, after.files, "the README edit plus the untracked file")
+        assertEquals(before.untracked + 1, after.untracked)
+    }
+
+    @Test
+    fun `stats leaves the main checkout out`() = runBlocking {
+        initRepo()
+        api.create(repo.toString(), CreateWorktreeRequestDto("feature/x"))
+        val root = repo.toRealPath()
+
+        // The main checkout holds the branch the others are compared against, so it has no base stats
+        // to report -- only its uncommitted counts, which dirty() answers for.
+        assertTrue(api.stats(repo.toString()).items.none { Path.of(it.path) == root })
+    }
+
     @Test
     fun `dirty counts commits missing from the upstream`() = runBlocking {
         initRepo()

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

@@ -8,6 +8,7 @@ import ai.kilocode.client.session.SessionSidePanelManager
 import ai.kilocode.client.telemetry.Telemetry
 import ai.kilocode.client.agentManager.worktree.GhStatusCoordinator
 import ai.kilocode.client.agentManager.worktree.KiloWorktreeService
+import ai.kilocode.client.agentManager.AgentManagerHost
 import ai.kilocode.client.agentManager.SidePanelKeys
 import ai.kilocode.client.agentManager.SidePanelMode
 import ai.kilocode.client.agentManager.applySidePanelMode
@@ -42,6 +43,9 @@ import kotlinx.coroutines.withContext
 import java.awt.BorderLayout
 import javax.swing.JPanel
 
+/** Registered id of the Kilo Code tool window (`kilo.jetbrains.frontend.xml`'s `<toolWindow id=...>`). */
+const val KILO_TOOL_WINDOW_ID = "Kilo Code"
+
 /**
  * Creates the Kilo Code tool window and delegates session content management.
  *
@@ -153,6 +157,19 @@ internal class KiloToolWindowSetupService(
                 agents()
                 agentManagerPanel.move(id, dir)
             }
+            // Same two flows, reachable from a worktree editor tab, which cannot see
+            // agentManagerPanel directly: see AgentManagerHost.
+            project.service<AgentManagerHost>().bind(
+                manager,
+                move = { id, dir, surface ->
+                    agents()
+                    agentManagerPanel.move(id, dir, surface)
+                },
+                newWorktree = {
+                    Telemetry.send("New Worktree Clicked", mapOf("surface" to "worktree_editor"))
+                    agentManagerPanel.configure(anchor = chat, onCreate = { agents() })
+                },
+            )
             val listener = object : ContentManagerListener {
                 override fun selectionChanged(event: ContentManagerEvent) {
                     if (event.operation != ContentManagerEvent.ContentOperation.add) return

packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/actions/MoveWorktreeSessionAction.kt

@@ -0,0 +1,23 @@
+package ai.kilocode.client.actions
+
+import ai.kilocode.client.agentManager.worktree.WorktreeSessionDataKeys
+import com.intellij.openapi.actionSystem.ActionUpdateThread
+import com.intellij.openapi.actionSystem.AnAction
+import com.intellij.openapi.actionSystem.AnActionEvent
+import com.intellij.openapi.project.DumbAware
+
+class MoveWorktreeSessionAction : AnAction(), DumbAware {
+    override fun getActionUpdateThread(): ActionUpdateThread = ActionUpdateThread.EDT
+
+    override fun update(e: AnActionEvent) {
+        val panel = e.getData(WorktreeSessionDataKeys.PANEL)
+        val item = e.getData(WorktreeSessionDataKeys.SESSION)
+        e.presentation.isEnabledAndVisible = panel != null && panel.canMove(item)
+    }
+
+    override fun actionPerformed(e: AnActionEvent) {
+        val panel = e.getData(WorktreeSessionDataKeys.PANEL) ?: return
+        val item = e.getData(WorktreeSessionDataKeys.SESSION) ?: return
+        if (panel.canMove(item)) panel.moveRow(item)
+    }
+}

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

@@ -0,0 +1,71 @@
+package ai.kilocode.client.agentManager
+
+import ai.kilocode.client.KILO_TOOL_WINDOW_ID
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.components.Service
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.Disposer
+import com.intellij.openapi.wm.ToolWindowManager
+import com.intellij.util.concurrency.annotations.RequiresEdt
+
+/**
+ * Project-level seam between a worktree session editor tab and the Agent Manager tool window panel,
+ * which the editor cannot reach directly (it lives inside [ai.kilocode.client.KiloToolWindowSetupService]).
+ * The tool window binds its two worktree flows here once it is created; an editor calls [move] /
+ * [newWorktree] the same way the chat branch dock does. When nothing is bound yet -- a tab restored
+ * before the tool window has been shown, or after a plugin reload -- the request is queued (the latest
+ * call wins) and the tool window is activated, which creates its content and flushes the queue via
+ * [bind].
+ */
+@Service(Service.Level.PROJECT)
+class AgentManagerHost(private val project: Project) {
+    private var onMove: ((String?, String, String) -> Unit)? = null
+    private var onNew: (() -> Unit)? = null
+    private var queued: (() -> Unit)? = null
+
+    /**
+     * Registers the tool window's worktree flows for the lifetime of [parent] (the tool window's
+     * disposable). A later [bind] call from a fresh tool window setup replaces the callbacks and
+     * flushes anything queued while none were bound.
+     */
+    @RequiresEdt
+    fun bind(parent: Disposable, move: (String?, String, String) -> Unit, newWorktree: () -> Unit) {
+        onMove = move
+        onNew = newWorktree
+        Disposer.register(parent) {
+            onMove = null
+            onNew = null
+        }
+        val pending = queued
+        queued = null
+        pending?.invoke()
+    }
+
+    /** Moves [sessionId] (or just the local changes in [directory] when null) into a new worktree. */
+    @RequiresEdt
+    fun move(sessionId: String?, directory: String, surface: String) {
+        val handler = onMove
+        if (handler != null) {
+            handler(sessionId, directory, surface)
+            return
+        }
+        queued = { onMove?.invoke(sessionId, directory, surface) }
+        activate()
+    }
+
+    /** Opens the New Worktree dialog. */
+    @RequiresEdt
+    fun newWorktree() {
+        val handler = onNew
+        if (handler != null) {
+            handler()
+            return
+        }
+        queued = { onNew?.invoke() }
+        activate()
+    }
+
+    private fun activate() {
+        ToolWindowManager.getInstance(project).getToolWindow(KILO_TOOL_WINDOW_ID)?.activate(null)
+    }
+}

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

@@ -245,7 +245,8 @@ class AgentManagerPanel(
         }
     }
 
-    internal fun move(sessionId: String?, directory: String) = controller.move(sessionId, directory)
+    internal fun move(sessionId: String?, directory: String, surface: String = "sidebar") =
+        controller.move(sessionId, directory, surface)
 
     private fun remove(item: WorktreeDto, force: Boolean) {
         controller.remove(item, force, onFailure = { result -> notifyFailed(item, result, force) })

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

@@ -242,9 +242,12 @@ class WorktreeController(
     /**
      * Copies working-tree changes into a new worktree. When [sessionId] is set, the source session is
      * also forked into the worktree; otherwise the opened worktree starts with a fresh session.
+     * [surface] is reported only on the "Continue in Worktree" telemetry event, so callers other than
+     * the sidebar chat dock (e.g. the worktree editor's session list or its own action toolbar) can be
+     * told apart.
      */
     @RequiresEdt
-    fun move(sessionId: String?, source: String = directory) {
+    fun move(sessionId: String?, source: String = directory, surface: String = "sidebar") {
         val key = sessionId ?: source
         if (!moves.add(key)) return
         val branch = suggestName()
@@ -278,7 +281,7 @@ class WorktreeController(
                                 onCreated?.invoke(worktree)
                                 telemetry(
                                     "Continue in Worktree",
-                                    mapOf("surface" to "sidebar", "session" to (sessionId != null).toString()),
+                                    mapOf("surface" to surface, "session" to (sessionId != null).toString()),
                                 )
                             }
                             MoveStage.ERROR -> failMove(key, temp, event.error, stage)

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

@@ -1,6 +1,7 @@
 package ai.kilocode.client.agentManager.worktree
 
 import ai.kilocode.client.KiloNotifications
+import ai.kilocode.client.agentManager.AgentManagerHost
 import ai.kilocode.client.app.KiloSessionService
 import ai.kilocode.client.app.KiloWorkspaceService
 import ai.kilocode.client.app.Workspace
@@ -72,8 +73,29 @@ open class WorktreeSessionEditorManager(
             project.service<KiloVfsManager>().updatePresentation(WorktreeSessionEditorKind.ID, worktreeSessionParams(updated))
         }
     },
+    // Whether this tab's directory is the repo's main working tree rather than a linked worktree.
+    // `git worktree list` answers this from any of the repo's worktrees, so no separate "repo root"
+    // is needed -- just the tab's own directory.
+    private val resolveBase: suspend (String) -> Boolean = { dir ->
+        service<KiloWorktreeService>().list(dir).worktrees
+            .firstOrNull { normalizeWorktreePath(it.path) == normalizeWorktreePath(dir) }
+            ?.main == true
+    },
+    private val moveHost: (String?, String, String) -> Unit = { id, dir, surface ->
+        project.service<AgentManagerHost>().move(id, dir, surface)
+    },
+    private val newWorktreeHost: () -> Unit = {
+        project.service<AgentManagerHost>().newWorktree()
+    },
 ) : SessionHost(project, worktree, create, resolve, status, timers, request) {
-    override val showsBranchDock: Boolean get() = false
+    // Both the branch dock and the New Worktree / Move to Worktree flows only make sense from the
+    // base checkout -- a linked worktree's own editor tab keeps today's plain session view. Resolved
+    // once in start(), before the first session opens; see resolveBase().
+    private var resolvedBase = false
+    private var baseResolved = false
+    override val showsBranchDock: Boolean get() = base()
+    override val supportsNewWorktree: Boolean get() = base()
+    override val supportsMoveToWorktree: Boolean get() = base()
     override val hostedInEditorTab: Boolean get() = true
     private val right = JPanel(BorderLayout())
     private val deleting = linkedSetOf<String>()
@@ -94,9 +116,39 @@ open class WorktreeSessionEditorManager(
         bindMigration()
     }
 
+    /** Whether this tab's directory is the repo's main working tree; see [resolveBase]. */
+    @RequiresEdt
+    open fun base(): Boolean = resolvedBase
+
+    @RequiresEdt
+    override fun newWorktree() {
+        if (base()) newWorktreeHost()
+    }
+
+    @RequiresEdt
+    override fun moveToWorktree(sessionId: String?, directory: String) {
+        if (base()) moveHost(sessionId, directory, "worktree_editor")
+    }
+
     @RequiresEdt
     fun start() {
         startedOnce = true
+        if (baseResolved) {
+            startSessions()
+            return
+        }
+        cs.launch {
+            val resolved = runCatching { resolveBase(worktree.directory) }.getOrDefault(false)
+            edt {
+                resolvedBase = resolved
+                baseResolved = true
+                startSessions()
+            }
+        }
+    }
+
+    @RequiresEdt
+    private fun startSessions() {
         list.reload {
             val target = session
             if (target != null) {

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

@@ -123,7 +123,7 @@ class WorktreeSessionEditorPanel @RequiresEdt constructor(
         onCell = { _, _ -> },
         onOpen = { row, focus -> open(row, focus) },
         menu = ActiveListMenu(WorktreeSessionDataKeys.SESSION, group, element = { row ->
-            (row as? SessionRow)?.session?.takeIf { canRename(it) || canDelete(it) }
+            (row as? SessionRow)?.session?.takeIf { canMove(it) || canRename(it) || canDelete(it) }
         }),
     )
     private val run = if (project != null && worktree.directory.isNotBlank()) {
@@ -221,6 +221,22 @@ class WorktreeSessionEditorPanel @RequiresEdt constructor(
     @RequiresEdt
     internal fun renameRow(item: SessionDto) = beginRename(item.id)
 
+    /**
+     * Only offered from the base checkout's tab (not a linked worktree's own tab, see
+     * [WorktreeSessionEditorManager.base]), for a real session that is not already being deleted, and
+     * hidden rather than disabled while the session is running -- matching the chat branch dock's
+     * Move to Worktree action.
+     */
+    @RequiresEdt
+    internal fun canMove(item: SessionDto?): Boolean =
+        manager.base() && canDelete(item) && manager.activity()[item?.id] != SessionActivityKind.RUNNING
+
+    @RequiresEdt
+    internal fun moveRow(item: SessionDto) {
+        if (!canMove(item)) return
+        manager.moveToWorktree(item.id, worktree.directory)
+    }
+
     @RequiresEdt
     private fun confirmDelete(ids: List<String>, cell: String? = null) {
         val active = ids.filter { it != SessionHost.NEW && it !in manager.deleting() }.distinct()