返回 DeepSeek-Reasonix
save_dag_test.go
根目录 / internal / agent / save_dag_test.go
1 package agent
2
3 import (
4 "errors"
5 "os"
6 "path/filepath"
7 "strings"
8 "sync"
9 "testing"
10
11 "reasonix/internal/fileutil"
12 "reasonix/internal/provider"
13 "reasonix/internal/store"
14 )
15
16 func dagSavedSession(t *testing.T, path string, contents ...string) *Session {
17 t.Helper()
18 s := NewSession("sys")
19 for i, c := range contents {
20 role := provider.RoleUser
21 if i%2 == 1 {
22 role = provider.RoleAssistant
23 }
24 s.Add(provider.Message{Role: role, Content: c})
25 }
26 if err := s.Save(path); err != nil {
27 t.Fatalf("save: %v", err)
28 }
29 return s
30 }
31
32 func dagEntryTypes(t *testing.T, path string) []string {
33 t.Helper()
34 b, err := os.ReadFile(store.SessionEventLog(path))
35 if err != nil {
36 t.Fatal(err)
37 }
38 var types []string
39 for line := range strings.SplitSeq(strings.TrimSpace(string(b)), "\n") {
40 _, rest, _ := strings.Cut(line, `"type":"`)
41 typ, _, _ := strings.Cut(rest, `"`)
42 types = append(types, typ)
43 }
44 return types
45 }
46
47 func assertNoTranscriptCopies(t *testing.T, path string) {
48 t.Helper()
49 entries, _ := os.ReadDir(filepath.Dir(path))
50 for _, entry := range entries {
51 if store.IsSessionTranscriptName(entry.Name()) && entry.Name() != filepath.Base(path) {
52 t.Fatalf("unexpected transcript copy %s", entry.Name())
53 }
54 }
55 }
56
57 func TestDAGSaveCreatesSchemaTwoLogAndAppendsDelta(t *testing.T) {
58 path := dagTestSession(t)
59 s := dagSavedSession(t, path, "q1", "a1")
60 probe, err := probeSessionEventLog(path)
61 if err != nil || !probe.dag {
62 t.Fatalf("probe = %+v err=%v", probe, err)
63 }
64 if got := dagEntryTypes(t, path); strings.Join(got, ",") != "log,writer,message,message,message" {
65 t.Fatalf("entries = %v", got)
66 }
67 ref, ok := s.Head()
68 if !ok || ref.HeadID != SessionMainHead || ref.LeafID != s.LeafID() || ref.LogGeneration != 1 {
69 t.Fatalf("head = %+v ok=%v", ref, ok)
70 }
71 if b, err := os.ReadFile(path); err != nil || strings.Count(string(b), "\n") != 3 {
72 t.Fatalf("checkpoint cache: %v %q", err, b)
73 }
74 idx, err := ReadSessionHeadIndex(path)
75 if err != nil || idx == nil || !idx.Current(path) || idx.MessageCount != 3 || idx.SelectedHead != SessionMainHead {
76 t.Fatalf("index = %+v err=%v", idx, err)
77 }
78 meta, _, err := LoadBranchMeta(path)
79 if err != nil || meta.HeadID != SessionMainHead || meta.LogSchema != 2 || meta.HeadCount != 1 || meta.Revision == 0 {
80 t.Fatalf("meta = %+v err=%v", meta, err)
81 }
82 s.Add(provider.Message{Role: provider.RoleUser, Content: "q2"})
83 if err := s.Save(path); err != nil {
84 t.Fatal(err)
85 }
86 if got := dagEntryTypes(t, path); strings.Join(got, ",") != "log,writer,message,message,message,message" {
87 t.Fatalf("entries after append = %v", got)
88 }
89 if err := s.Save(path); err != nil {
90 t.Fatal(err)
91 }
92 if got := len(dagEntryTypes(t, path)); got != 6 {
93 t.Fatalf("no-op save appended: %d entries", got)
94 }
95 loaded, err := LoadSession(path)
96 if err != nil || len(loaded.Messages) != 4 || loaded.LeafID() != s.LeafID() {
97 t.Fatalf("reload: err=%v len=%d", err, len(loaded.Messages))
98 }
99 assertNoTranscriptCopies(t, path)
100 }
101
102 func TestDAGSaveDisabledByEnvKeepsSchemaOne(t *testing.T) {
103 useSchemaOneLog(t)
104 path := dagTestSession(t)
105 dagSavedSession(t, path, "q1")
106 probe, err := probeSessionEventLog(path)
107 if err != nil || probe.dag || !probe.native {
108 t.Fatalf("probe = %+v err=%v", probe, err)
109 }
110 }
111
112 func TestDAGSaveLocalMetadataBecomesPatch(t *testing.T) {
113 path := dagTestSession(t)
114 s := dagSavedSession(t, path, "q1", "a1")
115 msgs := s.Snapshot()
116 msgs[1].Edited = true
117 msgs[1].WorkDurationMs = 42
118 s.ReplaceLocalMetadata(msgs)
119 if err := s.SaveRewrite(path); err != nil {
120 t.Fatal(err)
121 }
122 types := dagEntryTypes(t, path)
123 if types[len(types)-1] != sessionDAGTypePatch {
124 t.Fatalf("entries = %v", types)
125 }
126 loaded, err := LoadSession(path)
127 if err != nil || !loaded.Messages[1].Edited || loaded.Messages[1].WorkDurationMs != 42 || loaded.Messages[1].ID != msgs[1].ID {
128 t.Fatalf("reload = %+v err=%v", loaded.Messages[1], err)
129 }
130 if reasons := s.DrainContentRewriteReasons(); len(reasons) != 0 {
131 t.Fatalf("local metadata save queued cache reasons %v", reasons)
132 }
133 }
134
135 func TestDAGSaveSystemPromptRefreshKeepsLaterIDs(t *testing.T) {
136 path := dagTestSession(t)
137 s := dagSavedSession(t, path, "q1", "a1")
138 before := s.Snapshot()
139 s.SetLeadingSystemPrompt("sys-v2")
140 if err := s.SaveRewrite(path); err != nil {
141 t.Fatal(err)
142 }
143 types := dagEntryTypes(t, path)
144 if types[len(types)-1] != sessionDAGTypeSystem {
145 t.Fatalf("entries = %v", types)
146 }
147 loaded, err := LoadSession(path)
148 if err != nil || loaded.Messages[0].Content != "sys-v2" {
149 t.Fatalf("reload = %+v err=%v", loaded.Messages, err)
150 }
151 for i := range before {
152 if loaded.Messages[i].ID != before[i].ID {
153 t.Fatalf("message %d id changed across system refresh", i)
154 }
155 }
156 }
157
158 func TestDAGSaveTruncationRewindsWithoutErasingBytes(t *testing.T) {
159 path := dagTestSession(t)
160 s := dagSavedSession(t, path, "q1", "a1", "q2", "a2")
161 logBefore, _ := os.ReadFile(store.SessionEventLog(path))
162 msgs := s.Snapshot()
163 s.Rewrite(msgs[:3], "rewind_truncate")
164 if err := s.SaveRewrite(path); err != nil {
165 t.Fatal(err)
166 }
167 logAfter, _ := os.ReadFile(store.SessionEventLog(path))
168 if !strings.HasPrefix(string(logAfter), string(logBefore)) {
169 t.Fatal("rewind must not rewrite earlier bytes")
170 }
171 types := dagEntryTypes(t, path)
172 if types[len(types)-1] != sessionDAGTypeRewind {
173 t.Fatalf("entries = %v", types)
174 }
175 loaded, err := LoadSession(path)
176 if err != nil || len(loaded.Messages) != 3 || loaded.LeafID() != msgs[2].ID {
177 t.Fatalf("reload len=%d leaf=%q err=%v", len(loaded.Messages), loaded.LeafID(), err)
178 }
179 s.Add(provider.Message{Role: provider.RoleAssistant, Content: "a2-new"})
180 if err := s.Save(path); err != nil {
181 t.Fatal(err)
182 }
183 loaded, err = LoadSession(path)
184 if err != nil || strings.Join(dagContents(loaded.Messages), ",") != "sys,q1,a1,a2-new" {
185 t.Fatalf("after re-append: %v err=%v", dagContents(loaded.Messages), err)
186 }
187 }
188
189 func TestDAGSaveUpgradesSchemaOneLogOnlyUnderLease(t *testing.T) {
190 t.Setenv(SessionLogSchemaEnv, "v1")
191 path := dagTestSession(t)
192 v1 := dagSavedSession(t, path, "q1", "a1")
193 if err := os.Unsetenv(SessionLogSchemaEnv); err != nil {
194 t.Fatal(err)
195 }
196 loaded, err := LoadSession(path)
197 if err != nil {
198 t.Fatal(err)
199 }
200 loaded.Add(provider.Message{Role: provider.RoleUser, Content: "q2"})
201 if err := loaded.Save(path); err != nil {
202 t.Fatal(err)
203 }
204 if probe, _ := probeSessionEventLog(path); probe.dag {
205 t.Fatal("unleased writer must not upgrade an existing schema-1 log")
206 }
207 lease, err := TryAcquireSessionLease(path)
208 if err != nil {
209 t.Fatal(err)
210 }
211 defer lease.Release()
212 loaded.Add(provider.Message{Role: provider.RoleAssistant, Content: "a2"})
213 if err := loaded.Save(path); err != nil {
214 t.Fatal(err)
215 }
216 probe, _ := probeSessionEventLog(path)
217 if !probe.dag {
218 t.Fatal("lease holder must upgrade the schema-1 log on save")
219 }
220 again, err := LoadSession(path)
221 if err != nil || strings.Join(dagContents(again.Messages), ",") != "sys,q1,a1,q2,a2" {
222 t.Fatalf("after upgrade: %v err=%v", dagContents(again.Messages), err)
223 }
224 for i := range v1.Messages {
225 if again.Messages[i].ID != loaded.Messages[i].ID {
226 t.Fatalf("message %d id changed across upgrade", i)
227 }
228 }
229 if ref, ok := again.Head(); !ok || ref.HeadID != SessionMainHead {
230 t.Fatalf("head after upgrade = %+v ok=%v", ref, ok)
231 }
232 if st := dagReplay(t, path); st.upgradedFrom != sessionEventSchemaVersion {
233 t.Fatalf("upgradedFrom = %d", st.upgradedFrom)
234 }
235 }
236
237 func TestDAGSaveConcurrentWritersForkInsteadOfConflicting(t *testing.T) {
238 path := dagTestSession(t)
239 a := dagSavedSession(t, path, "q1", "a1")
240 b, err := LoadSession(path)
241 if err != nil {
242 t.Fatal(err)
243 }
244 a.Add(provider.Message{Role: provider.RoleUser, Content: "q2-from-a"})
245 if err := a.Save(path); err != nil {
246 t.Fatal(err)
247 }
248 b.Add(provider.Message{Role: provider.RoleUser, Content: "q2-from-b"})
249 if err := b.Save(path); err != nil {
250 t.Fatalf("second writer must not conflict: %v", err)
251 }
252 refA, _ := a.Head()
253 refB, _ := b.Head()
254 if refA.HeadID != SessionMainHead || refB.HeadID == SessionMainHead || refB.HeadID == "" {
255 t.Fatalf("heads a=%+v b=%+v", refA, refB)
256 }
257 events := b.DrainHeadEvents()
258 if len(events) != 1 || events[0].Kind != HeadEventForkedConcurrent || events[0].HeadID != refB.HeadID {
259 t.Fatalf("events = %+v", events)
260 }
261 heads, err := ListSessionHeads(path)
262 if err != nil || len(heads) != 2 || heads[1].Kind != HeadKindConcurrent || heads[1].MessageCount != 4 || heads[0].MessageCount != 4 {
263 t.Fatalf("heads = %+v err=%v", heads, err)
264 }
265 st := dagReplay(t, path)
266 if got := dagChain(st, SessionMainHead); strings.Join(got, ",") != "sys,q1,a1,q2-from-a" {
267 t.Fatalf("main chain %v", got)
268 }
269 if got := dagChain(st, refB.HeadID); strings.Join(got, ",") != "sys,q1,a1,q2-from-b" {
270 t.Fatalf("fork chain %v", got)
271 }
272 assertNoTranscriptCopies(t, path)
273 // Each writer keeps extending its own head afterwards.
274 a.Add(provider.Message{Role: provider.RoleAssistant, Content: "a2-from-a"})
275 b.Add(provider.Message{Role: provider.RoleAssistant, Content: "a2-from-b"})
276 if err := a.Save(path); err != nil {
277 t.Fatal(err)
278 }
279 if err := b.Save(path); err != nil {
280 t.Fatal(err)
281 }
282 if len(b.DrainHeadEvents()) != 0 {
283 t.Fatal("continuing on the fork must not fork again")
284 }
285 st = dagReplay(t, path)
286 if len(st.heads) != 2 || len(dagChain(st, SessionMainHead)) != 5 || len(dagChain(st, refB.HeadID)) != 5 {
287 t.Fatalf("heads=%d main=%d fork=%d", len(st.heads), len(dagChain(st, SessionMainHead)), len(dagChain(st, refB.HeadID)))
288 }
289 }
290
291 func TestDAGSaveBehindDiskReportsStalePrefix(t *testing.T) {
292 path := dagTestSession(t)
293 a := dagSavedSession(t, path, "q1", "a1")
294 b, err := LoadSession(path)
295 if err != nil {
296 t.Fatal(err)
297 }
298 a.Add(provider.Message{Role: provider.RoleUser, Content: "q2"})
299 if err := a.Save(path); err != nil {
300 t.Fatal(err)
301 }
302 err = b.Save(path)
303 if !errors.Is(err, ErrSessionSnapshotConflict) {
304 t.Fatalf("behind writer err = %v, want stale prefix conflict", err)
305 }
306 if kind, ok := SnapshotConflictKind(err); !ok || kind != SessionSnapshotConflictStalePrefix {
307 t.Fatalf("kind = %q ok=%v", kind, ok)
308 }
309 if st := dagReplay(t, path); len(st.heads) != 1 || len(dagChain(st, SessionMainHead)) != 4 {
310 t.Fatal("a behind writer must not append or fork")
311 }
312 assertNoTranscriptCopies(t, path)
313 }
314
315 func TestDAGSaveRedactionCompactErasesBytesUnderLease(t *testing.T) {
316 path := dagTestSession(t)
317 lease, err := TryAcquireSessionLease(path)
318 if err != nil {
319 t.Fatal(err)
320 }
321 defer lease.Release()
322 s := dagSavedSession(t, path, "q1 secret-token", "a1")
323 msgs := s.Snapshot()
324 ids := []string{msgs[0].ID, msgs[1].ID, msgs[2].ID}
325 msgs[1].Content = "q1 [redacted]"
326 s.Rewrite(msgs, "redact")
327 if err := s.SaveRewriteCompact(path); err != nil {
328 t.Fatal(err)
329 }
330 raw, _ := os.ReadFile(store.SessionEventLog(path))
331 if strings.Contains(string(raw), "secret-token") {
332 t.Fatal("redaction left the secret in the log")
333 }
334 st := dagReplay(t, path)
335 if st.generation != 2 {
336 t.Fatalf("generation = %d, want rotation", st.generation)
337 }
338 loaded, err := LoadSession(path)
339 if err != nil || loaded.Messages[1].Content != "q1 [redacted]" {
340 t.Fatalf("reload = %+v err=%v", loaded.Messages, err)
341 }
342 for i, id := range ids {
343 if loaded.Messages[i].ID != id {
344 t.Fatalf("message %d id changed by redaction", i)
345 }
346 }
347 if ref, _ := s.Head(); ref.LogGeneration != 2 {
348 t.Fatalf("session did not follow the rotation: %+v", ref)
349 }
350 }
351
352 func TestDAGSaveOversizeLogRotatesUnderLease(t *testing.T) {
353 path := dagTestSession(t)
354 lease, err := TryAcquireSessionLease(path)
355 if err != nil {
356 t.Fatal(err)
357 }
358 defer lease.Release()
359 big := strings.Repeat("x", 100<<10)
360 s := dagSavedSession(t, path, big+"1", big+"2", big+"3", big+"4", big+"5", big+"6")
361 msgs := s.Snapshot()
362 s.Rewrite(msgs[:2], "rewind_truncate")
363 if err := s.SaveRewrite(path); err != nil {
364 t.Fatal(err)
365 }
366 st := dagReplay(t, path)
367 if st.generation != 2 || len(st.nodes) != 2 {
368 t.Fatalf("generation=%d nodes=%d, want rotated log with only the live chain", st.generation, len(st.nodes))
369 }
370 if info, _ := os.Stat(store.SessionEventLog(path)); info.Size() > int64(len(big))*3 {
371 t.Fatalf("rotated log still %d bytes", info.Size())
372 }
373 loaded, err := LoadSession(path)
374 if err != nil || len(loaded.Messages) != 2 {
375 t.Fatalf("reload len=%d err=%v", len(loaded.Messages), err)
376 }
377 }
378
379 func TestDAGSaveCrashAtAppendRecoversOnNextSave(t *testing.T) {
380 path := dagTestSession(t)
381 s := dagSavedSession(t, path, "q1")
382 s.Add(provider.Message{Role: provider.RoleAssistant, Content: "a1"})
383 fileutil.CrashPoint = func(op, _ string) {
384 if op == "dag-append" {
385 panic("crash:dag-append")
386 }
387 }
388 func() {
389 defer func() {
390 if recover() == nil {
391 t.Fatal("crash point did not fire")
392 }
393 }()
394 _ = s.Save(path)
395 }()
396 fileutil.CrashPoint = nil
397 // The crash happened inside the locked save; a later save from the same
398 // session must still land exactly one copy of the message.
399 if err := s.Save(path); err != nil {
400 t.Fatalf("save after crash: %v", err)
401 }
402 loaded, err := LoadSession(path)
403 if err != nil || strings.Join(dagContents(loaded.Messages), ",") != "sys,q1,a1" {
404 t.Fatalf("after crash: %v err=%v", dagContents(loaded.Messages), err)
405 }
406 }
407
408 func TestDAGSaveConcurrentGoroutinesExtendTheirOwnHeads(t *testing.T) {
409 path := dagTestSession(t)
410 a := dagSavedSession(t, path, "q1", "a1")
411 b, err := LoadSession(path)
412 if err != nil {
413 t.Fatal(err)
414 }
415 const rounds = 15
416 var wg sync.WaitGroup
417 run := func(s *Session, tag string) {
418 defer wg.Done()
419 for i := range rounds {
420 s.Add(provider.Message{Role: provider.RoleUser, Content: tag + string(rune('a'+i))})
421 if err := s.Save(path); err != nil {
422 t.Errorf("%s save %d: %v", tag, i, err)
423 return
424 }
425 }
426 }
427 wg.Add(2)
428 go run(a, "A")
429 go run(b, "B")
430 wg.Wait()
431 st := dagReplay(t, path)
432 if st.damaged || len(st.heads) != 2 {
433 t.Fatalf("damaged=%v heads=%d", st.damaged, len(st.heads))
434 }
435 for _, id := range st.headOrder {
436 if got := len(dagChain(st, id)); got != 3+rounds {
437 t.Fatalf("head %s chain length %d", id, got)
438 }
439 }
440 assertNoTranscriptCopies(t, path)
441 }
442
443 func TestExportSessionSchemaOneWritesReadableSchemaOneSession(t *testing.T) {
444 path := dagTestSession(t)
445 s := dagSavedSession(t, path, "q1", "a1")
446 dst := filepath.Join(t.TempDir(), "export.jsonl")
447 if err := ExportSessionSchemaOne(path, dst); err != nil {
448 t.Fatal(err)
449 }
450 probe, err := probeSessionEventLog(dst)
451 if err != nil || !probe.native || probe.dag {
452 t.Fatalf("export probe = %+v err=%v", probe, err)
453 }
454 exported, err := LoadSession(dst)
455 if err != nil || strings.Join(dagContents(exported.Messages), ",") != strings.Join(dagContents(s.Messages), ",") {
456 t.Fatalf("export reload = %v err=%v", dagContents(exported.Messages), err)
457 }
458 if _, ok := exported.Head(); ok {
459 t.Fatal("exported session must be schema 1")
460 }
461 if err := ExportSessionSchemaOne(path, dst); err == nil {
462 t.Fatal("export must refuse to overwrite an existing destination")
463 }
464 }
465
466 func TestDAGSaveIndependentIdenticalTranscriptsConverge(t *testing.T) {
467 path := dagTestSession(t)
468 a := dagSavedSession(t, path, "q1", "a1")
469 b := NewSession("sys")
470 b.Add(provider.Message{Role: provider.RoleUser, Content: "q1"})
471 b.Add(provider.Message{Role: provider.RoleAssistant, Content: "a1"})
472 b.Add(provider.Message{Role: provider.RoleUser, Content: "q2"})
473 if err := b.Save(path); err != nil {
474 t.Fatal(err)
475 }
476 st := dagReplay(t, path)
477 if len(st.heads) != 1 || strings.Join(dagChain(st, SessionMainHead), ",") != "sys,q1,a1,q2" {
478 t.Fatalf("identical prefix must extend main: heads=%d chain=%v", len(st.heads), dagChain(st, SessionMainHead))
479 }
480 for i := range a.Messages {
481 if b.Messages[i].ID != a.Messages[i].ID {
482 t.Fatalf("message %d: independent writer did not adopt the persisted id", i)
483 }
484 }
485 if ref, _ := b.Head(); ref.HeadID != SessionMainHead || ref.LeafID != b.LeafID() {
486 t.Fatalf("b head = %+v", ref)
487 }
488 }
489
490 func TestDAGSaveUnrelatedWriterForksInsteadOfRewinding(t *testing.T) {
491 path := dagTestSession(t)
492 dagSavedSession(t, path, "q1", "a1", "q2", "a2")
493 b := NewSession("sys")
494 b.Add(provider.Message{Role: provider.RoleUser, Content: "q1"})
495 b.Add(provider.Message{Role: provider.RoleAssistant, Content: "a1"})
496 b.Add(provider.Message{Role: provider.RoleUser, Content: "q2-other"})
497 if err := b.Save(path); err != nil {
498 t.Fatal(err)
499 }
500 st := dagReplay(t, path)
501 ref, _ := b.Head()
502 if len(st.heads) != 2 || ref.HeadID == SessionMainHead {
503 t.Fatalf("unrelated writer must fork: heads=%d ref=%+v", len(st.heads), ref)
504 }
505 if got := dagChain(st, SessionMainHead); strings.Join(got, ",") != "sys,q1,a1,q2,a2" {
506 t.Fatalf("main was rewritten by an unrelated writer: %v", got)
507 }
508 if got := dagChain(st, ref.HeadID); strings.Join(got, ",") != "sys,q1,a1,q2-other" {
509 t.Fatalf("fork chain %v", got)
510 }
511 }
512
512 lines GO