fix(jetbrains): enable prompt undo redo

.changeset/jetbrains-prompt-mentions-undo.md

@@ -0,0 +1,5 @@
+---
+"@kilocode/kilo-jetbrains": patch
+---
+
+Fix prompt undo/redo, clean mid-token mention completion, and show unresolved file mentions in the JetBrains chat prompt.

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

@@ -2,14 +2,24 @@ package ai.kilocode.client.session.ui.editor
 
 import ai.kilocode.client.session.ui.prompt.PromptDataKeys
 import ai.kilocode.client.session.ui.prompt.SendPromptContext
+import com.intellij.ide.actions.UndoRedoAction
+import com.intellij.openapi.actionSystem.ActionManager
+import com.intellij.openapi.actionSystem.AnActionEvent
 import com.intellij.openapi.actionSystem.DataSink
+import com.intellij.openapi.actionSystem.IdeActions
+import com.intellij.openapi.actionSystem.PlatformCoreDataKeys
+import com.intellij.openapi.command.undo.UndoManager
+import com.intellij.openapi.editor.Editor
+import com.intellij.openapi.fileEditor.TextEditor
+import com.intellij.openapi.fileEditor.impl.text.TextEditorProvider
 import com.intellij.openapi.fileTypes.PlainTextFileType
+import com.intellij.openapi.fileTypes.PlainTextLanguage
+import com.intellij.openapi.project.DumbAwareAction
 import com.intellij.openapi.project.Project
 import com.intellij.ui.EditorTextField
 import com.intellij.ui.LanguageTextField
 import com.intellij.util.textCompletion.TextCompletionProvider
 import com.intellij.util.textCompletion.TextCompletionUtil
