feat(jetbrains): auto-collapse and preview reasoning blocks

Auto-collapse streaming reasoning when it finishes unless the user pinned the state, releasing the inner markdown editors to avoid leaks. Keep the in-place expanded body capped at 5 lines and add a hover preview popup capped at 10 lines for collapsed reasoning.

.changeset/jetbrains-reasoning-view.md

@@ -0,0 +1,5 @@
+---
+"@kilocode/kilo-jetbrains": patch
+---
+
+Auto-collapse JetBrains reasoning blocks when they finish streaming, keep manual expand/collapse choices, and preview collapsed reasoning on hover.

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

@@ -108,6 +108,7 @@ object SessionUiStyle {
         /** Reasoning block preview sizing. */
         object Reasoning {
             const val BODY_LINES = 5
+            const val POPUP_LINES = 10
             const val HEADER_VERTICAL_PADDING = 5
             const val BODY_VERTICAL_PADDING = 4
             const val BODY_HORIZONTAL_PADDING = 8

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

@@ -7,10 +7,13 @@ import ai.kilocode.client.session.SessionFileOpener
 import ai.kilocode.client.session.openSessionLink
 import ai.kilocode.client.session.model.Content
 import ai.kilocode.client.session.model.Reasoning
+import ai.kilocode.client.session.ui.popup.HeaderPopupBody
+import ai.kilocode.client.session.ui.popup.HeaderPopupRequest
 import ai.kilocode.client.session.ui.style.SessionEditorStyle
 import ai.kilocode.client.session.ui.selection.SessionSelection
 import ai.kilocode.client.session.ui.style.SessionUiStyle
 import ai.kilocode.client.session.views.base.SecondarySessionPartView
+import ai.kilocode.client.telemetry.Telemetry
 import ai.kilocode.client.ui.UiStyle
 import ai.kilocode.client.ui.md.MdView
 import ai.kilocode.client.ui.md.MdViewFactory
@@ -20,9 +23,11 @@ import com.intellij.ui.components.JBScrollPane
 import com.intellij.util.concurrency.annotations.RequiresEdt
 import com.intellij.util.ui.JBUI
 import java.awt.BorderLayout
+import java.awt.Container
 import java.awt.Dimension
 import java.awt.Font
 import java.awt.Rectangle
+import javax.swing.JComponent
 import javax.swing.JPanel
 import javax.swing.ScrollPaneConstants
 import javax.swing.Scrollable
@@ -63,6 +68,7 @@ class ReasoningView(
     private var done = reasoning.done
     private var registered = false
     private var following = false
+    private var pinned = false
 
     init {
         row.border = JBUI.Borders.empty(
@@ -100,6 +106,7 @@ class ReasoningView(
         var changed = false
         val next = content.content.toString()
         val follow = tailVisible()
+        val finishing = !done && content.done
         if (done != content.done) {
             done = content.done
             changed = true
@@ -112,10 +119,30 @@ class ReasoningView(
             }
             changed = true
         }
+        if (finishing && !pinned) {
+            changed = collapse() || changed
+            changed = releaseBody() || changed
+        }
         changed = sync() || changed
         if (changed) refresh()
     }
 
+    /** Detaches and disposes the markdown body so its editors are released when reasoning finishes. */
+    @RequiresEdt
+    private fun releaseBody(): Boolean {
+        if (!parts.bodyCreated()) return false
+        val detached = discardBody()
+        Disposer.dispose(parts.md(openFile, openUrl))
+        parts.reset()
+        registered = false
+        return detached
+    }
+
+    @RequiresEdt
+    override fun userToggled() {
+        pinned = true
+    }
+
     @RequiresEdt
     override fun appendDelta(delta: String) {
         if (delta.isEmpty()) return
@@ -150,6 +177,15 @@ class ReasoningView(
     @RequiresEdt
     internal fun bodyScrollBottom() = parts.scrollOrNull?.verticalScrollBar?.let { it.maximum - it.visibleAmount } ?: 0
 
+    @RequiresEdt
+    override fun headerPopup(): HeaderPopupRequest? {
+        if (isExpanded()) return null
+        val text = source.takeIf { it.isNotBlank() } ?: return null
+        return HeaderPopupRequest(row, build = { buildPopupBody(text) }) {
+            Telemetry.send("Header Popup Shown", mapOf("surface" to "session", "part" to "reasoning"))
+        }
+    }
+
     @RequiresEdt
     override fun applyStyle(style: SessionEditorStyle) {
         this.style = style
@@ -235,6 +271,38 @@ class ReasoningView(
         Disposer.register(this, md)
     }
 
+    @RequiresEdt
+    private fun buildPopupBody(text: String): HeaderPopupBody {
+        val md = MdViewFactory.create(style, null).apply {
+            addLinkListener { openSessionLink(it, openFile, openUrl) }
+        }
+        md.applyStyle(style)
+        md.font = style.smallEditorFont.deriveFont(Font.ITALIC)
+        md.codeFont = style.editorFamily
+        md.foreground = UiStyle.Colors.weak()
+        md.background = style.editorBackground
+        md.component.border = JBUI.Borders.empty()
+        md.set(text)
+        val panel = TrackPanel().apply {
+            isOpaque = true
+            background = style.editorBackground
+            border = JBUI.Borders.empty(
+                JBUI.scale(SessionUiStyle.View.Reasoning.BODY_VERTICAL_PADDING),
+                JBUI.scale(SessionUiStyle.View.Reasoning.BODY_HORIZONTAL_PADDING),
+            )
+            add(md.component, BorderLayout.CENTER)
+        }
+        val scroll = JBScrollPane(panel).apply {
+            border = JBUI.Borders.empty()
+            isOpaque = true
+            background = style.editorBackground
+            viewport.background = style.editorBackground
+            horizontalScrollBarPolicy = ScrollPaneConstants.HORIZONTAL_SCROLLBAR_NEVER
+            verticalScrollBarPolicy = ScrollPaneConstants.VERTICAL_SCROLLBAR_AS_NEEDED
+        }
+        return HeaderPopupBody(ReasoningPopupPanel(scroll, md), md, style.editorBackground)
+    }
+
     private fun bodyMaxHeight(): Int {
         if (!parts.bodyCreated()) return 0
         val md = md
@@ -280,6 +348,10 @@ class ReasoningParts(
 
     fun bodyCreated() = body != null
 
+    fun reset() {
+        body = null
+    }
+
     fun md(openFile: SessionFileOpener, openUrl: (String) -> Unit): MdView = body(openFile, openUrl).md
 
     fun scroll(openFile: SessionFileOpener, openUrl: (String) -> Unit): JBScrollPane = body(openFile, openUrl).scroll
@@ -344,3 +416,38 @@ class TrackPanel : JPanel(BorderLayout()), Scrollable {
         direction: Int,
     ) = visibleRect.height
 }
+
+private class ReasoningPopupPanel(
+    child: JComponent,
+    private val md: MdView,
+) : JPanel(BorderLayout()) {
+    init {
+        // Transparent so the balloon fill (editor background) shows uniformly behind the content.
+        isOpaque = false
+        add(child, BorderLayout.CENTER)
+    }
+
+    override fun getPreferredSize(): Dimension {
+        val size = super.getPreferredSize()
+        val cap = JBUI.scale(SessionUiStyle.View.Popup.MAX_WIDTH)
+        val width = maxOf(contentWidth(this), size.width).coerceAtMost(cap)
+        val height = md.component.getFontMetrics(md.font).height * SessionUiStyle.View.Reasoning.POPUP_LINES +
+            JBUI.scale(SessionUiStyle.View.Layout.BODY_EXTRA_HEIGHT)
+        return Dimension(width, minOf(size.height, height))
+    }
+}
+
+private fun contentWidth(root: Container): Int {
+    var max = 0
+    for (child in root.components) {
+        if (child is JBScrollPane) {
+            val view = child.viewport.view as? JComponent
+            val content = view?.preferredSize?.width ?: 0
+            val insets = child.insets
+            val viewport = child.viewportBorder?.getBorderInsets(child) ?: JBUI.emptyInsets()
+            max = maxOf(max, content + insets.left + insets.right + viewport.left + viewport.right)
+        }
+        if (child is Container) max = maxOf(max, contentWidth(child))
+    }
+    return max
+}

packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/views/base/AbstractSessionPartView.kt

@@ -67,10 +67,13 @@ abstract class AbstractSessionPartView(
         if (!expandable || !arrow.isVisible) return
         val changed = toggleLocal()
         if (!changed) return
+        userToggled()
         syncArrow()
         refresh()
     }
 
+    protected open fun userToggled() {}
+
     open fun expand(): Boolean {
         if (!expandable) return false
         if (isExpanded()) return false
@@ -89,6 +92,15 @@ abstract class AbstractSessionPartView(
 
     protected fun bodyComponent(): JComponent = body()
 
+    /** Detaches and forgets the cached body so the next expansion builds a fresh one. */
+    protected fun discardBody(): Boolean {
+        val item = body ?: return false
+        val attached = item.parent === this
+        if (attached) remove(item)
+        body = null
+        return attached
+    }
+
     private fun toggleLocal(): Boolean {
         val fn = resize ?: return toggleBody()
         val expanded = isExpanded()

packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/views/ReasoningViewStressTest.kt

@@ -33,7 +33,12 @@ class ReasoningViewStressTest : BasePlatformTestCase() {
         assertEquals(count, panel(view).componentCount)
 
         view.update(reasoning("r1", done = true, text = view.markdown() + "```"))
-        assertTrue(view.bodyVisible())
+        assertFalse(view.bodyVisible())
+        assertFalse(view.bodyCreated())
+        drainEdt()
+
+        assertEquals(base, EditorFactory.getInstance().allEditors.size)
+
         Disposer.dispose(view)
         drainEdt()
 

packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/views/ReasoningViewTest.kt

@@ -4,6 +4,8 @@ import ai.kilocode.client.session.model.Reasoning
 import ai.kilocode.client.session.ui.style.SessionEditorStyle
 import ai.kilocode.client.session.ui.style.SessionUiStyle
 import ai.kilocode.client.session.views.base.SecondarySessionPartView
+import com.intellij.openapi.editor.EditorFactory
+import com.intellij.openapi.util.Disposer
 import com.intellij.testFramework.fixtures.BasePlatformTestCase
 import com.intellij.ui.components.JBScrollPane
 import com.intellij.util.ui.JBUI
@@ -11,6 +13,7 @@ import com.intellij.util.ui.UIUtil
 import java.awt.Component
 import java.awt.Container
 import javax.swing.Icon
+import javax.swing.JComponent
 import javax.swing.JLabel
 import javax.swing.JPanel
 import javax.swing.ScrollPaneConstants
@@ -59,16 +62,37 @@ class ReasoningViewTest : BasePlatformTestCase() {
         assertEquals("one\ntwo\nthree\nfour", view.markdown())
     }
 
-    fun `test live reasoning stays expanded when marked done`() {
+    fun `test live reasoning auto-collapses and releases body when marked done`() {
         val view = ReasoningView(reasoning("p1", done = false, text = "one\ntwo\nthree\nfour"))
 
         assertTrue(view.isExpanded())
+        assertTrue(view.bodyCreated())
 
         view.update(reasoning("p1", done = true, text = "one\ntwo\nthree\nfour"))
 
+        assertFalse(view.isExpanded())
+        assertFalse(view.bodyVisible())
+        assertFalse(view.bodyCreated())
+
+        view.toggle()
+
         assertTrue(view.isExpanded())
         assertTrue(view.bodyVisible())
         assertTrue(view.bodyCreated())
+        assertEquals("one\ntwo\nthree\nfour", view.markdown())
+    }
+
+    fun `test auto-collapse releases streaming reasoning editors`() {
+        val base = EditorFactory.getInstance().allEditors.size
+
+        repeat(20) { i ->
+            val view = ReasoningView(reasoning("p$i", done = false, text = "```kotlin\nval x = $i\n```"))
+            popupEditors(view.md.component).forEach { it.getEditor(true) }
+            view.update(reasoning("p$i", done = true, text = "```kotlin\nval x = $i\n```"))
+        }
+        UIUtil.dispatchAllInvocationEvents()
+
+        assertEquals(base, EditorFactory.getInstance().allEditors.size)
     }
 
     fun `test manually expanded finished reasoning stays open on update`() {
@@ -91,6 +115,39 @@ class ReasoningViewTest : BasePlatformTestCase() {
         assertFalse(view.isExpanded())
     }
 
+    fun `test manual collapse during stream stays collapsed when marked done`() {
+        val view = ReasoningView(reasoning("p1", done = false, text = "one\ntwo"))
+
+        view.toggle()
+        view.update(reasoning("p1", done = true, text = "one\ntwo\nthree"))
+
+        assertFalse(view.isExpanded())
+        assertFalse(view.bodyVisible())
+        assertEquals("one\ntwo\nthree", view.markdown())
+    }
+
+    fun `test manual expand during stream stays open when marked done`() {
+        val view = ReasoningView(reasoning("p1", done = false, text = "one\ntwo"))
+
+        view.toggle()
+        view.toggle()
+        view.update(reasoning("p1", done = true, text = "one\ntwo\nthree"))
+
+        assertTrue(view.isExpanded())
+        assertTrue(view.bodyVisible())
+        assertEquals("one\ntwo\nthree", view.markdown())
+    }
+
+    fun `test unpinned appended reasoning auto-collapses when marked done`() {
+        val view = ReasoningView(reasoning("p1", done = false, text = "one"))
+
+        view.appendDelta("\ntwo")
+        view.update(reasoning("p1", done = true, text = "one\ntwo"))
+
+        assertFalse(view.isExpanded())
+        assertFalse(view.bodyVisible())
+    }
+
     fun `test collapsed reasoning stays collapsed on update`() {
         val view = ReasoningView(reasoning("p1", done = true, text = "one\ntwo"))
         view.update(reasoning("p1", done = true, text = "one\ntwo\nthree"))
@@ -200,15 +257,63 @@ class ReasoningViewTest : BasePlatformTestCase() {
         assertTrue(view.headerFont().size < style.editorSize)
     }
 
-    fun `test expanded reasoning body is capped to five rows`() {
+    fun `test expanded reasoning body is capped to configured rows`() {
         val view = ReasoningView(reasoning("p1", done = false, text = (1..20).joinToString("\n") { "line $it" }))
         val taller = ReasoningView(reasoning("p2", done = false, text = (1..200).joinToString("\n") { "line $it" }))
 
-        assertEquals(5, view.bodyMaxRows())
+        assertEquals(SessionUiStyle.View.Reasoning.BODY_LINES, view.bodyMaxRows())
         assertTrue(view.preferredSize.height > 0)
         assertEquals(view.preferredSize.height, taller.preferredSize.height)
     }
 
+    fun `test reasoning header popup is available only when collapsed with content`() {
+        val expanded = ReasoningView(reasoning("p1", done = false, text = "one"))
+        val blank = ReasoningView(reasoning("p2", done = true, text = ""))
+        val collapsed = ReasoningView(reasoning("p3", done = true, text = "one\ntwo"))
+
+        assertNull(expanded.headerPopup())
+        assertNull(blank.headerPopup())
+        assertNotNull(collapsed.headerPopup())
+
+        collapsed.toggle()
+
+        assertNull(collapsed.headerPopup())
+    }
+
+    fun `test reasoning header popup body is capped to reasoning preview rows`() {
+        val view = ReasoningView(reasoning("p1", done = true, text = (1..60).joinToString("\n") { "line $it" }))
+        val body = view.headerPopup()!!.build()
+
+        try {
+            val scroll = popupScrollPanes(body.component).first()
+            val panel = scroll.viewport.view as JPanel
+            val md = panel.components.filterIsInstance<JComponent>().single()
+            val line = md.getFontMetrics(md.font).height
+            val max = line * SessionUiStyle.View.Reasoning.POPUP_LINES +
+                JBUI.scale(SessionUiStyle.View.Layout.BODY_EXTRA_HEIGHT)
+
+            assertTrue(body.component.preferredSize.width in 1..JBUI.scale(SessionUiStyle.View.Popup.MAX_WIDTH))
+            assertTrue(body.component.preferredSize.height > 0)
+            assertTrue(body.component.preferredSize.height <= max)
+        } finally {
+            Disposer.dispose(body.disposable)
+        }
+    }
+
+    fun `test reasoning header popup editors are disposed after churn`() {
+        val base = EditorFactory.getInstance().allEditors.size
+        val view = ReasoningView(reasoning("p1", done = true, text = "```kotlin\nprintln(1)\n```"))
+
+        repeat(20) {
+            val body = view.headerPopup()!!.build()
+            popupEditors(body.component).forEach { it.getEditor(true) }
+            Disposer.dispose(body.disposable)
+        }
+        UIUtil.dispatchAllInvocationEvents()
+
+        assertEquals(base, EditorFactory.getInstance().allEditors.size)
+    }
+
     fun `test appended reasoning scrolls nested body to bottom`() {
         val view = ReasoningView(reasoning("p1", done = false, text = (1..20).joinToString("\n") { "line $it" }))
         view.setSize(300, 80)
@@ -302,6 +407,26 @@ class ReasoningViewTest : BasePlatformTestCase() {
         return found
     }
 
+    private fun popupEditors(root: JComponent): List<com.intellij.ui.EditorTextField> {
+        val found = mutableListOf<com.intellij.ui.EditorTextField>()
+        fun visit(component: JComponent) {
+            if (component is com.intellij.ui.EditorTextField) found.add(component)
+            component.components.filterIsInstance<JComponent>().forEach(::visit)
+        }
+        visit(root)
+        return found
+    }
+
+    private fun popupScrollPanes(root: JComponent): List<JBScrollPane> {
+        val found = mutableListOf<JBScrollPane>()
+        fun visit(component: JComponent) {
+            if (component is JBScrollPane) found.add(component)
+            component.components.filterIsInstance<JComponent>().forEach(::visit)
+        }
+        visit(root)
+        return found
+    }
+
     private fun collect(component: Component, found: MutableList<Icon>) {
         if (component is JLabel) component.icon?.let(found::add)
         if (component is Container) component.components.forEach { collect(it, found) }

packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/ui/UiStyleTest.kt

@@ -43,6 +43,9 @@ class UiStyleTest : BasePlatformTestCase() {
         assertTrue(JBUI.scale(SessionUiStyle.View.Layout.VERTICAL_PADDING) > 0)
         assertTrue(JBUI.scale(SessionUiStyle.View.Layout.HORIZONTAL_PADDING) > 0)
         assertTrue(SessionUiStyle.View.Tool.BODY_LINES > 0)
-        assertTrue(SessionUiStyle.View.Reasoning.BODY_LINES > 0)
+        assertEquals(5, SessionUiStyle.View.Reasoning.BODY_LINES)
+        assertEquals(10, SessionUiStyle.View.Reasoning.POPUP_LINES)
+        assertTrue(SessionUiStyle.View.Reasoning.POPUP_LINES != SessionUiStyle.View.Popup.MAX_LINES)
+        assertTrue(SessionUiStyle.View.Reasoning.POPUP_LINES != SessionUiStyle.View.Reasoning.BODY_LINES)
     }
 }