feat(jetbrains): add session copy controls

.changeset/jetbrains-hover-copy-overlay.md

@@ -0,0 +1,5 @@
+---
+"@kilocode/kilo-jetbrains": patch
+---
+
+Show a hover copy button for JetBrains session code and tool output blocks.

.changeset/jetbrains-message-toolbar.md

@@ -0,0 +1,5 @@
+---
+"@kilocode/kilo-jetbrains": patch
+---
+
+Add copy buttons below JetBrains session prompts and assistant responses.

.changeset/jetbrains-session-copy.md

@@ -0,0 +1,5 @@
+---
+"@kilocode/kilo-jetbrains": patch
+---
+
+Fix copying selected text from JetBrains session views.

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

@@ -28,6 +28,8 @@ import ai.kilocode.client.session.ui.attachment.attachmentParams
 import ai.kilocode.client.session.ui.attachment.ensureAttachmentEditorKind
 import ai.kilocode.client.session.ui.attachment.isEmbeddedAttachment
 import ai.kilocode.client.session.ui.header.SessionHeaderPanel
+import ai.kilocode.client.session.ui.selection.SessionContextMenu
+import ai.kilocode.client.session.ui.selection.SessionHoverCopyOverlay
 import ai.kilocode.client.session.ui.selection.SessionSelection
 import ai.kilocode.client.session.ui.style.SessionEditorStyle
 import ai.kilocode.client.session.ui.style.SessionEditorStyleTarget