-import com.intellij.openapi.fileTypes.PlainTextLanguage
 
 /**
  * A session-scoped [EditorTextField] for plain-text input.
@@ -26,7 +36,7 @@ import com.intellij.openapi.fileTypes.PlainTextLanguage
  * wrapping here.
  */
 internal open class SessionEditorTextField(
-    project: Project,
+    private val project: Project,
     private val ctx: SendPromptContext? = null,
     completion: TextCompletionProvider? = null,
 ) : EditorTextField(
@@ -41,8 +51,57 @@ internal open class SessionEditorTextField(
     project,
     PlainTextFileType.INSTANCE,
 ) {
+    private val undo = action("Kilo Session Undo", true)
+    private val redo = action("Kilo Session Redo", false)
+
+    init {
+        addSettingsProvider(::install)
+    }
+
     override fun uiDataSnapshot(sink: DataSink) {
         super.uiDataSnapshot(sink)
         ctx?.let { sink.set(PromptDataKeys.SEND, it) }
+        file()?.let { sink.set(PlatformCoreDataKeys.FILE_EDITOR, it) }
+    }
+
+    private fun install(editor: Editor) {
+        editor.contentComponent.putClientProperty(UndoRedoAction.IGNORE_SWING_UNDO_MANAGER, true)
+        // Workaround: global $Undo/$Redo can miss the synthetic FileEditor for this embedded
+        // EditorTextField. Bind the shortcuts locally until the platform data context targets it reliably.
+        val manager = ActionManager.getInstance()
+        manager.getAction(IdeActions.ACTION_UNDO)?.shortcutSet?.let {
+            undo.registerCustomShortcutSet(it, editor.contentComponent)
+        }
+        manager.getAction(IdeActions.ACTION_REDO)?.shortcutSet?.let {
+            redo.registerCustomShortcutSet(it, editor.contentComponent)
+        }
+    }
+
+    private fun action(text: String, undo: Boolean) = object : DumbAwareAction(text) {
+        override fun update(e: AnActionEvent) {
+            e.presentation.isEnabled = available(undo)
+        }
+
+        override fun actionPerformed(e: AnActionEvent) {
+            if (project.isDisposed) return
+            val file = file() ?: return
+            val manager = UndoManager.getInstance(project)
+            if (undo) {
+                if (manager.isUndoAvailable(file)) manager.undo(file)
+                return
+            }
+            if (manager.isRedoAvailable(file)) manager.redo(file)
+        }
+    }
+
+    private fun available(undo: Boolean): Boolean {
+        if (project.isDisposed) return false
+        val file = file() ?: return false
+        val manager = UndoManager.getInstance(project)
+        return if (undo) manager.isUndoAvailable(file) else manager.isRedoAvailable(file)
+    }
+
+    private fun file(): TextEditor? {
+        return getEditor(false)?.let(TextEditorProvider.getInstance()::getTextEditor)
     }
 }

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

@@ -18,6 +18,7 @@ import ai.kilocode.client.session.ui.prompt.SlashAction
 import ai.kilocode.client.testing.FakeWorkspaceRpcApi
 import ai.kilocode.rpc.dto.FileSearchResultDto
 import ai.kilocode.rpc.dto.WorkspaceFileDto
+import com.intellij.ide.actions.UndoRedoAction
 import com.intellij.icons.AllIcons
 import com.intellij.codeInsight.lookup.Lookup
 import com.intellij.codeInsight.lookup.LookupManager
@@ -25,6 +26,7 @@ import com.intellij.codeInsight.lookup.LookupPositionStrategy
 import com.intellij.codeInsight.lookup.impl.LookupImpl
 import com.intellij.notification.Notification
 import com.intellij.notification.Notifications
+import com.intellij.openapi.actionSystem.ActionManager
 import com.intellij.openapi.actionSystem.ActionPlaces
 import com.intellij.openapi.actionSystem.ActionUiKind
 import com.intellij.openapi.actionSystem.AnAction
@@ -32,17 +34,18 @@ import com.intellij.openapi.actionSystem.AnActionEvent
 import com.intellij.openapi.actionSystem.CommonDataKeys
 import com.intellij.openapi.actionSystem.DataContext
 import com.intellij.openapi.actionSystem.DataSink
+import com.intellij.openapi.actionSystem.IdeActions
 import com.intellij.openapi.actionSystem.PlatformCoreDataKeys
 import com.intellij.openapi.actionSystem.UiDataProvider
 import com.intellij.openapi.actionSystem.ex.ActionUtil
 import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.command.WriteCommandAction
+import com.intellij.openapi.command.undo.UndoManager
 import com.intellij.openapi.editor.DefaultLanguageHighlighterColors
 import com.intellij.openapi.editor.Editor
 import com.intellij.openapi.editor.EditorFactory
 import com.intellij.openapi.editor.actions.PasteAction
 import com.intellij.openapi.editor.colors.CodeInsightColors
-import com.intellij.openapi.command.undo.UndoManager
-import com.intellij.openapi.command.WriteCommandAction
 import com.intellij.openapi.fileEditor.TextEditor
 import com.intellij.openapi.keymap.KeymapUtil
 import com.intellij.testFramework.PlatformTestUtil
@@ -272,6 +275,44 @@ class PromptPanelTest : BasePlatformTestCase() {
         assertEquals("hello", editor.document.text)
     }
 
+    fun `test prompt editor platform undo redo actions target prompt editor`() {
+        val panel = PromptPanel(project = project, onSend = { _, _ -> }, onAbort = {}, onEnhance = { _, _ -> }, completion = completion())
+        val field = panel.defaultFocusedComponent as EditorTextField
+
+        realize(panel, 260, 400)
+        val editor = field.getEditor(false)!!
+        WriteCommandAction.runWriteCommandAction(project) {
+            editor.document.insertString(0, "hello")
+        }
+        assertSame(true, editor.contentComponent.getClientProperty(UndoRedoAction.IGNORE_SWING_UNDO_MANAGER))
+        val sink = TestSink()
+        (field as UiDataProvider).uiDataSnapshot(sink)
+        val file = sink.file as? TextEditor ?: error("missing file editor")
+        assertSame(editor.document, file.editor.document)
+        assertTrue("prompt file editor should have undo", UndoManager.getInstance(project).isUndoAvailable(file))
+
+        invokeAction(IdeActions.ACTION_UNDO, editor.contentComponent, file)
+        assertEquals("", editor.document.text)
+        invokeAction(IdeActions.ACTION_REDO, editor.contentComponent, file)
+        assertEquals("hello", editor.document.text)
+    }
+
+    fun `test prompt editor component undo redo shortcuts target prompt editor`() {
+        val panel = PromptPanel(project = project, onSend = { _, _ -> }, onAbort = {}, onEnhance = { _, _ -> }, completion = completion())
+        val field = panel.defaultFocusedComponent as EditorTextField
+
+        realize(panel, 260, 400)
+        val editor = field.getEditor(false)!!
+        WriteCommandAction.runWriteCommandAction(project) {
+            editor.document.insertString(0, "hello")
+        }
+
+        invokeComponentAction("Kilo Session Undo", editor)
+        assertEquals("", editor.document.text)
+        invokeComponentAction("Kilo Session Redo", editor)
+        assertEquals("hello", editor.document.text)
+    }
+
     fun `test prompt editor highlights missing mention as wrong reference`() {
         val panel = PromptPanel(project = project, onSend = { _, _ -> }, onAbort = {}, onEnhance = { _, _ -> }, completion = completion())
         val field = panel.defaultFocusedComponent as EditorTextField
@@ -973,6 +1014,34 @@ class PromptPanelTest : BasePlatformTestCase() {
         ActionUtil.performAction(action, event)
     }
 
+    private fun invokeComponentAction(text: String, editor: Editor) {
+        val action = ActionUtil.getActions(editor.contentComponent).first { item ->
+            item.templatePresentation.text == text
+        }
+        val event = event(action, editor)
+        ActionUtil.updateAction(action, event)
+        assertTrue("action $text should be enabled", event.presentation.isEnabled)
+        ActionUtil.performAction(action, event)
+        UIUtil.dispatchAllInvocationEvents()
+    }
+
+    private fun invokeAction(id: String, component: java.awt.Component, file: TextEditor) {
+        val action = ActionManager.getInstance().getAction(id) ?: error("missing action $id")
+        val ctx = DataContext { data ->
+            when (data) {
+                CommonDataKeys.PROJECT.name -> project
+                PlatformCoreDataKeys.CONTEXT_COMPONENT.name -> component
+                PlatformCoreDataKeys.FILE_EDITOR.name -> file
+                else -> null
+            }
+        }
+        val event = AnActionEvent.createEvent(action, ctx, null, ActionPlaces.UNKNOWN, ActionUiKind.NONE, null)
+        ActionUtil.updateAction(action, event)
+        assertTrue("action $id should be enabled", event.presentation.isEnabled)
+        ActionUtil.performAction(action, event)
+        UIUtil.dispatchAllInvocationEvents()
+    }
+
     private fun waitForLookupItems(editor: Editor): List<String> {
         repeat(50) {
             UIUtil.dispatchAllInvocationEvents()