fix(jetbrains): scope session copy context

2026-09-24 16:02:55 +08:00 · 2026-06-15 16:55:01 -04:00
parent ccd9890c46
commit b4864ebd43
19 changed files with 506 additions and 42 deletions
@@ -0,0 +1,5 @@
+---
+"@kilocode/kilo-jetbrains": patch
+---
+
+Fix copying selected text from JetBrains session views.
@@ -0,0 +1,85 @@
+# JetBrains Session View Copy Context Plan
+
+## Goal
+Make right-click Copy work for JetBrains session views without installing popup handlers on every child component.
+
+Use one session-level context-menu listener scoped to the session layered pane. At popup time, resolve the deepest component under the mouse inside the session UI, build the popup with that component's `DataContext`, and let the nearest view `UiDataProvider` provide `PlatformDataKeys.COPY_PROVIDER`.
+
+Do not add whole-session copy behavior in this pass. Copy support is view-scoped only.
+
+## Direction
+- Keep `SessionSelection` as the single source of selected text across selectable view surfaces.
+- Keep the standard IntelliJ Copy action only: `<reference ref="$Copy"/>`.
+- Do not install `PopupHandler.installPopupMenu(...)` on every markdown/prose/code/prompt component.
+- Install one popup dispatcher for the session root/layered pane from `SessionUi.buildUi()`.
+- Make session views expose copy data with `UiDataProvider`; the popup dispatcher should resolve and use the clicked component's normal IntelliJ data context.
+- Remove `SessionUi` root fallback copy provider if it is only providing copy for the entire session container. View roots and editor/text children should provide copy instead.
+
+## Current State To Adjust
+- `SessionSelection` currently owns both copy provider logic and per-component context-menu installation.
+- Several components currently call `selection.installContextMenu(...)` directly.
+- `SessionUi` currently implements `UiDataProvider` and delegates `COPY_PROVIDER` to `selection.provideCopy(sink)`.
+- `MdViewHybrid.RootPanel`, `MdViewHybrid.CodeField`, `PromptPanel`, and `SessionEditorTextField` already expose or can expose `selection.provideCopy(sink)`.
+- `PromptEditorTextField` already preserves `PromptDataKeys.SEND` through `SessionEditorTextField.uiDataSnapshot`.
+
+## Implementation Steps
+1. Split copy data from popup installation in `SessionSelection`.
+   - Keep the shared `TextCopyProvider` and `provideCopy(sink)`.
+   - Remove or stop using `installContextMenu(component)` for per-component popup installation.
+   - If a helper remains in this file, make it install only once on the session root/layered pane, not on view children.
+
+2. Add a session-scoped popup installer.
+   - Preferred shape: a small helper such as `SessionContextMenu.install(root: JComponent, parent: Disposable)` under `session/ui/selection/` or `session/ui/`.
+   - Scope it to the `SessionRootPanel`/`LayeredOverlayPanel` created in `SessionUi.buildUi()`.
+   - Because normal mouse listeners on an ancestor generally do not receive mouse events targeted at descendants, use a root-scoped AWT/event-queue listener unless a quick verification confirms the layered pane receives descendant popup events.
+   - Filter events to `MouseEvent` popup triggers whose source component is inside the session root and whose root is showing.
+   - Handle both press and release popup triggers for cross-platform behavior.
+   - Consume the event after showing the session popup to avoid duplicate native/component popups.
+
+3. Resolve the clicked component at popup time.
+   - Convert the event point from the source component into root coordinates.
+   - Use Swing's deepest-component lookup from the session root coordinates, with a fallback to the source component if needed.
+   - Prefer the deepest visible enabled component inside `root.content`/session UI; modal blocker content may still participate if it contains session views.
+   - Build the `DataContext` from the resolved component via `DataManager.getInstance().getDataContext(target)`.
+
+4. Build the popup from the target component's data context.
+   - Look up `Kilo.Session.ContextMenu` with `ActionManager.getInstance().getAction(...) as? ActionGroup`.
+   - Use a popup API that accepts an explicit `DataContext`, such as `JBPopupFactory.createActionGroupPopup(...)`, so Copy enablement is based on the clicked view/component rather than the layered-pane root.
+   - Show the popup at the original mouse location, converted to the component used as the popup anchor.
+   - Do not introduce a custom copy action.
+
+5. Keep copy providers view-scoped.
+   - `MdViewHybrid.RootPanel` should implement `UiDataProvider` and call `selection?.provideCopy(sink)`.
+   - `MdViewHybrid.CodeField.uiDataSnapshot` must keep `super.uiDataSnapshot(sink)` first, then add `selection?.provideCopy(sink)`.
+   - `PromptPanel` and prompt editor fields should expose `selection?.provideCopy(sink)` while preserving prompt-specific send data.
+   - Custom question editors should receive/register the shared selection and expose the same provider through `SessionEditorTextField`.
+   - Remove root/session-container fallback copy exposure from `SessionUi` if it is no longer needed for a specific focused view.
+
+6. Remove per-component popup installs.
+   - Delete calls such as `selection?.installContextMenu(...)` from markdown prose panes, code fields, fallback text areas, prompt shells/editors, scroll panes, and question custom editors.
+   - Keep registration calls like `selection?.register(...)`; those are still needed for actual selection tracking.
+
+7. Keep XML and bundle entries.
+   - Keep `Kilo.Session.ContextMenu` in `frontend/src/main/resources/kilo.jetbrains.frontend.xml`.
+   - The group should contain only `<reference ref="$Copy"/>`.
+   - Keep `action.Kilo.Session.ContextMenu.text=Session Actions` or equivalent bundle key.
+
+8. Update tests.
+   - Keep provider tests that assert view components expose `COPY_PROVIDER` directly: markdown root, code child, prompt editor, custom editor if covered.
+   - Adjust tests so `SessionUi` root is not the primary copy provider assertion.
+   - Add or update a test for single installer behavior: session root installs one popup/event dispatcher, and child components do not receive per-component popup client properties/listeners from `SessionSelection`.
+   - Add a focused test for target resolution if practical: given a nested view component under `SessionRootPanel`, resolving a popup point returns the nested component or its nearest useful data-provider ancestor.
+   - Keep descriptor coverage asserting `Kilo.Session.ContextMenu` exists and references only `$Copy`.
+
+9. Keep the changeset.
+   - Use the existing patch changeset text or keep it equivalent: `Fix copying selected text from JetBrains session views.`
+
+## Verification
+Run from `packages/kilo-jetbrains/` after implementation:
+- `./gradlew :frontend:test --tests '*SessionSelectionCopyTest' --tests '*PromptPanelTest' --tests '*MdViewHybridTest' --tests '*HistorySessionActionsTest'`
+- `./gradlew typecheck`
+
+## Notes
+- This stays in the JetBrains frontend module; no backend/RPC changes are needed.
+- No session transcript serialization or whole-session copy is part of this pass.
+- If the explicit-data-context popup API signature differs across platform versions, inspect IntelliJ source and use the nearest public `JBPopupFactory`/Action System API that accepts a `DataContext`. Avoid internal APIs.
@@ -28,6 +28,7 @@ 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.SessionSelection
 import ai.kilocode.client.session.ui.style.SessionEditorStyle
 import ai.kilocode.client.session.ui.style.SessionEditorStyleTarget
