返回 DeepSeek-Reasonix
checkpoint_test.go
根目录 / internal / checkpoint / checkpoint_test.go
1 package checkpoint
2
3 import (
4 "bytes"
5 "encoding/json"
6 "os"
7 "path/filepath"
8 "strings"
9 "testing"
10 "time"
11 "unicode/utf8"
12
13 "reasonix/internal/diff"
14 fileenc "reasonix/internal/fileutil/encoding"
15 )
16
17 func write(t *testing.T, p, s string) {
18 t.Helper()
19 if err := os.MkdirAll(filepath.Dir(p), 0o755); err != nil {
20 t.Fatal(err)
21 }
22 if err := os.WriteFile(p, []byte(s), 0o644); err != nil {
23 t.Fatal(err)
24 }
25 }
26 func read(t *testing.T, p string) string {
27 t.Helper()
28 b, err := os.ReadFile(p)
29 if err != nil {
30 t.Fatal(err)
31 }
32 return string(b)
33 }
34 func readBytes(t *testing.T, p string) []byte {
35 t.Helper()
36 b, err := os.ReadFile(p)
37 if err != nil {
38 t.Fatal(err)
39 }
40 return b
41 }
42
43 // Two turns edit a.txt and create b.txt; rewinding restores each file to its
44 // state at the start of the chosen turn (b.txt being deleted when it post-dates it).
45 func TestRestoreToStartOfTurn(t *testing.T) {
46 root := t.TempDir()
47 a := filepath.Join(root, "a.txt")
48 b := filepath.Join(root, "sub", "b.txt")
49 write(t, a, "v0")
50 s := New("", root)
51
52 s.Begin(0, "first", 0)
53 s.Snapshot(diff.Change{Path: a, Kind: diff.Modify, OldText: "v0"})
54 write(t, a, "v1") // the edit turn 0 made
55
56 s.Begin(1, "second", 2)
57 s.Snapshot(diff.Change{Path: a, Kind: diff.Modify, OldText: "v1"})
58 s.Snapshot(diff.Change{Path: b, Kind: diff.Create})
59 write(t, a, "v2")
60 write(t, b, "new")
61
62 // Rewind to the start of turn 1: a back to v1, b gone.
63 if _, _, err := s.RestoreCode(1); err != nil {
64 t.Fatal(err)
65 }
66 if got := read(t, a); got != "v1" {
67 t.Fatalf("a = %q, want v1", got)
68 }
69 if _, err := os.Stat(b); !os.IsNotExist(err) {
70 t.Fatalf("b should have been deleted, stat err=%v", err)
71 }
72 }
73
74 func TestRestoreToTurnZero(t *testing.T) {
75 root := t.TempDir()
76 a := filepath.Join(root, "a.txt")
77 write(t, a, "v0")
78 s := New("", root)
79 s.Begin(0, "first", 0)
80 s.Snapshot(diff.Change{Path: a, Kind: diff.Modify, OldText: "v0"})
81 write(t, a, "v1")
82 s.Begin(1, "second", 2)
83 s.Snapshot(diff.Change{Path: a, Kind: diff.Modify, OldText: "v1"})
84 write(t, a, "v2")
85
86 if _, _, err := s.RestoreCode(0); err != nil {
87 t.Fatal(err)
88 }
89 if got := read(t, a); got != "v0" {
90 t.Fatalf("a = %q, want v0 (earliest snapshot)", got)
91 }
92 }
93
94 func TestRestorePreservesGB18030Encoding(t *testing.T) {
95 root := t.TempDir()
96 a := filepath.Join(root, "gbk.txt")
97 original := "\u4f60\u597d\n\u65e7\u884c\n"
98 edited := "\u4f60\u597d\n\u65b0\u884c\n"
99 originalRaw := fileenc.Encode(original, fileenc.GB18030)
100 if err := os.WriteFile(a, originalRaw, 0o644); err != nil {
101 t.Fatal(err)
102 }
103
104 s := New("", root)
105 s.Begin(0, "edit gbk", 0)
106 s.Snapshot(diff.Change{Path: a, Kind: diff.Modify, OldText: original})
107 if err := os.WriteFile(a, fileenc.Encode(edited, fileenc.GB18030), 0o644); err != nil {
108 t.Fatal(err)
109 }
110
111 if _, _, err := s.RestoreCode(0); err != nil {
112 t.Fatal(err)
113 }
114 gotRaw := readBytes(t, a)
115 if utf8.Valid(gotRaw) {
116 t.Fatalf("restored GB18030 file became valid UTF-8 bytes: % x", gotRaw)
117 }
118 if !bytes.Equal(gotRaw, originalRaw) {
119 t.Fatalf("restored bytes = % x, want original GB18030 bytes % x", gotRaw, originalRaw)
120 }
121 }
122
123 func TestRestorePreservesGB18030EncodingAfterPersistence(t *testing.T) {
124 root := t.TempDir()
125 dir := filepath.Join(t.TempDir(), "sess.ckpt")
126 a := filepath.Join(root, "gbk.txt")
127 original := "\u4f60\u597d\n\u65e7\u884c\n"
128 edited := "\u4f60\u597d\n\u65b0\u884c\n"
129 originalRaw := fileenc.Encode(original, fileenc.GB18030)
130 if err := os.WriteFile(a, originalRaw, 0o644); err != nil {
131 t.Fatal(err)
132 }
133
134 s := New(dir, root)
135 s.Begin(0, "edit gbk", 0)
136 s.Snapshot(diff.Change{Path: a, Kind: diff.Modify, OldText: original})
137
138 resumed := New(dir, root)
139 if err := os.WriteFile(a, fileenc.Encode(edited, fileenc.GB18030), 0o644); err != nil {
140 t.Fatal(err)
141 }
142 if _, _, err := resumed.RestoreCode(0); err != nil {
143 t.Fatal(err)
144 }
145 if gotRaw := readBytes(t, a); !bytes.Equal(gotRaw, originalRaw) {
146 t.Fatalf("restored bytes after persistence = % x, want original GB18030 bytes % x", gotRaw, originalRaw)
147 }
148 }
149
150 func TestRestoreLegacySnapshotRequiresExplicitSafePath(t *testing.T) {
151 root := t.TempDir()
152 dir := filepath.Join(t.TempDir(), "sess.ckpt")
153 if err := os.MkdirAll(dir, 0o755); err != nil {
154 t.Fatal(err)
155 }
156 a := filepath.Join(root, "gbk.txt")
157 original := "\u4f60\u597d\n\u65e7\u884c\n"
158 edited := "\u4f60\u597d\n\u65b0\u884c\n"
159 if err := os.WriteFile(a, fileenc.Encode(edited, fileenc.GB18030), 0o644); err != nil {
160 t.Fatal(err)
161 }
162
163 legacy := Checkpoint{
164 Turn: 0,
165 Time: time.Now(),
166 Prompt: "legacy",
167 MsgIndex: 0,
168 Files: []FileSnap{{
169 Path: a,
170 Content: &original,
171 }},
172 }
173 b, err := json.Marshal(legacy)
174 if err != nil {
175 t.Fatal(err)
176 }
177 if err := os.WriteFile(filepath.Join(dir, "turn-0.json"), b, 0o644); err != nil {
178 t.Fatal(err)
179 }
180
181 resumed := New(dir, root)
182 if _, _, err := resumed.RestoreCode(0); err == nil {
183 t.Fatal("legacy restore must not silently overwrite an unverifiable file")
184 }
185 if got := string(fileenc.Decode(readBytes(t, a), fileenc.GB18030)); got != edited {
186 t.Fatalf("legacy refusal changed file to %q, want edited content preserved", got)
187 }
188 }
189
190 func TestSnapshotDedupsFirstTouchWins(t *testing.T) {
191 root := t.TempDir()
192 a := filepath.Join(root, "a.txt")
193 write(t, a, "orig")
194 s := New("", root)
195 s.Begin(0, "p", 0)
196 s.Snapshot(diff.Change{Path: a, Kind: diff.Modify, OldText: "orig"})
197 s.Snapshot(diff.Change{Path: a, Kind: diff.Modify, OldText: "edited-once"}) // ignored
198 write(t, a, "edited-twice")
199 if _, _, err := s.RestoreCode(0); err != nil {
200 t.Fatal(err)
201 }
202 if got := read(t, a); got != "orig" {
203 t.Fatalf("a = %q, want orig (first snapshot wins)", got)
204 }
205 }
206
207 func TestPersistV2RemainsReadableByLegacyBinary(t *testing.T) {
208 root := t.TempDir()
209 dir := filepath.Join(t.TempDir(), "sess.ckpt")
210 existing := filepath.Join(root, "existing.txt")
211 created := filepath.Join(root, "created.txt")
212 write(t, existing, "before")
213
214 s := New(dir, root)
215 s.Begin(0, "compat", 0)
216 s.CaptureBefore(existing, CaptureBeforeOpts{Source: CaptureBeforeMutation})
217 s.CaptureBefore(created, CaptureBeforeOpts{Source: CaptureBeforeMutation})
218
219 type legacyFile struct {
220 Path string `json:"path"`
221 Content *string `json:"content"`
222 Encoding json.RawMessage `json:"encoding,omitempty"`
223 }
224 type legacyCheckpoint struct {
225 Files []legacyFile `json:"files"`
226 }
227 var legacy legacyCheckpoint
228 b, err := os.ReadFile(filepath.Join(dir, "turn-0.json"))
229 if err != nil {
230 t.Fatal(err)
231 }
232 if err := json.Unmarshal(b, &legacy); err != nil {
233 t.Fatal(err)
234 }
235 byPath := map[string]*string{}
236 for _, file := range legacy.Files {
237 byPath[file.Path] = file.Content
238 }
239 if byPath[existing] == nil || *byPath[existing] != "before" {
240 t.Fatalf("legacy reader lost existing-file preimage: %#v", byPath[existing])
241 }
242 if content, ok := byPath[created]; !ok || content != nil {
243 t.Fatalf("legacy reader must keep created-file sentinel nil: present=%v content=%#v", ok, content)
244 }
245 }
246
247 func TestGCDoesNotDeleteSharedBlobStillReferencedByNewerCheckpoint(t *testing.T) {
248 root := t.TempDir()
249 dir := filepath.Join(t.TempDir(), "sess.ckpt")
250 a := filepath.Join(root, "a.txt")
251 b := filepath.Join(root, "b.txt")
252 write(t, a, "shared")
253 write(t, b, "shared")
254 s := New(dir, root)
255
256 s.Begin(0, "a", 0)
257 s.CaptureBefore(a, CaptureBeforeOpts{Source: CaptureBeforeMutation})
258 write(t, a, "a-edited")
259 s.CaptureAfter(a, CaptureAfterOpts{Seq: 1, Source: CaptureAfterMutation})
260 s.Begin(1, "b", 1)
261 s.CaptureBefore(b, CaptureBeforeOpts{Source: CaptureBeforeMutation})
262 write(t, b, "b-edited")
263 s.CaptureAfter(b, CaptureAfterOpts{Seq: 2, Source: CaptureAfterMutation})
264 s.Begin(2, "finalize", 2)
265
266 s.mu.Lock()
267 ref := s.done[1].Files[0].BlobRef
268 s.retainN = 1
269 s.gcLocked()
270 s.mu.Unlock()
271 if ref == "" || !s.blobs.Has(ref) {
272 t.Fatalf("shared blob %q was removed while the newer checkpoint still referenced it", ref)
273 }
274 plan, err := s.PrepareRewind(1, RewindCode, 1, 0, false)
275 if err != nil || !plan.CanFiles {
276 t.Fatalf("newer checkpoint became unrecoverable: plan=%+v err=%v", plan, err)
277 }
278 }
279
280 func TestExpiredV2PayloadRemainsSafeForLegacyReader(t *testing.T) {
281 root := t.TempDir()
282 dir := filepath.Join(t.TempDir(), "sess.ckpt")
283 content := "must not be interpreted as absent"
284 checkpoint := &Checkpoint{
285 SchemaVersion: SchemaV2,
286 Turn: 0,
287 Files: []FileSnap{{
288 Path: "a.txt", Content: &content, SHA256: Digest([]byte(content)), BlobRef: Digest([]byte(content)),
289 }},
290 }
291 store := New(dir, root)
292 if err := store.persist(checkpoint); err != nil {
293 t.Fatal(err)
294 }
295 store.mu.Lock()
296 err := store.expirePayloadLocked(checkpoint)
297 store.mu.Unlock()
298 if err != nil {
299 t.Fatal(err)
300 }
301
302 // A previous release only scans turn-*.json in the checkpoint root. If the
303 // expired checkpoint remains visible there, its content must never be nil:
304 // old RestoreCode interprets nil as "delete this file".
305 raw, err := os.ReadFile(filepath.Join(dir, "turn-0.json"))
306 if err == nil {
307 var legacy struct {
308 Files []struct {
309 Content *string `json:"content"`
310 } `json:"files"`
311 }
312 if err := json.Unmarshal(raw, &legacy); err != nil {
313 t.Fatal(err)
314 }
315 if len(legacy.Files) != 1 || legacy.Files[0].Content == nil {
316 t.Fatal("expired v2 payload tells a legacy reader to delete an existing file")
317 }
318 } else if !os.IsNotExist(err) {
319 t.Fatal(err)
320 }
321
322 reloaded := New(dir, root)
323 metas := reloaded.List()
324 if len(metas) != 1 || !metas[0].ExpiredFilePayload || metas[0].CanUndoFiles {
325 t.Fatalf("expired metadata was not preserved for the new reader: %+v", metas)
326 }
327 }
328
329 func TestBlobReadVerifiesContentAddress(t *testing.T) {
330 store := NewBlobStore(t.TempDir())
331 ref, err := store.Put([]byte("before"))
332 if err != nil {
333 t.Fatal(err)
334 }
335 if err := os.WriteFile(store.path(ref), []byte("corrupt"), 0o644); err != nil {
336 t.Fatal(err)
337 }
338 if got, err := store.Get(ref); err == nil {
339 t.Fatalf("content-addressed read accepted bytes %q that do not match %s", got, ref)
340 }
341 if store.Has(ref) {
342 t.Fatal("Has accepted a blob whose bytes do not match its content address")
343 }
344 if gotRef, err := store.Put([]byte("before")); err != nil || gotRef != ref {
345 t.Fatalf("Put did not repair corrupt blob: ref=%q err=%v", gotRef, err)
346 }
347 if got, err := store.Get(ref); err != nil || string(got) != "before" {
348 t.Fatalf("repaired blob = %q err=%v", got, err)
349 }
350 }
351
352 func TestRestoreRejectsPathEscape(t *testing.T) {
353 root := t.TempDir()
354 outside := filepath.Join(t.TempDir(), "evil.txt")
355 write(t, outside, "keep")
356 s := New("", root)
357 s.Begin(0, "p", 0)
358 s.Snapshot(diff.Change{Path: outside, Kind: diff.Modify, OldText: "hacked"})
359 if _, _, err := s.RestoreCode(0); err == nil {
360 t.Fatal("RestoreCode should reject a path outside the workspace")
361 }
362 if got := read(t, outside); got != "keep" {
363 t.Fatalf("outside file was modified: %q", got)
364 }
365 }
366
367 func TestPersistenceRoundTrip(t *testing.T) {
368 root := t.TempDir()
369 dir := filepath.Join(t.TempDir(), "sess.ckpt")
370 a := filepath.Join(root, "a.txt")
371
372 s := New(dir, root)
373 s.Begin(0, "hello", 1)
374 s.Snapshot(diff.Change{Path: a, Kind: diff.Modify, OldText: "v0"})
375 s.Begin(1, "world", 5)
376
377 // A fresh store over the same dir must see both turns and their boundaries.
378 s2 := New(dir, root)
379 metas := s2.List()
380 if len(metas) != 2 {
381 t.Fatalf("loaded %d checkpoints, want 2", len(metas))
382 }
383 if metas[0].Prompt != "hello" || metas[1].Prompt != "world" {
384 t.Fatalf("prompts = %q, %q", metas[0].Prompt, metas[1].Prompt)
385 }
386 // Boundaries must survive the round-trip so a resumed session can rewind/fork.
387 b := s2.Bounds()
388 if b[0] != 1 || b[1] != 5 {
389 t.Fatalf("bounds = %v, want {0:1, 1:5}", b)
390 }
391 if s2.NextTurn() != 2 {
392 t.Fatalf("NextTurn = %d, want 2", s2.NextTurn())
393 }
394 }
395
396 func TestListExposesCurrentTurnFiles(t *testing.T) {
397 root := t.TempDir()
398 a := filepath.Join(root, "a.txt")
399 write(t, a, "v0")
400 s := New("", root)
401 s.Begin(0, "edit current", 0)
402 s.Snapshot(diff.Change{Path: a, Kind: diff.Modify, OldText: "v0"})
403
404 metas := s.List()
405 if len(metas) != 1 {
406 t.Fatalf("metas = %d, want 1", len(metas))
407 }
408 if len(metas[0].Paths) != 1 || metas[0].Paths[0] != a {
409 t.Fatalf("current turn paths = %#v, want [%q]", metas[0].Paths, a)
410 }
411 }
412
413 func TestFileStateReturnsEarliestSnapshotAcrossPathForms(t *testing.T) {
414 root := t.TempDir()
415 path := filepath.Join(root, "nested", "file.txt")
416 s := New("", root)
417 s.Begin(0, "first", 0)
418 s.Snapshot(diff.Change{Path: path, Kind: diff.Modify, OldText: "original"})
419 s.Begin(1, "second", 2)
420 s.Snapshot(diff.Change{Path: filepath.Join("nested", "file.txt"), Kind: diff.Modify, OldText: "after first edit"})
421
422 state, ok := s.FileState(filepath.Join("nested", "file.txt"))
423 if !ok || state.Content == nil {
424 t.Fatalf("FileState = %+v, %v; want earliest content", state, ok)
425 }
426 if got := *state.Content; got != "original" {
427 t.Fatalf("FileState content = %q, want original", got)
428 }
429 if _, ok := s.FileState(filepath.Join("..", "outside.txt")); ok {
430 t.Fatal("FileState accepted a path outside the workspace")
431 }
432 }
433
434 func TestTruncateFromDropsFutureCheckpointsAndFiles(t *testing.T) {
435 root := t.TempDir()
436 dir := filepath.Join(t.TempDir(), "sess.ckpt")
437 a := filepath.Join(root, "a.txt")
438 write(t, a, "v0")
439 s := New(dir, root)
440 s.Begin(0, "first", 0)
441 s.Snapshot(diff.Change{Path: a, Kind: diff.Modify, OldText: "v0"})
442 s.Begin(1, "second", 2)
443 s.Snapshot(diff.Change{Path: a, Kind: diff.Modify, OldText: "v1"})
444 s.Begin(2, "third", 4)
445
446 if err := s.TruncateFrom(1); err != nil {
447 t.Fatal(err)
448 }
449
450 metas := s.List()
451 if len(metas) != 1 || metas[0].Turn != 0 {
452 t.Fatalf("metas after truncate = %+v, want only turn 0", metas)
453 }
454 if s.NextTurn() != 1 {
455 t.Fatalf("NextTurn after truncate = %d, want 1", s.NextTurn())
456 }
457 if _, err := os.Stat(filepath.Join(dir, "turn-1.json")); !os.IsNotExist(err) {
458 t.Fatalf("turn-1 checkpoint should be deleted, stat err=%v", err)
459 }
460 if _, err := os.Stat(filepath.Join(dir, "turn-2.json")); !os.IsNotExist(err) {
461 t.Fatalf("turn-2 checkpoint should be deleted, stat err=%v", err)
462 }
463 reloaded := New(dir, root)
464 if got := reloaded.List(); len(got) != 1 || got[0].Turn != 0 {
465 t.Fatalf("reloaded metas after truncate = %+v, want only turn 0", got)
466 }
467 }
468
469 func TestTruncateFromReportsPersistentDeleteFailure(t *testing.T) {
470 root := t.TempDir()
471 dir := filepath.Join(t.TempDir(), "sess.ckpt")
472 store := New(dir, root)
473 store.Begin(0, "first", 0)
474 store.Begin(1, "second", 2)
475 blocked := filepath.Join(dir, "turn-1.json")
476 if err := os.Remove(blocked); err != nil {
477 t.Fatal(err)
478 }
479 if err := os.Mkdir(blocked, 0o755); err != nil {
480 t.Fatal(err)
481 }
482 if err := os.WriteFile(filepath.Join(blocked, "keep"), []byte("x"), 0o644); err != nil {
483 t.Fatal(err)
484 }
485
486 if err := store.TruncateFrom(1); err == nil {
487 t.Fatal("truncate reported success despite a persistent checkpoint delete failure")
488 }
489 metas := store.List()
490 if len(metas) != 2 || metas[1].Turn != 1 {
491 t.Fatalf("failed truncate mutated in-memory checkpoints: %+v", metas)
492 }
493 }
494
495 func BenchmarkRestoreGB18030Encoding(b *testing.B) {
496 root := b.TempDir()
497 a := filepath.Join(root, "gbk.txt")
498 original := strings.Repeat("\u4f60\u597d\u4e16\u754c\n\u65e7\u884c\n", 8192)
499 edited := strings.Repeat("\u4f60\u597d\u4e16\u754c\n\u65b0\u884c\n", 8192)
500 originalRaw := fileenc.Encode(original, fileenc.GB18030)
501 editedRaw := fileenc.Encode(edited, fileenc.GB18030)
502 if err := os.WriteFile(a, originalRaw, 0o644); err != nil {
503 b.Fatal(err)
504 }
505
506 s := New("", root)
507 s.Begin(0, "edit gbk", 0)
508 s.Snapshot(diff.Change{Path: a, Kind: diff.Modify, OldText: original})
509
510 b.SetBytes(int64(len(originalRaw)))
511 b.ReportAllocs()
512 b.ResetTimer()
513 for i := 0; i < b.N; i++ {
514 if err := os.WriteFile(a, editedRaw, 0o644); err != nil {
515 b.Fatal(err)
516 }
517 if _, _, err := s.RestoreCode(0); err != nil {
518 b.Fatal(err)
519 }
520 }
521 }
522
523 func TestLazyDirectoryCreation(t *testing.T) {
524 root := t.TempDir()
525 dir := filepath.Join(t.TempDir(), "lazy-sess.ckpt")
526
527 s := New(dir, root)
528
529 if _, err := os.Stat(dir); !os.IsNotExist(err) {
530 t.Fatalf("directory should not exist yet: %v", err)
531 }
532
533 s.Begin(0, "lazy", 0)
534
535 if _, err := os.Stat(dir); err != nil {
536 t.Fatalf("directory should now exist: %v", err)
537 }
538 turnPath := filepath.Join(dir, "turn-0.json")
539 if _, err := os.Stat(turnPath); err != nil {
540 t.Fatalf("turn file should now exist: %v", err)
541 }
542 }
543
543 lines GO