| 1 | package agent |
| 2 | |
| 3 | import ( |
| 4 | "crypto/sha256" |
| 5 | "encoding/json" |
| 6 | "os" |
| 7 | "path/filepath" |
| 8 | "strings" |
| 9 | "testing" |
| 10 | |
| 11 | "reasonix/internal/provider" |
| 12 | ) |
| 13 | |
| 14 | // legacyDigestSessionMessages is the pre-fusion reference implementation: a |
| 15 | // full re-serialize pass over the final message slice. The incremental hasher |
| 16 | // fed during decode must reproduce it byte-for-byte. |
| 17 | func legacyDigestSessionMessages(msgs []provider.Message) ([sha256.Size]byte, error) { |
| 18 | h := sha256.New() |
| 19 | for _, m := range msgs { |
| 20 | m = messageForSessionIdentity(m) |
| 21 | b, err := json.Marshal(m) |
| 22 | if err != nil { |
| 23 | return [sha256.Size]byte{}, err |
| 24 | } |
| 25 | if _, err := h.Write(b); err != nil { |
| 26 | return [sha256.Size]byte{}, err |
| 27 | } |
| 28 | if _, err := h.Write([]byte{'\n'}); err != nil { |
| 29 | return [sha256.Size]byte{}, err |
| 30 | } |
| 31 | } |
| 32 | var out [sha256.Size]byte |
| 33 | copy(out[:], h.Sum(nil)) |
| 34 | return out, nil |
| 35 | } |
| 36 | |
| 37 | // representativeSessionMessages mixes the message shapes a long real session |
| 38 | // carries: system prompt, plain text, multimodal user input, reasoning with |
| 39 | // provider signatures, tool calls and their results, and display timestamps |
| 40 | // (which the digest must keep ignoring). |
| 41 | func representativeSessionMessages() []provider.Message { |
| 42 | return []provider.Message{ |
| 43 | {Role: provider.RoleSystem, Content: "sys"}, |
| 44 | {Role: provider.RoleUser, Content: "prompt one", CreatedAt: 1720000000000}, |
| 45 | {Role: provider.RoleAssistant, Content: "thinking about it", ReasoningContent: "chain of thought", ReasoningID: "rs_1", ReasoningStatus: "completed", ReasoningSignature: "sig"}, |
| 46 | {Role: provider.RoleUser, Content: "with image", Images: []string{"data:image/png;base64,aGVsbG8="}, CreatedAt: 1720000001000}, |
| 47 | {Role: provider.RoleAssistant, Content: "", ToolCalls: []provider.ToolCall{ |
| 48 | {ID: "call_1", Name: "read_file", Arguments: `{"path":"a.go"}`}, |
| 49 | }}, |
| 50 | {Role: provider.RoleTool, ToolCallID: "call_1", Name: "read_file", Content: "package main"}, |
| 51 | {Role: provider.RoleAssistant, Content: "final answer", WorkDurationMs: 42}, |
| 52 | } |
| 53 | } |
| 54 | |
| 55 | func writeLegacyJSONLSession(t *testing.T, path string, msgs []provider.Message) { |
| 56 | t.Helper() |
| 57 | f, err := os.OpenFile(path, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o600) |
| 58 | if err != nil { |
| 59 | t.Fatalf("open jsonl: %v", err) |
| 60 | } |
| 61 | enc := json.NewEncoder(f) |
| 62 | for _, m := range msgs { |
| 63 | if err := enc.Encode(m); err != nil { |
| 64 | f.Close() |
| 65 | t.Fatalf("encode jsonl: %v", err) |
| 66 | } |
| 67 | } |
| 68 | if err := f.Close(); err != nil { |
| 69 | t.Fatalf("close jsonl: %v", err) |
| 70 | } |
| 71 | } |
| 72 | |
| 73 | // loadAndDigestSessionMessages loads like LoadSession does and returns the |
| 74 | // decode-fused digest alongside the messages. |
| 75 | func loadAndDigestSessionMessages(path string) (msgs []provider.Message, fromEvents, damaged bool, digest [sha256.Size]byte, digestOK bool, err error) { |
| 76 | hasher := newSessionTranscriptHasher() |
| 77 | msgs, fromEvents, damaged, err = loadSessionMessagesWithLimits(path, defaultSessionReplayLimits, hasher) |
| 78 | digest, digestOK = hasher.sum() |
| 79 | return msgs, fromEvents, damaged, digest, digestOK, err |
| 80 | } |
| 81 | |
| 82 | func TestLoadSessionMessagesWithDigestMatchesLegacyEventLog(t *testing.T) { |
| 83 | path := filepath.Join(t.TempDir(), "session.jsonl") |
| 84 | s := NewSession("sys") |
| 85 | for _, m := range representativeSessionMessages()[1:] { |
| 86 | s.Add(m) |
| 87 | // One save per message forces a replace record plus a chain of append |
| 88 | // records, exercising the hasher across both event types. |
| 89 | if err := s.SaveSnapshot(path); err != nil { |
| 90 | t.Fatalf("SaveSnapshot: %v", err) |
| 91 | } |
| 92 | } |
| 93 | |
| 94 | msgs, fromEvents, damaged, digest, digestOK, err := loadAndDigestSessionMessages(path) |
| 95 | if err != nil { |
| 96 | t.Fatalf("load: %v", err) |
| 97 | } |
| 98 | if !fromEvents || damaged { |
| 99 | t.Fatalf("fromEvents=%v damaged=%v, want event-log replay without damage", fromEvents, damaged) |
| 100 | } |
| 101 | if !digestOK { |
| 102 | t.Fatal("digestOK = false, want true") |
| 103 | } |
| 104 | want, err := legacyDigestSessionMessages(msgs) |
| 105 | if err != nil { |
| 106 | t.Fatalf("legacy digest: %v", err) |
| 107 | } |
| 108 | if digest != want { |
| 109 | t.Fatalf("incremental digest %x != legacy digest %x", digest, want) |
| 110 | } |
| 111 | } |
| 112 | |
| 113 | func TestLoadSessionMessagesWithDigestMatchesLegacyJSONL(t *testing.T) { |
| 114 | path := filepath.Join(t.TempDir(), "session.jsonl") |
| 115 | writeLegacyJSONLSession(t, path, representativeSessionMessages()) |
| 116 | |
| 117 | msgs, fromEvents, damaged, digest, digestOK, err := loadAndDigestSessionMessages(path) |
| 118 | if err != nil { |
| 119 | t.Fatalf("load: %v", err) |
| 120 | } |
| 121 | if fromEvents || damaged { |
| 122 | t.Fatalf("fromEvents=%v damaged=%v, want plain jsonl load", fromEvents, damaged) |
| 123 | } |
| 124 | if !digestOK { |
| 125 | t.Fatal("digestOK = false, want true") |
| 126 | } |
| 127 | want, err := legacyDigestSessionMessages(msgs) |
| 128 | if err != nil { |
| 129 | t.Fatalf("legacy digest: %v", err) |
| 130 | } |
| 131 | if digest != want { |
| 132 | t.Fatalf("incremental digest %x != legacy digest %x", digest, want) |
| 133 | } |
| 134 | } |
| 135 | |
| 136 | func TestLoadSessionMessagesWithDigestCoversOnlyCleanPrefix(t *testing.T) { |
| 137 | path := filepath.Join(t.TempDir(), "session.jsonl") |
| 138 | sessionWithTurns(t, path, 2) |
| 139 | |
| 140 | logPath := SessionEventLogPath(path) |
| 141 | f, err := os.OpenFile(logPath, os.O_WRONLY|os.O_APPEND, 0o644) |
| 142 | if err != nil { |
| 143 | t.Fatalf("open log: %v", err) |
| 144 | } |
| 145 | if _, err := f.Write([]byte(`{"schema_version":1,"type":"append","message_index":5,"mess`)); err != nil { |
| 146 | t.Fatalf("write torn tail: %v", err) |
| 147 | } |
| 148 | f.Close() |
| 149 | |
| 150 | msgs, _, damaged, digest, digestOK, err := loadAndDigestSessionMessages(path) |
| 151 | if err != nil { |
| 152 | t.Fatalf("load: %v", err) |
| 153 | } |
| 154 | if !damaged { |
| 155 | t.Fatal("damaged = false, want true for torn tail") |
| 156 | } |
| 157 | if !digestOK { |
| 158 | t.Fatal("digestOK = false, want true") |
| 159 | } |
| 160 | // The digest must cover exactly the replayable prefix the caller received, |
| 161 | // never bytes from the torn record. |
| 162 | want, err := legacyDigestSessionMessages(msgs) |
| 163 | if err != nil { |
| 164 | t.Fatalf("legacy digest: %v", err) |
| 165 | } |
| 166 | if digest != want { |
| 167 | t.Fatalf("incremental digest %x != legacy digest %x", digest, want) |
| 168 | } |
| 169 | } |
| 170 | |
| 171 | func TestLoadSessionDigestFastPathKeepsPersistedBaseline(t *testing.T) { |
| 172 | path := filepath.Join(t.TempDir(), "session.jsonl") |
| 173 | s := NewSession("sys") |
| 174 | for _, m := range representativeSessionMessages()[1:] { |
| 175 | s.Add(m) |
| 176 | } |
| 177 | if err := s.SaveSnapshot(path); err != nil { |
| 178 | t.Fatalf("SaveSnapshot: %v", err) |
| 179 | } |
| 180 | |
| 181 | loaded, err := LoadSession(path) |
| 182 | if err != nil { |
| 183 | t.Fatalf("LoadSession: %v", err) |
| 184 | } |
| 185 | if loaded.normalizedDirty { |
| 186 | t.Fatal("normalizedDirty = true for a well-formed session, want false") |
| 187 | } |
| 188 | // The baseline anchored during load must equal the legacy digest of the |
| 189 | // loaded transcript; otherwise change detection (HasUnsavedChanges) would |
| 190 | // report phantom writes or miss real ones. |
| 191 | want, err := legacyDigestSessionMessages(loaded.Snapshot()) |
| 192 | if err != nil { |
| 193 | t.Fatalf("legacy digest: %v", err) |
| 194 | } |
| 195 | if !loaded.persisted.ok || loaded.persisted.digest != want { |
| 196 | t.Fatalf("persisted baseline digest %x (ok=%v) != legacy digest %x", loaded.persisted.digest, loaded.persisted.ok, want) |
| 197 | } |
| 198 | if loaded.HasUnsavedChanges(path) { |
| 199 | t.Fatal("HasUnsavedChanges = true right after a clean load, want false") |
| 200 | } |
| 201 | } |
| 202 | |
| 203 | func TestLoadSessionDigestRecomputedAfterNormalizationRepair(t *testing.T) { |
| 204 | path := filepath.Join(t.TempDir(), "session.jsonl") |
| 205 | // Dangling tool call: normalization fabricates a placeholder tool result, |
| 206 | // so the persisted baseline must be the digest of the repaired transcript. |
| 207 | writeLegacyJSONLSession(t, path, []provider.Message{ |
| 208 | {Role: provider.RoleSystem, Content: "sys"}, |
| 209 | {Role: provider.RoleUser, Content: "run it"}, |
| 210 | {Role: provider.RoleAssistant, Content: "", ToolCalls: []provider.ToolCall{ |
| 211 | {ID: "call_1", Name: "read_file", Arguments: `{"path":"a.go"}`}, |
| 212 | }}, |
| 213 | }) |
| 214 | |
| 215 | loaded, err := LoadSession(path) |
| 216 | if err != nil { |
| 217 | t.Fatalf("LoadSession: %v", err) |
| 218 | } |
| 219 | if !loaded.normalizedDirty { |
| 220 | t.Fatal("normalizedDirty = false, want true for a dangling tool call") |
| 221 | } |
| 222 | want, err := legacyDigestSessionMessages(loaded.Snapshot()) |
| 223 | if err != nil { |
| 224 | t.Fatalf("legacy digest: %v", err) |
| 225 | } |
| 226 | if !loaded.persisted.ok || loaded.persisted.digest != want { |
| 227 | t.Fatalf("persisted baseline digest %x (ok=%v) != legacy digest of repaired transcript %x", loaded.persisted.digest, loaded.persisted.ok, want) |
| 228 | } |
| 229 | rawWant, err := legacyDigestSessionMessages(loaded.rawMessages) |
| 230 | if err != nil { |
| 231 | t.Fatalf("legacy raw digest: %v", err) |
| 232 | } |
| 233 | if loaded.persisted.digest == rawWant { |
| 234 | t.Fatal("persisted baseline still hashes the pre-repair transcript") |
| 235 | } |
| 236 | } |
| 237 | |
| 238 | func TestLoadSessionDisplayMessagesDigestMatchesLegacy(t *testing.T) { |
| 239 | path := filepath.Join(t.TempDir(), "session.jsonl") |
| 240 | s := NewSession("sys") |
| 241 | for _, m := range representativeSessionMessages()[1:] { |
| 242 | s.Add(m) |
| 243 | } |
| 244 | if err := s.SaveSnapshot(path); err != nil { |
| 245 | t.Fatalf("SaveSnapshot: %v", err) |
| 246 | } |
| 247 | |
| 248 | msgs, state, clean, err := LoadSessionDisplayMessages(path) |
| 249 | if err != nil { |
| 250 | t.Fatalf("LoadSessionDisplayMessages: %v", err) |
| 251 | } |
| 252 | if !clean { |
| 253 | t.Fatal("clean = false, want true") |
| 254 | } |
| 255 | want, err := legacyDigestSessionMessages(msgs) |
| 256 | if err != nil { |
| 257 | t.Fatalf("legacy digest: %v", err) |
| 258 | } |
| 259 | if state.Digest != want || state.DigestHex != digestString(want) { |
| 260 | t.Fatalf("display digest %s != legacy digest %s", state.DigestHex, digestString(want)) |
| 261 | } |
| 262 | } |
| 263 | |
| 264 | // BenchmarkLoadSessionDigestFusion contrasts the old load shape (decode pass + |
| 265 | // separate full re-serialize digest pass) with the fused decode-time digest on |
| 266 | // a synthetic long session with realistic tool-output sizes. |
| 267 | func BenchmarkLoadSessionDigestFusion(b *testing.B) { |
| 268 | path := filepath.Join(b.TempDir(), "session.jsonl") |
| 269 | s := NewSession("sys") |
| 270 | bigResult := strings.Repeat("package main // line of tool output\n", 128) // ~4KB |
| 271 | for range 400 { |
| 272 | for _, m := range representativeSessionMessages()[1:] { |
| 273 | if m.Role == provider.RoleTool { |
| 274 | m.Content = bigResult |
| 275 | } |
| 276 | s.Add(m) |
| 277 | } |
| 278 | } |
| 279 | if err := s.SaveSnapshot(path); err != nil { |
| 280 | b.Fatalf("SaveSnapshot: %v", err) |
| 281 | } |
| 282 | |
| 283 | b.Run("digest_pass_only", func(b *testing.B) { |
| 284 | msgs, _, _, err := loadSessionMessages(path) |
| 285 | if err != nil { |
| 286 | b.Fatalf("loadSessionMessages: %v", err) |
| 287 | } |
| 288 | b.ResetTimer() |
| 289 | for b.Loop() { |
| 290 | if _, err := digestSessionMessages(msgs); err != nil { |
| 291 | b.Fatalf("digestSessionMessages: %v", err) |
| 292 | } |
| 293 | } |
| 294 | }) |
| 295 | b.Run("separate_digest_pass", func(b *testing.B) { |
| 296 | for b.Loop() { |
| 297 | msgs, _, _, err := loadSessionMessages(path) |
| 298 | if err != nil { |
| 299 | b.Fatalf("loadSessionMessages: %v", err) |
| 300 | } |
| 301 | if _, err := digestSessionMessages(msgs); err != nil { |
| 302 | b.Fatalf("digestSessionMessages: %v", err) |
| 303 | } |
| 304 | } |
| 305 | }) |
| 306 | b.Run("fused_decode_digest", func(b *testing.B) { |
| 307 | for b.Loop() { |
| 308 | if _, _, _, _, ok, err := loadAndDigestSessionMessages(path); err != nil || !ok { |
| 309 | b.Fatalf("fused load: ok=%v err=%v", ok, err) |
| 310 | } |
| 311 | } |
| 312 | }) |
| 313 | } |
| 314 |