From 10dbe8d0e4e8c1f29b7157a2eb796864b3129134 Mon Sep 17 00:00:00 2001 From: KazenDev Date: Sat, 12 Sep 2026 19:14:47 -0500 Subject: [PATCH] Add the /undo snapshot engine (part 1 of #944) Snapshot capture, restore, and the per-chat journal. No UI yet: the picker, the command registry entry, and the send-message hook come in the follow-up. Three data-loss paths closed, each with a test that fails without it: - a failed `git add` no longer hands back the tree of a stale index - every entry is anchored (`refs/freebuff/undo//`) so the cleanup job's prune cannot collect a snapshot the journal still lists - reverting checks the content, not just the tree: `git checkout` deletes a file whose blob it cannot read, and `ls-tree` still lists it The journal is the root set: refs are released when the entry is dropped, and `sweepAnchors` collects what a deleted chat left behind. Verified: 20 tests pass, the CLI typecheck has no new errors, and the full suite matches its baseline (+20 pass, same 61 pre-existing failures). --- cli/src/state/__tests__/undo-store.test.ts | 270 +++++++ cli/src/state/undo-store.ts | 406 ++++++++++ cli/src/utils/__tests__/undo-snapshot.test.ts | 257 ++++++ cli/src/utils/undo-snapshot.ts | 751 ++++++++++++++++++ 4 files changed, 1684 insertions(+) create mode 100644 cli/src/state/__tests__/undo-store.test.ts create mode 100644 cli/src/state/undo-store.ts create mode 100644 cli/src/utils/__tests__/undo-snapshot.test.ts create mode 100644 cli/src/utils/undo-snapshot.ts diff --git a/cli/src/state/__tests__/undo-store.test.ts b/cli/src/state/__tests__/undo-store.test.ts new file mode 100644 index 0000000000..207ab60098 --- /dev/null +++ b/cli/src/state/__tests__/undo-store.test.ts @@ -0,0 +1,270 @@ +import { afterAll, beforeAll, beforeEach, describe, expect, test } from 'bun:test' + +import { execFileSync } from 'node:child_process' +import { + existsSync, + mkdtempSync, + readdirSync, + readFileSync, + rmSync, + writeFileSync, +} from 'node:fs' +import os from 'node:os' +import path from 'node:path' + +import { getProjectDataDir, setCurrentChatId, setProjectRoot } from '../../project-files' +import { + anchorSnapshot, + listAnchors, + patchSnapshot, + releaseSnapshot, + setSnapshotDirOverrideForTesting, + trackSnapshot, +} from '../../utils/undo-snapshot' +import { + loadUndoState, + peekRedo, + peekUndo, + popRedo, + popUndo, + pushRedo, + pushUndo, + recordUndoEntry, + sweepAnchors, + undoToRecord, +} from '../undo-store' + +const CHAT_ID = 'undo-store-test-chat' + +let projectDir: string +let snapshotRoot: string + +beforeAll(() => { + projectDir = mkdtempSync(path.join(os.tmpdir(), 'undo-store-proj-')) + snapshotRoot = mkdtempSync(path.join(os.tmpdir(), 'undo-store-snapshots-')) + setProjectRoot(projectDir) + setCurrentChatId(CHAT_ID) + // The journal hands snapshots to undo-snapshot, which needs a git project. + setSnapshotDirOverrideForTesting(snapshotRoot) + execFileSync('git', ['init'], { cwd: projectDir, stdio: 'ignore' }) +}) + +// Each test starts from a clean journal for this chat. +beforeEach(() => { + const chatDir = path.join(getProjectDataDir(), 'chats', CHAT_ID) + rmSync(chatDir, { recursive: true, force: true }) + setCurrentChatId(CHAT_ID) +}) + +afterAll(() => { + const chatDir = path.join(getProjectDataDir(), 'chats', CHAT_ID) + rmSync(chatDir, { recursive: true, force: true }) + rmSync(projectDir, { recursive: true, force: true }) + setSnapshotDirOverrideForTesting(undefined) + rmSync(snapshotRoot, { recursive: true, force: true }) +}) + +/** Run the cleanup job's gc now, instead of waiting out its 7-day grace. */ +const pruneSnapshotRepo = (): void => { + const [entry] = readdirSync(snapshotRoot) + if (!entry) return + execFileSync( + 'git', + [ + '--git-dir', + path.join(snapshotRoot, entry), + '--work-tree', + projectDir, + 'gc', + '--quiet', + '--prune=now', + ], + { stdio: ['ignore', 'pipe', 'ignore'] }, + ) +} + +describe('recordUndoEntry', () => { + test('records an entry and ignores empty ones', () => { + recordUndoEntry(CHAT_ID, { + hashBefore: 'abc123', + files: ['a.txt', 'b.txt'], + message: 'fix the bug', + }) + const state = loadUndoState(CHAT_ID) + expect(state.undoStack).toHaveLength(1) + expect(state.undoStack[0]).toMatchObject({ + chatId: CHAT_ID, + hashBefore: 'abc123', + files: ['a.txt', 'b.txt'], + message: 'fix the bug', + }) + expect(peekUndo(CHAT_ID)?.hashBefore).toBe('abc123') + + // No hash or no files → nothing recorded. + recordUndoEntry(CHAT_ID, { hashBefore: '', files: [], message: 'x' }) + recordUndoEntry(CHAT_ID, { hashBefore: 'def', files: [], message: 'x' }) + expect(loadUndoState(CHAT_ID).undoStack).toHaveLength(1) + }) + + test('clears the redo stack when a new entry arrives', () => { + pushRedo(CHAT_ID, { + id: 'r1', + chatId: CHAT_ID, + hashBefore: 'old', + hashAfter: 'new', + files: ['a.txt'], + message: 'redo me', + createdAt: new Date().toISOString(), + }) + expect(peekRedo(CHAT_ID)).not.toBeNull() + + recordUndoEntry(CHAT_ID, { + hashBefore: 'xyz', + files: ['c.txt'], + message: 'new turn', + }) + expect(peekRedo(CHAT_ID)).toBeNull() + }) +}) + +describe('undo/redo stack operations', () => { + test('popUndo returns the most recent record and persists', () => { + recordUndoEntry(CHAT_ID, { + hashBefore: 'first', + files: ['one.txt'], + message: 'first turn', + }) + recordUndoEntry(CHAT_ID, { + hashBefore: 'second', + files: ['two.txt'], + message: 'second turn', + }) + + const record = popUndo(CHAT_ID) + expect(record?.hashBefore).toBe('second') + expect(peekUndo(CHAT_ID)?.hashBefore).toBe('first') + // Reload from disk to confirm the pop persisted. + expect(loadUndoState(CHAT_ID).undoStack).toHaveLength(1) + }) + + test('pushUndo restores a record and popRedo cycles', () => { + recordUndoEntry(CHAT_ID, { + hashBefore: 'cycle-hash', + files: ['x.txt'], + message: 'cycle', + }) + const record = popUndo(CHAT_ID)! + pushRedo(CHAT_ID, { ...record, hashAfter: 'after-state' }) + + const redoRecord = popRedo(CHAT_ID) + expect(redoRecord?.hashBefore).toBe('cycle-hash') + expect(redoRecord?.hashAfter).toBe('after-state') + + pushUndo(CHAT_ID, { ...redoRecord!, hashAfter: undefined }) + expect(peekUndo(CHAT_ID)?.hashBefore).toBe('cycle-hash') + }) + + test('returns null from empty stacks', () => { + setCurrentChatId('empty-chat') + expect(popUndo('empty-chat')).toBeNull() + expect(popRedo('empty-chat')).toBeNull() + setCurrentChatId(CHAT_ID) + }) +}) + +describe('a snapshot that is gone', () => { + test('leaves the worktree alone instead of deleting what it cannot restore', async () => { + const kept = path.join(projectDir, 'kept.txt') + writeFileSync(kept, 'the user wrote this\n') + const hash = await trackSnapshot(projectDir) + expect(hash).toBeTruthy() + + // The turn edits the file and the journal records what to revert. + writeFileSync(kept, 'the agent changed it\n') + recordUndoEntry(CHAT_ID, { + hashBefore: hash!, + files: ['kept.txt'], + message: 'a recorded turn', + }) + + // The snapshot is gone — pruned by the cleanup job, or a wiped config dir. + // `checkout` then fails, and "not in the snapshot" must not be read as + // "the turn created this file". + rmSync(snapshotRoot, { recursive: true, force: true }) + + const message = await undoToRecord(CHAT_ID, projectDir, peekUndo(CHAT_ID)!.id) + + expect(message).toBeNull() + expect(existsSync(kept)).toBe(true) + expect(readFileSync(kept, 'utf8')).toBe('the agent changed it\n') + }) +}) + +describe('snapshot anchors', () => { + test('a recorded turn stays restorable after the cleanup job runs', async () => { + const kept = path.join(projectDir, 'anchored.txt') + writeFileSync(kept, 'the user wrote this\n') + const hash = await trackSnapshot(projectDir) + expect(hash).toBeTruthy() + + // The turn edits the file, and the snapshot index moves on with it. + writeFileSync(kept, 'the agent changed it\n') + await patchSnapshot(projectDir, hash!) + recordUndoEntry(CHAT_ID, { + hashBefore: hash!, + files: ['anchored.txt'], + message: 'a recorded turn', + }) + pruneSnapshotRepo() + + const message = await undoToRecord(CHAT_ID, projectDir, peekUndo(CHAT_ID)!.id) + + expect(message).not.toBeNull() + expect(readFileSync(kept, 'utf8')).toBe('the user wrote this\n') + }) + + test('the sweep drops the anchors of a chat whose journal is gone', async () => { + writeFileSync(path.join(projectDir, 'swept.txt'), 'content\n') + const hash = await trackSnapshot(projectDir) + expect(anchorSnapshot(projectDir, 'abandoned-chat', hash!)).toBeTruthy() + expect(listAnchors(projectDir)).toContainEqual({ + key: 'abandoned-chat', + hash: hash!, + }) + + const { released } = sweepAnchors(projectDir) + + expect(released).toBeGreaterThan(0) + expect(listAnchors(projectDir)).not.toContainEqual({ + key: 'abandoned-chat', + hash: hash!, + }) + }) + + test('the sweep re-anchors an entry whose ref went missing', async () => { + writeFileSync(path.join(projectDir, 'reanchor.txt'), 'content\n') + const hash = await trackSnapshot(projectDir) + recordUndoEntry(CHAT_ID, { + hashBefore: hash!, + files: ['reanchor.txt'], + message: 'a recorded turn', + }) + // Pretend the release path ran early, or its write was lost. + expect(releaseSnapshot(projectDir, CHAT_ID, hash!)).toBe(true) + expect(listAnchors(projectDir)).not.toContainEqual({ + key: CHAT_ID, + hash: hash!, + }) + + sweepAnchors(projectDir) + + expect(listAnchors(projectDir)).toContainEqual({ key: CHAT_ID, hash: hash! }) + }) +}) + +describe('corrupt file handling', () => { + test('loads an empty state for a nonexistent chat', () => { + expect(loadUndoState('never-existed').undoStack).toEqual([]) + expect(loadUndoState('never-existed').redoStack).toEqual([]) + }) +}) diff --git a/cli/src/state/undo-store.ts b/cli/src/state/undo-store.ts new file mode 100644 index 0000000000..a04b78b475 --- /dev/null +++ b/cli/src/state/undo-store.ts @@ -0,0 +1,406 @@ +/** + * Per-chat undo/redo journal. + * + * Each assistant turn that changed files records an entry here: the snapshot + * hash captured before the turn plus the files it changed. `/undo` pops the + * most recent entry and reverts those files via the snapshot repo; `/redo` + * restores the state captured at undo time. Persisted as `undo.json` inside + * the chat's data directory so it survives restarts and follows the chat + * across `/history` resumes. + */ + +import { randomUUID } from 'node:crypto' +import { + existsSync, + mkdirSync, + readdirSync, + readFileSync, + writeFileSync, +} from 'node:fs' +import path from 'node:path' + +import { getProjectDataDir, tryGetProjectRoot } from '../project-files' +import { logger } from '../utils/logger' +import { + anchorSnapshot, + diffSnapshot, + isSnapshotAvailable, + listAnchors, + releaseSnapshot, + restoreSnapshot, + revertFiles, + trackSnapshot, +} from '../utils/undo-snapshot' + +export type UndoRecord = { + id: string + chatId: string + /** Snapshot (git tree hash) captured before the assistant turn. */ + hashBefore: string + /** Files the turn changed, relative to the project root. */ + files: string[] + /** The user message that started the turn. */ + message: string + createdAt: string + /** Snapshot captured at /undo time; set only while the record is redoable. */ + hashAfter?: string + /** + * Turns reverted by an undo action that jumped back past them. Stored on + * the redo record so /redo can restore both the files and the stacks. + */ + restored?: UndoRecord[] +} + +export type UndoState = { + undoStack: UndoRecord[] + redoStack: UndoRecord[] +} + +const MAX_UNDO_ENTRIES = 20 +const MAX_MESSAGE_CHARS = 120 + +/** Keep the journal small: first line of the prompt, truncated. */ +function truncateMessage(message: string): string { + const firstLine = (message.split('\n')[0] ?? '').trim() + if (firstLine.length <= MAX_MESSAGE_CHARS) return firstLine + return `${firstLine.slice(0, MAX_MESSAGE_CHARS - 1)}…` +} + +function chatDirFor(chatId: string): string { + return path.join(getProjectDataDir(), 'chats', chatId) +} + +function undoFilePath(chatId: string): string { + return path.join(chatDirFor(chatId), 'undo.json') +} + +export function loadUndoState(chatId: string): UndoState { + try { + const file = undoFilePath(chatId) + if (existsSync(file)) { + const parsed = JSON.parse(readFileSync(file, 'utf8')) as Partial + return { + undoStack: Array.isArray(parsed.undoStack) ? parsed.undoStack : [], + redoStack: Array.isArray(parsed.redoStack) ? parsed.redoStack : [], + } + } + } catch { + // Corrupt or unreadable — treat as empty rather than breaking commands. + } + return { undoStack: [], redoStack: [] } +} + +export function saveUndoState(chatId: string, state: UndoState): void { + try { + const dir = chatDirFor(chatId) + mkdirSync(dir, { recursive: true }) + writeFileSync(undoFilePath(chatId), JSON.stringify(state, null, 2)) + } catch { + // Best-effort; undo is a convenience feature. + } +} + +/** Every hash the journal would need in order to restore from either stack. */ +function referencedHashes(state: UndoState): Set { + const hashes = new Set() + for (const record of [...state.undoStack, ...state.redoStack]) { + if (record.hashBefore) hashes.add(record.hashBefore) + if (record.hashAfter) hashes.add(record.hashAfter) + for (const reverted of record.restored ?? []) { + if (reverted.hashBefore) hashes.add(reverted.hashBefore) + if (reverted.hashAfter) hashes.add(reverted.hashAfter) + } + } + return hashes +} + +/** + * Drop the anchors of the snapshots this chat no longer lists. Snapshots are + * held by refs (see undo-snapshot), so releasing them is what keeps the store + * from growing for the life of the project. + * + * Synchronous, like the rest of the journal's bookkeeping: the file it just + * wrote must not list an entry whose snapshot is still held by a stale ref, + * nor hold a ref no entry lists. + */ +function releaseDroppedAnchors( + chatId: string, + before: Set, + after: Set, +): void { + const projectRoot = tryGetProjectRoot() + if (!projectRoot) return + for (const hash of before) { + if (after.has(hash)) continue + releaseSnapshot(projectRoot, chatId, hash) + } +} + +/** Chat ids that have a data directory on disk, journalled or not. */ +function listJournalChatIds(): string[] { + try { + return readdirSync(path.join(getProjectDataDir(), 'chats'), { + withFileTypes: true, + }) + .filter((entry) => entry.isDirectory()) + .map((entry) => entry.name) + } catch { + // No chats yet, or the directory is unreadable. + return [] + } +} + +/** + * Make the anchors match the journals: hold what a journal lists, drop what + * none lists, and re-anchor an entry whose ref went missing. + * + * The journals are the root set. Releasing on drop covers the entries this + * store hands out, but it runs per chat and cannot know about a chat that was + * deleted: its `undo.json` disappears and the refs its turns created stay + * behind. Sweeping against the journals collects those, and repairs an anchor + * or a release that failed, which is what keeps the store bounded. + */ +export function sweepAnchors(projectRoot: string): { + released: number + anchored: number +} { + const anchors = listAnchors(projectRoot) + const chatIds = new Set(anchors.map((anchor) => anchor.key)) + for (const chatId of listJournalChatIds()) chatIds.add(chatId) + + let released = 0 + let anchored = 0 + for (const chatId of chatIds) { + const needed = referencedHashes(loadUndoState(chatId)) + for (const anchor of anchors) { + if (anchor.key !== chatId || needed.has(anchor.hash)) continue + if (releaseSnapshot(projectRoot, chatId, anchor.hash)) released += 1 + } + for (const hash of needed) { + if (anchors.some((a) => a.key === chatId && a.hash === hash)) continue + if (anchorSnapshot(projectRoot, chatId, hash)) anchored += 1 + } + } + return { released, anchored } +} + +/** The project root the sweep has already run for, so it runs once each. */ +let sweptProjectRoot: string | null = null + +/** + * Sweep once per project root per process, after the current work is done so + * the turn's own path never waits on maintenance. + */ +function sweepAnchorsInBackground(): void { + const projectRoot = tryGetProjectRoot() + if (!projectRoot || sweptProjectRoot === projectRoot) return + sweptProjectRoot = projectRoot + setTimeout(() => { + try { + const { released, anchored } = sweepAnchors(projectRoot) + if (released > 0 || anchored > 0) { + logger.debug( + { released, anchored }, + 'undo-store: swept the snapshot anchors', + ) + } + } catch (error) { + logger.debug({ error }, 'undo-store: anchor sweep failed') + } + }, 0) +} + +/** + * Record a completed turn. Clears the redo stack: new edits invalidate redo. + * No-ops when the snapshot is missing or nothing changed. + */ +export function recordUndoEntry( + chatId: string, + entry: { hashBefore: string; files: string[]; message: string }, +): void { + if (!entry.hashBefore || entry.files.length === 0) return + const state = loadUndoState(chatId) + const referenced = referencedHashes(state) + state.undoStack.push({ + id: randomUUID(), + chatId, + hashBefore: entry.hashBefore, + files: entry.files, + message: truncateMessage(entry.message), + createdAt: new Date().toISOString(), + }) + if (state.undoStack.length > MAX_UNDO_ENTRIES) { + state.undoStack.shift() + } + state.redoStack = [] + saveUndoState(chatId, state) + releaseDroppedAnchors(chatId, referenced, referencedHashes(state)) + // After the journal lists the turn, so the sweep can never see a snapshot + // no entry is holding yet. + const projectRoot = tryGetProjectRoot() + if (projectRoot) anchorSnapshot(projectRoot, chatId, entry.hashBefore) + sweepAnchorsInBackground() +} + +export function peekUndo(chatId: string): UndoRecord | null { + const stack = loadUndoState(chatId).undoStack + return stack[stack.length - 1] ?? null +} + +export function popUndo(chatId: string): UndoRecord | null { + const state = loadUndoState(chatId) + const record = state.undoStack.pop() ?? null + if (record) saveUndoState(chatId, state) + return record +} + +/** Push a record back onto the undo stack (used by /redo). Does not touch redo. */ +export function pushUndo(chatId: string, record: UndoRecord): void { + const state = loadUndoState(chatId) + const referenced = referencedHashes(state) + state.undoStack.push(record) + if (state.undoStack.length > MAX_UNDO_ENTRIES) { + state.undoStack.shift() + } + saveUndoState(chatId, state) + releaseDroppedAnchors(chatId, referenced, referencedHashes(state)) +} + +export function peekRedo(chatId: string): UndoRecord | null { + const stack = loadUndoState(chatId).redoStack + return stack[stack.length - 1] ?? null +} + +export function popRedo(chatId: string): UndoRecord | null { + const state = loadUndoState(chatId) + const record = state.redoStack.pop() ?? null + if (record) saveUndoState(chatId, state) + return record +} + +export function pushRedo(chatId: string, record: UndoRecord): void { + const state = loadUndoState(chatId) + state.redoStack.push(record) + saveUndoState(chatId, state) +} + +/** The undo stack for a chat, oldest turn first. */ +export function listUndoEntries(chatId: string): UndoRecord[] { + return loadUndoState(chatId).undoStack +} + +/** The redo stack for a chat, oldest action first. */ +export function listRedoEntries(chatId: string): UndoRecord[] { + return loadUndoState(chatId).redoStack +} + +/** + * Undo back to a specific recorded turn (the OpenCode model): reverts that + * turn's files AND everything the agent changed in newer turns, restoring the + * project to the snapshot captured before the selected turn. The reverted + * turns move onto the redo stack so /redo can restore the exact state that + * was left behind. Returns the confirmation message, or null when the + * snapshot store is unavailable (in which case nothing is changed). + */ +export async function undoToRecord( + chatId: string, + projectRoot: string, + recordId: string, +): Promise { + const state = loadUndoState(chatId) + const index = state.undoStack.findIndex((record) => record.id === recordId) + if (index === -1) return null + const record = state.undoStack[index]! + // The selected turn plus every newer one — all of it is reverted. + const affected = state.undoStack.slice(index) + const files = Array.from(new Set(affected.flatMap((r) => r.files))) + + // A snapshot can be gone — collected by the cleanup job, or lost with the + // config directory. Abort before touching anything: reverting against a + // snapshot it cannot read is how undoing deletes the files it meant to keep. + if (!(await isSnapshotAvailable(projectRoot, record.hashBefore))) return null + + // Capture the current state first so /redo can restore exactly what was + // undone. If the snapshot store is unavailable, abort without mutating the + // stacks (mirrors the old inline handler's behavior). + const hashAfter = await trackSnapshot(projectRoot) + if (!hashAfter) return null + // Held until the redo record that points at it is dropped. + anchorSnapshot(projectRoot, chatId, hashAfter) + + // Diff of what the affected turns changed (computed before reverting). + const diffStat = await diffSnapshot(projectRoot, record.hashBefore) + const { restored, deleted } = await revertFiles( + projectRoot, + record.hashBefore, + files, + ) + + state.undoStack = state.undoStack.slice(0, index) + state.redoStack.push({ ...record, hashAfter, restored: affected }) + saveUndoState(chatId, state) + + const undone = [ + ...restored.map((file) => ` ↺ ${file}`), + ...deleted.map((file) => ` 🗑 ${file} (deleted)`), + ].join('\n') + if (!undone) return 'Could not undo the selected change.' + const turns = affected.length + const heading = + turns === 1 + ? '**Undid the last change:**' + : `**Undid ${turns} change(s) back to: ${truncateMessage(record.message)}**` + return `${heading}\n${undone}${diffStat ? `\n\n${diffStat}` : ''}` +} + +/** + * Redo a specific undo action: restores the project to the state captured + * when that undo ran and moves the reverted turns back onto the undo stack. + * Redo actions newer than the selected one are invalidated by the jump. + * Returns the confirmation message, or null when the restore fails. + */ +export async function redoToRecord( + chatId: string, + projectRoot: string, + recordId: string, +): Promise { + const state = loadUndoState(chatId) + const index = state.redoStack.findIndex((record) => record.id === recordId) + if (index === -1) return null + const record = state.redoStack[index]! + if (!record.hashAfter) return null + + // Diff of what this redo restores (computed before the tree changes). + const diffStat = await diffSnapshot(projectRoot, record.hashAfter) + const ok = await restoreSnapshot(projectRoot, record.hashAfter) + if (!ok) return null + + // Drop newer redo actions — jumping back invalidates them. + const referenced = referencedHashes(state) + state.redoStack = state.redoStack.slice(0, index) + if (record.restored && record.restored.length > 0) { + state.undoStack.push(...record.restored) + } else { + state.undoStack.push({ ...record, hashAfter: undefined, restored: undefined }) + } + saveUndoState(chatId, state) + releaseDroppedAnchors(chatId, referenced, referencedHashes(state)) + + const turns = record.restored?.length ?? 1 + const files = record.restored + ? Array.from(new Set(record.restored.flatMap((r) => r.files))) + : record.files + // After the restore, files present in the tree came back; files missing + // were removed by the restored state (the agent had deleted them). + const fileLines = files.map((file) => + existsSync(path.join(projectRoot, file)) + ? ` ↺ ${file}` + : ` 🗑 ${file} (deleted)`, + ) + const fileWord = files.length === 1 ? 'file' : 'files' + const heading = + turns === 1 + ? `**Redid the last change (${files.length} ${fileWord}):**` + : `**Redid ${turns} change(s) (${files.length} ${fileWord}):**` + return `${heading}\n${fileLines.join('\n')}${diffStat ? `\n\n${diffStat}` : ''}` +} diff --git a/cli/src/utils/__tests__/undo-snapshot.test.ts b/cli/src/utils/__tests__/undo-snapshot.test.ts new file mode 100644 index 0000000000..03d43710fa --- /dev/null +++ b/cli/src/utils/__tests__/undo-snapshot.test.ts @@ -0,0 +1,257 @@ +import { afterAll, beforeAll, describe, expect, test } from 'bun:test' + +import { execFileSync } from 'node:child_process' +import { + chmodSync, + existsSync, + mkdtempSync, + readdirSync, + rmSync, + writeFileSync, +} from 'node:fs' +import os from 'node:os' +import path from 'node:path' + +import { + anchorSnapshot, + isSnapshotAvailable, + isUndoAvailable, + patchSnapshot, + releaseSnapshot, + restoreSnapshot, + revertFiles, + setSnapshotDirOverrideForTesting, + trackSnapshot, +} from '../undo-snapshot' + +let projectDir: string +let snapshotRoot: string +let plainDir: string + +/** The journal entry the snapshots belong to; anchors are keyed by it. */ +const SNAPSHOT_KEY = 'test-chat' + +const gitInProject = (args: string[]): string => + execFileSync('git', args, { + cwd: projectDir, + encoding: 'utf8', + // Keep expected stderr noise (e.g. the "ambiguous HEAD" probe) out of the + // test output. + stdio: ['ignore', 'pipe', 'ignore'], + }) + +const readProjectFile = (file: string): string => + execFileSync('cat', [path.join(projectDir, file)], { encoding: 'utf8' }) + +/** The snapshot repo the fixture has initialized under the override root. */ +const snapshotRepoDir = (): string => { + const [entry] = readdirSync(snapshotRoot) + if (!entry) throw new Error('no snapshot repo was initialized') + return path.join(snapshotRoot, entry) +} + +/** Run the cleanup job's gc now, instead of waiting out its 7-day grace. */ +const pruneSnapshotRepo = (): void => { + execFileSync( + 'git', + [ + '--git-dir', + snapshotRepoDir(), + '--work-tree', + projectDir, + 'gc', + '--quiet', + '--prune=now', + ], + { stdio: ['ignore', 'pipe', 'ignore'] }, + ) +} + +beforeAll(() => { + projectDir = mkdtempSync(path.join(os.tmpdir(), 'undo-snapshot-proj-')) + snapshotRoot = mkdtempSync(path.join(os.tmpdir(), 'undo-snapshot-store-')) + plainDir = mkdtempSync(path.join(os.tmpdir(), 'undo-snapshot-plain-')) + setSnapshotDirOverrideForTesting(snapshotRoot) + // A real (but commit-less) git repository. write-tree does not need commits + // or user identity, so this is all the fixture requires. + execFileSync('git', ['init'], { cwd: projectDir }) +}) + +afterAll(() => { + setSnapshotDirOverrideForTesting(undefined) + rmSync(projectDir, { recursive: true, force: true }) + rmSync(snapshotRoot, { recursive: true, force: true }) + rmSync(plainDir, { recursive: true, force: true }) +}) + +describe('isUndoAvailable', () => { + test('is true for a git repository', () => { + expect(isUndoAvailable(projectDir)).toBe(true) + }) + + test('is false for a directory without git', () => { + expect(isUndoAvailable(plainDir)).toBe(false) + }) +}) + +describe('trackSnapshot', () => { + test('returns null for a non-git directory', async () => { + expect(await trackSnapshot(plainDir)).toBeNull() + }) + + test('returns a hash and detects modified, added, and deleted files', async () => { + writeFileSync(path.join(projectDir, 'a.txt'), 'hello\n') + writeFileSync(path.join(projectDir, 'c.txt'), 'keep me\n') + const hash = await trackSnapshot(projectDir) + expect(hash).toBeTruthy() + + // Modify a tracked-in-snapshot file, add a new file, delete another. + writeFileSync(path.join(projectDir, 'a.txt'), 'hello world\n') + writeFileSync(path.join(projectDir, 'b.txt'), 'new file\n') + rmSync(path.join(projectDir, 'c.txt')) + + const changed = await patchSnapshot(projectDir, hash!) + expect(changed.sort()).toEqual(['a.txt', 'b.txt', 'c.txt']) + }) +}) + +describe('restoreSnapshot', () => { + test('overwrites the worktree with the tracked state', async () => { + writeFileSync(path.join(projectDir, 'restore.txt'), 'original\n') + const hash = await trackSnapshot(projectDir) + + writeFileSync(path.join(projectDir, 'restore.txt'), 'changed by agent\n') + + const ok = await restoreSnapshot(projectDir, hash!) + expect(ok).toBe(true) + expect(readProjectFile('restore.txt')).toBe('original\n') + }) +}) + +describe('revertFiles', () => { + test('restores a modified file and deletes a file created after the snapshot', async () => { + writeFileSync(path.join(projectDir, 'revert.txt'), 'v1\n') + const hash = await trackSnapshot(projectDir) + + writeFileSync(path.join(projectDir, 'revert.txt'), 'v2 by agent\n') + writeFileSync(path.join(projectDir, 'agent-created.txt'), 'agent made this\n') + + const { restored, deleted } = await revertFiles(projectDir, hash!, [ + 'revert.txt', + 'agent-created.txt', + ]) + // revert.txt existed in the snapshot and is restored; agent-created.txt did + // not exist in the snapshot (the agent created it) and is deleted. + expect(restored).toEqual(['revert.txt']) + expect(deleted).toEqual(['agent-created.txt']) + expect(readProjectFile('revert.txt')).toBe('v1\n') + expect(existsSync(path.join(projectDir, 'agent-created.txt'))).toBe(false) + }) +}) + +describe('a snapshot whose content is gone', () => { + test('leaves the worktree alone when a blob the tree lists is gone', async () => { + // Make sure the snapshot repo exists before reaching into it. + await trackSnapshot(projectDir) + + // A tree that references a blob which exists nowhere: the tree survives, + // its content does not. This is what borrowing the project's objects + // looks like once the project collects the object it was lending. + const missing = '0123456789abcdef0123456789abcdef01234567' + const tree = execFileSync( + 'git', + [ + '--git-dir', + snapshotRepoDir(), + '--work-tree', + projectDir, + 'mktree', + '--missing', + ], + { + cwd: projectDir, + encoding: 'utf8', + input: `100644 blob ${missing}\torphan.txt\n`, + }, + ).trim() + expect(tree).toBeTruthy() + + writeFileSync(path.join(projectDir, 'orphan.txt'), 'the agent wrote this\n') + + const { restored, deleted } = await revertFiles(projectDir, tree, [ + 'orphan.txt', + ]) + // `git checkout` deletes a worktree file whose blob it cannot read, so + // reverting against this tree would destroy the file instead of restoring + // it. Nothing may be restored, nothing deleted, and the file must survive. + expect(readProjectFile('orphan.txt')).toBe('the agent wrote this\n') + expect(restored).toEqual([]) + expect(deleted).toEqual([]) + expect(existsSync(path.join(projectDir, 'orphan.txt'))).toBe(true) + + // And it must not be advertised as restorable either. + expect(await isSnapshotAvailable(projectDir, tree)).toBe(false) + }) +}) + +describe('snapshot anchoring', () => { + test('git keeps the anchored snapshot and collects the unanchored one', async () => { + writeFileSync(path.join(projectDir, 'lifetime.txt'), 'before the turn\n') + const held = await trackSnapshot(projectDir) + expect(held).toBeTruthy() + expect(anchorSnapshot(projectDir, SNAPSHOT_KEY, held!)).toBeTruthy() + + // Two more turns move the snapshot index on, so nothing in git points at + // either tree any more. + writeFileSync(path.join(projectDir, 'lifetime.txt'), 'after the turn\n') + const unheld = await trackSnapshot(projectDir) + expect(unheld).toBeTruthy() + writeFileSync(path.join(projectDir, 'lifetime.txt'), 'and once more\n') + await trackSnapshot(projectDir) + pruneSnapshotRepo() + + // The anchored one survives; the one no journal holds is collected. + expect(await isSnapshotAvailable(projectDir, held!)).toBe(true) + expect(await isSnapshotAvailable(projectDir, unheld!)).toBe(false) + + // And the anchor is what holds it: drop it and git collects it too. + expect(releaseSnapshot(projectDir, SNAPSHOT_KEY, held!)).toBe(true) + pruneSnapshotRepo() + expect(await isSnapshotAvailable(projectDir, held!)).toBe(false) + expect(await restoreSnapshot(projectDir, held!)).toBe(false) + }) +}) + +describe('failed staging', () => { + test('a failed git add does not hand back a hash', async () => { + // A root user can read anything, so the fixture cannot fail the add there. + if (typeof process.getuid === 'function' && process.getuid() === 0) return + + // git refuses to index an unreadable file, but `write-tree` still returns + // the tree of the stale index: a hash that looks valid and is not. + const blocked = path.join(projectDir, 'blocked.txt') + writeFileSync(blocked, 'unreadable\n') + chmodSync(blocked, 0o000) + try { + expect(await trackSnapshot(projectDir)).toBeNull() + } finally { + chmodSync(blocked, 0o644) + rmSync(blocked, { force: true }) + } + }) +}) + +describe('isolation from the real repository', () => { + test('never stages or commits anything in the project git repo', async () => { + // Fresh fixture area. + writeFileSync(path.join(projectDir, 'isolated.txt'), 'content\n') + await trackSnapshot(projectDir) + writeFileSync(path.join(projectDir, 'isolated.txt'), 'content v2\n') + await trackSnapshot(projectDir) + + // The real repo must have no staged changes and no commits of ours. + const staged = gitInProject(['diff', '--cached', '--name-only']).trim() + expect(staged).toBe('') + expect(() => gitInProject(['rev-parse', 'HEAD'])).toThrow() + }) +}) diff --git a/cli/src/utils/undo-snapshot.ts b/cli/src/utils/undo-snapshot.ts new file mode 100644 index 0000000000..6818c0ceec --- /dev/null +++ b/cli/src/utils/undo-snapshot.ts @@ -0,0 +1,751 @@ +/** + * Undo snapshot store. + * + * Tracks the filesystem state before and after each assistant turn using a + * dedicated, hidden git repository per project — the same approach as + * OpenCode's snapshot service (`packages/opencode/src/snapshot/index.ts`). + * + * The snapshot repository never touches the project's real `.git`: every git + * invocation runs with `--git-dir --work-tree `. + * A snapshot is the tree hash produced by `git write-tree` after staging the + * project's changes, so restoring one is a matter of `git read-tree` + + * `git checkout-index`. The source repository's object database is reused via + * `objects/info/alternates` (with the source index copied) so the first + * snapshot stays fast even on huge repositories. + * + * All operations are best-effort: a failure disables undo for that turn + * rather than breaking the chat. + */ + +import { execFileSync, spawn } from 'node:child_process' +import { createHash } from 'node:crypto' +import { + copyFileSync, + existsSync, + mkdirSync, + readFileSync, + rmSync, + statSync, + writeFileSync, +} from 'node:fs' +import path from 'node:path' + +import { getConfigDir } from './config-dir' +import { findGitRoot } from './git' +import { logger } from './logger' + +/** Files larger than this are excluded from snapshots (same limit as OpenCode). */ +const MAX_SNAPSHOT_FILE_BYTES = 2 * 1024 * 1024 + +/** Absolute-path pathspec magic used with `--pathspec-from-file`. */ +const topLevelLiteral = (file: string): string => `:(top,literal)${file}` + +/** + * Refs that keep a snapshot alive. + * + * `write-tree` returns a bare tree and nothing points at it, so git considers + * it unreachable as soon as the snapshot index moves on with the next turn -- + * and the cleanup job's `gc --prune` collects it while the journal still lists + * the entry. Pointing a ref at it is how git's own docs say to keep an object + * around, and what other snapshot tools do. + * + * The ref is keyed by the journal entry's owner (the chat) so two chats can + * hold the same tree without one of them releasing the other's copy. One ref + * per live entry: `releaseSnapshot` drops it, which is what bounds the store. + */ +const ANCHOR_REF_PREFIX = 'refs/freebuff/undo/' + +const anchorRefFor = (key: string, hash: string): string => + `${ANCHOR_REF_PREFIX}${key}/${hash}` + +/** + * Identity and signing for the anchor commit, set explicitly so the user's own + * global git config (missing `user.email`, `commit.gpgsign` on) cannot fail it. + */ +const ANCHOR_GIT_CONFIG = [ + '-c', + 'user.name=freebuff', + '-c', + 'user.email=freebuff@localhost', + '-c', + 'commit.gpgsign=false', +] + +type GitResult = { + code: number + stdout: string + stderr: string +} + +async function runGit( + args: string[], + cwd: string, + options?: { input?: string; env?: Record }, +): Promise { + return new Promise((resolve) => { + const child = spawn('git', args, { + cwd, + env: options?.env ? { ...process.env, ...options.env } : process.env, + stdio: ['pipe', 'pipe', 'pipe'], + }) + let stdout = '' + let stderr = '' + child.stdout.setEncoding('utf8') + child.stderr.setEncoding('utf8') + child.stdout.on('data', (chunk: string) => { + stdout += chunk + }) + child.stderr.on('data', (chunk: string) => { + stderr += chunk + }) + child.on('error', (error) => { + resolve({ + code: 1, + stdout, + stderr: error instanceof Error ? error.message : String(error), + }) + }) + child.on('close', (code) => { + resolve({ code: code ?? 1, stdout, stderr }) + }) + if (options?.input) { + child.stdin.write(options.input) + } + child.stdin.end() + }) +} + +/** Serialize git operations per snapshot repo so concurrent tracks can't corrupt the index. */ +const locks = new Map>() + +/** Hourly GC keeps snapshot repos from growing without bound (prune 7 days). */ +const GC_INTERVAL_MS = 60 * 60 * 1000 +const GC_PRUNE = '7.days' +const lastGcByDir = new Map() + +function withLock(key: string, fn: () => Promise): Promise { + const previous = locks.get(key) ?? Promise.resolve() + const next = previous.then(fn, fn) + locks.set( + key, + next.catch(() => { + // Swallow so a failed op never wedges the chain for later ones. + }), + ) + return next +} + +/** Where snapshot repos live. Test-overridable via setSnapshotDirOverrideForTesting. */ +let snapshotDirOverride: string | undefined + +export function setSnapshotDirOverrideForTesting(dir: string | undefined): void { + snapshotDirOverride = dir +} + +function getSnapshotBaseDir(): string { + return snapshotDirOverride ?? path.join(getConfigDir(), 'undo-snapshots') +} + +function snapshotDirFor(projectRoot: string): string { + const key = createHash('sha1').update(projectRoot).digest('hex').slice(0, 12) + return path.join(getSnapshotBaseDir(), `${path.basename(projectRoot)}-${key}`) +} + +/** git args that point every command at the snapshot repo, not the real one. */ +const gitArgs = ( + snapshotDir: string, + projectRoot: string, + command: string[], +): string[] => ['--git-dir', snapshotDir, '--work-tree', projectRoot, ...command] + +/** Whether undo can work for this project at all (it must be a git repo). */ +export function isUndoAvailable(projectRoot: string): boolean { + return findGitRoot({ cwd: projectRoot }) !== null +} + +async function ensureInitialized( + snapshotDir: string, + projectRoot: string, +): Promise { + if (existsSync(snapshotDir)) return true + try { + mkdirSync(snapshotDir, { recursive: true }) + await runGit(['init'], projectRoot, { + env: { GIT_DIR: snapshotDir, GIT_WORK_TREE: projectRoot }, + }) + const configArgs = ['--git-dir', snapshotDir, 'config'] + await runGit([...configArgs, 'core.autocrlf', 'false'], projectRoot) + await runGit([...configArgs, 'core.longpaths', 'true'], projectRoot) + await runGit([...configArgs, 'feature.manyFiles', 'true'], projectRoot) + await runGit([...configArgs, 'index.version', '4'], projectRoot) + await runGit([...configArgs, 'core.untrackedCache', 'true'], projectRoot) + await seedFromSourceRepo(snapshotDir, projectRoot) + return true + } catch (error) { + logger.debug({ error }, 'undo-snapshot: failed to initialize snapshot repo') + return false + } +} + +/** + * Reuse the source repo's object database (and index) so already-hashed file + * content does not need re-hashing on the first snapshot. Best-effort. + */ +async function seedFromSourceRepo( + snapshotDir: string, + projectRoot: string, +): Promise { + try { + const common = await runGit( + ['rev-parse', '--path-format=absolute', '--git-common-dir'], + projectRoot, + ) + if (common.code !== 0) return + const source = common.stdout.trim() + if (!source || !existsSync(source)) return + + const sourceObjects = path.join(source, 'objects') + if (!existsSync(sourceObjects)) return + + const alternatesDir = path.join(snapshotDir, 'objects', 'info') + mkdirSync(alternatesDir, { recursive: true }) + writeFileSync(path.join(alternatesDir, 'alternates'), `${sourceObjects}\n`) + + const sourceIndex = path.join(source, 'index') + if (existsSync(sourceIndex)) { + copyFileSync(sourceIndex, path.join(snapshotDir, 'index')) + } + } catch (error) { + logger.debug({ error }, 'undo-snapshot: failed to seed from source repo') + } +} + +/** Mirror the source repo's info/exclude plus blocked (oversized) files into the snapshot repo. */ +function syncExcludes( + snapshotDir: string, + projectRoot: string, + gitRoot: string | null, + blocked: string[], +): void { + try { + const lines: string[] = [] + const sourceExclude = gitRoot + ? path.join(gitRoot, '.git', 'info', 'exclude') + : null + if (sourceExclude && existsSync(sourceExclude)) { + lines.push(...readFileSync(sourceExclude, 'utf8').split('\n')) + } + for (const file of blocked) { + lines.push(`/${file.replaceAll('\\', '/')}`) + } + mkdirSync(path.join(snapshotDir, 'info'), { recursive: true }) + writeFileSync( + path.join(snapshotDir, 'info', 'exclude'), + `${lines.filter((line) => line.trim()).join('\n')}\n`, + ) + } catch { + // Best-effort. + } +} + +/** + * Stage every changed, added, and deleted file (respecting the project's own + * ignore rules and skipping files over the size limit) into the snapshot + * index, so a following `write-tree` reflects the current state. + * @returns false when git refused to stage something, in which case the index + * no longer describes the worktree and must not be turned into a snapshot. + */ +async function stageChanges( + snapshotDir: string, + projectRoot: string, + gitRoot: string | null, +): Promise { + try { + const base = gitArgs(snapshotDir, projectRoot, []) + const [diff, others] = await Promise.all([ + runGit( + [...base, 'diff-files', '--name-only', '-z', '--', '.'], + projectRoot, + ), + runGit( + [ + ...base, + 'ls-files', + '--full-name', + '--others', + '--exclude-standard', + '-z', + '--', + '.', + ], + projectRoot, + ), + ]) + const changed = diff.stdout.split('\0').filter(Boolean) + const untracked = others.stdout.split('\0').filter(Boolean) + const all = Array.from(new Set([...changed, ...untracked])) + if (all.length === 0) return true + + const allowed: string[] = [] + const blocked: string[] = [] + for (const file of all) { + try { + const stat = statSync(path.join(projectRoot, file)) + if (stat.isFile() && stat.size > MAX_SNAPSHOT_FILE_BYTES) { + blocked.push(file) + } else { + allowed.push(file) + } + } catch { + // Deleted or unreadable — still stage the removal. + allowed.push(file) + } + } + + syncExcludes(snapshotDir, projectRoot, gitRoot, blocked) + + if (allowed.length === 0) return true + const pathspecs = `${allowed.map(topLevelLiteral).join('\0')}\0` + const staged = await runGit( + [ + ...base, + 'add', + '--all', + '--sparse', + '--pathspec-from-file=-', + '--pathspec-file-nul', + ], + projectRoot, + { input: pathspecs }, + ) + if (staged.code !== 0) { + // git refused a path (unreadable file, bad name, ...). The index keeps + // the state it had before, and `write-tree` would happily return the + // tree of that stale index: a hash that looks valid and is not. + logger.debug( + { stderr: staged.stderr }, + 'undo-snapshot: git add failed, snapshot would be stale', + ) + return false + } + return true + } catch (error) { + logger.debug({ error }, 'undo-snapshot: failed to stage changes') + return false + } +} + +/** + * Run git and return its trimmed stdout, or null when it failed. + * + * The journal's anchoring is synchronous on purpose. `saveUndoState` already + * writes its file synchronously, and an entry must never be listed before its + * snapshot is anchored -- the sweep reads the anchors and relies on that. + * These calls only touch refs and objects, never the index, so they do not + * contend with the staged operations above, which stay on the async path. + */ +function gitSync(args: string[], projectRoot: string): string | null { + try { + return execFileSync('git', args, { + cwd: projectRoot, + encoding: 'utf8', + stdio: ['ignore', 'pipe', 'ignore'], + }).trim() + } catch (error) { + logger.debug({ error }, 'undo-snapshot: git command failed') + return null + } +} + +/** + * Keep a snapshot out of git's reach by pointing a ref at it, so the cleanup + * job's prune cannot collect a snapshot the journal still lists. + * @returns the ref name, or null when it could not be created. + */ +export function anchorSnapshot( + projectRoot: string, + key: string, + hash: string, +): string | null { + if (!hash) return null + if (!findGitRoot({ cwd: projectRoot })) return null + const snapshotDir = snapshotDirFor(projectRoot) + const base = ['--git-dir', snapshotDir, '--work-tree', projectRoot] + const commit = gitSync( + [ + ...ANCHOR_GIT_CONFIG, + ...base, + 'commit-tree', + hash, + '-m', + 'freebuff undo snapshot', + ], + projectRoot, + ) + if (!commit) return null + const ref = anchorRefFor(key, hash) + // `update-ref` prints nothing on success, so compare against null rather + // than truthiness: an empty string is the happy path here. + return gitSync([...base, 'update-ref', ref, commit], projectRoot) !== null + ? ref + : null +} + +/** + * Drop a snapshot's anchor so the cleanup job can reclaim it. The journal + * calls this when it stops listing an entry (eviction, or the redo stack a new + * turn clears). + */ +export function releaseSnapshot( + projectRoot: string, + key: string, + hash: string, +): boolean { + if (!hash) return false + if (!findGitRoot({ cwd: projectRoot })) return false + const snapshotDir = snapshotDirFor(projectRoot) + return ( + gitSync( + [ + '--git-dir', + snapshotDir, + '--work-tree', + projectRoot, + 'update-ref', + '-d', + anchorRefFor(key, hash), + ], + projectRoot, + ) !== null + ) +} + +/** + * The anchors the snapshot repo is holding, as `{ key, hash }` pairs. The + * journal sweeps its entries against this list. + */ +export function listAnchors( + projectRoot: string, +): { key: string; hash: string }[] { + if (!findGitRoot({ cwd: projectRoot })) return [] + const snapshotDir = snapshotDirFor(projectRoot) + const refs = gitSync( + [ + '--git-dir', + snapshotDir, + '--work-tree', + projectRoot, + 'for-each-ref', + '--format=%(refname)', + ANCHOR_REF_PREFIX, + ], + projectRoot, + ) + if (!refs) return [] + return refs + .split('\n') + .map((line) => line.trim()) + .filter((line) => line.startsWith(ANCHOR_REF_PREFIX)) + .map((ref) => { + const rest = ref.slice(ANCHOR_REF_PREFIX.length) + const separator = rest.indexOf('/') + if (separator <= 0) return null + const key = rest.slice(0, separator) + const hash = rest.slice(separator + 1) + return key && hash && !hash.includes('/') ? { key, hash } : null + }) + .filter((entry): entry is { key: string; hash: string } => entry !== null) +} + +/** + * Capture a snapshot of the project's current state. + * + * The returned tree is not referenced by anything, so it only survives while + * its caller anchors it (`anchorSnapshot`). The cleanup job's prune has a + * grace period of days, which leaves room to anchor after the turn instead of + * on this path. + * + * @returns the snapshot hash, or null when unavailable/failed. + */ +export async function trackSnapshot(projectRoot: string): Promise { + const gitRoot = findGitRoot({ cwd: projectRoot }) + if (!gitRoot) return null + const snapshotDir = snapshotDirFor(projectRoot) + return withLock(snapshotDir, async () => { + try { + if (!(await ensureInitialized(snapshotDir, projectRoot))) return null + // A failed stage leaves the index stale, and `write-tree` would still + // return that stale tree. Refuse rather than hand back a hash that + // would later revert the worktree to the wrong state. + if (!(await stageChanges(snapshotDir, projectRoot, gitRoot))) return null + const result = await runGit( + [...gitArgs(snapshotDir, projectRoot, ['write-tree'])], + projectRoot, + ) + const hash = result.code === 0 ? result.stdout.trim() : '' + if (!hash) return null + void maybePrune(snapshotDir, projectRoot) + return hash + } catch (error) { + logger.debug({ error }, 'undo-snapshot: track failed') + return null + } + }) +} + +/** + * Run `git gc --prune=7.days` at most once per hour per snapshot repo. + * Fire-and-forget; never throws. + */ +async function maybePrune(snapshotDir: string, projectRoot: string): Promise { + const last = lastGcByDir.get(snapshotDir) ?? 0 + if (Date.now() - last < GC_INTERVAL_MS) return + lastGcByDir.set(snapshotDir, Date.now()) + const result = await runGit( + ['--git-dir', snapshotDir, 'gc', `--prune=${GC_PRUNE}`], + projectRoot, + ) + if (result.code !== 0) { + logger.debug({ stderr: result.stderr }, 'undo-snapshot: gc failed') + } +} + +/** + * List the files that changed since the given snapshot. + * @returns project-relative file paths (empty when nothing changed). + */ +export async function patchSnapshot( + projectRoot: string, + hash: string, +): Promise { + const gitRoot = findGitRoot({ cwd: projectRoot }) + if (!gitRoot) return [] + const snapshotDir = snapshotDirFor(projectRoot) + return withLock(snapshotDir, async () => { + try { + if (!(await stageChanges(snapshotDir, projectRoot, gitRoot))) return [] + const result = await runGit( + [ + ...gitArgs(snapshotDir, projectRoot, [ + 'diff', + '--cached', + '--no-ext-diff', + '--name-only', + hash, + '--', + '.', + ]), + ], + projectRoot, + ) + if (result.code !== 0) return [] + return result.stdout + .split('\n') + .map((line) => line.trim()) + .filter(Boolean) + } catch (error) { + logger.debug({ error }, 'undo-snapshot: patch failed') + return [] + } + }) +} + +/** Restore the entire worktree to a snapshot. Returns false on failure. */ +export async function restoreSnapshot( + projectRoot: string, + hash: string, +): Promise { + const gitRoot = findGitRoot({ cwd: projectRoot }) + if (!gitRoot) return false + const snapshotDir = snapshotDirFor(projectRoot) + return withLock(snapshotDir, async () => { + try { + const base = gitArgs(snapshotDir, projectRoot, []) + const read = await runGit([...base, 'read-tree', hash], projectRoot) + if (read.code !== 0) return false + const checkout = await runGit( + [...base, 'checkout-index', '-a', '-f'], + projectRoot, + ) + return checkout.code === 0 + } catch (error) { + logger.debug({ error }, 'undo-snapshot: restore failed') + return false + } + }) +} + +/** + * Whether a snapshot can still be restored from. + * + * A stored hash is a tree the journal kept by hand, and it can be gone: the + * cleanup job collected it before this module started anchoring snapshots, or + * the snapshot directory was wiped while `undo.json` stayed behind. Callers + * must ask before reverting. + * + * The tree surviving is not enough on its own. This module borrows the + * project's object database, so the project's own cleanup can collect a blob + * a kept tree still lists. Restoring from such a tree is not merely partial: + * `git checkout` removes the worktree file whose content it cannot read, so a + * snapshot with holes deletes files instead of restoring them. Ask about the + * content too. + * + * Assumes the caller already holds the snapshot repo's lock. + */ +async function snapshotIsComplete( + snapshotDir: string, + projectRoot: string, + hash: string, +): Promise { + const base = gitArgs(snapshotDir, projectRoot, []) + const tree = await runGit([...base, 'cat-file', '-e', hash], projectRoot) + if (tree.code !== 0) return false + // One command reports every object the tree references and cannot find, + // marking each with a leading `?`. + const objects = await runGit( + [...base, 'rev-list', '--objects', '--missing=print', hash], + projectRoot, + ) + if (objects.code !== 0) return false + return !objects.stdout + .split('\n') + .some((line) => line.trimStart().startsWith('?')) +} + +export async function isSnapshotAvailable( + projectRoot: string, + hash: string, +): Promise { + if (!hash) return false + const gitRoot = findGitRoot({ cwd: projectRoot }) + if (!gitRoot) return false + const snapshotDir = snapshotDirFor(projectRoot) + return withLock(snapshotDir, async () => { + try { + return await snapshotIsComplete(snapshotDir, projectRoot, hash) + } catch (error) { + logger.debug({ error }, 'undo-snapshot: availability check failed') + return false + } + }) +} + +/** + * Restore a specific set of files to a snapshot. Files that did not exist in + * the snapshot (created after it) are deleted. + * @returns the files that were restored and the files that were deleted. + */ +export async function revertFiles( + projectRoot: string, + hash: string, + files: string[], +): Promise<{ restored: string[]; deleted: string[] }> { + const empty = { restored: [] as string[], deleted: [] as string[] } + const gitRoot = findGitRoot({ cwd: projectRoot }) + if (!gitRoot || files.length === 0) return empty + const snapshotDir = snapshotDirFor(projectRoot) + return withLock(snapshotDir, async () => { + const result = { restored: [] as string[], deleted: [] as string[] } + try { + // Never start reverting against content that is not all there: the + // checkout below deletes a file whose blob it cannot read, so a snapshot + // with holes would destroy the worktree instead of restoring it. + if (!(await snapshotIsComplete(snapshotDir, projectRoot, hash))) { + logger.debug( + { hash }, + 'undo-snapshot: snapshot incomplete, leaving the worktree alone', + ) + return result + } + const base = gitArgs(snapshotDir, projectRoot, []) + const root = path.resolve(projectRoot) + for (const file of files) { + // Defense in depth: never touch a path that escapes the project. + const resolved = path.resolve(root, file) + if (resolved !== root && !resolved.startsWith(`${root}${path.sep}`)) { + continue + } + const checkout = await runGit( + [...base, 'checkout', hash, '--', file], + projectRoot, + ) + if (checkout.code === 0) { + result.restored.push(file) + continue + } + const tree = await runGit( + [...base, 'ls-tree', hash, '--', file], + projectRoot, + ) + if (tree.code !== 0) { + // The snapshot's tree itself is unreachable, so there is no way to + // tell whether this file was in it. Reading "we could not look" as + // "the turn created it" is how a missing snapshot deletes files the + // user had: leave the worktree alone instead. + logger.debug( + { file, hash }, + 'undo-snapshot: snapshot missing, leaving the file as it is', + ) + continue + } + if (tree.stdout.trim()) { + // It was in the snapshot, but bringing its content back did not work + // and the tree alone cannot say why (a collected object, most + // likely). Leave the file as it is and do not report a restore that + // did not happen: the caller's summary has to stay true. + logger.debug( + { file, hash }, + 'undo-snapshot: content unavailable, leaving the file as it is', + ) + continue + } + // Really was not in the snapshot — the turn created it. + try { + rmSync(resolved, { force: true }) + result.deleted.push(file) + } catch { + // Keep going with the remaining files. + } + } + return result + } catch (error) { + logger.debug({ error }, 'undo-snapshot: revert failed') + return result + } + }) +} + +/** + * Compact diff stat of everything that changed since a snapshot (used in the + * /undo result message). Empty string when there is nothing or on failure. + */ +export async function diffSnapshot( + projectRoot: string, + hash: string, +): Promise { + const gitRoot = findGitRoot({ cwd: projectRoot }) + if (!gitRoot) return '' + const snapshotDir = snapshotDirFor(projectRoot) + return withLock(snapshotDir, async () => { + try { + if (!(await stageChanges(snapshotDir, projectRoot, gitRoot))) return '' + const result = await runGit( + [ + ...gitArgs(snapshotDir, projectRoot, [ + 'diff', + '--cached', + '--no-ext-diff', + '--stat', + hash, + '--', + '.', + ]), + ], + projectRoot, + ) + return result.code === 0 ? result.stdout : '' + } catch { + return '' + } + }) +}