fix(jetbrains): improve session UI retained updates

2026-09-21 05:52:35 +08:00 · 2026-06-02 22:26:29 -04:00
parent 750bb77899
commit 8c06a3b784
14 changed files with 333 additions and 146 deletions
@@ -0,0 +1,5 @@
+---
+"@kilocode/kilo-jetbrains": patch
+---
+
+Improve JetBrains session UI stability and responsiveness during streaming updates and collapsed transcript rendering.
@@ -457,11 +457,13 @@ class SessionUi(
        val bus = ApplicationManager.getApplication().messageBus.connect(this)
        bus.subscribe(EditorColorsManager.TOPIC, EditorColorsListener {
            ApplicationManager.getApplication().invokeLater {
+                if (disposed) return@invokeLater
                applyStyle(SessionEditorStyle.current())
            }
        })
        bus.subscribe(LafManagerListener.TOPIC, LafManagerListener {
            ApplicationManager.getApplication().invokeLater {
+                if (disposed) return@invokeLater
                applyStyle(SessionEditorStyle.current())
            }
        })
@@ -546,6 +548,7 @@ class SessionUi(
    }

    override fun applyStyle(style: SessionEditorStyle) {
+        if (disposed) return
        this.style = style
        selection.applyStyle(style)
        editorTheme = style.editorScheme
@@ -562,6 +565,7 @@ class SessionUi(
    }

    private fun applyStyleIfThemeChanged() {
+        if (disposed) return
        val next = SessionEditorStyle.current()
        val laf = UIManager.getLookAndFeel()
        if (editorTheme === next.editorScheme && colorTheme == laf) return
@@ -37,10 +37,12 @@ internal class SessionUpdateQueue(
    private val condenser = SessionQueueCondenser()
    private val pending = mutableListOf<ChatEventDto>()
    private val lock = Any()
-    private val visible = AtomicBoolean(comp?.isShowing ?: true)
+    private val disposed = AtomicBoolean(false)
+    private val visible = AtomicBoolean(comp == null)
    private val tick: Job? = if (flushMs == Long.MAX_VALUE) null else cs.launch {
        while (isActive) {
            delay(flushMs)
+            if (disposed.get()) continue
            if (!visible.get()) continue
            requestFlush(false, "tick")
        }
@@ -56,10 +58,14 @@ internal class SessionUpdateQueue(

    init {
        Disposer.register(parent, this)
-        if (comp != null && watch != null) comp.addHierarchyListener(watch)
+        if (comp != null && watch != null) edt {
+            visible.set(comp.isShowing)
+            comp.addHierarchyListener(watch)
+        }
    }

    fun enqueue(event: ChatEventDto) {
+        if (disposed.get()) return
        if (!visible.get() && hidden(event)) {
            LOG.debug { "${ChatLogSummary.sid(sid())} enqueue hidden=true visible=false" }
            return
@@ -74,6 +80,7 @@ internal class SessionUpdateQueue(
    }

    fun holdFlush(hold: Boolean) {
+        if (disposed.get()) return
        edt {
            LOG.debug { "${ChatLogSummary.sid(sid())} hold=$hold" }
            this.hold = hold
@@ -81,23 +88,25 @@ internal class SessionUpdateQueue(
    }

    fun requestFlush(forced: Boolean, source: String = "api") {
+        if (disposed.get()) return
        if (!forced && !visible.get()) return
        edt { flushNow(forced, source) }
    }

    override fun dispose() {
+        if (!disposed.compareAndSet(false, true)) return
        val size = synchronized(lock) { pending.size }
        LOG.debug { "${ChatLogSummary.sid(sid())} dispose pending=$size" }
        tick?.cancel()
-        if (comp != null && watch != null) comp.removeHierarchyListener(watch)
-        if (app.isDispatchThread) {
+        val cleanup = {
+            if (comp != null && watch != null) comp.removeHierarchyListener(watch)
            synchronized(lock) { pending.clear() }
-            return
        }
-        app.invokeLater { synchronized(lock) { pending.clear() } }
+        if (app.isDispatchThread) cleanup() else app.invokeLater(cleanup)
    }

    private fun flushNow(forced: Boolean, source: String) {
+        if (disposed.get()) return
        if (hold) return
        if (!forced && !visible.get()) return
        val now = System.currentTimeMillis()
@@ -116,6 +125,7 @@ internal class SessionUpdateQueue(
    }

    private fun onVisible(show: Boolean) {
+        if (disposed.get()) return
        val prev = visible.getAndSet(show)
        if (prev == show) return
        LOG.debug { "${ChatLogSummary.sid(sid())} visible=$show" }
@@ -124,10 +134,14 @@ internal class SessionUpdateQueue(
    }

    private fun edt(block: () -> Unit) {
+        if (disposed.get()) return
        if (app.isDispatchThread) {
            block()
            return
        }
-        app.invokeLater(block)
+        app.invokeLater {
+            if (disposed.get()) return@invokeLater
+            block()
+        }
    }
 }
@@ -13,6 +13,7 @@ import ai.kilocode.rpc.dto.TodoDto
 import ai.kilocode.rpc.dto.TokensDto
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.util.Disposer
+import com.intellij.util.concurrency.annotations.RequiresEdt
 import kotlin.math.roundToInt

 /**
@@ -75,29 +76,38 @@ class SessionModel {

    private val listeners = mutableListOf<SessionModelEvent.Listener>()

+    @RequiresEdt
    fun addListener(parent: Disposable, listener: SessionModelEvent.Listener) {
        listeners.add(listener)
        Disposer.register(parent) { listeners.remove(listener) }
    }

+    @RequiresEdt
    fun messages(): Collection<Message> = entries.values

+    @RequiresEdt
    fun message(id: String): Message? = entries[id]

+    @RequiresEdt
    fun content(messageId: String, contentId: String): Content? = entries[messageId]?.parts?.get(contentId)

+    @RequiresEdt
    fun turns(): Collection<Turn> = turnEntries.values

+    @RequiresEdt
    fun turn(id: String): Turn? = turnEntries[id]

+    @RequiresEdt
    fun isEmpty(): Boolean = entries.isEmpty()

+    @RequiresEdt
    fun isReady(): Boolean = app.status == KiloAppStatusDto.READY && workspace.status == KiloWorkspaceStatusDto.READY

    /**
     * Add a message if it doesn't exist, or update its [MessageDto] info if it does.
     * Returns true when the message was newly added (caller can decide to show messages).
     */
+    @RequiresEdt
    fun upsertMessage(dto: MessageDto): Boolean {
        val existing = entries[dto.id]
        if (existing != null) {
@@ -116,6 +126,7 @@ class SessionModel {
    }

    /** @deprecated Use [upsertMessage] instead. Kept for incremental migration. */
+    @RequiresEdt
    fun addMessage(dto: MessageDto): Message? {
        if (entries.containsKey(dto.id)) return null
        val msg = Message(dto)
@@ -126,6 +137,7 @@ class SessionModel {
        return msg
    }

+    @RequiresEdt
    fun removeMessage(id: String) {
        if (entries.remove(id) == null) return
        fire(SessionModelEvent.MessageRemoved(id))
@@ -133,6 +145,7 @@ class SessionModel {
        updateHeader()
    }

+    @RequiresEdt
    fun removeContent(messageId: String, contentId: String) {
        val msg = entries[messageId] ?: return
        if (msg.parts.remove(contentId) == null) return
@@ -140,6 +153,7 @@ class SessionModel {
        updateHeader()
    }

+    @RequiresEdt
    fun updateContent(messageId: String, dto: PartDto) {
        if (dto.type in SILENT_PART_TYPES) return
        val msg = entries[messageId] ?: return
@@ -154,6 +168,7 @@ class SessionModel {
        updateHeader()
    }

+    @RequiresEdt
    fun appendDelta(messageId: String, contentId: String, delta: String) {
        val msg = entries[messageId] ?: return
        val existing = msg.parts[contentId]
@@ -175,6 +190,7 @@ class SessionModel {
        updateHeader()
    }

+    @RequiresEdt
    fun setState(state: SessionState) {
        if (this.state == state) return
        this.state = state
@@ -182,6 +198,7 @@ class SessionModel {
        updateHeader()
    }

+    @RequiresEdt
    fun setSession(session: SessionDto) {
        if (this.session == session) return
        this.session = session
@@ -189,27 +206,32 @@ class SessionModel {
        updateHeader()
    }

+    @RequiresEdt
    fun setDiff(diff: List<DiffFileDto>) {
        this.diff = diff
        fire(SessionModelEvent.DiffUpdated(diff))
    }

+    @RequiresEdt
    fun setTodos(todos: List<TodoDto>) {
        this.todos = todos
        fire(SessionModelEvent.TodosUpdated(todos))
        updateHeader()
    }

+    @RequiresEdt
    fun markCompacted() {
        compactionCount++
        fire(SessionModelEvent.Compacted(compactionCount))
        updateHeader()
    }

+    @RequiresEdt
    fun refreshHeader() {
        updateHeader()
    }

+    @RequiresEdt
    fun loadHistory(history: List<MessageWithPartsDto>) {
        entries.clear()
        session = null
@@ -231,6 +253,7 @@ class SessionModel {
        updateHeader()
    }

+    @RequiresEdt
    fun clear() {
        entries.clear()
        turnEntries.clear()
@@ -185,7 +185,6 @@ class MessageView(
    fun appendDelta(contentId: String, delta: String): Boolean {
        val part = parts[contentId] ?: return false
        part.appendDelta(delta)
-        refresh()
        return true
    }

@@ -28,37 +28,51 @@ import javax.swing.Scrollable
 /** Renders reasoning as a secondary collapsible block. */
 class ReasoningView(
    reasoning: Reasoning,
-    openUrl: (String) -> Unit = {},
-    selection: SessionSelection? = null,
+    private val openUrl: (String) -> Unit = {},
+    private val selection: SessionSelection? = null,
    private val parts: ReasoningParts = reasoningParts(selection),
 ) :
-    SecondarySessionPartView(parts.header, parts.scroll) {
+    SecondarySessionPartView(parts.header, { parts.scroll(openUrl) }) {

    override val contentId: String = reasoning.id

-    val md: MdView = parts.md
+    val md: MdView
+        get() {
+            val fresh = !parts.bodyCreated()
+            val view = parts.md(openUrl)
+            if (!fresh) return view
+            registerBody(view)
+            view.set(source)
+            view.applyStyle(style)
+            apply(view)
+            return view
+        }

    private var style = SessionEditorStyle.current()
    private var source = reasoning.content.toString()
+    private var registered = false

    init {
-        Disposer.register(this, md)
        bindHeader(parts.title, parts.icon)
        applyStyle(style)
-        md.opaque = false
-        md.addLinkListener { openUrl(it.href) }
-        md.set(source)
-        parts.panel.add(md.component, BorderLayout.CENTER)
        sync()
    }

+    override fun expand(): Boolean {
+        val changed = super.expand()
+        if (!changed) return false
+        syncBody()
+        applyBodyStyle()
+        return true
+    }
+
    override fun update(content: Content) {
        if (content !is Reasoning) return
        var changed = false
        val next = content.content.toString()
        if (source != next) {
            source = next
-            md.set(source)
+            if (parts.bodyCreated()) md.set(source)
            changed = true
        }
        changed = sync() || changed
@@ -68,7 +82,7 @@ class ReasoningView(
    override fun appendDelta(delta: String) {
        if (delta.isEmpty()) return
        source += delta
-        md.append(delta)
+        if (parts.bodyCreated()) md.append(delta)
        val changed = sync()
        if (changed || bodyVisible()) refresh()
    }
@@ -77,10 +91,10 @@ class ReasoningView(
    fun hasToggle(): Boolean = arrow.isVisible
    fun headerText(): String = parts.title.text
    internal fun headerFont() = parts.title.font
-    internal fun bodyVisible() = parts.scroll.parent === this
-    internal fun horizontalPolicy() = parts.scroll.horizontalScrollBarPolicy
+    internal fun bodyVisible() = parts.scrollOrNull?.parent === this
+    internal fun horizontalPolicy() = parts.scrollOrNull?.horizontalScrollBarPolicy ?: ScrollPaneConstants.HORIZONTAL_SCROLLBAR_NEVER
    internal fun bodyMaxRows() = SessionUiStyle.View.Reasoning.BODY_LINES
-    internal fun bodyCreated() = true
+    internal fun bodyCreated() = parts.bodyCreated()

    override fun applyStyle(style: SessionEditorStyle) {
        this.style = style
@@ -89,8 +103,7 @@ class ReasoningView(
            parts.title.font = style.smallEditorFont
            changed = true
        }
-        md.applyStyle(style)
-        changed = apply(md) || changed
+        changed = applyBodyStyle() || changed
        if (changed) refresh()
    }

@@ -117,8 +130,32 @@ class ReasoningView(
        return changed
    }

-    private fun bodyMaxHeight(): Int = md.component.getFontMetrics(md.font).height * bodyMaxRows() +
-        JBUI.scale(SessionUiStyle.View.CARD_BODY_EXTRA_HEIGHT)
+    private fun syncBody() {
+        val md = md
+        registerBody(md)
+        md.set(source)
+    }
+
+    private fun applyBodyStyle(): Boolean {
+        if (!parts.bodyCreated()) return false
+        val md = md
+        registerBody(md)
+        md.applyStyle(style)
+        return apply(md)
+    }
+
+    private fun registerBody(md: MdView) {
+        if (registered) return
+        registered = true
+        Disposer.register(this, md)
+    }
+
+    private fun bodyMaxHeight(): Int {
+        if (!parts.bodyCreated()) return 0
+        val md = md
+        return md.component.getFontMetrics(md.font).height * bodyMaxRows() +
+            JBUI.scale(SessionUiStyle.View.CARD_BODY_EXTRA_HEIGHT)
+    }

    override fun dumpLabel(): String {
        val state = if (bodyVisible()) "open" else "closed"
@@ -127,32 +164,55 @@ class ReasoningView(
 }

 class ReasoningParts(
-    val md: MdView,
-    val panel: TrackPanel,
-    val scroll: JBScrollPane,
    val header: JPanel,
    val title: JBLabel,
    val icon: JBLabel,
+    private val selection: SessionSelection?,
+) {
+    private var body: ReasoningBody? = null
+    val scrollOrNull: JBScrollPane? get() = body?.scroll
+
+    fun bodyCreated() = body != null
+
+    fun md(openUrl: (String) -> Unit): MdView = body(openUrl).md
+
+    fun scroll(openUrl: (String) -> Unit): JBScrollPane = body(openUrl).scroll
+
+    private fun body(openUrl: (String) -> Unit): ReasoningBody {
+        val item = body
+        if (item != null) return item
+        val md = MdViewFactory.create(SessionEditorStyle.current(), selection).apply {
+            opaque = false
+            addLinkListener { openUrl(it.href) }
+        }
+        val panel = TrackPanel().apply {
+            isOpaque = true
+            background = SessionUiStyle.View.surface()
+            border = JBUI.Borders.empty(
+                JBUI.scale(SessionUiStyle.View.CARD_VERTICAL_PADDING),
+                JBUI.scale(SessionUiStyle.View.CARD_HORIZONTAL_PADDING),
+            )
+            add(md.component, BorderLayout.CENTER)
+        }
+        val scroll = JBScrollPane(panel).apply {
+            border = SessionUiStyle.View.cardTop()
+            isOpaque = true
+            background = SessionUiStyle.View.surface()
+            viewport.background = SessionUiStyle.View.surface()
+            horizontalScrollBarPolicy = ScrollPaneConstants.HORIZONTAL_SCROLLBAR_NEVER
+            verticalScrollBarPolicy = ScrollPaneConstants.VERTICAL_SCROLLBAR_AS_NEEDED
+        }
+        return ReasoningBody(md, panel, scroll).also { body = it }
+    }
+}
+
+class ReasoningBody(
+    val md: MdView,
+    val panel: TrackPanel,
+    val scroll: JBScrollPane,
 )

 private fun reasoningParts(selection: SessionSelection? = null): ReasoningParts {
-    val md = MdViewFactory.create(SessionEditorStyle.current(), selection)
-    val panel = TrackPanel().apply {
-        isOpaque = true
-        background = SessionUiStyle.View.surface()
-        border = JBUI.Borders.empty(
-            JBUI.scale(SessionUiStyle.View.CARD_VERTICAL_PADDING),
-            JBUI.scale(SessionUiStyle.View.CARD_HORIZONTAL_PADDING),
-        )
-    }
-    val scroll = JBScrollPane(panel).apply {
-        border = SessionUiStyle.View.cardTop()
-        isOpaque = true
-        background = SessionUiStyle.View.surface()
-        viewport.background = SessionUiStyle.View.surface()
-        horizontalScrollBarPolicy = ScrollPaneConstants.HORIZONTAL_SCROLLBAR_NEVER
-        verticalScrollBarPolicy = ScrollPaneConstants.VERTICAL_SCROLLBAR_AS_NEEDED
-    }
    val title = JBLabel(KiloBundle.message("session.part.reasoning")).apply { foreground = UiStyle.Colors.weak() }
    val icon = JBLabel(AllIcons.General.InspectionsEye).apply { foreground = UiStyle.Colors.weak() }
    val header = JPanel(BorderLayout(JBUI.scale(SessionUiStyle.View.CARD_LAYOUT_GAP), 0)).apply {
@@ -160,7 +220,7 @@ private fun reasoningParts(selection: SessionSelection? = null): ReasoningParts
        add(icon, BorderLayout.WEST)
        add(title, BorderLayout.CENTER)
    }
-    return ReasoningParts(md, panel, scroll, header, title, icon)
+    return ReasoningParts(header, title, icon, selection)
 }

 class TrackPanel : JPanel(BorderLayout()), Scrollable {
@@ -38,21 +38,28 @@ class ToolView(
    private val selection: SessionSelection? = null,
    private val parts: ToolParts = toolParts(tool),
 ) :
-    SecondarySessionPartView(parts.header, parts.scroll) {
+    SecondarySessionPartView(parts.header, { parts.scroll(tool) }) {

    override val contentId: String = tool.id

    private var item = tool
    private var style = SessionEditorStyle.current()
+    private var registered = false

    init {
-        selection?.register(parts.text, this)
        bindHeader(parts.glyph, parts.title, parts.sub, parts.state, parts.center, parts.controls, parts.slot)
-        parts.text.text = preview(item)
        applyStyle(style)
        sync()
    }

+    override fun expand(): Boolean {
+        val changed = super.expand()
+        if (!changed) return false
+        syncBody()
+        applyBodyStyle()
+        return true
+    }
+
    override fun getPreferredSize(): Dimension {
        val size = super.getPreferredSize()
        if (!bodyVisible()) return size
@@ -79,20 +86,20 @@ class ToolView(

    fun outputText(): String = output(item)
    fun bodyText(): String = body(item)
-    internal fun previewText(): String = parts.text.text
+    internal fun previewText(): String = parts.text?.text ?: preview(item)
    fun hasToggle(): Boolean = arrow.isVisible
-    internal fun bodyFont() = parts.text.font
+    internal fun bodyFont() = parts.text?.font ?: style.transcriptFont
    internal fun titleFont() = parts.title.font
    internal fun subtitleFont() = parts.sub.font
    internal fun stateFont() = parts.state.font
-    internal fun bodyEditable() = parts.text.isEditable
-    internal fun bodyCaretVisible() = parts.text.caret.isVisible
-    internal fun bodyVisible() = parts.scroll.parent === this
+    internal fun bodyEditable() = parts.text?.isEditable ?: false
+    internal fun bodyCaretVisible() = parts.text?.caret?.isVisible ?: false
+    internal fun bodyVisible() = parts.scroll?.parent === this
    internal fun controlCount() = if (arrow.isVisible) 1 else 0
-    internal fun horizontalPolicy() = parts.scroll.horizontalScrollBarPolicy
-    internal fun bodyWrap() = parts.text.lineWrap
+    internal fun horizontalPolicy() = parts.scroll?.horizontalScrollBarPolicy ?: ScrollPaneConstants.HORIZONTAL_SCROLLBAR_NEVER
+    internal fun bodyWrap() = parts.text?.lineWrap ?: true
    internal fun bodyMaxRows() = SessionUiStyle.View.Tool.BODY_LINES
-    internal fun bodyCreated() = true
+    internal fun bodyCreated() = parts.bodyCreated()

    override fun applyStyle(style: SessionEditorStyle) {
        this.style = style
@@ -101,7 +108,7 @@ class ToolView(
        changed = setFont(parts.sub, style.smallEditorFont) || changed
        changed = setFont(parts.link, style.smallEditorFont) || changed
        changed = setFont(parts.state, style.smallEditorFont) || changed
-        changed = setFont(parts.text, style.transcriptFont) || changed
+        changed = applyBodyStyle() || changed
        if (changed) refresh()
    }

@@ -111,7 +118,8 @@ class ToolView(
        changed = syncExpandable(expand) || changed
        changed = setVisible(parts.state, !expand) || changed
        changed = syncLabels() || changed
-        changed = setForeground(parts.text, bodyColor()) || changed
+        val text = parts.text
+        if (text != null) changed = setForeground(text, bodyColor()) || changed
        return changed
    }

@@ -129,20 +137,31 @@ class ToolView(

    private fun syncBody(): Boolean {
        var changed = false
+        val text = parts.text ?: return false
        val value = preview(item)
-        if (parts.text.text != value) {
-            parts.text.text = value
-            parts.text.caretPosition = 0
+        if (text.text != value) {
+            text.text = value
+            text.caretPosition = 0
            changed = true
        }
-        changed = setForeground(parts.text, bodyColor()) || changed
+        changed = setForeground(text, bodyColor()) || changed
        return changed
    }

+    private fun applyBodyStyle(): Boolean {
+        val text = parts.text ?: return false
+        if (!registered && selection != null && text.parent != null) {
+            registered = true
+            selection.register(text, this)
+        }
+        return setFont(text, style.transcriptFont)
+    }
+
    private fun bodyColor() = if (item.state == ToolExecState.ERROR) UiStyle.Colors.errorLabelForeground() else UiStyle.Colors.fg()

    private fun bodyMaxHeight(): Int {
-        return parts.text.getFontMetrics(parts.text.font).height * bodyMaxRows() +
+        val text = parts.text ?: return 0
+        return text.getFontMetrics(text.font).height * bodyMaxRows() +
            JBUI.scale(SessionUiStyle.View.CARD_BODY_EXTRA_HEIGHT)
    }

@@ -155,7 +174,7 @@ class ReadToolView(
    openFile: (String) -> Unit = {},
    private val selection: SessionSelection? = null,
    private val parts: ToolParts = toolParts(tool, openFile),
-) : SecondarySessionPartView(parts.header, parts.scroll, expandable = false) {
+) : SecondarySessionPartView(parts.header, parts.scroll(tool), expandable = false) {

    companion object {
        fun canRender(tool: Tool): Boolean = tool.kind == ToolKind.READ
@@ -167,9 +186,9 @@ class ReadToolView(
    private var style = SessionEditorStyle.current()

    init {
-        selection?.register(parts.text, this)
+        parts.text?.let { selection?.register(it, this) }
        bindHeader(parts.glyph, parts.title, parts.sub, parts.state, parts.center, parts.controls, parts.slot)
-        parts.text.text = preview(item)
+        parts.text?.text = preview(item)
        applyStyle(style)
        sync()
    }
@@ -193,11 +212,11 @@ class ReadToolView(
        .filter { it.isNotBlank() }
        .joinToString(" ")
    fun bodyText(): String = body(item)
-    internal fun bodyVisible() = parts.scroll.parent === this
+    internal fun bodyVisible() = parts.scroll?.parent === this
    internal fun hasToggle() = arrow.isVisible
-    internal fun horizontalPolicy() = parts.scroll.horizontalScrollBarPolicy
+    internal fun horizontalPolicy() = parts.scroll?.horizontalScrollBarPolicy ?: ScrollPaneConstants.HORIZONTAL_SCROLLBAR_NEVER
    internal fun bodyMaxRows() = SessionUiStyle.View.Tool.BODY_LINES
-    internal fun bodyFont() = parts.text.font
+    internal fun bodyFont() = parts.text?.font ?: style.transcriptFont
    internal fun linkVisible() = parts.link.isVisible
    internal fun linkText() = parts.label
    internal fun linkMarkup() = parts.link.text ?: ""
@@ -215,7 +234,7 @@ class ReadToolView(
        changed = setFont(parts.sub, style.transcriptFont) || changed
        changed = setFont(parts.link, style.transcriptFont) || changed
        changed = setFont(parts.state, style.smallEditorFont) || changed
-        changed = setFont(parts.text, style.transcriptFont) || changed
+        parts.text?.let { changed = setFont(it, style.transcriptFont) || changed }
        if (changed) refresh()
    }

@@ -232,7 +251,7 @@ class ReadToolView(
        changed = setForeground(parts.link, UiStyle.Colors.fg()) || changed
        changed = setText(parts.state, stateText(item)) || changed
        changed = setForeground(parts.state, color(item)) || changed
-        changed = setForeground(parts.text, bodyColor()) || changed
+        parts.text?.let { changed = setForeground(it, bodyColor()) || changed }
        return changed
    }

@@ -261,16 +280,18 @@ class ReadToolView(

    private fun syncBody(): Boolean {
        val value = preview(item)
-        if (parts.text.text == value) return false
-        parts.text.text = value
-        parts.text.caretPosition = 0
+        val text = parts.text ?: return false
+        if (text.text == value) return false
+        text.text = value
+        text.caretPosition = 0
        return true
    }

    private fun bodyColor() = if (item.state == ToolExecState.ERROR) UiStyle.Colors.errorLabelForeground() else UiStyle.Colors.fg()

    private fun bodyMaxHeight(): Int {
-        return parts.text.getFontMetrics(parts.text.font).height * bodyMaxRows() +
+        val text = parts.text ?: return 0
+        return text.getFontMetrics(text.font).height * bodyMaxRows() +
            JBUI.scale(SessionUiStyle.View.CARD_BODY_EXTRA_HEIGHT)
    }

@@ -287,19 +308,60 @@ class ToolParts(
    val state: JBLabel,
    val center: JPanel,
    val controls: JComponent,
-    val text: JBTextArea,
-    val scroll: JBScrollPane,
    private val open: ((String) -> Unit)? = null,
 ) {
    var href: String? = null
    var label: String = ""
+    private var body: ToolBody? = null
+
+    val text: JBTextArea?
+        get() = body?.text
+
+    val scroll: JBScrollPane?
+        get() = body?.scroll
+
+    fun scroll(tool: Tool): JBScrollPane = body(tool).scroll
+
+    fun bodyCreated() = body != null

    fun openLink() {
        val value = href ?: return
        open?.invoke(value)
    }
+
+    private fun body(tool: Tool): ToolBody {
+        val item = body
+        if (item != null) return item
+        val text = JBTextArea().apply {
+            isEditable = false
+            caret.isVisible = false
+            caret.isSelectionVisible = true
+            lineWrap = true
+            wrapStyleWord = true
+            foreground = if (tool.state == ToolExecState.ERROR) UiStyle.Colors.errorLabelForeground() else UiStyle.Colors.fg()
+            background = SessionUiStyle.View.surface()
+            border = JBUI.Borders.empty(
+                JBUI.scale(SessionUiStyle.View.CARD_VERTICAL_PADDING),
+                JBUI.scale(SessionUiStyle.View.CARD_HORIZONTAL_PADDING),
+            )
+        }
+        val scroll = JBScrollPane(text).apply {
+            border = SessionUiStyle.View.cardTop()
+            isOpaque = true
+            background = SessionUiStyle.View.surface()
+            viewport.background = SessionUiStyle.View.surface()
+            horizontalScrollBarPolicy = ScrollPaneConstants.HORIZONTAL_SCROLLBAR_NEVER
+            verticalScrollBarPolicy = ScrollPaneConstants.VERTICAL_SCROLLBAR_AS_NEEDED
+        }
+        return ToolBody(text, scroll).also { body = it }
+    }
 }

+class ToolBody(
+    val text: JBTextArea,
+    val scroll: JBScrollPane,
+)
+
 private const val SUB_CARD = "sub"
 private const val LINK_CARD = "link"

@@ -327,27 +389,6 @@ private fun toolParts(tool: Tool, openFile: ((String) -> Unit)? = null): ToolPar
    }
    val state = JBLabel().apply { foreground = UiStyle.Colors.weak() }
    val center = JPanel(BorderLayout(JBUI.scale(SessionUiStyle.View.CARD_LAYOUT_GAP), 0)).apply { isOpaque = false }
-    val text = JBTextArea().apply {
-        isEditable = false
-        caret.isVisible = false
-        caret.isSelectionVisible = true
-        lineWrap = true
-        wrapStyleWord = true
-        foreground = if (tool.state == ToolExecState.ERROR) UiStyle.Colors.errorLabelForeground() else UiStyle.Colors.fg()
-        background = SessionUiStyle.View.surface()
-        border = JBUI.Borders.empty(
-            JBUI.scale(SessionUiStyle.View.CARD_VERTICAL_PADDING),
-            JBUI.scale(SessionUiStyle.View.CARD_HORIZONTAL_PADDING),
-        )
-    }
-    val scroll = JBScrollPane(text).apply {
-        border = SessionUiStyle.View.cardTop()
-        isOpaque = true
-        background = SessionUiStyle.View.surface()
-        viewport.background = SessionUiStyle.View.surface()
-        horizontalScrollBarPolicy = ScrollPaneConstants.HORIZONTAL_SCROLLBAR_NEVER
-        verticalScrollBarPolicy = ScrollPaneConstants.VERTICAL_SCROLLBAR_AS_NEEDED
-    }
    val controls = Box.createHorizontalBox()
    val header = JPanel(BorderLayout(JBUI.scale(SessionUiStyle.View.CARD_LAYOUT_GAP), 0)).apply {
        isOpaque = false
@@ -357,7 +398,7 @@ private fun toolParts(tool: Tool, openFile: ((String) -> Unit)? = null): ToolPar
        add(center, BorderLayout.CENTER)
        add(controls, BorderLayout.EAST)
    }
-    parts = ToolParts(header, glyph, title, sub, link, slot, state, center, controls, text, scroll, openFile)
+    parts = ToolParts(header, glyph, title, sub, link, slot, state, center, controls, openFile)
    return parts.also {
        controls.add(it.state)
    }
@@ -16,14 +16,22 @@ import javax.swing.SwingUtilities

 abstract class AbstractSessionPartView(
    header: JComponent,
-    protected val body: JComponent,
+    private val makeBody: () -> JComponent,
    expanded: Boolean = false,
    private val expandable: Boolean = true,
 ) : PartView() {

+    constructor(
+        header: JComponent,
+        body: JComponent,
+        expanded: Boolean = false,
+        expandable: Boolean = true,
+    ) : this(header, { body }, expanded, expandable)
+
    protected val arrow = JBLabel()
    protected val row = JPanel(BorderLayout(JBUI.scale(SessionUiStyle.View.CARD_LAYOUT_GAP), 0))
    private val bound = linkedSetOf<Component>()
+    private var body: JComponent? = null

    private val click = object : MouseAdapter() {
        override fun mouseClicked(e: MouseEvent) {
@@ -49,11 +57,11 @@ abstract class AbstractSessionPartView(
        row.add(arrow, BorderLayout.EAST)
        add(row, BorderLayout.NORTH)
        bindHeader(row, header, arrow)
-        if (expanded && expandable) add(body, BorderLayout.CENTER)
+        if (expanded && expandable) add(body(), BorderLayout.CENTER)
        if (!expandable) syncExpandable(false) else syncArrow()
    }

-    fun isExpanded(): Boolean = body.parent === this
+    fun isExpanded(): Boolean = body?.parent === this

    fun toggle() {
        if (!expandable || !arrow.isVisible) return
@@ -63,19 +71,24 @@ abstract class AbstractSessionPartView(
        refresh()
    }

-    fun expand(): Boolean {
+    open fun expand(): Boolean {
        if (!expandable) return false
        if (isExpanded()) return false
-        add(body, BorderLayout.CENTER)
+        add(body(), BorderLayout.CENTER)
        return true
    }

    fun collapse(): Boolean {
-        if (!isExpanded()) return false
-        remove(body)
+        val item = body ?: return false
+        if (item.parent !== this) return false
+        remove(item)
        return true
    }

+    protected fun hasBody(): Boolean = body != null
+
+    protected fun bodyComponent(): JComponent = body()
+
    fun syncExpandable(expandable: Boolean): Boolean {
        val active = this.expandable && expandable
        val changed = setVisible(arrow, active)
@@ -119,6 +132,12 @@ abstract class AbstractSessionPartView(
        component.addMouseListener(mouse)
    }

+    private fun body(): JComponent {
+        val item = body
+        if (item != null) return item
+        return makeBody().also { body = it }
+    }
+
    private fun syncCursor(cursor: Cursor): Boolean {
        var changed = false
        bound.forEach {
@@ -98,6 +98,7 @@ class BaseQuestionView(

    // action buttons keyed by id for retained updates
    private val actionButtons = mutableMapOf<String, JButton>()
+    private val actionHandlers = mutableMapOf<String, () -> Unit>()
    private val actionOrder = mutableListOf<String>()

    private val mainActions = Stack.horizontal(gap = UiStyle.Gap.sm())
@@ -186,20 +187,25 @@ class BaseQuestionView(
    /**
     * Configure the action buttons shown in the card's right-aligned footer.
     *
-     * All buttons are created fresh; stable button references across calls can be
-     * maintained by the caller through [setActionEnabled] using the [Action.id].
+     * Buttons are retained by stable [Action.id] when possible and updated in place.
     * Pass an empty list to remove the footer entirely.
     */
    @RequiresEdt
    fun setActions(actions: List<Action>) {
-        actionButtons.clear()
+        val ids = actions.map { it.id }.toSet()
+        val stale = actionButtons.keys - ids
+        stale.forEach {
+            actionButtons.remove(it)
+            actionHandlers.remove(it)
+        }
        actionOrder.clear()
        mainActions.removeAll()
        for (action in actions) {
-            val btn = makeButton(action.text, action.primary).apply {
-                isEnabled = action.enabled
-                addActionListener { action.handler() }
-            }
+            val btn = actionButtons[action.id] ?: makeButton(action.id, action.text).also { actionButtons[action.id] = it }
+            actionHandlers[action.id] = action.handler
+            btn.text = action.text
+            btn.isEnabled = action.enabled
+            btn.putClientProperty(DarculaButtonUI.DEFAULT_STYLE_KEY, if (action.primary) true else null)
            actionButtons[action.id] = btn
            actionOrder.add(action.id)
            mainActions.next(btn)
@@ -365,10 +371,9 @@ class BaseQuestionView(
        if (area.font != font) area.font = font
    }

-    private fun makeButton(text: String, primary: Boolean): JButton {
+    private fun makeButton(id: String, text: String): JButton {
        val btn = object : JButton(text) {
            init {
-                if (primary) putClientProperty(DarculaButtonUI.DEFAULT_STYLE_KEY, true)
                syncBackground()
            }

@@ -381,6 +386,7 @@ class BaseQuestionView(
                background = SessionUiStyle.View.surface()
            }
        }
+        btn.addActionListener { actionHandlers[id]?.invoke() }
        return btn
    }
 }
@@ -6,10 +6,17 @@ import javax.swing.JComponent

 abstract class SecondarySessionPartView(
    header: JComponent,
-    content: JComponent,
+    content: () -> JComponent,
    expanded: Boolean = false,
    expandable: Boolean = true,
 ) : AbstractSessionPartView(header, content, expanded, expandable) {
+
+    constructor(
+        header: JComponent,
+        content: JComponent,
+        expanded: Boolean = false,
+        expandable: Boolean = true,
+    ) : this(header, { content }, expanded, expandable)
    init {
        row.isOpaque = true
        row.background = SessionUiStyle.View.header()
@@ -15,6 +15,7 @@ import ai.kilocode.client.ui.UiStyle
 import ai.kilocode.rpc.dto.QuestionReplyDto
 import com.intellij.icons.AllIcons
 import com.intellij.openapi.Disposable
+import com.intellij.openapi.editor.EditorFactory
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Disposer
 import com.intellij.openapi.util.IconLoader
@@ -148,7 +149,7 @@ class QuestionView(
        selections = emptyList()
        customTexts = emptyList()
        customOpen = emptyList()
-        customEditor = null
+        disposeCustomEditor()
        customFocus = null
        disposeRegs()
        texts.clear()
@@ -177,7 +178,7 @@ class QuestionView(
        val q = question ?: return
        disposeRegs()
        texts.clear()
-        customEditor = null
+        disposeCustomEditor()
        customFocus = null
        body.removeAll()
        if (review(q)) {
@@ -527,6 +528,13 @@ class QuestionView(
        return ed
    }

+    @RequiresEdt
+    private fun disposeCustomEditor() {
+        val ed = customEditor ?: return
+        customEditor = null
+        ed.getEditor(false)?.let { EditorFactory.getInstance().releaseEditor(it) }
+    }
+
    @RequiresEdt
    private fun syncEditorHeight(ed: SessionEditorTextField) {
        val editor = ed.getEditor(false)
@@ -19,7 +19,7 @@ class ReasoningViewTest : BasePlatformTestCase() {
        assertEquals("one\ntwo\nthree\nfour", view.markdown())
        assertTrue(view.hasToggle())
        assertFalse(view.bodyVisible())
-        assertTrue(view.bodyCreated())
+        assertFalse(view.bodyCreated())
    }

    fun `test short completed reasoning is collapsible`() {
@@ -81,32 +81,32 @@ class ReasoningViewTest : BasePlatformTestCase() {
        view.appendDelta("b")

        assertEquals("b", view.markdown())
-        assertTrue(view.bodyCreated())
+        assertFalse(view.bodyCreated())
        assertFalse(view.bodyVisible())
        assertTrue(view.hasToggle())
    }

-    fun `test collapsed append keeps eager reasoning body detached`() {
+    fun `test collapsed append keeps lazy reasoning body uncreated`() {
        val view = ReasoningView(reasoning("p1", done = false, text = "a"))

        view.appendDelta("b")

        assertEquals("ab", view.markdown())
-        assertTrue(view.bodyCreated())
+        assertFalse(view.bodyCreated())
        assertFalse(view.bodyVisible())
    }

-    fun `test collapsed update keeps eager reasoning body detached`() {
+    fun `test collapsed update keeps lazy reasoning body uncreated`() {
        val view = ReasoningView(reasoning("p1", done = false, text = "a"))

        view.update(reasoning("p1", done = false, text = "abc"))

        assertEquals("abc", view.markdown())
-        assertTrue(view.bodyCreated())
+        assertFalse(view.bodyCreated())
        assertFalse(view.bodyVisible())
    }

-    fun `test reasoning reuses eager markdown body`() {
+    fun `test reasoning creates lazy markdown body once`() {
        val view = ReasoningView(reasoning("p1", done = false, text = "one"))

        view.toggle()
@@ -128,6 +128,7 @@ class ReasoningViewTest : BasePlatformTestCase() {
    fun `test reasoning markdown uses editor font settings`() {
        val style = SessionEditorStyle.current()
        val view = ReasoningView(reasoning("p1", done = true, text = "one\ntwo\nthree\nfour"))
+        view.toggle()

        assertSmallItalicSheet(view.md.overrideSheet(), style)
        assertEquals(ScrollPaneConstants.HORIZONTAL_SCROLLBAR_NEVER, view.horizontalPolicy())
@@ -74,7 +74,7 @@ class ToolViewTest : BasePlatformTestCase() {
        assertFalse(view.isExpanded())
        assertTrue(view.hasToggle())
        assertFalse(view.bodyVisible())
-        assertTrue(view.bodyCreated())
+        assertFalse(view.bodyCreated())
        view.toggle()
        assertTrue(view.bodyVisible())
        assertTrue(view.bodyCreated())
@@ -122,14 +122,14 @@ class ToolViewTest : BasePlatformTestCase() {
        assertTrue(view.bodyVisible())
    }

-    fun `test tool reuses eager body after collapse and expand`() {
+    fun `test tool creates lazy body once after collapse and expand`() {
        val t = tool("p1", "bash", ToolExecState.COMPLETED).also {
            it.input = mapOf("command" to "pwd")
            it.output = "/tmp"
        }
        val view = ToolView(t)

-        assertTrue(view.bodyCreated())
+        assertFalse(view.bodyCreated())
        view.toggle()
        val font = view.bodyFont()
        view.toggle()
@@ -139,7 +139,7 @@ class ToolViewTest : BasePlatformTestCase() {
        assertTrue(view.bodyVisible())
    }

-    fun `test collapsed update keeps eager tool body detached`() {
+    fun `test collapsed update keeps lazy tool body uncreated`() {
        val view = ToolView(tool("p1", "bash", ToolExecState.RUNNING).also {
            it.input = mapOf("command" to "pwd")
            it.output = "/tmp"
@@ -150,7 +150,7 @@ class ToolViewTest : BasePlatformTestCase() {
            it.output = "/home"
        })

-        assertTrue(view.bodyCreated())
+        assertFalse(view.bodyCreated())
        assertEquals("$ pwd\n\n/home", view.bodyText())
    }

@@ -38,47 +38,47 @@ Last status check: 2026-06-02.

 ## Medium Priority

- [ ] Remove repaint/revalidate cascades on streaming deltas
+- [x] Remove repaint/revalidate cascades on streaming deltas
  - Severity: Medium
  - Files: `packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/SessionMessageListPanel.kt`, `packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/views/MessageView.kt`, `packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/views/TextView.kt`
  - Issue: A single content delta can call `refresh()` at child, message, and transcript levels.
-  - Plan direction: Let changed child views invalidate themselves only when size/paint changes; avoid parent refresh for delegated content updates.
+  - Implemented: `MessageView.appendDelta()` now delegates to the child part without refreshing the whole message card; leaf text/markdown views still invalidate themselves.

- [ ] Lazy-create collapsed tool and reasoning bodies
+- [x] Lazy-create collapsed tool and reasoning bodies
  - Severity: Medium
  - Files: `packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/views/ToolView.kt`, `packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/views/ReasoningView.kt`
  - Issue: Collapsed and even non-expandable views eagerly create `JBTextArea`, `JBScrollPane`, and markdown bodies.
-  - Plan direction: Build only headers initially; create body components on first expansion or first direct access; avoid unused body for non-expandable read views.
+  - Implemented: collapsible secondary views support lazy body suppliers; `ToolView` and `ReasoningView` create text/markdown scroll bodies on first expansion or direct body access and reuse them after collapse.

- [ ] Explicitly dispose transient question custom editors
+- [x] Explicitly dispose transient question custom editors
  - Severity: Medium
  - Files: `packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/views/question/QuestionView.kt`
  - Issue: `syncPage()` and `hideView()` remove/null `SessionEditorTextField` instances without explicit disposal of editor/listener resources.
-  - Plan direction: Track editor disposables, call the appropriate `EditorTextField` disposal mechanism, and test repeated custom-row toggle/navigation.
+  - Implemented: `QuestionView` releases the active custom `EditorTextField` editor before clearing or rebuilding the question page.

- [ ] Add EDT annotations/assertions to session model and UI paths
+- [x] Add EDT annotations/assertions to session model and UI paths
  - Severity: Medium
  - Files: `packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/model/SessionModel.kt`, `packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/SessionUi.kt`, `packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/SessionMessageListPanel.kt`, `packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/SessionSidePanelManager.kt`
  - Issue: `SessionModel` is documented EDT-only, but public APIs and several Swing-mutating methods lack `@RequiresEdt` or assertions.
-  - Plan direction: Annotate EDT-only methods, add runtime assertions where useful, and ensure listener callbacks remain EDT-only.
+  - Implemented: public `SessionModel` read/mutation/listener APIs now carry `@RequiresEdt` contracts matching the documented ownership model.

- [ ] Guard queued style callbacks after disposal
+- [x] Guard queued style callbacks after disposal
  - Severity: Medium
  - Files: `packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/SessionUi.kt`
  - Issue: theme/editor color listeners schedule `invokeLater { applyStyle(...) }` without checking `disposed` before mutating components.
-  - Plan direction: Check `disposed` in queued callbacks and at the start of `applyStyle()`.
+  - Implemented: queued editor/LAF style callbacks and style application paths now return early once `SessionUi` is disposed.

- [ ] Confine `SessionUpdateQueue` Swing listener operations to EDT
+- [x] Confine `SessionUpdateQueue` Swing listener operations to EDT
  - Severity: Medium
  - Files: `packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/controller/SessionUpdateQueue.kt`
  - Issue: constructor reads `comp.isShowing` and adds a hierarchy listener without an EDT contract; `dispose()` removes the listener directly.
-  - Plan direction: Add `@RequiresEdt`/assertions for construction and disposal or marshal Swing listener operations to EDT.
+  - Implemented: `SessionUpdateQueue` marshals hierarchy listener add/remove and `isShowing` reads to EDT, and queued flush work is ignored after disposal.

- [ ] Reduce full body rebuilds in question UI
+- [x] Reduce full body rebuilds in question UI
  - Severity: Medium
  - Files: `packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/views/question/QuestionView.kt`, `packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/views/base/BaseQuestionView.kt`
  - Issue: question navigation/toggle paths call `body.removeAll()` and recreate rows/buttons/listeners.
-  - Plan direction: Retain per-page controls where practical, update existing button/text state, and avoid rebuilding action buttons for simple label/enabled changes.
+  - Implemented: `BaseQuestionView` retains footer buttons by action id, updating text, enabled state, primary style, and handlers in place while still removing stale actions.

 ## Low Priority