From f8b17e2f67f9f52295be09da1e510991669e66c4 Mon Sep 17 00:00:00 2001 From: Marius Date: Fri, 12 Jun 2026 19:10:29 +0200 Subject: [PATCH] fix(jetbrains): clean up prompt enhancement lifecycle (#11176) --- .../session/controller/SessionController.kt | 25 +++++++++++--- .../client/session/ui/prompt/PromptPanel.kt | 5 ++- .../session/controller/PromptEnhancerTest.kt | 22 +++++++++--- .../client/session/ui/PromptPanelTest.kt | 34 +++++++++++++++++++ 4 files changed, 76 insertions(+), 10 deletions(-) diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/controller/SessionController.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/controller/SessionController.kt index b869c15a46..940437682e 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/controller/SessionController.kt +++ b/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) -> 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) -> 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(CancellationException("Session controller disposed")) + callbacks.forEach { it(result) } } - cs.cancel() } override fun toString(): String { diff --git a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/prompt/PromptPanel.kt b/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/prompt/PromptPanel.kt index 83a159e372..e154d67ccf 100644 --- a/packages/kilo-jetbrains/frontend/src/main/kotlin/ai/kilocode/client/session/ui/prompt/PromptPanel.kt +++ b/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 { diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/controller/PromptEnhancerTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/controller/PromptEnhancerTest.kt index b3de87418d..4c07869a7a 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/controller/PromptEnhancerTest.kt +++ b/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() + val results = mutableListOf>() 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) } } diff --git a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/ui/PromptPanelTest.kt b/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/ui/PromptPanelTest.kt index f9afc6d88e..2edb90556d 100644 --- a/packages/kilo-jetbrains/frontend/src/test/kotlin/ai/kilocode/client/session/ui/PromptPanelTest.kt +++ b/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() + 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) -> 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++ })