fix(jetbrains): show revert progress inline

.changeset/jetbrains-revert-loading.md

@@ -0,0 +1,5 @@
+---
+"@kilocode/kilo-jetbrains": patch
+---
+
+Show rollback/redo progress inline (on the message and redo controls) with a cancel action instead of a full-screen loading overlay.

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

@@ -243,6 +243,7 @@ class SessionUi(
         is SessionState.Idle,
         is SessionState.Loading,
         is SessionState.Busy,
+        is SessionState.Reverting,
         is SessionState.Retry,
         is SessionState.Offline,
         is SessionState.Error -> null
@@ -361,7 +362,8 @@ class SessionUi(
             repo = workspace.directory,
             resize = { anchor, fn -> scroll.preserve(anchor, fn) },
             revert = ::revert,
-            banner = RevertBanner(controller.model, controller::redo, controller::redoAll, focus),
+            cancelRevert = controller::cancelRevert,
+            banner = RevertBanner(controller.model, controller::redo, controller::redoAll, controller::cancelRevert, focus),
         ).also {
             it.onHover = { view, on -> if (on) popup.show(view) else popup.notifyExit(view) }
         }

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

@@ -23,6 +23,7 @@ import ai.kilocode.client.session.model.Text
 import ai.kilocode.client.plugin.KiloPluginSettings
 import ai.kilocode.client.session.SessionRef
 import ai.kilocode.client.telemetry.Telemetry
+import ai.kilocode.client.util.UiTimer
 import ai.kilocode.client.util.UiTimerSource
 import ai.kilocode.client.util.UiTimers
 import ai.kilocode.rpc.dto.ChatEventDto
@@ -89,6 +90,7 @@ class SessionController(
   private val flushMs: Long = EVENT_FLUSH_MS,
   private val condense: Boolean = true,
   private val displayMs: Long = DISPLAY_DELAY_MS,
+  private val revertTimeoutMs: Long = REVERT_TIMEOUT_MS,
   private val open: (SessionRef) -> Unit = {},
   private val beforeUpdate: () -> Boolean = { false },
   private val afterUpdate: (Boolean) -> Unit = {},
@@ -102,6 +104,7 @@ class SessionController(
     private data class OrganizationTarget(val org: String?)
     private data class Followup(val dir: String, val time: Long)
     private data class Pref(val agent: String?, val model: String?, val variants: List<String>, val variant: String?, val reset: Boolean)
+    private data class RevertOp(val key: Long)
     private data class Dispatch(
         val kind: String,
         val source: String,
@@ -115,6 +118,7 @@ class SessionController(
         private val LOG = KiloLog.create(SessionController::class.java)
         internal const val RECENT_LIMIT = 5
         internal const val DISPLAY_DELAY_MS = 1_000L
+        internal const val REVERT_TIMEOUT_MS = 30_000L
         private const val FOLLOWUP_TTL_MS = 30_000L
         private const val FOLLOWUP_NEW_SESSION = "Start new session"
     }
@@ -147,6 +151,10 @@ class SessionController(
     private var tool: String? = null
     private var eventJob: Job? = null
     private var drainJob: Job? = null
+    private var revertJob: Job? = null
+    private var revertOp: RevertOp? = null
+    private var revertSeq = 0L
+    private var revertWatchdog: UiTimer? = null
     private var creating: CompletableDeferred<String?>? = null
     private val childJobs: MutableMap<String, Job> = mutableMapOf()
     private val childIds: MutableSet<String> = mutableSetOf()
@@ -271,6 +279,7 @@ class SessionController(
 
     private fun dispatch(data: Dispatch, send: suspend (String) -> Unit) {
         assertEdt()
+        if (revertOp != null) return
         val props = data.props + if (data.kind == "command") slashProps() else emptyMap()
         capture("Conversation Send Clicked", sessionProps(sid ?: ref?.key) + mapOf(
             "source" to data.source,
@@ -343,6 +352,10 @@ class SessionController(
     fun abort() {
         assertEdt()
         LOG.debug { "${ChatLogSummary.sid(sid ?: ref?.key ?: "pending")} kind=abort" }
+        if (revertOp != null) {
+            cancelRevert()
+            return
+        }
         val id = sid ?: return
         capture("Session Stop Clicked", sessionProps(id))
         cs.launch {
@@ -379,7 +392,7 @@ class SessionController(
     fun compact() {
         assertEdt()
         val id = sid ?: return
-        if (model.state.isBusy()) return
+        if (revertOp != null || model.state.isBusy()) return
         if (model.isEmpty()) return
         val parsed = model.model?.let(::parseModel) ?: return
         val sel = ModelSelectionDto(parsed.first, parsed.second)
@@ -411,12 +424,19 @@ class SessionController(
             )
             return
         }
-        val busy = model.state.isBusy()
+        val state = model.state
+        if (revertOp != null) return
+        val busy = state.isBusy()
         LOG.info(
             "${ChatLogSummary.sid(id)} kind=revert clicked=true message=$message " +
                 "part=${part ?: "none"} busy=$busy",
         )
-        cs.launch {
+        val op = beginReverting(
+            KiloBundle.message("session.status.rollingback"),
+            SessionState.Reverting.Kind.ROLLBACK,
+            message,
+        ) ?: return
+        revertJob = cs.launch {
             try {
                 if (busy) {
                     LOG.info("${ChatLogSummary.sid(id)} kind=revert abort=true reason=busy")
@@ -427,24 +447,37 @@ class SessionController(
                 capture("Session Rollback", sessionProps(id))
                 synchronizeFromDisk(id, "revert")
                 LOG.info("${ChatLogSummary.sid(id)} kind=revert ok=true")
+                edt { clearReverting(op) }
+            } catch (e: CancellationException) {
+                throw e
             } catch (e: Exception) {
                 capture("Session Error", sessionProps(id) + mapOf("context" to "revert", "errorClass" to e::class.java.name))
                 LOG.warn("${ChatLogSummary.sid(id)} kind=revert dir=${ChatLogSummary.dir(directory)} failed message=${e.message}", e)
+                edt { failReverting(op, e) }
             }
         }
     }
 
     fun unrevert() {
         assertEdt()
         val id = sid ?: return
-        cs.launch {
+        if (revertOp != null) return
+        val op = beginReverting(
+            KiloBundle.message("session.status.redoing"),
+            SessionState.Reverting.Kind.REDO,
+        ) ?: return
+        revertJob = cs.launch {
             try {
                 sessions.unrevert(id, directory)
                 capture("Session Unrevert", sessionProps(id))
                 synchronizeFromDisk(id, "unrevert")
+                edt { clearReverting(op) }
+            } catch (e: CancellationException) {
+                throw e
             } catch (e: Exception) {
                 capture("Session Error", sessionProps(id) + mapOf("context" to "unrevert", "errorClass" to e::class.java.name))
                 LOG.warn("${ChatLogSummary.sid(id)} kind=unrevert dir=${ChatLogSummary.dir(directory)} failed message=${e.message}", e)
+                edt { failReverting(op, e) }
             }
         }
     }
@@ -466,7 +499,36 @@ class SessionController(
             return
         }
         sid?.let { capture("Session Redo", sessionProps(it)) }
-        revert(next.info.id)
+        redoTo(next.info.id)
+    }
+
+    private fun redoTo(message: String) {
+        assertEdt()
+        val id = sid ?: return
+        val state = model.state
+        if (revertOp != null) return
+        val busy = state.isBusy()
+        val op = beginReverting(
+            KiloBundle.message("session.status.redoing"),
+            SessionState.Reverting.Kind.REDO,
+            message,
+        ) ?: return
+        revertJob = cs.launch {
+            try {
+                if (busy) {
+                    LOG.info("${ChatLogSummary.sid(id)} kind=redo abort=true reason=busy")
+                    sessions.abort(id, directory)
+                    LOG.info("${ChatLogSummary.sid(id)} kind=redo abort=true ok=true")
+                }
+                sessions.revert(id, directory, message, null)
+                synchronizeFromDisk(id, "redo")
+                edt { clearReverting(op) }
+            } catch (e: CancellationException) {
+                throw e
+            } catch (e: Exception) {
+                edt { failReverting(op, e) }
+            }
+        }
     }
 
     fun redoAll() {
@@ -475,6 +537,17 @@ class SessionController(
         unrevert()
     }
 
+    fun cancelRevert() {
+        assertEdt()
+        if (revertOp == null) return
+        LOG.info("${ChatLogSummary.sid(sid ?: "?")} kind=revert cancelRequested=true")
+        sid?.let { capture("Session Revert Cancel Requested", sessionProps(it)) }
+        val state = model.state
+        if (state is SessionState.Reverting) {
+            model.setState(state.copy(text = KiloBundle.message("session.status.operation.finishing")))
+        }
+    }
+
     private fun synchronizeFromDisk(id: String, kind: String) {
         ApplicationManager.getApplication().invokeLater {
             runCatching {
@@ -1393,10 +1466,11 @@ class SessionController(
     }
 
     private fun status(dto: SessionStatusDto) {
+        if (revertOp != null) return
         val state = when (dto.type) {
             "idle" -> {
                 val current = model.state
-                if (current is SessionState.LoginRequired) return
+                if (current is SessionState.LoginRequired || current is SessionState.Reverting) return
                 SessionState.Idle
             }
             "busy" -> {
@@ -1419,14 +1493,68 @@ class SessionController(
         model.setState(state)
     }
 
+    private fun beginReverting(text: String, kind: SessionState.Reverting.Kind, message: String? = null): RevertOp? {
+        assertEdt()
+        if (revertOp != null) return null
+        val op = RevertOp(++revertSeq)
+        revertOp = op
+        model.setState(SessionState.Reverting(text, kind, message))
+        startRevertWatchdog(op)
+        return op
+    }
+
+    private fun startRevertWatchdog(op: RevertOp) {
+        assertEdt()
+        stopRevertWatchdog()
+        val ms = revertTimeoutMs.coerceIn(1, Int.MAX_VALUE.toLong()).toInt()
+        revertWatchdog = timers.timer(ms, repeats = false) { onRevertTimeout(op) }.also { it.start() }
+    }
+
+    private fun stopRevertWatchdog() {
+        revertWatchdog?.stop()
+        revertWatchdog = null
+    }
+
+    private fun onRevertTimeout(op: RevertOp) {
+        assertEdt()
+        if (revertOp?.key != op.key) return
+        LOG.warn("${ChatLogSummary.sid(sid ?: "?")} kind=revert timeout=true after=${revertTimeoutMs}ms")
+        stopRevertWatchdog()
+        sid?.let { capture("Session Revert Timeout", sessionProps(it)) }
+        val state = model.state
+        if (state is SessionState.Reverting) {
+            model.setState(state.copy(text = KiloBundle.message("session.error.revert.timeout")))
+        }
+    }
+
+    private fun clearReverting(op: RevertOp) {
+        assertEdt()
+        if (revertOp?.key != op.key) return
+        stopRevertWatchdog()
+        revertJob = null
+        revertOp = null
+        if (model.state is SessionState.Reverting) model.setState(SessionState.Idle)
+    }
+
+    private fun failReverting(op: RevertOp, e: Exception) {
+        assertEdt()
+        if (revertOp?.key != op.key) return
+        stopRevertWatchdog()
+        revertJob = null
+        revertOp = null
+        model.setState(SessionState.Error(e.message ?: KiloBundle.message("session.error.unknown")))
+    }
+
     private fun idle() {
+        if (revertOp != null) return
         // Treat session.idle as an explicit signal to return to Idle.
         // Only apply if we're not in a more specific non-terminal state.
         val current = model.state
         if (current !is SessionState.Error
             && current !is SessionState.AwaitingPermission
             && current !is SessionState.AwaitingQuestion
             && current !is SessionState.LoginRequired
+            && current !is SessionState.Reverting
         ) {
             model.setState(SessionState.Idle)
         }
@@ -2019,6 +2147,11 @@ class SessionController(
             cancelSubscriptions()
             drainJob?.cancel()
             drainJob = null
+            revertWatchdog?.stop()
+            revertWatchdog = null
+            revertJob?.cancel()
+            revertJob = null
+            revertOp = null
             val callbacks = enhancements.values.toList()
             enhancements.clear()
             cs.cancel()
@@ -2054,6 +2187,10 @@ class SessionController(
                 out.add("[busy]")
                 out.add("[${state.text.toDumpText()}]")
             }
+            is SessionState.Reverting -> {
+                out.add("[reverting]")
+                out.add("[${state.text.toDumpText()}]")
+            }
             is SessionState.AwaitingQuestion -> out.add("[awaiting-question]")
             is SessionState.AwaitingPermission -> out.add("[awaiting-permission]")
             is SessionState.Retry -> {

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

@@ -8,6 +8,10 @@ sealed class SessionState {
 
     data class Busy(val text: String) : SessionState()
 
+    data class Reverting(val text: String, val kind: Kind, val message: String? = null) : SessionState() {
+        enum class Kind { ROLLBACK, REDO }
+    }
+
     data class AwaitingQuestion(val question: Question) : SessionState()
 
     data class AwaitingPermission(val permission: Permission) : SessionState()

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

@@ -2,6 +2,7 @@ package ai.kilocode.client.session.ui
 
 import ai.kilocode.client.plugin.KiloBundle
 import ai.kilocode.client.session.model.SessionModel
+import ai.kilocode.client.session.model.SessionState
 import ai.kilocode.client.session.ui.style.SessionEditorStyle
 import ai.kilocode.client.session.ui.style.SessionEditorStyleTarget
 import ai.kilocode.client.session.views.base.BaseQuestionView
@@ -21,6 +22,7 @@ class RevertBanner(
     private val model: SessionModel,
     private val redoAction: () -> Unit,
     private val redoAllAction: () -> Unit,
+    private val cancelAction: () -> Unit,
     focus: (() -> Unit)? = null,
 ) : BorderLayoutPanel(), SessionView, SessionEditorStyleTarget {
     override val sessionViewKind = SessionView.Kind.Default
@@ -32,6 +34,7 @@ class RevertBanner(
     private val files = Stack.vertical(UiStyle.Gap.xs())
 
     private val rows = LinkedHashMap<String, Row>()
+    private var progress: RevertProgress? = null
 
     private val hint = JBLabel(KiloBundle.message("revert.banner.hint")).apply {
         font = JBFont.small()
@@ -81,8 +84,28 @@ class RevertBanner(
         repaint()
     }
 
+    @RequiresEdt
+    fun setReverting(state: SessionState) {
+        val busy = state is SessionState.Reverting && state.kind == SessionState.Reverting.Kind.REDO
+        if (busy) {
+            card.setActionEnabled("redo", false)
+            card.setActionEnabled("all", false)
+            val node = progress ?: RevertProgress(cancelAction).also {
+                it.applyStyle(SessionEditorStyle.current())
+                progress = it
+            }
+            node.setText(state.text)
+            card.setActionLeft(node)
+            return
+        }
+        card.setActionLeft(null)
+        card.setActionEnabled("redo", true)
+        card.setActionEnabled("all", true)
+    }
+
     override fun applyStyle(style: SessionEditorStyle) {
         card.applyStyle(style)
+        progress?.applyStyle(style)
         hint.foreground = UIUtil.getLabelForeground()
         notice.foreground = UIUtil.getContextHelpForeground()
         rows.values.forEach { it.applyStyle() }