@@ -142,6 +144,7 @@ class SessionUi(
     private lateinit var root: SessionRootPanel
     private lateinit var account: SessionAccountOverlay
     private lateinit var drop: SessionDropOverlay
+    private lateinit var overlay: SessionHoverCopyOverlay
     private val hide = timers.timer(HIDE_MS, repeats = false) {
         if (disposed || !this::drop.isInitialized) return@timer
         drop.setActive(false)
@@ -171,7 +174,7 @@ class SessionUi(
     private var modalFocus: (() -> JComponent)? = null
     private var style = SessionEditorStyle.current()
     private val selection = SessionSelection()
-    private val copy = object : TextCopyProvider() {
+    private val provider = object : TextCopyProvider() {
         override fun getActionUpdateThread() = ActionUpdateThread.EDT
 
         override fun getTextLinesToCopy(): Collection<String>? {
@@ -216,7 +219,7 @@ class SessionUi(
     internal fun currentStyle() = style
 
     override fun uiDataSnapshot(sink: DataSink) {
-        sink[PlatformDataKeys.COPY_PROVIDER] = copy
+        sink[PlatformDataKeys.COPY_PROVIDER] = provider
     }
 
     @RequiresEdt
@@ -253,6 +256,7 @@ class SessionUi(
 
     private fun buildUi() {
         root = SessionRootPanel()
+        SessionContextMenu.install(root, this)
 
         migrationOverlay = MigrationOverlayPanel().apply {
             onSkip = { migration.skip() }
@@ -283,6 +287,11 @@ class SessionUi(
             )
         }
 
+        overlay = SessionHoverCopyOverlay(root, this)
+        root.addOverlay(overlay) { pane, child ->
+            overlay.bounds(pane, child)
+        }
+
         sessionContent = JPanel(BorderLayout())
 
         blankBody = JPanel(BorderLayout()).apply {
@@ -324,10 +333,12 @@ class SessionUi(
         header = SessionHeaderPanel(controller, this)
 
         scroll = SessionScroll(root, sessionContent, messageBody, blankBody)
+        scroll.onScroll = overlay::clear
         connection = ConnectionPanel(this, controller)
 
         prompt = PromptPanel(
             project = project,
+            selection = selection,
             onSend = { text, files -> sendPrompt(text, files) },
             onAbort = { controller.abort() },
             onEnhance = controller::enhancePrompt,

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

@@ -44,6 +44,7 @@ internal class SessionScroll(
     internal val bar: JScrollBar get() = component.verticalScrollBar
     internal val jump: JBLabel
     val view: JComponent? get() = component.viewport.view as? JComponent
+    var onScroll: (() -> Unit)? = null
 
     private var style = SessionEditorStyle.current()
     private var tail = true
@@ -322,6 +323,7 @@ internal class SessionScroll(
         val moved = bar.value != value
         val down = bar.value > value
         syncValue()
+        if (moved) onScroll?.invoke()
         if (auto || opening) {
             updateJump()
             return

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

@@ -90,25 +90,34 @@ class SessionMessageListPanel(
 
                 is SessionModelEvent.ContentAdded -> {
                     msgToView[event.messageId]?.upsertPart(event.content)
+                    msgToTurn[event.messageId]?.syncCopyToolbars()
                     refresh()
                 }
 
                 is SessionModelEvent.ContentUpdated -> {
                     msgToView[event.messageId]?.upsertPart(event.content)
+                    msgToTurn[event.messageId]?.syncCopyToolbars()
                     refresh()
                 }
 
                 is SessionModelEvent.ContentRemoved -> {
                     msgToView[event.messageId]?.removePart(event.contentId)
+                    msgToTurn[event.messageId]?.syncCopyToolbars()
                     refresh()
                 }
 
                 is SessionModelEvent.ContentDelta -> {
                     if (event.created) return@addListener
                     val handled = msgToView[event.messageId]?.appendDelta(event.contentId, event.delta) == true
-                    if (handled) return@addListener
+                    if (handled) {
+                        msgToTurn[event.messageId]?.syncCopyToolbars()
+                        return@addListener
+                    }
                     val content = model.content(event.messageId, event.contentId)
-                    if (content != null) msgToView[event.messageId]?.upsertPart(content)
+                    if (content != null) {
+                        msgToView[event.messageId]?.upsertPart(content)
+                        msgToTurn[event.messageId]?.syncCopyToolbars()
+                    }
                 }
 
                 is SessionModelEvent.HistoryLoaded -> rebuild()
@@ -191,6 +200,7 @@ class SessionMessageListPanel(
             val mv = tv.addMessage(msg)
             register(msgId, tv, mv)
         }
+        tv.syncCopyToolbars()
         add(tv)
         anchorFooter()
         refresh()
@@ -216,6 +226,7 @@ class SessionMessageListPanel(
             val mv = tv.addMessage(msg)
             register(id, tv, mv)
         }
+        tv.syncCopyToolbars()
 
         refresh()
     }
@@ -248,6 +259,7 @@ class SessionMessageListPanel(
                 val mv = tv.addMessage(msg)
                 register(msgId, tv, mv)
             }
+            tv.syncCopyToolbars()
             add(tv)
         }
 

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

@@ -2,6 +2,7 @@ package ai.kilocode.client.session.ui.editor
 
 import ai.kilocode.client.session.ui.prompt.PromptDataKeys
 import ai.kilocode.client.session.ui.prompt.SendPromptContext
+import ai.kilocode.client.session.ui.selection.SessionSelection
 import com.intellij.openapi.actionSystem.DataSink
 import com.intellij.openapi.fileTypes.PlainTextFileType
 import com.intellij.openapi.project.Project
@@ -24,9 +25,11 @@ import com.intellij.ui.EditorTextField
 internal open class SessionEditorTextField(
     project: Project,
     private val ctx: SendPromptContext? = null,
+    private val selection: SessionSelection? = null,
 ) : EditorTextField(project, PlainTextFileType.INSTANCE) {
     override fun uiDataSnapshot(sink: DataSink) {
         super.uiDataSnapshot(sink)
+        selection?.provideCopy(sink) { text }
         ctx?.let { sink.set(PromptDataKeys.SEND, it) }
     }
 }

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

@@ -1,9 +1,11 @@
 package ai.kilocode.client.session.ui.prompt
 
 import ai.kilocode.client.session.ui.editor.SessionEditorTextField
+import ai.kilocode.client.session.ui.selection.SessionSelection
 import com.intellij.openapi.project.Project
 
 internal class PromptEditorTextField(
     project: Project,
     ctx: SendPromptContext,
-) : SessionEditorTextField(project, ctx)
+    selection: SessionSelection? = null,
+) : SessionEditorTextField(project, ctx, selection)

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

@@ -13,6 +13,7 @@ import ai.kilocode.client.session.ui.style.SessionEditorStyleTarget
 import ai.kilocode.client.session.ui.style.SessionUiStyle
 import ai.kilocode.client.session.ui.mode.ModePicker
 import ai.kilocode.client.session.ui.model.ModelPicker
+import ai.kilocode.client.session.ui.selection.SessionSelection
 import ai.kilocode.client.ui.HoverIcon
 import ai.kilocode.client.ui.UiStyle
 import ai.kilocode.client.ui.iconButton
@@ -81,7 +82,8 @@ class PromptPanel(
     private val onSend: (String, List<PromptPartDto>) -> Unit,
     private val onAbort: () -> Unit,
     private val onEnhance: (String, (Result<String>) -> Unit) -> Unit,
-) : BorderLayoutPanel(), SessionEditorStyleTarget, SendPromptContext {
+    private val selection: SessionSelection? = null,
+) : BorderLayoutPanel(), SessionEditorStyleTarget, SendPromptContext, UiDataProvider {
 
     companion object {
         private val LOG = KiloLog.create(PromptPanel::class.java)
@@ -124,7 +126,7 @@ class PromptPanel(
         }
     }
 
-    private val editor = PromptEditorTextField(project, this).apply {
+    private val editor = PromptEditorTextField(project, this, selection).apply {
         border = JBUI.Borders.empty()
         setFontInheritedFromLAF(false)
         setPlaceholder(placeholder())
@@ -210,6 +212,7 @@ class PromptPanel(
 
     init {
         applyStyle(style)
+        selection?.register(editor)
         editor.text = ""
         editor.addDocumentListener(object : DocumentListener {
             override fun documentChanged(e: DocumentEvent) {
@@ -317,6 +320,10 @@ class PromptPanel(
 
     internal val defaultFocusedComponent: JComponent get() = editor
 
+    override fun uiDataSnapshot(sink: DataSink) {
+        selection?.provideCopy(sink) { editor.text }
+    }
+
     @RequiresEdt
     override fun applyStyle(style: SessionEditorStyle) {
         this.style = style

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

@@ -0,0 +1,63 @@
+package ai.kilocode.client.session.ui.selection
+
+import com.intellij.ide.DataManager
+import com.intellij.openapi.Disposable
+import com.intellij.openapi.actionSystem.ActionGroup
+import com.intellij.openapi.actionSystem.ActionManager
+import com.intellij.openapi.ui.popup.JBPopupFactory
+import com.intellij.openapi.util.Disposer
+import com.intellij.ui.awt.RelativePoint
+import com.intellij.util.concurrency.annotations.RequiresEdt
+import java.awt.AWTEvent
+import java.awt.Component
+import java.awt.Point
+import java.awt.Toolkit
+import java.awt.event.AWTEventListener
+import java.awt.event.MouseEvent
+import javax.swing.JComponent
+import javax.swing.SwingUtilities
+
+internal object SessionContextMenu {
+    private val KEY = Any()
+    private const val ID = "Kilo.Session.ContextMenu"
+
+    @RequiresEdt
+    fun install(root: JComponent, parent: Disposable) {
+        if (root.getClientProperty(KEY) == true) return
+        val listener = AWTEventListener { event ->
+            val mouse = event as? MouseEvent ?: return@AWTEventListener
+            if (!mouse.isPopupTrigger) return@AWTEventListener
+            if (mouse.id != MouseEvent.MOUSE_PRESSED && mouse.id != MouseEvent.MOUSE_RELEASED) return@AWTEventListener
+            show(root, mouse)
+        }
+        Toolkit.getDefaultToolkit().addAWTEventListener(listener, AWTEvent.MOUSE_EVENT_MASK)
+        root.putClientProperty(KEY, true)
+        Disposer.register(parent) {
+            Toolkit.getDefaultToolkit().removeAWTEventListener(listener)
+            root.putClientProperty(KEY, null)
+        }
+    }
+
+    @RequiresEdt
+    internal fun target(root: JComponent, src: Component, point: Point): Component? =
+        SessionTargetResolver.target(root, src, point)
+
+    @RequiresEdt
+    private fun show(root: JComponent, event: MouseEvent) {
+        if (!root.isShowing) return
+        val src = event.component ?: return
+        val target = target(root, src, event.point) ?: return
+        val group = ActionManager.getInstance().getAction(ID) as? ActionGroup ?: return
+        val ctx = DataManager.getInstance().getDataContext(target)
+        val popup = JBPopupFactory.getInstance().createActionGroupPopup(
+            null,
+            group,
+            ctx,
+            JBPopupFactory.ActionSelectionAid.SPEEDSEARCH,
+            true,
+        )
+        val point = SwingUtilities.convertPoint(src, event.point, target)
+        popup.show(RelativePoint(target, point))
+        event.consume()
+    }
+}

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

@@ -0,0 +1,56 @@
+package ai.kilocode.client.session.ui.selection
+
+import ai.kilocode.client.plugin.KiloBundle
+import ai.kilocode.client.ui.HoverIcon
+import com.intellij.icons.AllIcons
+import com.intellij.openapi.ide.CopyPasteManager
+import com.intellij.openapi.ui.popup.Balloon
+import com.intellij.openapi.ui.popup.JBPopupFactory
+import com.intellij.ui.awt.RelativePoint
+import com.intellij.util.concurrency.annotations.RequiresEdt
+import java.awt.Cursor
+import java.awt.Point
+import java.awt.datatransfer.StringSelection
+import java.awt.event.MouseAdapter
+import java.awt.event.MouseEvent
+
+internal class SessionCopyButton(
+    fill: Boolean = false,
+    private val text: () -> String?,
+) {
+    private var balloon: Balloon? = null
+    val button = HoverIcon(fill = fill).apply {
+        icon = AllIcons.Actions.Copy
+        cursor = Cursor.getPredefinedCursor(Cursor.HAND_CURSOR)
+        toolTipText = KiloBundle.message("session.copy.hover")
+    }
+
+    init {
+        button.addActionListener { copy() }
+        button.addMouseListener(object : MouseAdapter() {
+            override fun mouseExited(e: MouseEvent) {
+                dismiss()
+            }
+        })
+    }
+
+    @RequiresEdt
+    fun dismiss() {
+        balloon?.hide()
+        balloon = null
+    }
+
+    @RequiresEdt
+    fun copy() {
+        val value = text()?.takeIf { it.isNotEmpty() } ?: return
+        CopyPasteManager.getInstance().setContents(StringSelection(value))
+        dismiss()
+        balloon = JBPopupFactory.getInstance()
+            .createHtmlTextBalloonBuilder(KiloBundle.message("session.copy.copied"), null, null, null)
+            .createBalloon()
+            .also { item ->
+                item.setAnimationEnabled(false)
+                item.show(RelativePoint(button, Point(button.width / 2, 0)), Balloon.Position.above)
+            }
+    }
+}

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

@@ -0,0 +1,11 @@
+package ai.kilocode.client.session.ui.selection
+
+import com.intellij.util.concurrency.annotations.RequiresEdt
+import javax.swing.JComponent
+
+internal interface SessionCopyTarget {
+    val copyAnchor: JComponent
+
+    @RequiresEdt
+    fun copyText(): String?
+}