fix(cli): keep session reverts atomic

.changeset/atomic-session-revert.md

@@ -0,0 +1,6 @@
+---
+"@kilocode/cli": patch
+"kilo-code": patch
+---
+
+Keep conversations and workspace files unchanged when a checkpoint cannot be fully restored.

packages/opencode/src/kilocode/session/revert.ts

@@ -0,0 +1,34 @@
+import fs from "node:fs/promises"
+import path from "node:path"
+import { Effect } from "effect"
+import { Instance } from "@/kilocode/instance"
+import type { Snapshot } from "@/snapshot"
+
+export namespace KiloSessionRevert {
+  export const apply = Effect.fn("KiloSessionRevert.apply")(function* (
+    snap: Snapshot.Interface,
+    patches: Snapshot.Patch[],
+    baseline?: string,
+  ) {
+    return yield* snap.revert(patches).pipe(
+      Effect.catchCause((cause) => {
+        if (!baseline) return Effect.failCause(cause)
+        return restore(snap, baseline).pipe(Effect.andThen(Effect.failCause(cause)))
+      }),
+    )
+  })
+
+  export const restore = Effect.fn("KiloSessionRevert.restore")(function* (snap: Snapshot.Interface, hash: string) {
+    const current = yield* snap.track()
+    const removed = current
+      ? (yield* snap.diffFull(hash, current)).flatMap((file) =>
+          file.status === "added" && file.file ? [path.resolve(Instance.directory, file.file)] : [],
+        )
+      : []
+
+    yield* snap.restore(hash)
+    yield* Effect.forEach(removed, (file) => Effect.promise(() => fs.rm(file, { force: true, recursive: true })), {
+      discard: true,
+    })
+  })
+}

packages/opencode/src/session/revert.ts

@@ -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,14 +82,19 @@ 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 - do not mutate files without a durable compensation snapshot
+      if (patches.some((patch) => patch.files.length > 0) && !rev.snapshot) {
+        return yield* Effect.die(new Error("Cannot rewind files because the current workspace snapshot is unavailable"))
+      }
+      if (session.revert?.snapshot) yield* KiloSessionRevert.restore(snap, session.revert.snapshot)
+      // kilocode_change end
 
       // 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 })
       // kilocode_change end
 
-      yield* snap.revert(patches)
+      yield* KiloSessionRevert.apply(snap, patches, rev.snapshot) // kilocode_change
       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 })
@@ -118,7 +124,7 @@ 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)
+      if (session.revert.snapshot) yield* KiloSessionRevert.restore(snap, session.revert.snapshot) // kilocode_change
       yield* sessions.clearRevert(input.sessionID)
       return yield* sessions.get(input.sessionID).pipe(Effect.orDie)
     })

packages/opencode/src/snapshot/index.ts

@@ -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,

packages/opencode/test/kilocode/session/revert.test.ts

@@ -1,5 +1,5 @@
 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"
@@ -242,4 +242,114 @@ describe("workspace revert status", () => {
       { git: true },
     ),
   )
+
+  it.live(
+    "keeps the conversation and workspace unchanged when a checkpoint cannot be fully restored",
+    provideTmpdirInstance(
+      (dir) =>
+        Effect.gen(function* () {
+          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")
+          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,
+          })
+
+          yield* Effect.promise(() => fs.chmod(protectedFile, 0o444))
+          yield* Effect.promise(() => fs.chmod(locked, 0o555))
+          const outcome = yield* revert.revert({ sessionID: session.id, messageID: user.id }).pipe(
+            Effect.exit,
+            Effect.ensuring(
+              Effect.promise(async () => {
+                await fs.chmod(locked, 0o755)
+                await fs.chmod(protectedFile, 0o644)
+              }),
+            ),
+          )
+          const current = yield* sessions.get(session.id)
+          const actual = {
+            failed: Exit.isFailure(outcome),
+            reverted: current.revert !== undefined,
+            protected: yield* Effect.promise(() => fs.readFile(protectedFile, "utf8")),
+            writable: yield* Effect.promise(() => fs.readFile(writableFile, "utf8")),
+          }
+
+          expect(actual).toEqual({
+            failed: true,
+            reverted: false,
+            protected: "after",
+            writable: "after",
+          })
+        }),
+      { git: true },
+    ),
+    30_000,
+  )
 })