Skip to content

Commit 10dbe8d

Browse files
committed
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/<chat>/<hash>`) 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).
1 parent 746b20d commit 10dbe8d

4 files changed

Lines changed: 1684 additions & 0 deletions

File tree

Lines changed: 270 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,270 @@
1+
import { afterAll, beforeAll, beforeEach, describe, expect, test } from 'bun:test'
2+
3+
import { execFileSync } from 'node:child_process'
4+
import {
5+
existsSync,
6+
mkdtempSync,
7+
readdirSync,
8+
readFileSync,
9+
rmSync,
10+
writeFileSync,
11+
} from 'node:fs'
12+
import os from 'node:os'
13+
import path from 'node:path'
14+
15+
import { getProjectDataDir, setCurrentChatId, setProjectRoot } from '../../project-files'
16+
import {
17+
anchorSnapshot,
18+
listAnchors,
19+
patchSnapshot,
20+
releaseSnapshot,
21+
setSnapshotDirOverrideForTesting,
22+
trackSnapshot,
23+
} from '../../utils/undo-snapshot'
24+
import {
25+
loadUndoState,
26+
peekRedo,
27+
peekUndo,
28+
popRedo,
29+
popUndo,
30+
pushRedo,
31+
pushUndo,
32+
recordUndoEntry,
33+
sweepAnchors,
34+
undoToRecord,
35+
} from '../undo-store'
36+
37+
const CHAT_ID = 'undo-store-test-chat'
38+
39+
let projectDir: string
40+
let snapshotRoot: string
41+
42+
beforeAll(() => {
43+
projectDir = mkdtempSync(path.join(os.tmpdir(), 'undo-store-proj-'))
44+
snapshotRoot = mkdtempSync(path.join(os.tmpdir(), 'undo-store-snapshots-'))
45+
setProjectRoot(projectDir)
46+
setCurrentChatId(CHAT_ID)
47+
// The journal hands snapshots to undo-snapshot, which needs a git project.
48+
setSnapshotDirOverrideForTesting(snapshotRoot)
49+
execFileSync('git', ['init'], { cwd: projectDir, stdio: 'ignore' })
50+
})
51+
52+
// Each test starts from a clean journal for this chat.
53+
beforeEach(() => {
54+
const chatDir = path.join(getProjectDataDir(), 'chats', CHAT_ID)
55+
rmSync(chatDir, { recursive: true, force: true })
56+
setCurrentChatId(CHAT_ID)
57+
})
58+
59+
afterAll(() => {
60+
const chatDir = path.join(getProjectDataDir(), 'chats', CHAT_ID)
61+
rmSync(chatDir, { recursive: true, force: true })
62+
rmSync(projectDir, { recursive: true, force: true })
63+
setSnapshotDirOverrideForTesting(undefined)
64+
rmSync(snapshotRoot, { recursive: true, force: true })
65+
})
66+
67+
/** Run the cleanup job's gc now, instead of waiting out its 7-day grace. */
68+
const pruneSnapshotRepo = (): void => {
69+
const [entry] = readdirSync(snapshotRoot)
70+
if (!entry) return
71+
execFileSync(
72+
'git',
73+
[
74+
'--git-dir',
75+
path.join(snapshotRoot, entry),
76+
'--work-tree',
77+
projectDir,
78+
'gc',
79+
'--quiet',
80+
'--prune=now',
81+
],
82+
{ stdio: ['ignore', 'pipe', 'ignore'] },
83+
)
84+
}
85+
86+
describe('recordUndoEntry', () => {
87+
test('records an entry and ignores empty ones', () => {
88+
recordUndoEntry(CHAT_ID, {
89+
hashBefore: 'abc123',
90+
files: ['a.txt', 'b.txt'],
91+
message: 'fix the bug',
92+
})
93+
const state = loadUndoState(CHAT_ID)
94+
expect(state.undoStack).toHaveLength(1)
95+
expect(state.undoStack[0]).toMatchObject({
96+
chatId: CHAT_ID,
97+
hashBefore: 'abc123',
98+
files: ['a.txt', 'b.txt'],
99+
message: 'fix the bug',
100+
})
101+
expect(peekUndo(CHAT_ID)?.hashBefore).toBe('abc123')
102+
103+
// No hash or no files → nothing recorded.
104+
recordUndoEntry(CHAT_ID, { hashBefore: '', files: [], message: 'x' })
105+
recordUndoEntry(CHAT_ID, { hashBefore: 'def', files: [], message: 'x' })
106+
expect(loadUndoState(CHAT_ID).undoStack).toHaveLength(1)
107+
})
108+
109+
test('clears the redo stack when a new entry arrives', () => {
110+
pushRedo(CHAT_ID, {
111+
id: 'r1',
112+
chatId: CHAT_ID,
113+
hashBefore: 'old',
114+
hashAfter: 'new',
115+
files: ['a.txt'],
116+
message: 'redo me',
117+
createdAt: new Date().toISOString(),
118+
})
119+
expect(peekRedo(CHAT_ID)).not.toBeNull()
120+
121+
recordUndoEntry(CHAT_ID, {
122+
hashBefore: 'xyz',
123+
files: ['c.txt'],
124+
message: 'new turn',
125+
})
126+
expect(peekRedo(CHAT_ID)).toBeNull()
127+
})
128+
})
129+
130+
describe('undo/redo stack operations', () => {
131+
test('popUndo returns the most recent record and persists', () => {
132+
recordUndoEntry(CHAT_ID, {
133+
hashBefore: 'first',
134+
files: ['one.txt'],
135+
message: 'first turn',
136+
})
137+
recordUndoEntry(CHAT_ID, {
138+
hashBefore: 'second',
139+
files: ['two.txt'],
140+
message: 'second turn',
141+
})
142+
143+
const record = popUndo(CHAT_ID)
144+
expect(record?.hashBefore).toBe('second')
145+
expect(peekUndo(CHAT_ID)?.hashBefore).toBe('first')
146+
// Reload from disk to confirm the pop persisted.
147+
expect(loadUndoState(CHAT_ID).undoStack).toHaveLength(1)
148+
})
149+
150+
test('pushUndo restores a record and popRedo cycles', () => {
151+
recordUndoEntry(CHAT_ID, {
152+
hashBefore: 'cycle-hash',
153+
files: ['x.txt'],
154+
message: 'cycle',
155+
})
156+
const record = popUndo(CHAT_ID)!
157+
pushRedo(CHAT_ID, { ...record, hashAfter: 'after-state' })
158+
159+
const redoRecord = popRedo(CHAT_ID)
160+
expect(redoRecord?.hashBefore).toBe('cycle-hash')
161+
expect(redoRecord?.hashAfter).toBe('after-state')
162+
163+
pushUndo(CHAT_ID, { ...redoRecord!, hashAfter: undefined })
164+
expect(peekUndo(CHAT_ID)?.hashBefore).toBe('cycle-hash')
165+
})
166+
167+
test('returns null from empty stacks', () => {
168+
setCurrentChatId('empty-chat')
169+
expect(popUndo('empty-chat')).toBeNull()
170+
expect(popRedo('empty-chat')).toBeNull()
171+
setCurrentChatId(CHAT_ID)
172+
})
173+
})
174+
175+
describe('a snapshot that is gone', () => {
176+
test('leaves the worktree alone instead of deleting what it cannot restore', async () => {
177+
const kept = path.join(projectDir, 'kept.txt')
178+
writeFileSync(kept, 'the user wrote this\n')
179+
const hash = await trackSnapshot(projectDir)
180+
expect(hash).toBeTruthy()
181+
182+
// The turn edits the file and the journal records what to revert.
183+
writeFileSync(kept, 'the agent changed it\n')
184+
recordUndoEntry(CHAT_ID, {
185+
hashBefore: hash!,
186+
files: ['kept.txt'],
187+
message: 'a recorded turn',
188+
})
189+
190+
// The snapshot is gone — pruned by the cleanup job, or a wiped config dir.
191+
// `checkout` then fails, and "not in the snapshot" must not be read as
192+
// "the turn created this file".
193+
rmSync(snapshotRoot, { recursive: true, force: true })
194+
195+
const message = await undoToRecord(CHAT_ID, projectDir, peekUndo(CHAT_ID)!.id)
196+
197+
expect(message).toBeNull()
198+
expect(existsSync(kept)).toBe(true)
199+
expect(readFileSync(kept, 'utf8')).toBe('the agent changed it\n')
200+
})
201+
})
202+
203+
describe('snapshot anchors', () => {
204+
test('a recorded turn stays restorable after the cleanup job runs', async () => {
205+
const kept = path.join(projectDir, 'anchored.txt')
206+
writeFileSync(kept, 'the user wrote this\n')
207+
const hash = await trackSnapshot(projectDir)
208+
expect(hash).toBeTruthy()
209+
210+
// The turn edits the file, and the snapshot index moves on with it.
211+
writeFileSync(kept, 'the agent changed it\n')
212+
await patchSnapshot(projectDir, hash!)
213+
recordUndoEntry(CHAT_ID, {
214+
hashBefore: hash!,
215+
files: ['anchored.txt'],
216+
message: 'a recorded turn',
217+
})
218+
pruneSnapshotRepo()
219+
220+
const message = await undoToRecord(CHAT_ID, projectDir, peekUndo(CHAT_ID)!.id)
221+
222+
expect(message).not.toBeNull()
223+
expect(readFileSync(kept, 'utf8')).toBe('the user wrote this\n')
224+
})
225+
226+
test('the sweep drops the anchors of a chat whose journal is gone', async () => {
227+
writeFileSync(path.join(projectDir, 'swept.txt'), 'content\n')
228+
const hash = await trackSnapshot(projectDir)
229+
expect(anchorSnapshot(projectDir, 'abandoned-chat', hash!)).toBeTruthy()
230+
expect(listAnchors(projectDir)).toContainEqual({
231+
key: 'abandoned-chat',
232+
hash: hash!,
233+
})
234+
235+
const { released } = sweepAnchors(projectDir)
236+
237+
expect(released).toBeGreaterThan(0)
238+
expect(listAnchors(projectDir)).not.toContainEqual({
239+
key: 'abandoned-chat',
240+
hash: hash!,
241+
})
242+
})
243+
244+
test('the sweep re-anchors an entry whose ref went missing', async () => {
245+
writeFileSync(path.join(projectDir, 'reanchor.txt'), 'content\n')
246+
const hash = await trackSnapshot(projectDir)
247+
recordUndoEntry(CHAT_ID, {
248+
hashBefore: hash!,
249+
files: ['reanchor.txt'],
250+
message: 'a recorded turn',
251+
})
252+
// Pretend the release path ran early, or its write was lost.
253+
expect(releaseSnapshot(projectDir, CHAT_ID, hash!)).toBe(true)
254+
expect(listAnchors(projectDir)).not.toContainEqual({
255+
key: CHAT_ID,
256+
hash: hash!,
257+
})
258+
259+
sweepAnchors(projectDir)
260+
261+
expect(listAnchors(projectDir)).toContainEqual({ key: CHAT_ID, hash: hash! })
262+
})
263+
})
264+
265+
describe('corrupt file handling', () => {
266+
test('loads an empty state for a nonexistent chat', () => {
267+
expect(loadUndoState('never-existed').undoStack).toEqual([])
268+
expect(loadUndoState('never-existed').redoStack).toEqual([])
269+
})
270+
})

0 commit comments

Comments
 (0)