fix(jetbrains): load reverted session diffs

.changeset/jetbrains-reverted-session-diff-list.md

@@ -0,0 +1,5 @@
+---
+"@kilocode/kilo-jetbrains": patch
+---
+
+Show the rolled-back file list when reopening a session whose last message was reverted.

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

@@ -1011,6 +1011,7 @@ class SessionController(
                     }
                 }
                 recoverPending(id)
+                seedRevertDiff(id)
                 runEdt {
                     if (disposed) return@runEdt
                     if (sid != id) return@runEdt
@@ -1060,6 +1061,7 @@ class SessionController(
                     }
                 }
                 recoverPending(session.id)
+                seedRevertDiff(session.id)
                 runEdt {
                     if (disposed) return@runEdt
                     subscribeEvents()
@@ -1091,6 +1093,25 @@ class SessionController(
         }
     }
 
+    /**
+     * Seed [SessionModel.diff] when opening a reverted session. The rolled-back file list in
+     * [ai.kilocode.client.session.ui.RevertBanner] falls back to `model.diff` when the CLI does not
+     * attach a diff to the revert marker. On a live revert a `session.diff` event seeds that; on
+     * reload nothing does, so fetch the persisted session diff once here. Skipped for sessions
+     * without a revert or once a diff is already present (e.g. a concurrent `session.diff` event).
+     */
+    private suspend fun seedRevertDiff(id: String) {
+        var fetch = false
+        runEdt { fetch = !disposed && sid == id && model.revert() != null && model.diff.isEmpty() }
+        if (!fetch) return
+        val diffs = runCatching { sessions.diff(id, directory) }.getOrNull()?.takeIf { it.isNotEmpty() } ?: return
+        runEdt {
+            if (disposed || sid != id) return@runEdt
+            if (model.revert() == null || model.diff.isNotEmpty()) return@runEdt
+            updateModel { model.setDiff(diffs) }
+        }
+    }
+
     private fun startSessionLoading(token: SessionLoadState.Loading) {
         assertEdt()
         setSessionLoadState(token)

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

@@ -0,0 +1,38 @@
+package ai.kilocode.client.session.controller
+
+import ai.kilocode.rpc.dto.DiffFileDto
+import ai.kilocode.rpc.dto.MessageWithPartsDto
+import ai.kilocode.rpc.dto.SessionRevertDto
+
+/**
+ * A session reverted in a previous run has no `session.diff` event to replay on open, and the
+ * pinned CLI may not attach a diff to the revert marker. [SessionController.seedRevertDiff] fetches
+ * the persisted session diff so the reverted-files banner has something to render.
+ */
+class RevertDiffLoadingTest : SessionControllerTestBase() {
+
+    fun `test opening a reverted session seeds model diff from the diff rpc`() {
+        rpc.session = session("ses_test").copy(revert = SessionRevertDto(messageID = "msg1", snapshot = "snap"))
+        rpc.history.add(MessageWithPartsDto(msg("msg1", "ses_test", "user"), emptyList()))
+        rpc.history.add(MessageWithPartsDto(msg("msg2", "ses_test", "assistant"), emptyList()))
+        rpc.diffs["ses_test"] = mutableListOf(DiffFileDto("src/A.kt", 2, 1, "@@ patch"))
+
+        val c = controller("ses_test")
+        flush()
+
+        assertEquals("msg1", c.model.revert()?.messageID)
+        assertEquals(listOf("src/A.kt"), c.model.diff.map { it.file })
+    }
+
+    fun `test opening a session without a revert does not fetch the diff`() {
+        rpc.session = session("ses_test")
+        rpc.history.add(MessageWithPartsDto(msg("msg1", "ses_test", "user"), emptyList()))
+        rpc.diffs["ses_test"] = mutableListOf(DiffFileDto("src/A.kt", 2, 1, "@@ patch"))
+
+        val c = controller("ses_test")
+        flush()
+
+        assertNull(c.model.revert())
+        assertTrue(c.model.diff.isEmpty())
+    }
+}