fix(jetbrains): improve diff performance and session layout

.changeset/jetbrains-attachment-button-sync.md

@@ -0,0 +1,5 @@
+---
+"@kilocode/kilo-jetbrains": patch
+---
+
+Keep the JetBrains prompt send/stop button in sync when attachments are added or removed while a session is busy.

.changeset/jetbrains-diff-view-fixes.md

@@ -0,0 +1,5 @@
+---
+"@kilocode/kilo-jetbrains": patch
+---
+
+Fix JetBrains diff views to show compact workspace-relative file paths and keep added-file content visible in large branch diffs.

.changeset/jetbrains-session-load-crop.md

@@ -0,0 +1,5 @@
+---
+"@kilocode/kilo-jetbrains": patch
+---
+
+Fix JetBrains chat transcripts rendering cropped when opening existing sessions.

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

@@ -81,6 +81,7 @@ class KiloWorkspaceRpcApiImpl internal constructor(
         private val GLOBAL = MODERN + LEGACY + "config.json"
         private val LOCAL_DIRS = listOf(".kilo", ".kilocode", ".opencode")
         private const val DIFF_CAP = 200_000
+        private const val BRANCH_DIFF_CAP = 8 * 1024 * 1024
         private const val LARGE_FILE = 2 * 1024 * 1024L
         private val JSON = Json { ignoreUnknownKeys = true }
         private val CONFIG = """{
@@ -271,9 +272,9 @@ class KiloWorkspaceRpcApiImpl internal constructor(
         val files = stats.map { DiffFileDto(it.path, it.additions, it.deletions, "", status[it.path] ?: "modified") } +
             untrackedPaths.map { untracked(base, it, withPatch = false) }
         if (!patches) return@withContext files
-        // Fetch patches lazily and stop once the running total reaches DIFF_CAP, so a branch with
+        // Fetch patches lazily and stop once the running total reaches BRANCH_DIFF_CAP, so a branch with
         // hundreds of changed files doesn't spawn a git subprocess (or read a file) per entry.
-        capDiff(files, DIFF_CAP) { file ->
+        capDiff(files, BRANCH_DIFF_CAP) { file ->
             if (file.status == "untracked") untracked(base, file.file, withPatch = true).patch.orEmpty()
             else fileDiff(base, anc, file.file)
         }

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

@@ -75,6 +75,16 @@ class BranchDiffTest {
         assertEquals("untracked", diff[1].status)
     }
 
+    @Test
+    fun `capDiff keeps all patches under a large budget`() {
+        val files = (1..110).map { i -> stat("src/File$i.kt") }
+
+        val diff = capDiff(files, cap = 8 * 1024 * 1024) { file -> "patch-${file.file}\n".repeat(30) }
+
+        assertEquals(files.map { it.file }, diff.map { it.file })
+        assertEquals(files.map { file -> "patch-${file.file}\n".repeat(30) }, diff.map { it.patch })
+    }
+
     @Test
     fun `parses git numstat output`() {
         val stats = parseNumstat("1\t2\tsrc/A.kt\n0\t3\tsrc/B.kt\n-\t-\tbin.png\n")

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

@@ -8,6 +8,7 @@ import com.intellij.diff.requests.SimpleDiffRequest
 import com.intellij.diff.util.DiffUserDataKeys
 import com.intellij.openapi.fileTypes.FileTypeManager
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.vcs.FileStatus
 
 internal fun diffRequest(
     project: Project,
@@ -18,15 +19,20 @@ internal fun diffRequest(
     val sides = DiffPatchReconstruct.sides(dto)
     val type = FileTypeManager.getInstance().getFileTypeByFileName(dto.file)
     val factory = DiffContentFactory.getInstance()
+    val status = fileStatus(dto)
+    val patch = dto.patch?.takeIf { it.isNotBlank() }
+    val fallback = patch ?: KiloBundle.message("diff.editor.patch.unavailable")
     val left = when {
         DiffPatchReconstruct.added(dto.patch) -> factory.createEmpty()
         sides.renderable -> factory.create(project, sides.before, type)
+        status == FileStatus.DELETED -> factory.create(project, fallback, type)
         else -> factory.createEmpty()
     }
     val right = when {
         DiffPatchReconstruct.deleted(dto.patch) -> factory.createEmpty()
         sides.renderable -> factory.create(project, sides.after, type)
-        else -> factory.create(project, dto.patch ?: KiloBundle.message("diff.editor.patch.unavailable"), type)
+        status == FileStatus.DELETED -> factory.createEmpty()
+        else -> factory.create(project, fallback, type)
     }
     return SimpleDiffRequest(diffTitle(dto.file, branch), left, right, labels.first, labels.second).also {
         it.putUserData(DiffUserDataKeys.FORCE_READ_ONLY, true)

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

@@ -49,6 +49,7 @@ import com.intellij.ui.components.JBScrollPane
 import com.intellij.ui.treeStructure.Tree
 import com.intellij.util.concurrency.annotations.RequiresEdt
 import com.intellij.util.ui.JBUI
+import com.intellij.util.ui.tree.TreeUtil
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.Job
@@ -67,6 +68,8 @@ import javax.swing.JPanel
 import javax.swing.ScrollPaneConstants
 import javax.swing.JViewport
 import javax.swing.JTree
+import javax.swing.event.TreeExpansionEvent
+import javax.swing.event.TreeExpansionListener
 import javax.swing.tree.DefaultMutableTreeNode
 import javax.swing.tree.DefaultTreeModel
 import javax.swing.tree.TreeCellRenderer
@@ -98,10 +101,11 @@ internal class DiffEditorView(
     private val load: ((DiffEditorData) -> Unit) -> Job,
     private val replace: (DiffEditorData) -> Unit,
 ) : Disposable {
+    private val start = normalize(initial)
     private val disposed = AtomicBoolean(false)
     private val outdated = AtomicBoolean(false)
     private val refreshing = AtomicBoolean(false)
-    private val tree = buildFileTree(initial)
+    private val tree = buildFileTree(start)
     private val badge = DiffStatBadge(0, 0, inset = UiStyle.Gap.pad())
     private val splitter = OnePixelSplitter(false, 0.25f)
     private val select = Debouncer<Int>(scope, parent) { show(it) }
@@ -114,12 +118,12 @@ internal class DiffEditorView(
         add(banner, BorderLayout.NORTH)
         add(splitter, BorderLayout.CENTER)
     }
-    private var files = initial
+    private var files = start
     private var branch = branch
     private var syncing = false
-    private var requested: String? = initial.firstOrNull()?.file
+    private var requested: String? = start.firstOrNull()?.file
     private var refreshJob: Job? = null
-    private var processor = processor(initial, selected(initial.firstOrNull()?.file))
+    private var processor = processor(start, selected(start.firstOrNull()?.file))
     private val openFileAction = object : DumbAwareAction(
         KiloBundle.message("diff.editor.openFile"),
         KiloBundle.message("diff.editor.openFile"),
@@ -155,11 +159,11 @@ internal class DiffEditorView(
         // disposes the old processor on each refresh, and registering under parent would leak a
         // removal hook (holding the dead processor) for every refresh across the editor's lifetime.
         processor.addListener(DiffRequestProcessorListener { syncTree() }, processor)
-        splitter.firstComponent = buildTreePanel(tree, initial, badge, processor.component, ::refresh)
+        splitter.firstComponent = buildTreePanel(tree, start, badge, processor.component, ::refresh)
         splitter.secondComponent = processor.component
         processor.updateRequest()
-        applyBadge(initial)
-        select(initial.firstOrNull()?.file)
+        applyBadge(start)
+        select(start.firstOrNull()?.file)
         listen()
     }
 
@@ -174,24 +178,25 @@ internal class DiffEditorView(
 
     @RequiresEdt
     fun applyFiles(next: List<DiffFileDto>, nextBranch: String? = branch) {
-        if (same(files, next) && branch == nextBranch) return
+        val items = normalize(next)
+        if (same(files, items) && branch == nextBranch) return
         val path = selectedFile()?.file ?: activePath() ?: files.firstOrNull()?.file
-        val index = selected(path, next)
+        val index = selected(path, items)
         val old = processor
-        files = next
+        files = items
         branch = nextBranch
-        requested = next.getOrNull(index)?.file
-        tree.model = buildFileModel(next)
+        requested = items.getOrNull(index)?.file
+        tree.model = buildFileModel(items)
         expandAll(tree)
-        processor = processor(next, index)
+        processor = processor(items, index)
         Disposer.register(parent, processor)
         processor.addListener(DiffRequestProcessorListener { syncTree() }, processor)
-        splitter.firstComponent = buildTreePanel(tree, next, badge, processor.component, ::refresh)
+        splitter.firstComponent = buildTreePanel(tree, items, badge, processor.component, ::refresh)
         splitter.secondComponent = processor.component
         processor.updateRequest()
         Disposer.dispose(old)
-        applyBadge(next)
-        select(next.getOrNull(index)?.file)
+        applyBadge(items)
+        select(items.getOrNull(index)?.file)
         root.revalidate()
         root.repaint()
     }
@@ -360,6 +365,22 @@ internal class DiffEditorView(
 
     private fun same(a: List<DiffFileDto>, b: List<DiffFileDto>): Boolean = a == b
 
+    private fun normalize(files: List<DiffFileDto>): List<DiffFileDto> = files.map { file -> file.copy(file = display(file.file)) }
+
+    private fun display(file: String): String {
+        val dir = params["directory"] ?: return file
+        val root = clean(dir) ?: return file
+        return try {
+            val raw = Path.of(file)
+            if (!raw.isAbsolute) return file
+            val path = raw.normalize()
+            if (!path.startsWith(root)) return file
+            root.relativize(path).toString().replace('\\', '/')
+        } catch (_: InvalidPathException) {
+            file
+        }
+    }
+
     private fun clean(dir: String): Path? = try {
         Path.of(dir).normalize()
     } catch (_: InvalidPathException) {
@@ -421,17 +442,51 @@ private fun buildFileTree(files: List<DiffFileDto>): Tree {
         val node = path.lastPathComponent as? DefaultMutableTreeNode
         (node?.userObject as? Node)?.name.orEmpty()
     }
+    // A folder row hides its rolled-up badge while expanded, so its preferred width depends on
+    // expansion state. JTree only invalidates cached path bounds on model changes, not on
+    // expand/collapse, so a collapsed folder would keep its narrower expanded-state bounds and the
+    // re-shown badge would squeeze the name until an unrelated re-measure. Invalidate the layout
+    // cache on a user toggle so the row re-measures. invalidateCacheAndRepaint is UI-scoped (whole
+    // tree), so [bulkToggle] suppresses this during expand/collapse-all and invalidates once at the
+    // end — otherwise a bulk op would re-measure the whole tree per row. Registered after the initial
+    // expandAll (already fully expanded, nothing to re-measure).
     expandAll(tree)
+    tree.addTreeExpansionListener(object : TreeExpansionListener {
+        override fun treeExpanded(event: TreeExpansionEvent) = onToggle()
+        override fun treeCollapsed(event: TreeExpansionEvent) = onToggle()
+        private fun onToggle() { if (!tree.bulk) TreeUtil.invalidateCacheAndRepaint(tree.ui) }
+    })
     return tree
 }
 
 private fun buildFileModel(files: List<DiffFileDto>): DefaultTreeModel {
     val root = DefaultMutableTreeNode(Node("", "", true, null))
     for (file in files) addFile(root, file)
+    compact(root)
     updateStats(root)
     return DefaultTreeModel(root)
 }
 
+// Collapse chains of single-child directories into one node (e.g. "pkg/ui/list") so the tree
+// doesn't nest through directories that never branch, mirroring the IDE's compact directories.
+private fun compact(node: DefaultMutableTreeNode) {
+    val item = node.userObject as? Node
+    if (item != null && item.file == null && item.path.isNotEmpty()) {
+        while (node.childCount == 1) {
+            val parent = node.userObject as? Node ?: break
+            val child = node.getChildAt(0) as? DefaultMutableTreeNode ?: break
+            val kid = child.userObject as? Node ?: break
+            if (kid.file != null) break
+            node.userObject = Node("${parent.name}/${kid.name}", kid.path, true, null)
+            node.removeAllChildren()
+            while (child.childCount > 0) node.add(child.getChildAt(0) as DefaultMutableTreeNode)
+        }
+    }
+    for (i in 0 until node.childCount) {
+        compact(node.getChildAt(i) as? DefaultMutableTreeNode ?: continue)
+    }
+}
+
 private fun buildTreePanel(tree: Tree, files: List<DiffFileDto>, badge: DiffStatBadge, target: JComponent, refresh: () -> Unit): JComponent {
     val toolbar = ActionManager.getInstance().createActionToolbar(
         ActionPlaces.TOOLBAR,
@@ -464,7 +519,7 @@ private fun buildTreePanel(tree: Tree, files: List<DiffFileDto>, badge: DiffStat
             JBScrollPane(tree).apply {
                 border = JBUI.Borders.empty()
                 viewportBorder = JBUI.Borders.empty()
-                horizontalScrollBarPolicy = ScrollPaneConstants.HORIZONTAL_SCROLLBAR_NEVER
+                horizontalScrollBarPolicy = ScrollPaneConstants.HORIZONTAL_SCROLLBAR_AS_NEEDED
             },
             BorderLayout.CENTER,
         )
@@ -483,18 +538,34 @@ private fun fileCount(count: Int): String = KiloBundle.message(
     count,
 )
 
-private fun expandAll(tree: Tree) {
+private fun expandAll(tree: Tree) = bulkToggle(tree) {
     var i = 0
     while (i < tree.rowCount) {
         tree.expandRow(i)
         i += 1
     }
 }
 
-private fun collapseAll(tree: Tree) {
+private fun collapseAll(tree: Tree) = bulkToggle(tree) {
     for (i in tree.rowCount - 1 downTo 0) tree.collapseRow(i)
 }
 
+/**
+ * Run a bulk expand/collapse without firing the per-row layout-cache invalidation. Each toggle would
+ * otherwise invalidate the whole tree (invalidateCacheAndRepaint is UI-scoped), making a bulk op
+ * O(rows^2) to re-measure. Suppress the toggle listener for the loop and invalidate once at the end.
+ */
+private fun bulkToggle(tree: Tree, action: () -> Unit) {
+    val diff = tree as? DiffTree
+    diff?.bulk = true
+    try {
+        action()
+    } finally {
+        diff?.bulk = false
+    }
+    TreeUtil.invalidateCacheAndRepaint(tree.ui)
+}
+
 private fun addFile(root: DefaultMutableTreeNode, file: DiffFileDto) {
     var node = root
     val parts = file.file.split('/').filter { it.isNotBlank() }
@@ -549,6 +620,9 @@ private class Node(val name: String, val path: String, val dir: Boolean, val fil
 }
 
 private class DiffTree(model: TreeModel) : Tree(model) {
+    /** Set while a bulk expand/collapse runs so the toggle listener skips its per-row invalidation. */
+    var bulk = false
+
     override fun getBackground(): Color = JBUI.CurrentTheme.ToolWindow.background()
 
     override fun getScrollableTracksViewportHeight(): Boolean {
@@ -584,9 +658,12 @@ private class Renderer : JPanel(BorderLayout()), TreeCellRenderer {
         val name = item?.name?.ifBlank { item.path }.orEmpty()
         val color = item?.file?.let(::fileStatus)?.color
         if (color == null) text.append(name) else text.append(name, SimpleTextAttributes(SimpleTextAttributes.STYLE_PLAIN, color))
-        val changed = item != null && (item.additions != 0 || item.deletions != 0)
-        badge.isVisible = changed
-        if (changed) badge.update(item.additions, item.deletions)
+        // A folder's badge rolls up its descendants' stats, which is only meaningful while the
+        // folder is collapsed. Once expanded the child rows carry their own badges, so hide the
+        // folder aggregate to avoid duplicating the numbers. Leaf files always show their badge.
+        val show = item != null && (item.additions != 0 || item.deletions != 0) && !(item.dir && expanded)
+        badge.isVisible = show
+        if (show) badge.update(item.additions, item.deletions)
         return this
     }
 }

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

@@ -389,6 +389,7 @@ class SessionUi(
         header = SessionHeaderPanel(controller, this) { openBranchChanges() }
 
         scroll = SessionScroll(root, sessionContent, messageBody, blankBody)
+        messageBody.onReflow = { on -> if (on && !opening) scroll.followTail() }
         scroll.onScroll = {
             overlay.clear()
             popup.hideAll()
@@ -673,6 +674,7 @@ class SessionUi(
         if (width <= 0 || height <= 0) return
         if (body(controller.model.state) !== messageBody) return
         pending = false
+        messageBody.reflow()
         scroll.openBottom {
             opening = false
         }

packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/ModifiedFilesView.kt

@@ -53,6 +53,7 @@ class ModifiedFilesView private constructor(
 
     init {
         body.parent = this
+        body.overflow = ::openDiffViewer
         parts.diff.addActionListener { openDiffViewer() }
         isVisible = false
         bindHeader(parts.glyph, parts.title, parts.count, parts.panel.left, parts.bars, parts.anchor)
@@ -149,7 +150,10 @@ class ModifiedFilesView private constructor(
     @RequiresEdt
     private fun buildPopup(files: List<EditFileChange>): HeaderPopupBody {
         val owner = Disposer.newDisposable("Modified files popup body")
-        val popup = PatchBody(selection, openFile, POPUP_OPTS).also { it.parent = owner }
+        val popup = PatchBody(selection, openFile, POPUP_OPTS).also {
+            it.parent = owner
+            it.overflow = ::openDiffViewer
+        }
         val panel = popup.mountFiles(files)
         popup.applyStyle(style)
         return HeaderPopupBody(panel, owner, style.editorBackground, SessionUiStyle.View.Popup.WIDE_MAX_WIDTH)

packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/SessionLayout.kt

@@ -91,6 +91,11 @@ class SessionLayout(
         cache.remove(comp)
     }
 
+    /** Drop every cached measurement so the next layout pass re-measures all children. */
+    fun forgetAll() {
+        cache.clear()
+    }
+
     private fun measure(comp: Component, width: Int): Int {
         val hit = cache[comp]
         if (comp.isValid && hit?.width == width) return hit.height