@@ -48,11 +49,6 @@ import ai.kilocode.rpc.dto.PromptPartDto
 import com.intellij.util.ui.JBUI
 import ai.kilocode.log.KiloLog
 import com.intellij.ide.BrowserUtil
-import com.intellij.ide.TextCopyProvider
-import com.intellij.openapi.actionSystem.ActionUpdateThread
-import com.intellij.openapi.actionSystem.DataSink
-import com.intellij.openapi.actionSystem.PlatformDataKeys
-import com.intellij.openapi.actionSystem.UiDataProvider
 import com.intellij.ide.ui.LafManagerListener
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.application.ApplicationManager
@@ -97,7 +93,7 @@ class SessionUi(
    private val manager: SessionManager? = null,
    private val workspaces: KiloWorkspaceService = service(),
    private val migration: MigrationUiController = service<KiloMigrationService>(),
-) : JPanel(BorderLayout()), Disposable, SessionEditorStyleTarget, UiDataProvider {
+) : JPanel(BorderLayout()), Disposable, SessionEditorStyleTarget {

    companion object {
        private val LOG = KiloLog.create(SessionUi::class.java)
@@ -170,14 +166,6 @@ class SessionUi(
    private var modalFocus: (() -> JComponent)? = null
    private var style = SessionEditorStyle.current()
    private val selection = SessionSelection()
-    private val copy = object : TextCopyProvider() {
-        override fun getActionUpdateThread() = ActionUpdateThread.EDT
-
-        override fun getTextLinesToCopy(): Collection<String>? {
-            val text = selection.selectedText()?.takeIf { it.isNotEmpty() } ?: return null
-            return listOf(text)
-        }
-    }
    private var editorTheme = style.editorScheme
    private var colorTheme = UIManager.getLookAndFeel()
    private var disposed = false
@@ -214,10 +202,6 @@ class SessionUi(

    internal fun currentStyle() = style

-    override fun uiDataSnapshot(sink: DataSink) {
-        sink[PlatformDataKeys.COPY_PROVIDER] = copy
-    }
-
    @RequiresEdt
    internal fun activityKind(): SessionActivityKind? = when (val state = controller.model.state) {
        is SessionState.Idle,
@@ -252,6 +236,7 @@ class SessionUi(

    private fun buildUi() {
        root = SessionRootPanel()
+        SessionContextMenu.install(root, this)

        migrationOverlay = MigrationOverlayPanel().apply {
            onSkip = { migration.skip() }
@@ -327,6 +312,7 @@ class SessionUi(

        prompt = PromptPanel(
            project = project,
+            selection = selection,
            onSend = { text, files -> sendPrompt(text, files) },
            onAbort = { controller.abort() },
            onEnhance = controller::enhancePrompt,
@@ -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) }
    }
 }
@@ -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)
@@ -12,6 +12,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
@@ -77,7 +78,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)
@@ -114,7 +116,7 @@ class PromptPanel(
    private var attachment = true
    private var submitting = false

-    private val editor = PromptEditorTextField(project, this).apply {
+    private val editor = PromptEditorTextField(project, this, selection).apply {
        border = JBUI.Borders.empty()
        setFontInheritedFromLAF(false)
        setPlaceholder(placeholder())
@@ -197,6 +199,7 @@ class PromptPanel(

    init {
        applyStyle(style)
+        selection?.register(editor)
        editor.text = ""
        editor.addDocumentListener(object : DocumentListener {
            override fun documentChanged(e: DocumentEvent) {
@@ -303,6 +306,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
@@ -0,0 +1,81 @@
+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.actionSystem.ActionPlaces
+import com.intellij.openapi.actionSystem.UiDataProvider
+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? {
+        if (!inside(root, src)) return null
+        val pt = SwingUtilities.convertPoint(src, point, root)
+        if (!root.contains(pt)) return null
+        val deep = SwingUtilities.getDeepestComponentAt(root, pt.x, pt.y)?.takeIf { inside(root, it) } ?: src
+        return provider(root, deep) ?: deep
+    }
+
+    @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()
+    }
+
+    private fun provider(root: JComponent, comp: Component): Component? {
+        var current: Component? = comp
+        while (current != null && inside(root, current)) {
+            if (current is UiDataProvider) return current
+            current = current.parent
+        }
+        return null
+    }
+
+    private fun inside(root: JComponent, comp: Component): Boolean = comp === root || SwingUtilities.isDescendingFrom(comp, root)
+}
@@ -1,7 +1,11 @@
 package ai.kilocode.client.session.ui.selection

 import ai.kilocode.client.session.ui.style.SessionEditorStyle
+import com.intellij.ide.TextCopyProvider
 import com.intellij.openapi.Disposable
+import com.intellij.openapi.actionSystem.ActionUpdateThread
+import com.intellij.openapi.actionSystem.DataSink
+import com.intellij.openapi.actionSystem.PlatformDataKeys
 import com.intellij.openapi.editor.colors.EditorColors
 import com.intellij.openapi.editor.event.SelectionEvent
 import com.intellij.openapi.editor.event.SelectionListener
@@ -23,6 +27,14 @@ class SessionSelection : Disposable {
    private var style: SessionEditorStyle? = null
    private var clearing = false
    private var disposed = false
+    private val copy = object : TextCopyProvider() {
+        override fun getActionUpdateThread() = ActionUpdateThread.EDT
+
+        override fun getTextLinesToCopy(): Collection<String>? {
+            val text = selectedText()?.takeIf { it.isNotEmpty() } ?: return null
+            return listOf(text)
+        }
+    }

    @RequiresEdt
    fun selectedText(): String? {
@@ -37,6 +49,24 @@ class SessionSelection : Disposable {
        return item.selectedText()?.takeIf { it.isNotEmpty() }
    }

+    @RequiresEdt
+    fun provideCopy(sink: DataSink, content: (() -> String?)? = null) {
+        if (content == null) {
+            sink.set(PlatformDataKeys.COPY_PROVIDER, copy)
+            return
+        }
+        sink.set(PlatformDataKeys.COPY_PROVIDER, object : TextCopyProvider() {
+            override fun getActionUpdateThread() = ActionUpdateThread.EDT
+
+            override fun getTextLinesToCopy(): Collection<String>? {
+                val text = selectedText()?.takeIf { it.isNotEmpty() }
+                    ?: content()?.takeIf { it.isNotEmpty() }
+                    ?: return null
+                return listOf(text)
+            }
+        })
+    }
+
    @RequiresEdt
    fun register(component: JTextComponent, parent: Disposable? = null): Disposable {
        val item = TextItem(component)
@@ -488,7 +488,7 @@ class QuestionView(
     */
    @RequiresEdt
    private fun buildCustomEditor(): SessionEditorTextField {
-        val ed = SessionEditorTextField(project)
+        val ed = SessionEditorTextField(project, selection = selection)
        ed.border = JBUI.Borders.empty()
        ed.setFontInheritedFromLAF(false)
        ed.setPlaceholder(KiloBundle.message("session.question.custom.placeholder"))
@@ -506,6 +506,7 @@ class QuestionView(
            ex.settings.isAdditionalPageAtBottom = false
            ex.scrollPane.horizontalScrollBarPolicy = ScrollPaneConstants.HORIZONTAL_SCROLLBAR_NEVER
        }
+        selection?.register(ed)?.let(regs::add)
        ed.font = style.editorFont
        ed.background = style.editorScheme.defaultBackground

@@ -13,6 +13,8 @@ import ai.kilocode.client.ui.md.MdCodeBlockFactory
 import ai.kilocode.client.ui.md.MdCodeBlockOptions
 import ai.kilocode.client.ui.md.MdViewFactory
 import ai.kilocode.client.ui.md.hybrid.MdTerminal
+import com.intellij.openapi.actionSystem.DataSink
+import com.intellij.openapi.actionSystem.UiDataProvider
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.util.Disposer
 import com.intellij.ui.EditorTextField
@@ -26,10 +28,10 @@ import javax.swing.ScrollPaneConstants

 class ShellToolView(
    tool: Tool,
-    selection: SessionSelection? = null,
+    private val selection: SessionSelection? = null,
    private val parts: ToolParts = toolParts(tool),
    private val holder: ShellHolder = ShellHolder(tool, selection),
-) : SecondarySessionPartView(parts.header, { holder.body().panel }) {
+) : SecondarySessionPartView(parts.header, { holder.body().panel }), UiDataProvider {

    override val contentId: String = tool.id

@@ -43,6 +45,12 @@ class ShellToolView(
        sync()
    }

+    override fun uiDataSnapshot(sink: DataSink) {
+        selection?.provideCopy(sink) { holder.shell?.markdown() ?: fallbackText() }
+    }
+
+    private fun fallbackText() = ShellContent(item).body
+
    @RequiresEdt
    override fun expand(): Boolean {
        val changed = super.expand()
@@ -16,6 +16,8 @@ import ai.kilocode.client.ui.layout.VAlign
 import ai.kilocode.client.ui.layout.align
 import ai.kilocode.cli.KiloCliParser
 import ai.kilocode.log.KiloLog
+import com.intellij.openapi.actionSystem.DataSink
+import com.intellij.openapi.actionSystem.UiDataProvider
 import com.intellij.openapi.Disposable
 import com.intellij.openapi.editor.EditorFactory
 import com.intellij.openapi.fileTypes.PlainTextFileType
@@ -164,10 +166,14 @@ class ToolBody private constructor(
    fun register(selection: SessionSelection, parent: Disposable) {
        val field = ed
        if (field != null) {
+            (field as? ToolField)?.selection = selection
            selection.register(field, parent)
            return
        }
-        area?.let { selection.register(it, parent) }
+        area?.let {
+            (it as? ToolArea)?.selection = selection
+            selection.register(it, parent)
+        }
    }

    @RequiresEdt
@@ -233,7 +239,7 @@ class ToolBody private constructor(
            return body
        }

-        private fun area(tool: Tool, wrap: Boolean) = JBTextArea().apply {
+        private fun area(tool: Tool, wrap: Boolean) = ToolArea().apply {
            isEditable = false
            caret.isVisible = false
            caret.isSelectionVisible = true
@@ -272,6 +278,14 @@ class ToolBody private constructor(
    }
 }

+private class ToolArea : JBTextArea(), UiDataProvider {
+    var selection: SessionSelection? = null
+
+    override fun uiDataSnapshot(sink: DataSink) {
+        selection?.provideCopy(sink) { text }
+    }
+}
+
 private class ToolField(value: String, private var style: SessionEditorStyle) : EditorTextField(
    EditorFactory.getInstance().createDocument(value.trimEnd('\n')),
    ProjectManager.getInstance().defaultProject,
@@ -279,6 +293,8 @@ private class ToolField(value: String, private var style: SessionEditorStyle) :
    true,
    false,
 ) {
+    var selection: SessionSelection? = null
+
    init {
        setFontInheritedFromLAF(false)
        font = style.editorFont
@@ -296,6 +312,11 @@ private class ToolField(value: String, private var style: SessionEditorStyle) :
            ed.scrollPane.verticalScrollBarPolicy = ScrollPaneConstants.VERTICAL_SCROLLBAR_NEVER
        }
    }
+
+    override fun uiDataSnapshot(sink: DataSink) {
+        super.uiDataSnapshot(sink)
+        selection?.provideCopy(sink) { text }
+    }
 }

 private const val SUB_CARD = "sub"
@@ -8,6 +8,8 @@ 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 ai.kilocode.client.ui.UiStyle
+import com.intellij.openapi.actionSystem.DataSink
+import com.intellij.openapi.actionSystem.UiDataProvider
 import com.intellij.openapi.util.Disposer
 import com.intellij.util.concurrency.annotations.RequiresEdt
 import com.intellij.util.ui.JBUI
@@ -19,7 +21,7 @@ class ToolView(
    tool: Tool,
    private val selection: SessionSelection? = null,
    private val parts: ToolParts = toolParts(tool, mode = ToolBodyMode.EDITOR),
-) : SecondarySessionPartView(parts.header, { parts.scroll(tool) }) {
+) : SecondarySessionPartView(parts.header, { parts.scroll(tool) }), UiDataProvider {

    override val contentId: String = tool.id

@@ -34,6 +36,12 @@ class ToolView(
        sync()
    }

+    override fun uiDataSnapshot(sink: DataSink) {
+        selection?.provideCopy(sink) { parts.content?.text ?: fallbackText() }
+    }
+
+    private fun fallbackText() = listOf(commandText(), outputText()).filter { it.isNotBlank() }.joinToString("\n\n")
+
    @RequiresEdt
    override fun expand(): Boolean {
        val changed = super.expand()
@@ -11,6 +11,8 @@ import ai.kilocode.client.ui.md.MdView
 import ai.kilocode.log.KiloLog
 import com.intellij.execution.ui.ConsoleViewContentType
 import com.intellij.openapi.Disposable
+import com.intellij.openapi.actionSystem.DataSink
+import com.intellij.openapi.actionSystem.UiDataProvider
 import com.intellij.openapi.editor.EditorFactory
 import com.intellij.openapi.editor.markup.HighlighterLayer
 import com.intellij.openapi.editor.markup.HighlighterTargetArea
@@ -93,7 +95,7 @@ internal open class MdViewHybrid(
    private var tableBorderOverride: Color? = null
    private var opaqueState = true

-    private val root = JPanel().apply {
+    private val root = RootPanel().apply {
        layout = BoxLayout(this, BoxLayout.Y_AXIS)
        isOpaque = true
        background = opts().background
@@ -407,7 +409,7 @@ internal open class MdViewHybrid(

    private fun htmlBlock(body: String, disposable: Disposable): JBHtmlPane {
        val opts = opts()
-        return JBHtmlPane(
+        return object : JBHtmlPane(
            JBHtmlPaneStyleConfiguration {
                enableInlineCodeBackground = true
                enableCodeBlocksBackground = true
@@ -415,7 +417,11 @@ internal open class MdViewHybrid(
            JBHtmlPaneConfiguration {
                customStyleSheetProvider { sheet() }
            },
-        ).apply {
+        ), UiDataProvider {
+            override fun uiDataSnapshot(sink: DataSink) {
+                selection?.provideCopy(sink) { document.getText(0, document.length).trim() }
+            }
+        }.apply {
            isEditable = false
            isOpaque = opts.opaque
            background = opts.background
@@ -630,6 +636,17 @@ internal open class MdViewHybrid(
                ed.scrollPane.verticalScrollBarPolicy = ScrollPaneConstants.VERTICAL_SCROLLBAR_NEVER
            }
        }
+
+        override fun uiDataSnapshot(sink: DataSink) {
+            super.uiDataSnapshot(sink)
+            selection?.provideCopy(sink) { text }
+        }
+    }
+
+    private inner class RootPanel : JPanel(), UiDataProvider {
+        override fun uiDataSnapshot(sink: DataSink) {
+            selection?.provideCopy(sink) { markdown() }
+        }
    }

    private fun shellDisplay(term: Term, mode: Mode): ShellDisplay {
@@ -153,5 +153,9 @@
            <reference ref="Kilo.Session.Rename"/>
            <reference ref="Kilo.Session.Delete"/>
        </group>
+
+        <group id="Kilo.Session.ContextMenu">
+            <reference ref="$Copy"/>
+        </group>
    </actions>
 </idea-plugin>
@@ -318,6 +318,7 @@ action.Kilo.Session.Rename.description=Rename the selected session
 action.Kilo.Session.Delete.text=Delete
 action.Kilo.Session.Delete.description=Delete the selected session(s)
 action.Kilo.History.ContextMenu.text=History Actions
+action.Kilo.Session.ContextMenu.text=Session Actions

 # Migration wizard
 migration.migrate.title=Migrate Your Settings
@@ -414,9 +414,11 @@ class HistorySessionActionsTest : BasePlatformTestCase() {
        assertTrue(xml.contains("id=\"Kilo.Session.Rename\""))
        assertTrue(xml.contains("id=\"Kilo.Session.Delete\""))
        assertTrue(xml.contains("id=\"Kilo.History.ContextMenu\""))
+        assertTrue(xml.contains("id=\"Kilo.Session.ContextMenu\""))
        assertTrue(xml.contains("ref=\"Kilo.Session.Open\""))
        assertTrue(xml.contains("ref=\"Kilo.Session.Rename\""))
        assertTrue(xml.contains("ref=\"Kilo.Session.Delete\""))
+        assertTrue(xml.contains("ref=\"${'$'}Copy\""))
    }

    // ------ Helpers ------
@@ -11,17 +11,25 @@ import ai.kilocode.client.session.ui.prompt.PromptAttachmentPasteHandler
 import ai.kilocode.client.session.ui.prompt.PromptAttachmentPasteProvider
 import ai.kilocode.client.session.ui.prompt.PromptDataKeys
 import ai.kilocode.client.session.ui.prompt.PromptPanel
+import ai.kilocode.client.session.ui.selection.SessionSelection
 import com.intellij.icons.AllIcons
+import com.intellij.ide.CopyProvider
 import com.intellij.notification.Notification
 import com.intellij.notification.Notifications
 import com.intellij.openapi.actionSystem.CommonDataKeys
 import com.intellij.openapi.actionSystem.DataContext
+import com.intellij.openapi.actionSystem.DataKey
+import com.intellij.openapi.actionSystem.DataMap
+import com.intellij.openapi.actionSystem.DataProvider
 import com.intellij.openapi.actionSystem.DataSink
+import com.intellij.openapi.actionSystem.DataSnapshotProvider
+import com.intellij.openapi.actionSystem.PlatformDataKeys
 import com.intellij.openapi.actionSystem.UiDataProvider
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.editor.EditorFactory
 import com.intellij.openapi.editor.actions.PasteAction
+import com.intellij.openapi.ide.CopyPasteManager
 import com.intellij.openapi.keymap.KeymapUtil
 import com.intellij.testFramework.PlatformTestUtil
 import com.intellij.testFramework.fixtures.BasePlatformTestCase
@@ -44,6 +52,7 @@ import java.io.File
 import java.util.Base64
 import javax.imageio.ImageIO
 import javax.swing.JButton
+import javax.swing.JPanel
 import javax.swing.ImageIcon
 import javax.swing.SwingUtilities

@@ -113,6 +122,49 @@ class PromptPanelTest : BasePlatformTestCase() {
        assertEquals(min, editor.preferredSize.height)
    }

+    fun `test prompt editor exposes selection copy provider`() {
+        val selection = SessionSelection()
+        val panel = PromptPanel(project = project, selection = selection, onSend = { _, _ -> }, onAbort = {}, onEnhance = { _, _ -> })
+        val editor = panel.defaultFocusedComponent as EditorTextField
+        val host = JPanel()
+        host.add(panel)
+        host.addNotify()
+        try {
+            editor.text = "alpha prompt"
+
+            editor.getEditor(true)!!.selectionModel.setSelection(0, 5)
+            val sink = TestSink()
+            (editor as UiDataProvider).uiDataSnapshot(sink)
+            sink.copy!!.performCopy(DataContext.EMPTY_CONTEXT)
+
+            assertEquals("alpha", CopyPasteManager.getInstance().getContents(DataFlavor.stringFlavor))
+        } finally {
+            editor.getEditor(false)?.let(EditorFactory.getInstance()::releaseEditor)
+            selection.dispose()
+        }
+    }
+
+    fun `test prompt editor copies full content without selection`() {
+        val selection = SessionSelection()
+        val panel = PromptPanel(project = project, selection = selection, onSend = { _, _ -> }, onAbort = {}, onEnhance = { _, _ -> })
+        val editor = panel.defaultFocusedComponent as EditorTextField
+        val host = JPanel()
+        host.add(panel)
+        host.addNotify()
+        try {
+            editor.text = "alpha prompt"
+
+            val sink = TestSink()
+            (editor as UiDataProvider).uiDataSnapshot(sink)
+            sink.copy!!.performCopy(DataContext.EMPTY_CONTEXT)
+
+            assertEquals("alpha prompt", CopyPasteManager.getInstance().getContents(DataFlavor.stringFlavor))
+        } finally {
+            editor.getEditor(false)?.let(EditorFactory.getInstance()::releaseEditor)
+            selection.dispose()
+        }
+    }
+
    fun `test attachment only prompt can send`() {
        var sent = false
        val panel = PromptPanel(project, { text, files ->
@@ -673,9 +725,11 @@ class PromptPanelTest : BasePlatformTestCase() {

    private class TestSink : DataSink {
        var send: Any? = null
+        var copy: CopyProvider? = null

-        override fun <T : Any> set(key: com.intellij.openapi.actionSystem.DataKey<T>, data: T?) {
+        override fun <T : Any> set(key: DataKey<T>, data: T?) {
            if (key == PromptDataKeys.SEND) send = data
+            if (key == PlatformDataKeys.COPY_PROVIDER) copy = data as? CopyProvider
        }

        override fun <T : Any> setNull(key: com.intellij.openapi.actionSystem.DataKey<T>) {
@@ -685,20 +739,20 @@ class PromptPanelTest : BasePlatformTestCase() {
        }

        override fun <T : Any> lazyValue(
-            key: com.intellij.openapi.actionSystem.DataKey<T>,
-            data: (com.intellij.openapi.actionSystem.DataMap) -> T?,
+            key: DataKey<T>,
+            data: (DataMap) -> T?,
        ) {
        }

-        override fun uiDataSnapshot(provider: com.intellij.openapi.actionSystem.UiDataProvider) {
+        override fun uiDataSnapshot(provider: UiDataProvider) {
            provider.uiDataSnapshot(this)
        }

-        override fun dataSnapshot(provider: com.intellij.openapi.actionSystem.DataSnapshotProvider) {
+        override fun dataSnapshot(provider: DataSnapshotProvider) {
            provider.dataSnapshot(this)
        }

-        override fun uiDataSnapshot(provider: com.intellij.openapi.actionSystem.DataProvider) {
+        override fun uiDataSnapshot(provider: DataProvider) {
        }
    }

@@ -1,6 +1,7 @@
 package ai.kilocode.client.session.ui

 import ai.kilocode.client.session.SessionUiTestBase
+import ai.kilocode.client.session.ui.selection.SessionContextMenu
 import ai.kilocode.client.session.views.tool.ShellToolView
 import ai.kilocode.client.session.views.tool.ToolView
 import ai.kilocode.rpc.dto.ChatEventDto
@@ -15,17 +16,23 @@ import com.intellij.openapi.actionSystem.DataSnapshotProvider
 import com.intellij.openapi.actionSystem.PlatformDataKeys
 import com.intellij.openapi.actionSystem.UiDataProvider
 import com.intellij.openapi.ide.CopyPasteManager
+import com.intellij.ui.EditorTextField
+import com.intellij.ui.components.JBScrollPane
+import java.awt.Component
 import java.awt.Container
 import java.awt.datatransfer.DataFlavor
+import java.awt.Point
+import javax.swing.JComponent
+import javax.swing.JPanel
 import javax.swing.text.JTextComponent

@Suppress("UnstableApiUsage")
 class SessionSelectionCopyTest : SessionUiTestBase() {
-    fun `test session ui exposes copy provider when selection exists`() {
+    fun `test transcript view exposes copy provider when selection exists`() {
        val area = showTool("alpha output")

        select(area, "alpha")
-        val provider = copyProvider()
+        val provider = copyProvider(area)

        assertNotNull(provider)
        assertTrue(provider!!.isCopyEnabled(DataContext.EMPTY_CONTEXT))
@@ -35,23 +42,65 @@ class SessionSelectionCopyTest : SessionUiTestBase() {
        val area = showTool("alpha output")

        select(area, "alpha")
-        copyProvider()!!.performCopy(DataContext.EMPTY_CONTEXT)
+        copyProvider(area)!!.performCopy(DataContext.EMPTY_CONTEXT)

        assertEquals("alpha", CopyPasteManager.getInstance().getContents(DataFlavor.stringFlavor))
    }

+    fun `test copy provider writes full component text without selection`() {
+        val area = showTool("alpha output")
+
+        copyProvider(area)!!.performCopy(DataContext.EMPTY_CONTEXT)
+
+        assertEquals("alpha output", CopyPasteManager.getInstance().getContents(DataFlavor.stringFlavor))
+    }
+
    fun `test selecting another transcript component changes copied text`() {
        val one = showTool("alpha output", id = "tool_a")
        val two = showTool("bravo output", id = "tool_b")

        select(one, "alpha")
        select(two, "bravo")
-        copyProvider()!!.performCopy(DataContext.EMPTY_CONTEXT)
+        copyProvider(two)!!.performCopy(DataContext.EMPTY_CONTEXT)

        assertTrue(one.selectedText.isNullOrEmpty())
        assertEquals("bravo", CopyPasteManager.getInstance().getContents(DataFlavor.stringFlavor))
    }

+    fun `test code block child context exposes session copy provider`() {
+        showText("```text\nalpha code\n```")
+        val field = textEditors(ui).first { it.text.contains("alpha code") }
+        val editor = field.getEditor(true)!!
+
+        editor.selectionModel.setSelection(0, 5)
+        val provider = copyProvider(field as UiDataProvider)
+        provider!!.performCopy(DataContext.EMPTY_CONTEXT)
+
+        assertEquals("alpha", CopyPasteManager.getInstance().getContents(DataFlavor.stringFlavor))
+    }
+
+    fun `test code block child copies full content without selection`() {
+        showText("```text\nalpha code\n```")
+        val field = textEditors(ui).first { it.text.contains("alpha code") }
+
+        copyProvider(field as UiDataProvider)!!.performCopy(DataContext.EMPTY_CONTEXT)
+
+        assertEquals("alpha code", CopyPasteManager.getInstance().getContents(DataFlavor.stringFlavor))
+    }
+
+    fun `test session context menu resolves deepest component`() {
+        val root = JPanel(null)
+        val mid = JPanel(null)
+        val child = JPanel(null)
+        root.setBounds(0, 0, 100, 100)
+        mid.setBounds(10, 10, 80, 80)
+        child.setBounds(5, 5, 20, 20)
+        root.add(mid)
+        mid.add(child)
+
+        assertSame(child, SessionContextMenu.target(root, root, Point(20, 20)))
+    }
+
    private fun select(area: JTextComponent, text: String) {
        val start = area.text.indexOf(text)
        assertTrue(start >= 0)
@@ -79,6 +128,13 @@ class SessionSelectionCopyTest : SessionUiTestBase() {
        return textComponent(text)
    }

+    private fun showText(text: String) {
+        if (controller().id == null) showMessages()
+        emit(ChatEventDto.MessageUpdated("ses_test", message("msg_text")))
+        emit(ChatEventDto.PartUpdated("ses_test", part("part_text", "msg_text", "text", text)))
+        layout()
+    }
+
    private fun toolViews(root: Container): List<Container> {
        val out = mutableListOf<Container>()
        if (root is ShellToolView || root is ToolView) out.add(root)
@@ -94,12 +150,46 @@ class SessionSelectionCopyTest : SessionUiTestBase() {
        else -> false
    }

-    private fun copyProvider(): CopyProvider? {
+    private fun copyProvider(provider: UiDataProvider): CopyProvider? {
        val sink = CopySink()
-        (ui as UiDataProvider).uiDataSnapshot(sink)
+        provider.uiDataSnapshot(sink)
        return sink.copy
    }

+    private fun copyProvider(component: Component): CopyProvider? {
+        (component as? UiDataProvider)?.let(::copyProvider)?.let { return it }
+        ancestors(component).filterIsInstance<UiDataProvider>().firstNotNullOfOrNull(::copyProvider)?.let { return it }
+        val point = Point((component.width / 2).coerceAtLeast(0), (component.height / 2).coerceAtLeast(0))
+        val target = SessionContextMenu.target(ui as JComponent, component, point) ?: component
+        ancestors(target).filterIsInstance<UiDataProvider>().firstNotNullOfOrNull(::copyProvider)?.let { return it }
+        return providers(ui).firstNotNullOfOrNull(::copyProvider)
+    }
+
+    private fun providers(root: Component): Sequence<UiDataProvider> = sequence {
+        if (root is UiDataProvider) yield(root)
+        if (root is Container) {
+            for (child in root.components) yieldAll(providers(child))
+        }
+    }
+
+    private fun ancestors(component: Component): Sequence<Component> = sequence {
+        var comp: Component? = component
+        while (comp != null) {
+            yield(comp)
+            comp = comp.parent
+        }
+    }
+
+    private fun textEditors(root: Container): List<EditorTextField> {
+        val out = mutableListOf<EditorTextField>()
+        if (root is EditorTextField) out.add(root)
+        for (child in root.components) {
+            if (child is JBScrollPane) (child.viewport.view as? EditorTextField)?.let(out::add)
+            if (child is Container) out.addAll(textEditors(child))
+        }
+        return out.distinct()
+    }
+
    private fun textComponent(needle: String): JTextComponent = textComponents(ui)
        .first { it.text.contains(needle) }

@@ -1,9 +1,19 @@
 package ai.kilocode.client.ui.md

 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 com.intellij.ide.CopyProvider
 import com.intellij.execution.process.ProcessOutputTypes
 import com.intellij.execution.ui.ConsoleViewContentType
+import com.intellij.openapi.actionSystem.DataContext
+import com.intellij.openapi.actionSystem.DataKey
+import com.intellij.openapi.actionSystem.DataMap
+import com.intellij.openapi.actionSystem.DataProvider
+import com.intellij.openapi.actionSystem.DataSink
+import com.intellij.openapi.actionSystem.DataSnapshotProvider
+import com.intellij.openapi.actionSystem.PlatformDataKeys
+import com.intellij.openapi.actionSystem.UiDataProvider
 import com.intellij.openapi.editor.DefaultLanguageHighlighterColors
 import com.intellij.openapi.editor.HighlighterColors
 import com.intellij.openapi.editor.colors.CodeInsightColors
@@ -15,6 +25,7 @@ import com.intellij.openapi.fileTypes.FileTypeRegistry
 import com.intellij.openapi.fileTypes.PlainTextFileType
 import com.intellij.openapi.fileTypes.UnknownFileType
 import com.intellij.openapi.editor.markup.TextAttributes
+import com.intellij.openapi.ide.CopyPasteManager
 import com.intellij.openapi.util.Disposer
 import com.intellij.testFramework.fixtures.BasePlatformTestCase
 import com.intellij.ui.EditorTextField
@@ -28,6 +39,7 @@ import java.awt.Font
 import javax.swing.Box
 import javax.swing.JPanel
 import javax.swing.ScrollPaneConstants
+import java.awt.datatransfer.DataFlavor

@Suppress("UnstableApiUsage")
 class MdViewHybridTest : BasePlatformTestCase() {
@@ -793,6 +805,33 @@ class MdViewHybridTest : BasePlatformTestCase() {
        assertEquals("https://example.com", received.single().href)
    }

+    fun `test markdown root and code child expose selection copy provider`() {
+        Disposer.dispose(view)
+        disposed = true
+        val selection = SessionSelection()
+        val local = MdViewFactory.hybrid(selection = selection)
+        try {
+            local.set("```text\nalpha code\n```")
+            val field = (local.component as JPanel).components
+                .filterIsInstance<JBScrollPane>()
+                .mapNotNull { it.viewport.view as? EditorTextField }
+                .single()
+            field.getEditor(true)!!.selectionModel.setSelection(0, 5)
+
+            val root = CopySink()
+            (local.component as UiDataProvider).uiDataSnapshot(root)
+            val child = CopySink()
+            (field as UiDataProvider).uiDataSnapshot(child)
+            child.copy!!.performCopy(DataContext.EMPTY_CONTEXT)
+
+            assertNotNull(root.copy)
+            assertEquals("alpha", CopyPasteManager.getInstance().getContents(DataFlavor.stringFlavor))
+        } finally {
+            Disposer.dispose(local)
+            selection.dispose()
+        }
+    }
+
    private fun scrolls(): List<JBScrollPane> = (view.component as JPanel).components.filterIsInstance<JBScrollPane>()

    private fun htmls(): List<JBHtmlPane> = (view.component as JPanel).components.filterIsInstance<JBHtmlPane>()
@@ -820,6 +859,26 @@ class MdViewHybridTest : BasePlatformTestCase() {
        UIUtil.dispatchAllInvocationEvents()
    }

+    private class CopySink : DataSink {
+        var copy: CopyProvider? = null
+
+        override fun <T : Any> set(key: DataKey<T>, data: T?) {
+            if (key == PlatformDataKeys.COPY_PROVIDER) copy = data as? CopyProvider
+        }
+
+        override fun <T : Any> setNull(key: DataKey<T>) {}
+
+        override fun <T : Any> lazyNull(key: DataKey<T>) {}
+
+        override fun <T : Any> lazyValue(key: DataKey<T>, data: (DataMap) -> T?) {}
+
+        override fun uiDataSnapshot(provider: UiDataProvider) = provider.uiDataSnapshot(this)
+
+        override fun dataSnapshot(provider: DataSnapshotProvider) = provider.dataSnapshot(this)
+
+        override fun uiDataSnapshot(provider: DataProvider) {}
+    }
+
    private fun customStyle(): SessionEditorStyle {
        val scheme = EditorColorsManager.getInstance().globalScheme.clone() as EditorColorsScheme
        scheme.setAttributes(