Merge pull request #12587 from Kilo-Org/reproduce-and-fix-revert-file-rollback-bug

fix(cli): keep session reverts atomic
2026-08-28 11:05:31 +08:00 · 2026-07-28 11:58:14 +02:00
parent d41e0cdc04 fc04f4f6ff
commit b20986b366
5 changed files with 360 additions and 36 deletions
@@ -0,0 +1,6 @@
+---
+"@kilocode/cli": patch
+"kilo-code": patch
+---
+
+Keep conversations and workspace files unchanged when a checkpoint cannot be fully restored.
@@ -0,0 +1,51 @@
+import { Cause, Effect } from "effect"
+import type { MessageV2 } from "@/session/message-v2"
+import type { Session } from "@/session/session"
+import type { Snapshot } from "@/snapshot"
+
+export namespace KiloSessionRevert {
+  const rollback = <E>(snap: Snapshot.Interface, hash: string, files: string[], cause: Cause.Cause<E>) =>
+    restore(snap, hash, files).pipe(
+      Effect.matchCauseEffect({
+        onFailure: (next) => Effect.failCause(Cause.combine(cause, next)),
+        onSuccess: () => Effect.failCause(cause),
+      }),
+    )
+
+  export function files(messages: MessageV2.WithParts[], rev: NonNullable<Session.Info["revert"]>) {
+    const result: string[] = []
+    let active = false
+    for (const msg of messages) {
+      for (const part of msg.parts) {
+        if (active && part.type === "patch") result.push(...part.files)
+        if (active || msg.info.id !== rev.messageID) continue
+        if (rev.partID && part.id !== rev.partID) continue
+        active = true
+      }
+    }
+    return [...new Set(result)]
+  }
+
+  export const apply = Effect.fn("KiloSessionRevert.apply")(function* <A, E, R>(
+    snap: Snapshot.Interface,
+    baseline: string | undefined,
+    files: string[],
+    effect: Effect.Effect<A, E, R>,
+  ) {
+    return yield* effect.pipe(
+      Effect.catchCause((cause) => {
+        if (!baseline || files.length === 0) return Effect.failCause(cause)
+        return rollback(snap, baseline, files, cause)
+      }),
+    )
+  })
+
+  export const restore = Effect.fn("KiloSessionRevert.restore")(function* (
+    snap: Snapshot.Interface,
+    hash: string,
+    files: string[],
+  ) {
+    if (files.length === 0) return
+    yield* snap.revert([{ hash, files }])
+  })
+}
@@ -10,6 +10,7 @@ import { MessageV2 } from "./message-v2"
 import { SessionID, MessageID, PartID } from "./schema"
 import { SessionRunState } from "./run-state"
 import { SessionSummary } from "./summary"
