fix(jetbrains): harden agent manager session lifecycle

2026-10-29 18:52:45 +08:00 · 2026-08-10 12:21:49 -04:00
parent 5ba298792d
commit 745ee1ecce
11 changed files with 473 additions and 114 deletions
@@ -5,6 +5,7 @@ import ai.kilocode.client.agentManager.worktree.ConfigureWorktreeDialog
 import ai.kilocode.client.agentManager.worktree.WorktreeController
 import ai.kilocode.client.agentManager.worktree.WorktreeDataKeys
 import ai.kilocode.client.agentManager.worktree.WorktreeIcons
+import ai.kilocode.client.agentManager.worktree.WorktreeStatusBinding
 import ai.kilocode.client.agentManager.worktree.WorktreeStatusService
 import ai.kilocode.client.agentManager.worktree.WorktreeNameCache
 import ai.kilocode.client.agentManager.worktree.WorktreeEditorMatchers
@@ -22,6 +23,7 @@ import ai.kilocode.client.diff.ensureDiffEditorKind
 import ai.kilocode.client.plugin.KiloBundle
 import ai.kilocode.client.session.SessionActivityKind
 import ai.kilocode.client.ui.UiStyle
+import ai.kilocode.client.util.bindTheme
 import ai.kilocode.client.ui.list.ActiveList
 import ai.kilocode.client.ui.list.ActiveListBadge
 import ai.kilocode.client.ui.list.ActiveListConfig
@@ -40,7 +42,6 @@ import ai.kilocode.rpc.dto.WorktreeStatsDto
 import com.intellij.icons.AllIcons
 import com.intellij.ide.BrowserUtil
 import com.intellij.ide.DeleteProvider
-import com.intellij.ide.ui.LafManagerListener
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.actionSystem.ActionManager
 import com.intellij.openapi.actionSystem.ActionGroup
@@ -52,7 +53,6 @@ import com.intellij.openapi.actionSystem.DataSink
 import com.intellij.openapi.actionSystem.DefaultActionGroup
 import com.intellij.openapi.actionSystem.PlatformDataKeys
 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
@@ -67,13 +67,6 @@ import java.awt.Color
 import javax.swing.event.ListDataEvent
 import javax.swing.event.ListDataListener
 import javax.swing.JComponent
