diff --git a/internal/state/sessionstore/save.go b/internal/state/sessionstore/save.go index bdb1769157..e7197f5909 100644 --- a/internal/state/sessionstore/save.go +++ b/internal/state/sessionstore/save.go @@ -99,6 +99,7 @@ type sessionPersistState struct { // and heal a stale ledger. saveVerified bool ok bool + disk persistedDisk } type sessionSaveMode int @@ -272,7 +273,7 @@ func (s *Session) saveLocked(path string, mode sessionSaveMode) error { // stalest capture written last would then read the newer transcript it // lost the race to as a bogus stale-prefix conflict. msgs, version, rewriteVersion := s.SnapshotWithVersion() - digest, contentBytes, err := digestAndSizeSessionMessages(msgs) + digest, contentBytes, prefixDigest, err := digestSizeAndPrefix(context.Background(), msgs, s.diskPrefixLen(path)) if err != nil { return err } @@ -292,7 +293,7 @@ func (s *Session) saveLocked(path string, mode sessionSaveMode) error { } repairLog := false ownedRewrite := mode == sessionSaveRewrite || mode == sessionSaveRewriteCompact - decision, err := s.checkSnapshotWrite(path, msgs, digest, version, ownedRewrite) + decision, err := s.snapshotDecision(path, msgs, digest, prefixDigest, version, ownedRewrite) if err != nil { return err } @@ -1119,28 +1120,8 @@ func digestAndSizeSessionMessages(msgs []provider.Message) ([sha256.Size]byte, i } func digestAndSizeSessionMessagesContext(ctx context.Context, msgs []provider.Message) ([sha256.Size]byte, int64, error) { - h := sha256.New() - size := int64(0) - for _, m := range msgs { - if err := ctx.Err(); err != nil { - return [sha256.Size]byte{}, 0, err - } - m = messageForSessionIdentity(m) - b, err := json.Marshal(m) - if err != nil { - return [sha256.Size]byte{}, 0, err - } - if _, err := h.Write(b); err != nil { - return [sha256.Size]byte{}, 0, err - } - if _, err := h.Write([]byte{'\n'}); err != nil { - return [sha256.Size]byte{}, 0, err - } - size += int64(len(b)) + 1 - } - var out [sha256.Size]byte - copy(out[:], h.Sum(nil)) - return out, size, nil + digest, size, _, err := digestSizeAndPrefix(ctx, msgs, -1) + return digest, size, err } func messagesHavePrefix(full, prefix []provider.Message) bool { diff --git a/internal/state/sessionstore/save_fast_baseline.go b/internal/state/sessionstore/save_fast_baseline.go new file mode 100644 index 0000000000..5cd8e96456 --- /dev/null +++ b/internal/state/sessionstore/save_fast_baseline.go @@ -0,0 +1,125 @@ +package sessionstore + +import ( + "context" + "crypto/sha256" + "encoding/json" + "os" + + "reasonix/internal/contract/provider" + "reasonix/internal/state/store" +) + +type fileStampOnDisk struct { + size, mod int64 + here bool +} + +// baselineStamp prints the four files a save is judged against: the +// checkpoint, the event log, the ledger sidecar and the damage marker. +type baselineStamp [4]fileStampOnDisk + +func stampBaseline(path string) baselineStamp { + var out baselineStamp + for i, name := range []string{path, store.SessionEventLog(path), store.SessionMeta(path), store.SessionEventLogDamaged(path)} { + if info, err := os.Stat(name); err == nil { + out[i] = fileStampOnDisk{size: info.Size(), mod: info.ModTime().UnixNano(), here: true} + } + } + return out +} + +// persistedDisk is what this runtime knows about disk right after its own +// completed save: the files' print and how many messages they hold. +type persistedDisk struct { + stamp baselineStamp + n int + armed bool +} + +// armDiskBaseline records the print of the files this save just wrote. It runs +// under the save locks, so no cooperating writer can land between the write and +// the print. +func (s *Session) armDiskBaseline(path string, digest [sha256.Size]byte, n int) { + stamp := stampBaseline(path) + s.mu.Lock() + defer s.mu.Unlock() + if !s.persisted.ok || !s.persisted.saveVerified || s.persisted.path != canonicalSessionSavePath(path) || s.persisted.digest != digest { + return + } + s.persisted.disk = persistedDisk{stamp: stamp, n: n, armed: true} +} + +// diskPrefixLen is how many leading messages the next save must account for to +// use fastSnapshotDecision, or -1 when this runtime has no print to trust. +func (s *Session) diskPrefixLen(path string) int { + st := s.persistState(path) + if !st.ok || !st.saveVerified || !st.revisionKnown || !st.disk.armed { + return -1 + } + return st.disk.n +} + +// fastSnapshotDecision answers checkSnapshotWrite without reading the +// transcript back. The files are exactly as this runtime left them, so disk +// holds the persisted messages, and prefix — the digest of next's first n +// messages, taken in the pass that digested all of it — says whether next still +// begins with them. Anything else falls to the full comparison. +func (s *Session) fastSnapshotDecision(path string, next []provider.Message, prefix [sha256.Size]byte) (snapshotWriteDecision, bool) { + st := s.persistState(path) + if !st.ok || !st.saveVerified || !st.revisionKnown || !st.disk.armed { + return snapshotWriteDecision{}, false + } + n := st.disk.n + if n > len(next) || prefix != st.digest || stampBaseline(path) != st.disk.stamp { + return snapshotWriteDecision{}, false + } + for _, m := range next { + if len(m.Images) > 0 || len(m.UnavailableImages) > 0 { + return snapshotWriteDecision{}, false + } + } + return snapshotWriteDecision{ + revision: st.revision, + upToDate: len(next) == n, + appendOnly: len(next) > n, + appendFrom: n, + }, true +} + +// digestSizeAndPrefix is digestAndSizeSessionMessagesContext that also reports +// the digest of the first k messages, accumulated in the same pass. k < 0 asks +// for none. +func digestSizeAndPrefix(ctx context.Context, msgs []provider.Message, k int) (digest [sha256.Size]byte, size int64, prefix [sha256.Size]byte, err error) { + h := sha256.New() + if k == 0 { + copy(prefix[:], h.Sum(nil)) + } + for i, m := range msgs { + if err := ctx.Err(); err != nil { + return digest, 0, prefix, err + } + m = messageForSessionIdentity(m) + b, err := json.Marshal(m) + if err != nil { + return digest, 0, prefix, err + } + h.Write(b) + h.Write([]byte{'\n'}) + size += int64(len(b)) + 1 + if i+1 == k { + copy(prefix[:], h.Sum(nil)) + } + } + copy(digest[:], h.Sum(nil)) + return digest, size, prefix, nil +} + +// snapshotDecision is checkSnapshotWrite, answered from the print when the +// files are as this runtime left them. +func (s *Session) snapshotDecision(path string, msgs []provider.Message, digest, prefix [sha256.Size]byte, version uint64, ownedRewrite bool) (snapshotWriteDecision, error) { + if decision, ok := s.fastSnapshotDecision(path, msgs, prefix); ok { + return decision, nil + } + return s.checkSnapshotWrite(path, msgs, digest, version, ownedRewrite) +} diff --git a/internal/state/sessionstore/save_fast_baseline_test.go b/internal/state/sessionstore/save_fast_baseline_test.go new file mode 100644 index 0000000000..870dd2cb7c --- /dev/null +++ b/internal/state/sessionstore/save_fast_baseline_test.go @@ -0,0 +1,86 @@ +package sessionstore + +import ( + "fmt" + "path/filepath" + "runtime" + "strings" + "testing" + + "reasonix/internal/contract/provider" +) + +func seededSession(t *testing.T, path string, n int) *Session { + t.Helper() + s := NewSession("sys") + for i := range n / 2 { + s.Add(provider.Message{Role: provider.RoleUser, Content: fmt.Sprintf("turn %d: %s", i, strings.Repeat("please keep going. ", 20))}) + s.Add(provider.Message{Role: provider.RoleAssistant, Content: strings.Repeat("a plain english sentence about the work. ", 40)}) + } + if err := s.Save(path); err != nil { + t.Fatalf("seed save: %v", err) + } + return s +} + +func allocatedBy(f func()) uint64 { + var before, after runtime.MemStats + runtime.GC() + runtime.ReadMemStats(&before) + f() + runtime.ReadMemStats(&after) + return after.TotalAlloc - before.TotalAlloc +} + +// A runtime that wrote the last save and finds the files untouched has nothing +// to reload: one more message costs about one encoding pass, not a replay of the +// event log plus a decode and re-encode of what disk already holds. +func TestSavingOneMoreMessageDoesNotReplayTheLog(t *testing.T) { + path := filepath.Join(t.TempDir(), "s.jsonl") + s := seededSession(t, path, 2000) + if err := s.Save(path); err != nil { + t.Fatal(err) + } + _, size, err := digestAndSizeSessionMessages(s.Snapshot()) + if err != nil { + t.Fatal(err) + } + s.Add(provider.Message{Role: provider.RoleUser, Content: "next"}) + got := allocatedBy(func() { + if err := s.Save(path); err != nil { + t.Fatal(err) + } + }) + if limit := uint64(size) * 4; got > limit { + t.Fatalf("one-message save allocated %d bytes for a %d byte transcript (limit %d)", got, size, limit) + } +} + +func TestSavingAfterAnotherRuntimeWroteStillSeesIt(t *testing.T) { + path := filepath.Join(t.TempDir(), "s.jsonl") + a := seededSession(t, path, 40) + b, err := LoadSession(path) + if err != nil { + t.Fatal(err) + } + b.Add(provider.Message{Role: provider.RoleUser, Content: "written by the other runtime"}) + if err := b.Save(path); err != nil { + t.Fatalf("other runtime save: %v", err) + } + a.Add(provider.Message{Role: provider.RoleUser, Content: "written by the stale runtime"}) + if err := a.Save(path); err == nil { + after, lerr := LoadSession(path) + if lerr != nil { + t.Fatal(lerr) + } + var kept int + for _, m := range after.Snapshot() { + if strings.HasPrefix(m.Content, "written by the other runtime") { + kept++ + } + } + if kept != 1 { + t.Fatalf("the stale runtime's save dropped the other runtime's message (kept %d)", kept) + } + } +} diff --git a/internal/state/sessionstore/save_listing_projection.go b/internal/state/sessionstore/save_listing_projection.go index 0996e58b02..156140d4fd 100644 --- a/internal/state/sessionstore/save_listing_projection.go +++ b/internal/state/sessionstore/save_listing_projection.go @@ -10,6 +10,7 @@ import ( func (s *Session) markPersistedWithListing(path string, digest [sha256.Size]byte, version uint64, revision int64, rewriteVersion int, msgs []provider.Message) { s.markPersisted(path, digest, version, revision, rewriteVersion) persistSessionListingProjection(path, msgs) + s.armDiskBaseline(path, digest, len(msgs)) } func persistSessionListingProjection(path string, msgs []provider.Message) {