fix(jetbrains): show message toolbars on hover

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

@@ -4,8 +4,12 @@ import com.intellij.util.concurrency.annotations.RequiresEdt
 import javax.swing.JComponent
 
 internal interface SessionCopyTarget {
+    val copyEligible: Boolean get() = true
+
     val copyAnchor: JComponent
 
+    val copyToolbar: JComponent? get() = null
+
     @RequiresEdt
     fun copyText(): String?
 }

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

@@ -22,6 +22,7 @@ internal class SessionHoverCopyOverlay(
     private var target: SessionCopyTarget? = null
     private val copy = SessionCopyButton(fill = true) { target?.copyText() }
     private val button = copy.button
+    private var child: JComponent = button
 
     init {
         isVisible = false
@@ -62,14 +63,14 @@ internal class SessionHoverCopyOverlay(
     }
 
     override fun doLayout() {
-        button.setBounds(0, 0, width, height)
+        child.setBounds(0, 0, width, height)
     }
 
-    override fun getPreferredSize() = button.preferredSize
+    override fun getPreferredSize() = child.preferredSize
 
-    override fun getMinimumSize() = button.minimumSize
+    override fun getMinimumSize() = child.minimumSize
 
-    override fun getMaximumSize() = button.maximumSize
+    override fun getMaximumSize() = child.maximumSize
 
     @RequiresEdt
     private fun sync(event: MouseEvent) {
@@ -88,6 +89,7 @@ internal class SessionHoverCopyOverlay(
     private fun show(item: SessionCopyTarget) {
         if (target === item && isVisible) return
         target = item
+        use(item.copyToolbar ?: button)
         isVisible = true
         parent?.doLayout()
         revalidate()
@@ -116,6 +118,7 @@ internal class SessionHoverCopyOverlay(
     @RequiresEdt
     private fun conceal() {
         copy.dismiss()
+        use(button)
         if (target == null && !isVisible) return
         target = null
         isVisible = false
@@ -123,6 +126,14 @@ internal class SessionHoverCopyOverlay(
         repaint()
     }
 
+    @RequiresEdt
+    private fun use(comp: JComponent) {
+        if (child === comp && comp.parent === this) return
+        removeAll()
+        child = comp
+        add(child)
+    }
+
     override fun dispose() {
         clear()
     }

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

@@ -64,11 +64,13 @@ internal object SessionTargetResolver {
     @RequiresEdt
     private fun copy(root: JComponent, comp: Component): SessionCopyTarget? {
         var current: Component? = comp
-        var target: SessionCopyTarget? = null
+        val targets = mutableListOf<SessionCopyTarget>()
         while (current != null && inside(root, current)) {
-            if (current is SessionCopyTarget) target = current
+            if (current is SessionCopyTarget && current.copyEligible) targets.add(current)
             current = current.parent
         }
-        return target
+        val toolbar = targets.indexOfFirst { it.copyToolbar != null }
+        if (toolbar > 0) return targets.take(toolbar).firstOrNull { it.copyToolbar == null }
+        return targets.lastOrNull()
     }
 }

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

@@ -10,6 +10,8 @@ import ai.kilocode.client.plugin.KiloBundle
 import com.intellij.util.concurrency.annotations.RequiresEdt
 import com.intellij.util.ui.JBUI
 import java.awt.BorderLayout
+import java.awt.Dimension
+import javax.swing.JComponent
 import javax.swing.JPanel
 
 internal class MessageToolbar(
@@ -65,6 +67,18 @@ internal class MessageToolbar(
     @RequiresEdt
     fun copyButton() = button
 
+    fun placeholder(): JComponent = object : JPanel() {
+        init {
+            isOpaque = false
+        }
+
+        override fun getPreferredSize(): Dimension = this@MessageToolbar.preferredSize
+
+        override fun getMinimumSize(): Dimension = this@MessageToolbar.minimumSize
+
+        override fun getMaximumSize(): Dimension = this@MessageToolbar.maximumSize
+    }
+
     override fun removeNotify() {
         copy.dismiss()
         super.removeNotify()

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

@@ -12,6 +12,7 @@ import ai.kilocode.client.session.model.ToolCallRef
 import ai.kilocode.client.session.model.ToolExecState
 import ai.kilocode.client.session.ui.SessionView
 import ai.kilocode.client.session.ui.style.SessionEditorStyle
+import ai.kilocode.client.session.ui.selection.SessionCopyTarget
 import ai.kilocode.client.session.ui.selection.SessionSelection
 import ai.kilocode.client.session.ui.style.SessionEditorStyleTarget
 import ai.kilocode.client.session.views.base.PartView
@@ -78,6 +79,7 @@ class MessageView(
     private var prompt: PromptView? = null
     private var promptBox: JPanel? = null
     private var promptToolbar: MessageToolbar? = null
+    private var promptToolbarPlaceholder: JComponent? = null
 
     init {
         isOpaque = false
@@ -293,6 +295,7 @@ class MessageView(
         prompt = null
         promptBox = null
         promptToolbar = null
+        promptToolbarPlaceholder = null
         for ((_, content) in msg.parts) {
             if (content is StepFinish) continue
             if (isHidden(content)) continue
@@ -396,6 +399,7 @@ class MessageView(
         prompt = null
         promptBox = null
         promptToolbar = null
+        promptToolbarPlaceholder = null
         hidden = null
     }
 
@@ -452,19 +456,27 @@ class MessageView(
         prompt = view
         val bar = promptToolbar ?: MessageToolbar(
             { prompt?.copyMarkdown(trim = false) },
-            BorderLayout.LINE_START,
+            BorderLayout.LINE_END,
             revert?.let { fn -> { fn(msg.info.id) } },
         ).also { promptToolbar = it }
+        val placeholder = promptToolbarPlaceholder ?: bar.placeholder().also { promptToolbarPlaceholder = it }
         val box = JPanel(BorderLayout()).also {
             it.isOpaque = false
             it.add(view, BorderLayout.CENTER)
             promptBox = it
         }
         bar.setActive(true)
-        return JPanel(BorderLayout()).also {
-            it.isOpaque = false
-            it.add(box, BorderLayout.CENTER)
-            it.add(bar, BorderLayout.SOUTH)
+        return object : JPanel(BorderLayout()), SessionCopyTarget {
+            override val copyAnchor: JComponent get() = placeholder
+            override val copyToolbar: JComponent get() = bar
+
+            override fun copyText(): String? = prompt?.copyMarkdown(trim = false)
+
+            init {
+                isOpaque = false
+                add(box, BorderLayout.CENTER)
+                add(placeholder, BorderLayout.SOUTH)
+            }
         }
     }
 

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

@@ -7,6 +7,7 @@ import ai.kilocode.client.session.openSessionLink
 import ai.kilocode.client.session.model.Content
 import ai.kilocode.client.session.model.Text
 import ai.kilocode.client.session.ui.style.SessionEditorStyle
+import ai.kilocode.client.session.ui.selection.SessionCopyTarget
 import ai.kilocode.client.session.ui.selection.SessionSelection
 import ai.kilocode.client.session.views.base.PartView
 import ai.kilocode.client.ui.md.MdView
@@ -15,6 +16,7 @@ import com.intellij.openapi.util.Disposer
 import com.intellij.util.concurrency.annotations.RequiresEdt
 import java.awt.BorderLayout
 import javax.swing.JButton
+import javax.swing.JComponent
 
 /**
  * Renders a [Text] part as markdown using [MdView].
@@ -27,7 +29,7 @@ open class TextView(
     private val openFile: SessionFileOpener = { _, _ -> },
     private val openUrl: (String) -> Unit = {},
     selection: SessionSelection? = null,
-) : PartView() {
+) : PartView(), SessionCopyTarget {
 
     override val contentId: String = text.id
 
@@ -37,6 +39,13 @@ open class TextView(
         text = { copyText() },
         tooltip = KiloBundle.message("session.copy.response"),
     )
+    private val placeholder = toolbar.placeholder()
+
+    override val copyEligible: Boolean get() = hasCopyToolbar()
+
+    override val copyAnchor: JComponent get() = placeholder
+
+    override val copyToolbar: JComponent? get() = toolbar.takeIf { hasCopyToolbar() }
 
     init {
         layout = BorderLayout()
@@ -46,7 +55,7 @@ open class TextView(
         md.addLinkListener { onLink(it) }
         applyStyle(SessionEditorStyle.current())
         add(md.component, BorderLayout.CENTER)
-        add(toolbar, BorderLayout.SOUTH)
+        add(placeholder, BorderLayout.SOUTH)
         if (text.content.isNotEmpty()) md.set(text.content.toString())
         syncToolbar()
     }
@@ -121,11 +130,15 @@ open class TextView(
 
     @RequiresEdt
     private fun syncToolbar() {
-        toolbar.sync(copyText()?.isNotEmpty() == true)
+        val on = copyText()?.isNotEmpty() == true
+        toolbar.sync(on)
+        if (placeholder.isVisible == on) return
+        placeholder.isVisible = on
+        refresh()
     }
 
     @RequiresEdt
-    private fun copyText(): String? {
+    override fun copyText(): String? {
         val item = mode ?: return null
         return copyMarkdown(item.trim)
     }

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

@@ -18,6 +18,7 @@ import ai.kilocode.client.session.views.base.BaseQuestionView
 import ai.kilocode.client.session.views.permission.PermissionView
 import ai.kilocode.client.session.views.question.QuestionResultView
 import ai.kilocode.client.session.views.question.QuestionView
+import ai.kilocode.client.session.ui.selection.SessionCopyTarget
 import ai.kilocode.client.session.views.MessageToolbar
 import ai.kilocode.client.session.views.MessageView
 import ai.kilocode.client.session.views.TextView
@@ -237,10 +238,17 @@ class SessionMessageListPanelTest : BasePlatformTestCase() {
 
         val view = panel.findMessage("u1")!!.part("p1") as TextView
         val message = panel.findMessage("u1")!!
-        assertNotNull(find<MessageToolbar>(message))
+        val target = components(message).filterIsInstance<SessionCopyTarget>().single { it.copyToolbar != null }
+        val toolbar = target.copyToolbar as MessageToolbar
+        val placeholder = target.copyAnchor
+
+        assertFalse(components(message).filterIsInstance<MessageToolbar>().any { it === toolbar })
+        assertNull(toolbar.parent)
         assertFalse(view.hasCopyToolbar())
-        assertEquals(BorderLayout.LINE_START, message.promptToolbarAlignment())
+        assertEquals(BorderLayout.LINE_END, message.promptToolbarAlignment())
+        assertEquals(BorderLayout.LINE_END, toolbar.alignment())
         assertTrue(message.promptToolbarActive())
+        assertEquals(toolbar.preferredSize, placeholder.preferredSize)
     }
 
     fun `test user prompt toolbar omits rollback when revert handler is absent`() {
@@ -248,8 +256,9 @@ class SessionMessageListPanelTest : BasePlatformTestCase() {
         model.updateContent("u1", part("p1", "u1", "text", text = "hello"))
 
         val message = panel.findMessage("u1")!!
+        val toolbar = components(message).filterIsInstance<SessionCopyTarget>().single { it.copyToolbar != null }.copyToolbar as MessageToolbar
 
-        assertFalse(components(message).filterIsInstance<JButton>().any { it.toolTipText == KiloBundle.message("revert.message.rollback") })
+        assertFalse(components(toolbar).filterIsInstance<JButton>().any { it.toolTipText == KiloBundle.message("revert.message.rollback") })
     }
 
     fun `test user prompt toolbar shows rollback when revert handler is present`() {
@@ -259,7 +268,8 @@ class SessionMessageListPanelTest : BasePlatformTestCase() {
         model.updateContent("u1", part("p1", "u1", "text", text = "hello"))
 
         val message = panel.findMessage("u1")!!
-        val rollback = components(message)
+        val toolbar = components(message).filterIsInstance<SessionCopyTarget>().single { it.copyToolbar != null }.copyToolbar as MessageToolbar
+        val rollback = components(toolbar)
             .filterIsInstance<JButton>()
             .first { it.toolTipText == KiloBundle.message("revert.message.rollback") }
         rollback.doClick()

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

@@ -5,6 +5,9 @@ import ai.kilocode.client.session.ui.selection.SessionCopyTarget
 import ai.kilocode.client.session.ui.selection.SessionContextMenu
 import ai.kilocode.client.session.ui.selection.SessionHoverCopyOverlay
 import ai.kilocode.client.session.ui.selection.SessionTargetResolver
+import ai.kilocode.client.session.views.MessageToolbar
+import ai.kilocode.client.session.views.MessageView
+import ai.kilocode.client.session.views.TextView
 import ai.kilocode.client.session.views.tool.ShellToolView
 import ai.kilocode.client.session.views.tool.ToolView
 import ai.kilocode.client.test.CopyProviderSink
@@ -173,6 +176,57 @@ class SessionSelectionCopyTest : SessionUiTestBase() {
         assertFalse(overlay.isVisible)
     }
 
+    fun `test user prompt toolbar appears in hover overlay`() {
+        showMessages()
+        emit(ChatEventDto.MessageUpdated("ses_test", message("u1")), flush = false)
+        emit(ChatEventDto.PartUpdated("ses_test", part("p1", "u1", "text", "hello")))
+        val root = find<SessionRootPanel>(ui)
+        val message = find<MessageView>(ui)
+        val target = components(message).filterIsInstance<SessionCopyTarget>().single { it.copyToolbar != null }
+        val overlay = find<SessionHoverCopyOverlay>(ui)
+        val anchor = target.copyAnchor
+        val toolbar = target.copyToolbar as MessageToolbar
+
+        assertFalse(components(ui).any { it === toolbar })
+        assertNull(toolbar.parent)
+
+        show(overlay, target)
+        layout()
+
+        assertTrue(overlay.isVisible)
+        assertSame(root.overlay, overlay.parent)
+        val shown = overlay.components.single() as MessageToolbar
+        assertSame(toolbar, shown)
+        assertEquals(anchor.preferredSize, shown.preferredSize)
+
+        shown.copyButton().doClick()
+
+        assertEquals("hello", CopyPasteManager.getInstance().getContents(DataFlavor.stringFlavor))
+    }
+
+    fun `test assistant response toolbar appears in hover overlay`() {
+        showText(" response ")
+        val root = find<SessionRootPanel>(ui)
+        val target = components(ui).filterIsInstance<TextView>().first { it.markdown().contains("response") }
+        val overlay = find<SessionHoverCopyOverlay>(ui)
+        val toolbar = target.copyToolbar as MessageToolbar
+
+        assertTrue(target.hasCopyToolbar())
+        assertFalse(components(ui).any { it === toolbar })
+        assertNull(toolbar.parent)
+
+        show(overlay, target)
+        layout()
+
+        assertTrue(overlay.isVisible)
+        assertSame(root.overlay, overlay.parent)
+        val shown = overlay.components.single() as MessageToolbar
+        assertSame(toolbar, shown)
+        shown.copyButton().doClick()
+
+        assertEquals("response", CopyPasteManager.getInstance().getContents(DataFlavor.stringFlavor))
+    }
+
     fun `test hover copy overlay ignores mouse events after disposal`() {
         val root = ShowingPanel()
         val parent = Disposer.newDisposable("overlay-test")
@@ -229,6 +283,12 @@ class SessionSelectionCopyTest : SessionUiTestBase() {
 
     private fun rgb(color: java.awt.Color): Int = color.rgb and RGB_MASK
 
+    private fun show(overlay: SessionHoverCopyOverlay, target: SessionCopyTarget) {
+        val method = SessionHoverCopyOverlay::class.java.getDeclaredMethod("show", SessionCopyTarget::class.java)
+        method.isAccessible = true
+        method.invoke(overlay, target)
+    }
+
     fun `test hover overlay keeps current target while pointer remains inside anchor`() {
         val overlay = find<SessionHoverCopyOverlay>(ui)
         val target = TargetPanel("alpha")
@@ -267,7 +327,7 @@ class SessionSelectionCopyTest : SessionUiTestBase() {
 
     private fun showText(text: String) {
         if (controller().id == null) showMessages()
-        emit(ChatEventDto.MessageUpdated("ses_test", message("msg_text")))
+        emit(ChatEventDto.MessageUpdated("ses_test", message("msg_text").copy(role = "assistant")))
         emit(ChatEventDto.PartUpdated("ses_test", part("part_text", "msg_text", "text", text)))
         layout()
     }
@@ -309,6 +369,13 @@ class SessionSelectionCopyTest : SessionUiTestBase() {
         }
     }
 
+    private fun components(root: Component): Sequence<Component> = sequence {
+        yield(root)
+        if (root is Container) {
+            for (child in root.components) yieldAll(components(child))
+        }
+    }
+
     private fun ancestors(component: Component): Sequence<Component> = sequence {
         var comp: Component? = component
         while (comp != null) {

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

@@ -120,9 +120,13 @@ class TextViewTest : BasePlatformTestCase() {
 
         val layout = view.layout as BorderLayout
         assertSame(view.md.component, layout.getLayoutComponent(BorderLayout.CENTER))
-        val bar = layout.getLayoutComponent(BorderLayout.SOUTH) as MessageToolbar
+        val bar = view.copyToolbar as MessageToolbar
+        val placeholder = layout.getLayoutComponent(BorderLayout.SOUTH)
         assertEquals(BorderLayout.LINE_END, bar.alignment())
         assertTrue(components(bar).contains(view.copyButton()))
+        assertSame(view.copyAnchor, placeholder)
+        assertNull(bar.parent)
+        assertEquals(bar.preferredSize, placeholder.preferredSize)
         assertEquals(UiStyle.Gap.xs(), bar.insets.top)
         assertTrue(view.hasCopyToolbar())
     }
@@ -159,14 +163,14 @@ class TextViewTest : BasePlatformTestCase() {
         val view = TextView(Text("p1").also { it.content.append(" first ") })
         view.setCopyToolbar(true)
         val comp = view.md.component
-        val bar = (view.layout as BorderLayout).getLayoutComponent(BorderLayout.SOUTH)
+        val bar = view.copyToolbar
 
         view.update(Text("p1").also { it.content.append(" second ") })
         view.appendDelta(" third ")
         view.copyButton().doClick()
 
         assertSame(comp, view.md.component)
-        assertSame(bar, (view.layout as BorderLayout).getLayoutComponent(BorderLayout.SOUTH))
+        assertSame(bar, view.copyToolbar)
         assertEquals("second  third", clipboard())
     }