-import javax.swing.SwingUtilities
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.SupervisorJob
-import kotlinx.coroutines.cancel
-import kotlinx.coroutines.flow.collectLatest
-import kotlinx.coroutines.launch

 /**
 * Agent Manager panel: a git-worktree list with search and a delete action revealed on selection,
@@ -103,10 +96,8 @@ class AgentManagerPanel(
        }),
    )
    private var selected: String? = null
-    private val cs = CoroutineScope(SupervisorJob() + Dispatchers.Default)
    private var stats: Map<String, WorktreeStatsDto> = emptyMap()
    private var prs: Map<String, WorktreePrDto> = emptyMap()
-    private var status: AutoCloseable? = null

    init {
        Disposer.register(parent, this)
@@ -116,7 +107,7 @@ class AgentManagerPanel(
        list.installPopup(group)
        sync()
        bindModel()
-        bindTheme()
+        bindTheme(this, this)
        controller.onSelect = { key ->
            // Focus the list so the freshly created worktree renders as an active selection rather
            // than the inactive highlight it would get while focus stays on the toolbar.
@@ -312,15 +303,6 @@ class AgentManagerPanel(
        KiloNotifications.error(project, title, result.error)
    }

-    private fun bindTheme() {
-        val bus = ApplicationManager.getApplication().messageBus.connect(this)
-        bus.subscribe(LafManagerListener.TOPIC, LafManagerListener {
-            ApplicationManager.getApplication().invokeLater {
-                SwingUtilities.updateComponentTreeUI(this)
-            }
-        })
-    }
-
    private fun bindEditorSelection() {
        val target = project ?: return
        target.service<WorktreeEditorMatchers>().register(WorktreeSessionEditorMatcher)
@@ -391,32 +373,12 @@ class AgentManagerPanel(

    private fun bindStatus() {
        val target = project ?: return
-        val service = target.service<WorktreeStatusService>()
-        status = service.attach()
-        service.refreshStats()
-        service.refreshPr()
-        cs.launch {
-            service.stats.collectLatest { value ->
-                edtIfAlive {
-                    stats = value
-                    sync()
-                }
-            }
-        }
-        cs.launch {
-            service.pr.collectLatest { value ->
-                edtIfAlive {
-                    prs = value
-                    sync()
-                }
-            }
-        }
-    }
-
-    private fun edtIfAlive(block: () -> Unit) {
-        ApplicationManager.getApplication().invokeLater {
-            if ((project == null || !project.isDisposed) && !Disposer.isDisposed(this)) block()
-        }
+        WorktreeStatusBinding(
+            target,
+            this,
+            onStats = { value -> stats = value; sync() },
+            onPr = { value -> prs = value; sync() },
+        )
    }

    override fun dispose() {
@@ -424,8 +386,6 @@ class AgentManagerPanel(
        controller.onCreateFailure = null
        controller.onRemoveSuccess = null
        controller.onActivityChanged = null
-        status?.close()
-        cs.cancel()
    }

    override fun uiDataSnapshot(sink: DataSink) {
@@ -166,8 +166,7 @@ open class WorktreeSessionEditorManager(
    @RequiresEdt
    private fun adoptTitle(): String? {
        val live = titles()
-        return (0 until list.model.size)
-            .map { list.model.getElementAt(it) }
+        return list.sessions()
            .sortedBy { it.time.created }
            .firstNotNullOfOrNull { s ->
                (live[s.id]?.takeIf { it.isNotBlank() } ?: s.title.takeIf { it.isNotBlank() })
@@ -228,15 +227,14 @@ open class WorktreeSessionEditorManager(

    @RequiresEdt
    private fun latest(): SessionDto? {
-        return (0 until list.model.size)
-            .map { list.model.getElementAt(it) }
+        return list.sessions()
            .filter { it.id !in deleting }
            .maxByOrNull { it.time.updated }
    }

    @RequiresEdt
    private fun next(key: String?): SessionDto? {
-        val rows = HistoryTime.sorted((0 until list.model.size).map { LocalHistoryItem(list.model.getElementAt(it)) })
+        val rows = HistoryTime.sorted(list.sessions().map { LocalHistoryItem(it) })
            .map { it.session }
        val idx = rows.indexOfFirst { it.id == key }
        if (idx < 0) return rows.firstOrNull { it.id !in deleting }
@@ -246,9 +244,7 @@ open class WorktreeSessionEditorManager(

    @RequiresEdt
    private fun title(id: String): String {
-        return (0 until list.model.size)
-            .map { list.model.getElementAt(it) }
-            .firstOrNull { it.id == id }
+        return list.session(id)
            ?.title
            ?.takeIf { it.isNotBlank() }
            ?: KiloBundle.message("worktree.session.untitled")
@@ -11,6 +11,7 @@ import ai.kilocode.client.session.SessionRef
 import ai.kilocode.client.session.history.HistorySection
 import ai.kilocode.client.session.history.HistoryTime
 import ai.kilocode.client.session.history.LocalHistoryItem
+import ai.kilocode.client.util.bindTheme
 import ai.kilocode.client.ui.list.ActiveList
 import ai.kilocode.client.ui.list.ActiveListBadge
 import ai.kilocode.client.ui.list.ActiveListConfig
@@ -27,7 +28,6 @@ import ai.kilocode.rpc.dto.SessionDto
 import ai.kilocode.rpc.dto.WorktreePrDto
 import ai.kilocode.rpc.dto.WorktreeStatsDto
 import com.intellij.icons.AllIcons
-import com.intellij.ide.ui.LafManagerListener
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.actionSystem.ActionManager
 import com.intellij.openapi.actionSystem.ActionPlaces
@@ -38,7 +38,6 @@ import com.intellij.openapi.actionSystem.AnActionEvent
 import com.intellij.openapi.actionSystem.DataSink
 import com.intellij.openapi.actionSystem.DefaultActionGroup
 import com.intellij.openapi.actionSystem.UiDataProvider
-import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.components.service
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Disposer
@@ -54,15 +53,8 @@ import java.awt.Color
 import javax.swing.JComponent
 import javax.swing.ListSelectionModel
 import javax.swing.JPanel
-import javax.swing.SwingUtilities
 import javax.swing.event.ListDataEvent
 import javax.swing.event.ListDataListener
-import kotlinx.coroutines.CoroutineScope
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.SupervisorJob
-import kotlinx.coroutines.cancel
-import kotlinx.coroutines.flow.collectLatest
-import kotlinx.coroutines.launch

 class WorktreeSessionEditorPanel(
    parent: Disposable,
@@ -96,8 +88,6 @@ class WorktreeSessionEditorPanel(
    )
    private val statsView = WorktreeStatsView(::openBranchDiff)
    private var started = false
-    private val cs = CoroutineScope(SupervisorJob() + Dispatchers.Default)
-    private var status: AutoCloseable? = null
    private var stats: WorktreeStatsDto? = null
    private var pr: WorktreePrDto? = null

@@ -115,7 +105,6 @@ class WorktreeSessionEditorPanel(
        splitter.secondComponent = manager.component
        addToCenter(splitter)
        bindModel()
-        bindTheme()
        manager.onPresent = { key -> select(key) }
        manager.onListChanged = {
            sync()
@@ -132,6 +121,7 @@ class WorktreeSessionEditorPanel(
            }
        }
        bindStatus()
+        bindTheme(this, this)
        sync()
    }

@@ -262,7 +252,7 @@ class WorktreeSessionEditorPanel(
        val titles = manager.titles()
        val deleting = manager.deleting()
        if (pending || key == SessionHost.NEW) rows += NewRow
-        rows += HistoryTime.sorted((0 until controller.model.size).map { LocalHistoryItem(controller.model.getElementAt(it)) })
+        rows += HistoryTime.sorted(controller.sessions().map { LocalHistoryItem(it) })
            .map { SessionRow(it.session, kinds[it.id], deleting = it.id in deleting, live = titles[it.id]) }
        list.update(rows, ActiveListSelection.PreserveNoScroll)
        select(if (pending) SessionHost.NEW else key)
@@ -275,11 +265,7 @@ class WorktreeSessionEditorPanel(
    }

    @RequiresEdt
-    private fun item(key: String): SessionDto? {
-        return (0 until controller.model.size)
-            .map { controller.model.getElementAt(it) }
-            .firstOrNull { it.id == key }
-    }
+    private fun item(key: String): SessionDto? = controller.session(key)

    @RequiresEdt
    private fun title(key: String): String {
@@ -301,41 +287,15 @@ class WorktreeSessionEditorPanel(
        Disposer.register(this) { controller.model.removeListDataListener(listener) }
    }

-    private fun bindTheme() {
-        val bus = ApplicationManager.getApplication().messageBus.connect(this)
-        bus.subscribe(LafManagerListener.TOPIC, LafManagerListener {
-            ApplicationManager.getApplication().invokeLater {
-                SwingUtilities.updateComponentTreeUI(this)
-            }
-        })
-    }
-
    private fun bindStatus() {
        val target = project ?: return
-        val service = target.service<WorktreeStatusService>()
-        status = service.attach()
-        cs.launch {
-            service.stats.collectLatest { value ->
-                edtIfAlive {
-                    stats = value[normalizeWorktreePath(worktree.directory)]
-                    statsView.update(stats, pr)
-                }
-            }
-        }
-        cs.launch {
-            service.pr.collectLatest { value ->
-                edtIfAlive {
-                    pr = value[normalizeWorktreePath(worktree.directory)]
-                    statsView.update(stats, pr)
-                }
-            }
-        }
-    }
-
-    private fun edtIfAlive(block: () -> Unit) {
-        ApplicationManager.getApplication().invokeLater {
-            if ((project == null || !project.isDisposed) && !Disposer.isDisposed(this)) block()
-        }
+        val key = normalizeWorktreePath(worktree.directory)
+        WorktreeStatusBinding(
+            target,
+            this,
+            onStats = { value -> stats = value[key]; statsView.update(stats, pr) },
+            onPr = { value -> pr = value[key]; statsView.update(stats, pr) },
+        )
    }

    override fun uiDataSnapshot(sink: DataSink) {
@@ -348,8 +308,6 @@ class WorktreeSessionEditorPanel(
    override fun dispose() {
        manager.onPresent = null
        manager.onListChanged = null
-        status?.close()
-        cs.cancel()
    }

    private inner class NewAction : AnAction(
@@ -18,10 +18,18 @@ class WorktreeSessionListController(
 ) {
    val model = CollectionListModel<SessionDto>()

+    /** Snapshot of the listed sessions, in model order. */
+    @RequiresEdt
+    fun sessions(): List<SessionDto> = (0 until model.size).map { model.getElementAt(it) }
+
+    /** The listed session with [id], or null when it is not in the model. */
+    @RequiresEdt
+    fun session(id: String): SessionDto? = sessions().firstOrNull { it.id == id }
+
    fun reload(done: (() -> Unit)? = null) {
        cs.launch {
            try {
-                val result = service.list(dir)
+                val result = service.sessionsFor(dir)
                edt {
                    model.replaceAll(result.sessions)
                    capture("Worktree Session List Loaded", mapOf("count" to result.sessions.size.toString()))
@@ -0,0 +1,46 @@
+package ai.kilocode.client.agentManager.worktree
+
+import ai.kilocode.client.util.edt
+import ai.kilocode.rpc.dto.WorktreePrDto
+import ai.kilocode.rpc.dto.WorktreeStatsDto
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.components.service
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.Disposer
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.SupervisorJob
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.flow.collectLatest
+import kotlinx.coroutines.launch
+
+/**
+ * Shared wiring to the project [WorktreeStatusService] for the surfaces that show worktree
+ * stats/PR badges (the Agent Manager list and the worktree session editor). Owns the attach
+ * handle and a coroutine scope, forwards each flow value to [onStats]/[onPr] on the EDT, and
+ * skips delivery once [parent] or the project is disposed. Cleanup is tied to [parent].
+ */
+internal class WorktreeStatusBinding(
+    private val project: Project,
+    private val parent: Disposable,
+    private val onStats: (Map<String, WorktreeStatsDto>) -> Unit,
+    private val onPr: (Map<String, WorktreePrDto>) -> Unit,
+) {
+    private val cs = CoroutineScope(SupervisorJob() + Dispatchers.Default)
+    private val handle: AutoCloseable
+
+    init {
+        val service = project.service<WorktreeStatusService>()
+        handle = service.attach()
+        cs.launch { service.stats.collectLatest { value -> alive { onStats(value) } } }
+        cs.launch { service.pr.collectLatest { value -> alive { onPr(value) } } }
+        Disposer.register(parent) { close() }
+    }
+
+    private fun alive(block: () -> Unit) = edt({ !project.isDisposed && !Disposer.isDisposed(parent) }, block)
+
+    private fun close() {
+        handle.close()
+        cs.cancel()
+    }
+}
@@ -57,13 +57,13 @@ class WorktreeStatusService internal constructor(
    }

    fun refreshStats() {
-        if (project.isDisposed) return
+        if (project.isDisposed || refs == 0) return
        val timer = debounce ?: timers.timer(STATS_DEBOUNCE, repeats = false) { loadStats() }.also { debounce = it }
        timer.restart()
    }

    fun refreshPr(force: Boolean = false) {
-        if (project.isDisposed) return
+        if (project.isDisposed || refs == 0) return
        val now = timers.now()
        if (!force && now - lastPr < PR_THROTTLE) return
        lastPr = now
@@ -81,6 +81,7 @@ class WorktreeStatusService internal constructor(
        debounce?.stop()
        statsTimer?.stop()
        prTimer?.stop()
+        debounce = null
        statsTimer = null
        prTimer = null
    }
@@ -108,7 +109,12 @@ class WorktreeStatusService internal constructor(
    }

    private fun notify(value: GhAvailability) {
-        if (notified || value == GhAvailability.OK) return
+        // Re-arm once gh recovers so a later regression notifies again instead of staying silent.
+        if (value == GhAvailability.OK) {
+            notified = false
+            return
+        }
+        if (notified) return
        notified = true
        if (value == GhAvailability.MISSING) {
            KiloNotifications.suggestion(
@@ -60,6 +60,10 @@ class KiloSessionService internal constructor(
        private val LOG = KiloLog.create(KiloSessionService::class.java)
    }

+    // Reflects the sessions from the most recent tracking [list]/[renameSession] call, which is
+    // scoped to a single directory. It is NOT a per-workspace source of truth: a caller listing a
+    // different directory (e.g. an Agent Manager worktree tab) overwrites it. Directory-scoped
+    // callers must consume the return value of [list]/[sessionsFor], never this flow.
    private val _sessions = MutableStateFlow<List<SessionDto>>(emptyList())
    val sessions: StateFlow<List<SessionDto>> = _sessions.asStateFlow()

@@ -108,6 +112,13 @@ class KiloSessionService internal constructor(
        return result
    }

+    /**
+     * List sessions for [dir] without touching the shared [sessions] flow. Use this for
+     * directory-scoped views (e.g. Agent Manager worktree tabs) that maintain their own model, so a
+     * background refresh does not clobber the primary workspace's [sessions] snapshot.
+     */
+    suspend fun sessionsFor(dir: String): SessionListDto = call { list(dir) }
+
    /** Load recent sessions for the current worktree family. */
    suspend fun recent(dir: String, limit: Int): List<SessionDto> =
        call { recent(dir, limit) }.sessions
@@ -0,0 +1,21 @@
+package ai.kilocode.client.util
+
+import com.intellij.ide.ui.LafManagerListener
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.application.ApplicationManager
+import javax.swing.JComponent
+import javax.swing.SwingUtilities
+
+/**
+ * Refresh [roots] on Look-and-Feel changes. The subscription is tied to [parent], so it is removed
+ * when the owning component is disposed. Replaces the copy/pasted `LafManagerListener` blocks that
+ * each panel used to carry.
+ */
+internal fun bindTheme(parent: Disposable, vararg roots: JComponent) {
+    val bus = ApplicationManager.getApplication().messageBus.connect(parent)
+    bus.subscribe(LafManagerListener.TOPIC, LafManagerListener {
+        ApplicationManager.getApplication().invokeLater {
+            roots.forEach { SwingUtilities.updateComponentTreeUI(it) }
+        }
+    })
+}
@@ -0,0 +1,152 @@
+package ai.kilocode.client.agentManager.worktree
+
+import ai.kilocode.client.app.KiloSessionService
+import ai.kilocode.client.testing.FakeSessionRpcApi
+import ai.kilocode.client.testing.TestCoroutines
+import ai.kilocode.client.util.edtWait
+import ai.kilocode.rpc.dto.SessionDto
+import ai.kilocode.rpc.dto.SessionTimeDto
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.testFramework.fixtures.BasePlatformTestCase
+import com.intellij.util.ui.UIUtil
+
+@Suppress("UnstableApiUsage")
+class WorktreeSessionListControllerTest : BasePlatformTestCase() {
+    private lateinit var coroutines: TestCoroutines
+    private lateinit var rpc: FakeSessionRpcApi
+    private lateinit var sessions: KiloSessionService
+    private lateinit var controller: WorktreeSessionListController
+
+    private val dir = "/repo/.kilo/worktrees/feature-x"
+
+    override fun setUp() {
+        super.setUp()
+        coroutines = TestCoroutines()
+        rpc = FakeSessionRpcApi()
+        sessions = KiloSessionService(project, coroutines.scope, rpc)
+        controller = WorktreeSessionListController(sessions, dir, coroutines.scope, telemetry = { _, _ -> })
+    }
+
+    override fun tearDown() {
+        try {
+            coroutines.close(::pump)
+        } finally {
+            super.tearDown()
+        }
+    }
+
+    fun `test reload populates the model from the directory listing`() {
+        rpc.listed += session("ses_1", "One")
+        rpc.listed += session("ses_2", "Two")
+
+        controller.reload()
+        drain()
+
+        assertEquals(2, controller.model.size)
+        assertEquals(setOf("ses_1", "ses_2"), controller.sessions().map { it.id }.toSet())
+    }
+
+    fun `test reload does not clobber the shared sessions flow`() {
+        // The shared flow reflects a different directory; a worktree reload must not overwrite it.
+        rpc.listed += session("ses_main", "Main")
+        kotlinx.coroutines.runBlocking(kotlinx.coroutines.Dispatchers.Default) { sessions.list("/repo") }
+        assertEquals(listOf("ses_main"), sessions.sessions.value.map { it.id })
+
+        rpc.listed.clear()
+        rpc.listed += session("ses_wt", "Worktree")
+        controller.reload()
+        drain()
+
+        assertEquals(listOf("ses_wt"), controller.sessions().map { it.id })
+        assertEquals(listOf("ses_main"), sessions.sessions.value.map { it.id })
+    }
+
+    fun `test create prepends the new session and keeps existing rows`() {
+        rpc.listed += session("existing", "Existing")
+        controller.reload()
+        drain()
+
+        var created: SessionDto? = null
+        controller.create { created = it }
+        drain()
+
+        assertEquals("ses_test", created?.id)
+        assertEquals("ses_test", controller.model.getElementAt(0).id)
+        assertTrue(controller.sessions().any { it.id == "existing" })
+    }
+
+    fun `test delete removes the session and reports success`() {
+        rpc.listed += session("ses_1", "One")
+        rpc.listed += session("ses_2", "Two")
+        controller.reload()
+        drain()
+
+        var ok: Boolean? = null
+        controller.delete("ses_1") { success, _ -> ok = success }
+        drain()
+
+        assertEquals(true, ok)
+        assertNull(controller.session("ses_1"))
+        assertNotNull(controller.session("ses_2"))
+    }
+
+    fun `test delete failure reports the error and reloads`() {
+        rpc.listed += session("ses_1", "One")
+        controller.reload()
+        drain()
+        rpc.deleteThrows = RuntimeException("boom")
+
+        var err: String? = null
+        var ok: Boolean? = null
+        controller.delete("ses_1") { success, message -> ok = success; err = message }
+        drain()
+
+        assertEquals(false, ok)
+        assertEquals("boom", err)
+        // The failed delete never removed ses_1 from the backend, so the reload restores it.
+        assertNotNull(controller.session("ses_1"))
+    }
+
+    fun `test rename applies optimistically and keeps the server title on success`() {
+        rpc.listed += session("ses_1", "Old")
+        controller.reload()
+        drain()
+
+        edtWait { controller.rename("ses_1", "New") { _, _ -> } }
+        // Optimistic title is visible before the RPC resolves.
+        assertEquals("New", controller.session("ses_1")?.title)
+
+        drain()
+        assertEquals("New", controller.session("ses_1")?.title)
+        assertEquals(Triple("ses_1", dir, "New"), rpc.renames.single())
+    }
+
+    fun `test rename failure reverts to the prior title`() {
+        rpc.listed += session("ses_1", "Old")
+        controller.reload()
+        drain()
+        rpc.renameThrows = RuntimeException("nope")
+
+        var err: String? = null
+        edtWait { controller.rename("ses_1", "New") { _, message -> err = message } }
+        drain()
+
+        assertEquals("nope", err)
+        assertEquals("Old", controller.session("ses_1")?.title)
+    }
+
+    private fun session(id: String, title: String) = SessionDto(
+        id = id,
+        projectID = "prj",
+        directory = dir,
+        title = title,
+        version = "1",
+        time = SessionTimeDto(created = 1.0, updated = 2.0),
+    )
+
+    private fun drain() = coroutines.drain(::pump)
+
+    private fun pump() {
+        ApplicationManager.getApplication().invokeAndWait { UIUtil.dispatchAllInvocationEvents() }
+    }
+}
@@ -0,0 +1,130 @@
+package ai.kilocode.client.agentManager.worktree
+
+import ai.kilocode.client.testing.FakeWorktreeRpcApi
+import ai.kilocode.client.testing.TestCoroutines
+import ai.kilocode.client.testing.TestUiTimers
+import ai.kilocode.rpc.dto.GhAvailability
+import ai.kilocode.rpc.dto.GhState
+import ai.kilocode.rpc.dto.WorktreePrDto
+import ai.kilocode.rpc.dto.WorktreePrListDto
+import ai.kilocode.rpc.dto.WorktreeStatsDto
+import ai.kilocode.rpc.dto.WorktreeStatsListDto
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.testFramework.fixtures.BasePlatformTestCase
+import com.intellij.testFramework.replaceService
+import com.intellij.util.ui.UIUtil
+
+@Suppress("UnstableApiUsage")
+class WorktreeStatusServiceTest : BasePlatformTestCase() {
+    private lateinit var coroutines: TestCoroutines
+    private lateinit var rpc: FakeWorktreeRpcApi
+    private lateinit var timers: TestUiTimers
+    private lateinit var service: WorktreeStatusService
+
+    override fun setUp() {
+        super.setUp()
+        coroutines = TestCoroutines()
+        rpc = FakeWorktreeRpcApi()
+        ApplicationManager.getApplication()
+            .replaceService(KiloWorktreeService::class.java, KiloWorktreeService(coroutines.scope, rpc), testRootDisposable)
+        timers = TestUiTimers()
+        service = WorktreeStatusService(project, coroutines.scope, timers)
+    }
+
+    override fun tearDown() {
+        try {
+            coroutines.close(::pump)
+        } finally {
+            super.tearDown()
+        }
+    }
+
+    fun `test attach loads pr immediately and stats after debounce`() {
+        val path = "${project.basePath}/.kilo/worktrees/feature-x"
+        rpc.statsResult = WorktreeStatsListDto(listOf(WorktreeStatsDto(path, additions = 4)))
+        rpc.prResult = WorktreePrListDto(GhAvailability.OK, listOf(WorktreePrDto(path, 3, GhState.OPEN, "https://pr/3")))
+        val key = normalizeWorktreePath(path)
+
+        val handle = service.attach()
+        drain()
+        assertEquals(3, service.pr.value[key]?.number)
+
+        timers.advanceBy(300)
+        drain()
+        assertEquals(4, service.stats.value[key]?.additions)
+        handle.close()
+    }
+
+    fun `test refresh is ignored after the last handle closes`() {
+        val path = "${project.basePath}/.kilo/worktrees/feature-x"
+        val key = normalizeWorktreePath(path)
+        rpc.statsResult = WorktreeStatsListDto(listOf(WorktreeStatsDto(path, additions = 1)))
+        val handle = service.attach()
+        timers.advanceBy(300)
+        drain()
+        assertEquals(1, service.stats.value[key]?.additions)
+
+        handle.close()
+        rpc.statsResult = WorktreeStatsListDto(listOf(WorktreeStatsDto(path, additions = 5)))
+        service.refreshStats()
+        timers.advanceBy(300)
+        drain()
+
+        assertEquals(1, service.stats.value[key]?.additions)
+    }
+
+    fun `test polling keeps loading while a handle remains`() {
+        val path = "${project.basePath}/.kilo/worktrees/feature-x"
+        val key = normalizeWorktreePath(path)
+        rpc.statsResult = WorktreeStatsListDto(listOf(WorktreeStatsDto(path, additions = 1)))
+        val first = service.attach()
+        val second = service.attach()
+        timers.advanceBy(300)
+        drain()
+        assertEquals(1, service.stats.value[key]?.additions)
+
+        first.close()
+        rpc.statsResult = WorktreeStatsListDto(listOf(WorktreeStatsDto(path, additions = 7)))
+        // The 30s poll fires refreshStats, which reschedules the 300ms debounce.
+        timers.advanceBy(30_000)
+        timers.advanceBy(300)
+        drain()
+
+        assertEquals(7, service.stats.value[key]?.additions)
+        second.close()
+    }
+
+    fun `test pr refresh throttles non-forced calls but honors force`() {
+        val path = "${project.basePath}/.kilo/worktrees/feature-x"
+        val key = normalizeWorktreePath(path)
+        rpc.prResult = WorktreePrListDto(GhAvailability.OK, listOf(WorktreePrDto(path, 1, GhState.OPEN, "https://pr/1")))
+        val handle = service.attach()
+        drain()
+        assertEquals(1, service.pr.value[key]?.number)
+
+        rpc.prResult = WorktreePrListDto(GhAvailability.OK, listOf(WorktreePrDto(path, 2, GhState.OPEN, "https://pr/2")))
+        service.refreshPr()
+        drain()
+        assertEquals(1, service.pr.value[key]?.number)
+
+        service.refreshPr(force = true)
+        drain()
+        assertEquals(2, service.pr.value[key]?.number)
+        handle.close()
+    }
+
+    fun `test gh availability propagates from pr status`() {
+        rpc.prResult = WorktreePrListDto(GhAvailability.MISSING)
+        val handle = service.attach()
+        drain()
+
+        assertEquals(GhAvailability.MISSING, service.gh.value)
+        handle.close()
+    }
+
+    private fun drain() = coroutines.drain(::pump)
+
+    private fun pump() {
+        ApplicationManager.getApplication().invokeAndWait { UIUtil.dispatchAllInvocationEvents() }
+    }
+}
@@ -0,0 +1,71 @@
+package ai.kilocode.client.vfs
+
+import ai.kilocode.client.util.edtWait
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.components.service
+import com.intellij.openapi.fileEditor.FileEditorManager
+import com.intellij.openapi.project.Project
+import com.intellij.testFramework.fixtures.BasePlatformTestCase
+import com.intellij.util.ui.components.BorderLayoutPanel
+import javax.swing.JComponent
+
+@Suppress("UnstableApiUsage")
+class KiloVfsManagerTest : BasePlatformTestCase() {
+    private val kind = "test-close"
+
+    override fun setUp() {
+        super.setUp()
+        service<KiloEditorKindRegistry>().register(TestKind)
+    }
+
+    override fun tearDown() {
+        try {
+            service<KiloEditorKindRegistry>().unregister(kind)
+            KiloVirtualFileSystem.getInstance().clear()
+        } finally {
+            super.tearDown()
+        }
+    }
+
+    fun `test close closes the open editor and releases the cache`() {
+        val params = mapOf("path" to "/repo/wt")
+        val vfs = project.service<KiloVfsManager>()
+        val manager = FileEditorManager.getInstance(project)
+
+        edtWait { assertTrue(vfs.open(kind, params)) }
+        assertTrue(manager.openFiles.any { it is KiloVirtualFile })
+
+        edtWait { vfs.close(kind, params) }
+
+        assertTrue(manager.openFiles.none { it is KiloVirtualFile })
+        assertNull(KiloVirtualFileSystem.getInstance().cached(KiloPath(kind, params)))
+    }
+
+    fun `test close matches by canonical params after the cache was released`() {
+        // Two params so canonicalization sorts them; open and close use different insertion orders.
+        val opened = linkedMapOf("path" to "/repo/wt", "extra" to "1")
+        val closed = linkedMapOf("extra" to "1", "path" to "/repo/wt")
+        val vfs = project.service<KiloVfsManager>()
+        val manager = FileEditorManager.getInstance(project)
+
+        edtWait { assertTrue(vfs.open(kind, opened)) }
+        assertTrue(manager.openFiles.any { it is KiloVirtualFile })
+
+        // Drop the VFS cache entry so close() must fall back to scanning the open editors.
+        KiloVirtualFileSystem.getInstance().release(KiloPath(kind, opened))
+        assertNull(KiloVirtualFileSystem.getInstance().cached(KiloPath(kind, opened)))
+
+        edtWait { vfs.close(kind, closed) }
+
+        assertTrue(manager.openFiles.none { it is KiloVirtualFile })
+    }
+
+    private object TestKind : KiloEditorKind {
+        override val id: String = "test-close"
+
+        override fun title(params: Map<String, String>): String = "Test"
+
+        override fun createContent(project: Project, file: KiloVirtualFile, parent: Disposable): JComponent =
+            BorderLayoutPanel()
+    }
+}