+import { KiloSessionRevert } from "@/kilocode/session/revert" // kilocode_change

 export const RevertInput = Schema.Struct({
  sessionID: SessionID,
@@ -81,35 +82,45 @@ export const layer = Layer.effect(
          : "unavailable"
      // kilocode_change end
      rev.snapshot = session.revert?.snapshot ?? (yield* snap.track())
-      if (session.revert?.snapshot) yield* snap.restore(session.revert.snapshot)
+      // kilocode_change start - keep the entire workspace transition atomic
+      const prior = session.revert ? KiloSessionRevert.files(all, session.revert) : []
+      const files = [...new Set([...prior, ...patches.flatMap((patch) => patch.files)])]
+      const baseline = session.revert?.snapshot && files.length > 0 ? yield* snap.track() : rev.snapshot
+      if (files.length > 0 && !baseline) {
+        return yield* Effect.die(new Error("Cannot rewind files because the current workspace snapshot is unavailable"))
+      }
+      yield* KiloSessionRevert.apply(
+        snap,
+        baseline,
+        files,
+        Effect.gen(function* () {
+          if (session.revert?.snapshot) yield* KiloSessionRevert.restore(snap, session.revert.snapshot, prior)

-      // kilocode_change start - compute diffs BEFORE reverting files so the diff
-      // reflects changes being undone (files on disk still have AI modifications)
-      const diffs = yield* summary.computeDiff({ messages: range })
+          // Compute the user-facing diff while files still contain the changes being undone.
+          const diffs = yield* summary.computeDiff({ messages: range })
+          yield* snap.revert(patches)
+          if (rev.snapshot) rev.diff = yield* snap.diff(rev.snapshot)
+          yield* storage.write(["session_diff", input.sessionID], diffs).pipe(Effect.ignore)
+          yield* events.publish(Session.Event.Diff, { sessionID: input.sessionID, diff: diffs })
+          const summaryDiffs: Snapshot.SummaryFileDiff[] = diffs.map((d) => ({
+            file: d.file,
+            additions: d.additions,
+            deletions: d.deletions,
+            status: d.status,
+          }))
+          yield* sessions.setRevert({
+            sessionID: input.sessionID,
+            revert: rev,
+            summary: {
+              additions: diffs.reduce((sum, x) => sum + x.additions, 0),
+              deletions: diffs.reduce((sum, x) => sum + x.deletions, 0),
+              files: diffs.length,
+              diffs: summaryDiffs,
+            },
+          })
+        }),
+      )
      // kilocode_change end
-
-      yield* snap.revert(patches)
-      if (rev.snapshot) rev.diff = yield* snap.diff(rev.snapshot)
-      yield* storage.write(["session_diff", input.sessionID], diffs).pipe(Effect.ignore)
-      yield* events.publish(Session.Event.Diff, { sessionID: input.sessionID, diff: diffs })
-      // kilocode_change start
-      const summaryDiffs: Snapshot.SummaryFileDiff[] = diffs.map((d) => ({
-        file: d.file,
-        additions: d.additions,
-        deletions: d.deletions,
-        status: d.status,
-      }))
-      // kilocode_change end
-      yield* sessions.setRevert({
-        sessionID: input.sessionID,
-        revert: rev,
-        summary: {
-          additions: diffs.reduce((sum, x) => sum + x.additions, 0),
-          deletions: diffs.reduce((sum, x) => sum + x.deletions, 0),
-          files: diffs.length,
-          diffs: summaryDiffs, // kilocode_change
-        },
-      })
      return yield* sessions.get(input.sessionID).pipe(Effect.orDie)
    })

@@ -118,8 +129,25 @@ export const layer = Layer.effect(
      yield* state.assertNotBusy(input.sessionID)
      const session = yield* sessions.get(input.sessionID).pipe(Effect.orDie)
      if (!session.revert) return session
-      if (session.revert.snapshot) yield* snap.restore(session.revert.snapshot)
-      yield* sessions.clearRevert(input.sessionID)
+      // kilocode_change start - preserve the reverted workspace if redo cannot complete
+      const all = yield* sessions.messages({ sessionID: input.sessionID }).pipe(Effect.orDie)
+      const files = KiloSessionRevert.files(all, session.revert)
+      const baseline = files.length > 0 ? yield* snap.track() : undefined
+      if (files.length > 0 && !baseline) {
+        return yield* Effect.die(
+          new Error("Cannot restore files because the current workspace snapshot is unavailable"),
+        )
+      }
+      yield* KiloSessionRevert.apply(
+        snap,
+        baseline,
+        files,
+        Effect.gen(function* () {
+          if (session.revert?.snapshot) yield* KiloSessionRevert.restore(snap, session.revert.snapshot, files)
+          yield* sessions.clearRevert(input.sessionID)
+        }),
+      )
+      // kilocode_change end
      return yield* sessions.get(input.sessionID).pipe(Effect.orDie)
    })

@@ -208,7 +208,10 @@ export const layer: Layer.Layer<Service, never, Requirements> =

          const exists = (file: string) => fs.exists(file).pipe(Effect.orDie)
          const read = (file: string) => fs.readFileString(file).pipe(Effect.catch(() => Effect.succeed("")))
-          const remove = (file: string) => fs.remove(file).pipe(Effect.catch(() => Effect.void))
+          // kilocode_change start - restoration must fail if deletion fails
+          const remove = (file: string) =>
+            fs.remove(file, { force: true }).pipe(Effect.orDie)
+          // kilocode_change end
          // kilocode_change start - serialize snapshot repositories across CLI and extension processes
          const locked = <A, R>(fx: Effect.Effect<A, never, R>) =>
            lock(state.gitdir).withPermits(1)(flock.withLock(fx, `snapshot:${state.gitdir}`).pipe(Effect.orDie))
