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