diff --git a/.changeset/atomic-session-revert.md b/.changeset/atomic-session-revert.md new file mode 100644 index 0000000000..ed922d1a6e --- /dev/null +++ b/.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. diff --git a/packages/opencode/src/kilocode/session/revert.ts b/packages/opencode/src/kilocode/session/revert.ts new file mode 100644 index 0000000000..64ba8f98a0 --- /dev/null +++ b/packages/opencode/src/kilocode/session/revert.ts @@ -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 = (snap: Snapshot.Interface, hash: string, files: string[], cause: Cause.Cause) => + 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) { + 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* ( + snap: Snapshot.Interface, + baseline: string | undefined, + files: string[], + effect: Effect.Effect, + ) { + 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 }]) + }) +} diff --git a/packages/opencode/src/session/revert.ts b/packages/opencode/src/session/revert.ts index 4281cacca4..a2d6801175 100644 --- a/packages/opencode/src/session/revert.ts +++ b/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,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) }) diff --git a/packages/opencode/src/snapshot/index.ts b/packages/opencode/src/snapshot/index.ts index 5b974dfd60..a1d7480978 100644 --- a/packages/opencode/src/snapshot/index.ts +++ b/packages/opencode/src/snapshot/index.ts @@ -208,7 +208,10 @@ export const layer: Layer.Layer = 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 = (fx: Effect.Effect) => lock(state.gitdir).withPermits(1)(flock.withLock(fx, `snapshot:${state.gitdir}`).pipe(Effect.orDie)) @@ -471,13 +474,14 @@ export const layer: Layer.Layer = 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 = 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() for (const item of patches) { @@ -508,13 +520,20 @@ export const layer: Layer.Layer = 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, diff --git a/packages/opencode/test/kilocode/session/revert.test.ts b/packages/opencode/test/kilocode/session/revert.test.ts index 0299328419..a2274bb55c 100644 --- a/packages/opencode/test/kilocode/session/revert.test.ts +++ b/packages/opencode/test/kilocode/session/revert.test.ts @@ -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, + ) })