@@ -471,13 +474,14 @@ export const layer: Layer.Layer<Service, never, Requirements> =
                    exitCode: checkout.code,
                    stderr: checkout.stderr,
                  })
-                  return
+                  return yield* Effect.die(new Error(`Failed to restore snapshot ${snapshot}`)) // kilocode_change
                }
                yield* Effect.logError("failed to restore snapshot", {
                  snapshot,
                  exitCode: result.code,
                  stderr: result.stderr,
                })
+                return yield* Effect.die(new Error(`Failed to restore snapshot ${snapshot}`)) // kilocode_change
              }),
            )
          })
@@ -485,6 +489,14 @@ export const layer: Layer.Layer<Service, never, Requirements> =
          const revert = Effect.fnUntraced(function* (patches: Patch[]) {
            return yield* locked(
              Effect.gen(function* () {
+                // kilocode_change start - validate every checkpoint before mutating workspace files
+                for (const hash of new Set(patches.filter((item) => item.files.length > 0).map((item) => item.hash))) {
+                  const tree = yield* git([...core, ...args(["cat-file", "-e", `${hash}^{tree}`])], {
+                    cwd: state.worktree,
+                  })
+                  if (tree.code !== 0) return yield* Effect.die(new Error(`Snapshot ${hash} is unavailable`))
+                }
+                // kilocode_change end
                const ops: { hash: string; file: string; rel: string }[] = []
                const seen = new Set<string>()
                for (const item of patches) {
@@ -508,13 +520,20 @@ export const layer: Layer.Layer<Service, never, Requirements> =
                  const tree = yield* git([...core, ...args(["ls-tree", op.hash, "--", op.rel])], {
                    cwd: state.worktree,
                  })
-                  if (tree.code === 0 && tree.text.trim()) {
-                    yield* Effect.logInfo("file existed in snapshot but checkout failed, keeping", {
+                  // kilocode_change start - never report success for a file that Git could not restore
+                  if (tree.code !== 0) {
+                    return yield* Effect.die(new Error(`Snapshot ${op.hash} is unavailable`))
+                  }
+                  if (tree.text.trim()) {
+                    yield* Effect.logError("file existed in snapshot but checkout failed", {
                      file: op.file,
                      hash: op.hash,
+                      exitCode: result.code,
+                      stderr: result.stderr,
                    })
-                    return
+                    return yield* Effect.die(new Error(`Failed to restore ${op.file} from snapshot ${op.hash}`))
                  }
+                  // kilocode_change end
                  yield* Effect.logInfo("file did not exist in snapshot, deleting", {
                    file: op.file,
                    hash: op.hash,
@@ -1,16 +1,17 @@
 import { describe, expect } from "bun:test"
-import { Effect, Layer } from "effect"
+import { Effect, Exit, Layer } from "effect"
 import fs from "node:fs/promises"
 import path from "node:path"
 import { CrossSpawnSpawner } from "@opencode-ai/core/cross-spawn-spawner"
 import { ProviderV2 } from "@opencode-ai/core/provider"
 import { ModelV2 } from "@opencode-ai/core/model"
 import { MessageV2 } from "@/session/message-v2"
+import { KiloSessionRevert } from "@/kilocode/session/revert"
 import { SessionRevert } from "@/session/revert"
 import { MessageID, PartID } from "@/session/schema"
 import { Session } from "@/session/session"
 import { Snapshot } from "@/snapshot"
-import { provideTmpdirInstance } from "../../fixture/fixture"
+import { provideInstance, provideTmpdirInstance } from "../../fixture/fixture"
 import { testEffect } from "../../lib/effect"

 const env = Layer.mergeAll(
@@ -20,6 +21,96 @@ const env = Layer.mergeAll(
  CrossSpawnSpawner.defaultLayer,
 )
 const it = testEffect(env)
+const guarded = process.platform === "win32" ? it.live.skip : it.live
+
+const setup = Effect.fnUntraced(function* (dir: string, deleted = false) {
+  const sessions = yield* Session.Service
+  const revert = yield* SessionRevert.Service
+  const snapshot = yield* Snapshot.Service
+  const session = yield* sessions.create({})
+  const locked = path.join(dir, "locked")
+  const protectedFile = path.join(locked, "protected.txt")
+  const writableFile = path.join(dir, "writable.txt")
+  const providerID = ProviderV2.ID.make("test")
+  yield* Effect.promise(() => fs.mkdir(locked))
+  yield* Effect.promise(() => fs.writeFile(protectedFile, "before"))
+  yield* Effect.promise(() => fs.writeFile(writableFile, "before"))
+  const user = yield* sessions.updateMessage({
+    id: MessageID.ascending(),
+    sessionID: session.id,
+    role: "user",
+    agent: "default",
+    model: { providerID, modelID: ModelV2.ID.make("test") },
+    time: { created: Date.now() },
+  })
+  yield* sessions.updatePart({
+    id: PartID.ascending(),
+    messageID: user.id,
+    sessionID: session.id,
+    type: "text",
+    text: "change both files",
+  })
+  const assistant = yield* sessions.updateMessage({
+    id: MessageID.ascending(),
+    sessionID: session.id,
+    role: "assistant",
+    parentID: user.id,
+    mode: "default",
+    agent: "default",
+    path: { cwd: dir, root: dir },
+    cost: 0,
+    tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
+    modelID: ModelV2.ID.make("test"),
+    providerID,
+    time: { created: Date.now() },
+    finish: "end_turn",
+  })
+  const before = yield* snapshot.track()
+  if (!before) throw new Error("expected snapshot")
+  if (deleted) yield* Effect.promise(() => fs.rm(protectedFile))
+  if (!deleted) yield* Effect.promise(() => fs.writeFile(protectedFile, "after"))
+  yield* Effect.promise(() => fs.writeFile(writableFile, "after"))
+  const after = yield* snapshot.track()
+  if (!after) throw new Error("expected snapshot")
+  const patch = yield* snapshot.patch(before)
+  yield* sessions.updatePart({
+    id: PartID.ascending(),
+    messageID: assistant.id,
+    sessionID: session.id,
+    type: "step-start",
+    snapshot: before,
+  })
+  yield* sessions.updatePart({
+    id: PartID.ascending(),
+    messageID: assistant.id,
+    sessionID: session.id,
+    type: "step-finish",
+    reason: "stop",
+    snapshot: after,
+    cost: 0,
+    tokens: { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } },
+  })
+  yield* sessions.updatePart({
+    id: PartID.ascending(),
+    messageID: assistant.id,
+    sessionID: session.id,
+    type: "patch",
+    hash: patch.hash,
+    files: patch.files,
+  })
+  return {
+    sessions,
+    revert,
+    snapshot,
+    session,
+    user,
+    after,
+    patch,
+    locked,
+    protected: protectedFile,
+    writable: writableFile,
+  }
+})

 describe("partial assistant revert", () => {
  it.live(
@@ -242,4 +333,133 @@ describe("workspace revert status", () => {
      { git: true },
    ),
  )
+
+  guarded(
+    "keeps the conversation and workspace unchanged when a checkpoint cannot be fully restored",
+    provideTmpdirInstance(
+      (dir) =>
+        Effect.gen(function* () {
+          const item = yield* setup(dir)
+          yield* Effect.promise(() => fs.chmod(item.protected, 0o444))
+          yield* Effect.promise(() => fs.chmod(item.locked, 0o555))
+          const outcome = yield* item.revert.revert({ sessionID: item.session.id, messageID: item.user.id }).pipe(
+            Effect.exit,
+            Effect.ensuring(
+              Effect.promise(async () => {
+                await fs.chmod(item.locked, 0o755)
+                await fs.chmod(item.protected, 0o644)
+              }),
+            ),
+          )
+          const current = yield* item.sessions.get(item.session.id)
+          const actual = {
+            failed: Exit.isFailure(outcome),
+            reverted: current.revert !== undefined,
+            protected: yield* Effect.promise(() => fs.readFile(item.protected, "utf8")),
+            writable: yield* Effect.promise(() => fs.readFile(item.writable, "utf8")),
+          }
+
+          expect(actual).toEqual({
+            failed: true,
+            reverted: false,
+            protected: "after",
+            writable: "after",
+          })
+        }),
+      { git: true },
+    ),
+    30_000,
+  )
+
+  guarded(
+    "keeps the reverted state when unrevert cannot fully restore files",
+    provideTmpdirInstance(
+      (dir) =>
+        Effect.gen(function* () {
+          const item = yield* setup(dir)
+          yield* item.revert.revert({ sessionID: item.session.id, messageID: item.user.id })
+          yield* Effect.promise(() => fs.chmod(item.protected, 0o444))
+          yield* Effect.promise(() => fs.chmod(item.locked, 0o555))
+          const outcome = yield* item.revert.unrevert({ sessionID: item.session.id }).pipe(
+            Effect.exit,
+            Effect.ensuring(
+              Effect.promise(async () => {
+                await fs.chmod(item.locked, 0o755)
+                await fs.chmod(item.protected, 0o644)
+              }),
+            ),
+          )
+          const current = yield* item.sessions.get(item.session.id)
+
+          expect({
+            failed: Exit.isFailure(outcome),
+            reverted: current.revert !== undefined,
+            protected: yield* Effect.promise(() => fs.readFile(item.protected, "utf8")),
+            writable: yield* Effect.promise(() => fs.readFile(item.writable, "utf8")),
+          }).toEqual({ failed: true, reverted: true, protected: "before", writable: "before" })
+        }),
+      { git: true },
+    ),
+    30_000,
+  )
+
+  guarded(
+    "keeps the prior revert when replacing its checkpoint cannot restore files",
+    provideTmpdirInstance(
+      (dir) =>
+        Effect.gen(function* () {
+          const item = yield* setup(dir)
+          yield* item.revert.revert({ sessionID: item.session.id, messageID: item.user.id })
+          yield* Effect.promise(() => fs.chmod(item.protected, 0o444))
+          yield* Effect.promise(() => fs.chmod(item.locked, 0o555))
+          const outcome = yield* item.revert.revert({ sessionID: item.session.id, messageID: item.user.id }).pipe(
+            Effect.exit,
+            Effect.ensuring(
+              Effect.promise(async () => {
+                await fs.chmod(item.locked, 0o755)
+                await fs.chmod(item.protected, 0o644)
+              }),
+            ),
+          )
+          const current = yield* item.sessions.get(item.session.id)
+
+          expect({
+            failed: Exit.isFailure(outcome),
+            reverted: current.revert !== undefined,
+            protected: yield* Effect.promise(() => fs.readFile(item.protected, "utf8")),
+            writable: yield* Effect.promise(() => fs.readFile(item.writable, "utf8")),
+          }).toEqual({ failed: true, reverted: true, protected: "before", writable: "before" })
+        }),
+      { git: true },
+    ),
+    30_000,
+  )
+
+  it.live(
+    "unreverts deleted files from a session rooted in a worktree subdirectory",
+    provideTmpdirInstance(
+      (root) =>
+        Effect.gen(function* () {
+          const dir = path.join(root, "nested")
+          yield* Effect.promise(() => fs.mkdir(dir))
+          const item = yield* setup(dir, true)
+          yield* item.revert.revert({ sessionID: item.session.id, messageID: item.user.id })
+          expect(yield* Effect.promise(() => fs.readFile(item.protected, "utf8"))).toBe("before")
+
+          yield* KiloSessionRevert.restore(item.snapshot, item.after, item.patch.files).pipe(provideInstance(dir))
+
+          expect(
+            yield* Effect.promise(() =>
+              fs.stat(item.protected).then(
+                () => true,
+                () => false,
+              ),
+            ),
+          ).toBe(false)
+          expect(yield* Effect.promise(() => fs.readFile(item.writable, "utf8"))).toBe("after")
+        }),
+      { git: true },
+    ),
+    30_000,
+  )
 })