fix(jetbrains): clean up prompt enhancement lifecycle (#11176)

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

@@ -125,6 +125,8 @@ class SessionController(
     ) { sid ?: ref?.key ?: "pending" }
 
     private var disposed = false
+    private var enhancement = 0L
+    private val enhancements = mutableMapOf<Long, (Result<String>) -> Unit>()
     private var partType: String? = null
     private var tool: String? = null
     private var eventJob: Job? = null
@@ -198,23 +200,31 @@ class SessionController(
 
     fun enhancePrompt(text: String, complete: (Result<String>) -> Unit) {
         assertEdt()
+        if (disposed) {
+            complete(Result.failure(CancellationException("Session controller disposed")))
+            return
+        }
+        val id = ++enhancement
+        enhancements[id] = complete
         capture("Prompt Enhance Clicked", mapOf("textLength" to bucket(text)))
         cs.launch {
             val result = try {
                 Result.success(sessions.enhancePrompt(directory, text))
             } catch (e: CancellationException) {
-                throw e
+                Result.failure(e)
             } catch (e: Exception) {
                 Result.failure(e)
             }
             edt {
-                if (disposed) return@edt
+                val callback = enhancements.remove(id) ?: return@edt
                 result.onSuccess {
                     capture("Prompt Enhanced", mapOf("textLength" to bucket(text)))
                 }.onFailure { e ->
-                    capture("Session Error", mapOf("context" to "enhance-prompt", "errorClass" to e::class.java.name))
+                    if (e !is CancellationException) {
+                        capture("Session Error", mapOf("context" to "enhance-prompt", "errorClass" to e::class.java.name))
+                    }
                 }
-                complete(result)
+                callback(result)
             }
         }
     }
@@ -1740,13 +1750,18 @@ class SessionController(
 
     override fun dispose() {
         runEdt {
+            if (disposed) return@runEdt
             disposed = true
             connectionDelay.dispose()
             cancelSubscriptions()
             drainJob?.cancel()
             drainJob = null
+            val callbacks = enhancements.values.toList()
+            enhancements.clear()
+            cs.cancel()
+            val result = Result.failure<String>(CancellationException("Session controller disposed"))
+            callbacks.forEach { it(result) }
         }
-        cs.cancel()
     }
 
     override fun toString(): String {

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

@@ -43,6 +43,7 @@ import com.intellij.util.ui.JBUI
 import com.intellij.util.ui.UIUtil
 import com.intellij.util.ui.components.BorderLayoutPanel
 import com.intellij.util.messages.MessageBusConnection
+import kotlinx.coroutines.CancellationException
 import java.awt.BorderLayout
 import java.awt.Cursor
 import java.awt.Graphics
@@ -150,6 +151,7 @@ class PromptPanel(
         addActionListener { onAutoApproveToggle(!autoApprove) }
     }
 
+    private val enhancingIcon = AnimatedIcon.Default()
     private val enhance = HoverIcon().apply {
         icon = WAND_ICON
         toolTipText = KiloBundle.message("prompt.action.enhance")
@@ -326,6 +328,7 @@ class PromptPanel(
             editor.text = it
             focus()
         }.onFailure {
+            if (it is CancellationException) return@onFailure
             KiloNotifications.error(
                 project,
                 KiloBundle.message("prompt.action.enhance.failed"),
@@ -344,7 +347,7 @@ class PromptPanel(
     @RequiresEdt
     private fun syncEnhance() {
         enhance.isEnabled = ready && !busy && !enhancing
-        enhance.icon = if (enhancing) AnimatedIcon.Default() else WAND_ICON
+        enhance.icon = if (enhancing) enhancingIcon else WAND_ICON
         enhance.toolTipText = if (enhancing) {
             KiloBundle.message("prompt.action.enhance.loading")
         } else {

packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/controller/PromptEnhancerTest.kt

@@ -1,6 +1,7 @@
 package ai.kilocode.client.session.controller
 
 import com.intellij.openapi.application.ApplicationManager
+import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.runBlocking
 
@@ -36,18 +37,31 @@ class PromptEnhancerTest : SessionControllerTestBase() {
         assertSame(before, edt { controller.model.state })
     }
 
-    fun `test enhance prompt ignores completion after disposal`() {
+    fun `test enhance prompt cancels pending completions on disposal`() {
         val controller = controller()
         val gate = CompletableDeferred<Unit>()
+        val results = mutableListOf<Result<String>>()
         rpc.enhanceGate = gate
-        var completed = false
 
-        edt { controller.enhancePrompt("make a plan") { completed = true } }
+        edt {
+            controller.enhancePrompt("make a plan") {
+                assertTrue(ApplicationManager.getApplication().isDispatchThread)
+                results.add(it)
+            }
+            controller.enhancePrompt("rewrite a plan") {
+                assertTrue(ApplicationManager.getApplication().isDispatchThread)
+                results.add(it)
+            }
+        }
         settle()
         controller.dispose()
+
+        assertEquals(2, results.size)
+        assertTrue(results.all { it.exceptionOrNull() is CancellationException })
+
         runBlocking { gate.complete(Unit) }
         settle()
 
-        assertFalse(completed)
+        assertEquals(2, results.size)
     }
 }

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

@@ -5,13 +5,17 @@ import ai.kilocode.client.plugin.KiloBundle
 import ai.kilocode.client.session.ui.prompt.PromptDataKeys
 import ai.kilocode.client.session.ui.prompt.PromptPanel
 import com.intellij.icons.AllIcons
+import com.intellij.notification.Notification
+import com.intellij.notification.Notifications
 import com.intellij.openapi.actionSystem.DataSink
 import com.intellij.openapi.actionSystem.UiDataProvider
+import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.keymap.KeymapUtil
 import com.intellij.testFramework.fixtures.BasePlatformTestCase
 import com.intellij.ui.AnimatedIcon
 import com.intellij.ui.EditorTextField
 import com.intellij.util.ui.EmptyIcon
+import kotlinx.coroutines.CancellationException
 import java.awt.Container
 import javax.swing.JButton
 import javax.swing.SwingUtilities
@@ -230,6 +234,11 @@ class PromptPanelTest : BasePlatformTestCase() {
         assertEquals("make a plan", seen)
         assertFalse(enhance.isEnabled)
         assertTrue(enhance.icon is AnimatedIcon)
+        val icon = enhance.icon
+
+        panel.setReady(true)
+
+        assertSame(icon, enhance.icon)
 
         complete!!(Result.success("Use a focused implementation plan"))
 
@@ -254,6 +263,31 @@ class PromptPanelTest : BasePlatformTestCase() {
         assertTrue(enhance.isEnabled)
     }
 
+    fun `test cancelled enhancement restores button without notification`() {
+        val notes = mutableListOf<Notification>()
+        val listener = object : Notifications {
+            override fun notify(notification: Notification) {
+                notes.add(notification)
+            }
+        }
+        ApplicationManager.getApplication().messageBus.connect(testRootDisposable).subscribe(Notifications.TOPIC, listener)
+        project.messageBus.connect(testRootDisposable).subscribe(Notifications.TOPIC, listener)
+        var complete: ((Result<String>) -> Unit)? = null
+        val panel = PromptPanel(project, {}, {}, { _, done -> complete = done })
+        val editor = panel.defaultFocusedComponent as EditorTextField
+        val enhance = enhanceButton(panel)
+        panel.setReady(true)
+        editor.text = "keep this draft"
+
+        enhance.doClick()
+        complete!!(Result.failure(CancellationException("disposed")))
+
+        assertEquals("keep this draft", editor.text)
+        assertTrue(enhance.isEnabled)
+        assertFalse(enhance.icon is AnimatedIcon)
+        assertTrue(notes.isEmpty())
+    }
+
     fun `test empty enhancement inserts explanation without request`() {
         var requests = 0
         val panel = PromptPanel(project, {}, {}, { _, _ -> requests++ })