git.giacomolli.uk - onechat/blob - memory_test.go

   1 package main
   2
   3 import (
   4         "context"
   5         "fmt"
   6         "log"
   7         "maps"
   8         "os"
   9         "path/filepath"
  10         "slices"
  11         "strings"
  12         "testing"
  13         "time"
  14
  15         "github.com/stretchr/testify/require"
  16 )
  17
  18 func TestMemoryRestartLoadsTheSameView(t *testing.T) {
  19         root := t.TempDir()
  20         clock := newTestClock()
  21         memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
  22         require.NoError(t, err)
  23
  24         memory.chat.high = 200
  25         memory.chat.low = 100
  26         memory.cview.high = 80
  27         memory.cview.low = 40
  28         for n := range 100 {
  29                 _, err = memory.Append("user", fmt.Sprint(n))
  30                 require.NoError(t, err)
  31         }
  32
  33         view := memory.View()
  34         context := memory.Context(0, 100)
  35         saved, err := os.ReadFile(filepath.Join(root, "view.json"))
  36         require.NoError(t, err)
  37         require.Less(t, len(memory.chat.lines), 100)
  38         require.NoError(t, memory.Close())
  39
  40         memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
  41         require.NoError(t, err)
  42
  43         defer func() {
  44                 require.NoError(t, memory.Close())
  45         }()
  46
  47         require.Equal(t, view, memory.View())
  48         require.Equal(t, context, memory.Context(0, 100))
  49
  50         loaded, err := os.ReadFile(filepath.Join(root, "view.json"))
  51         require.NoError(t, err)
  52         require.Equal(t, saved, loaded)
  53 }
  54
  55 func TestMemoryOpensAfterACrashBeforeTheViewSave(t *testing.T) {
  56         root := t.TempDir()
  57         clock := newTestClock()
  58         memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
  59         require.NoError(t, err)
  60
  61         _, err = memory.Append("user", "first")
  62         require.NoError(t, err)
  63
  64         saved, err := os.ReadFile(filepath.Join(root, "view.json"))
  65         require.NoError(t, err)
  66
  67         _, err = memory.Append("user", "second")
  68         require.NoError(t, err)
  69
  70         view := memory.View()
  71         require.NoError(t, memory.Close())
  72
  73         // The log and the tree hold message 1, but the view does not.
  74         require.NoError(t, os.WriteFile(filepath.Join(root, "view.json"), saved, 0o600))
  75
  76         memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
  77         require.NoError(t, err)
  78         require.Equal(t, view, memory.View())
  79         require.NoError(t, memory.Close())
  80 }
  81
  82 func TestMemoryKeepsTheCompactionViewCoarser(t *testing.T) {
  83         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer(), discardLog())
  84         require.NoError(t, err)
  85
  86         defer func() {
  87                 require.NoError(t, memory.Close())
  88         }()
  89
  90         memory.chat.high = 200
  91         memory.chat.low = 100
  92         memory.cview.high = 80
  93         memory.cview.low = 40
  94         for n := range 64 {
  95                 _, err = memory.Append("user", fmt.Sprint(n%10))
  96                 require.NoError(t, err)
  97
  98                 coarse := 0
  99                 for _, line := range memory.chat.lines {
 100                         outer := memory.cview.lines[coarse]
 101                         for nodeEnd(outer.level, outer.index) <= line.index<<line.level {
 102                                 coarse++
 103                                 outer = memory.cview.lines[coarse]
 104                         }
 105
 106                         require.GreaterOrEqual(t, outer.level, line.level, "message %d", n)
 107                 }
 108         }
 109 }
 110
 111 func TestMemoryRunsAtMostEightBuilds(t *testing.T) {
 112         summarizer := newFakeSummarizer()
 113         summarizer.gate = make(chan struct{})
 114         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
 115         require.NoError(t, err)
 116
 117         defer func() {
 118                 require.NoError(t, memory.Close())
 119         }()
 120
 121         for n := range 12 {
 122                 _, err = memory.Append("echo", memoryLongText(n))
 123                 require.NoError(t, err)
 124         }
 125
 126         memory.mutex.Lock()
 127         running := slices.Collect(maps.Keys(memory.queue.running))
 128         memory.mutex.Unlock()
 129         slices.SortFunc(running, func(a, b NodeKey) int {
 130                 return a.index - b.index
 131         })
 132         require.Equal(t, []NodeKey{{0, 0}, {0, 1}, {0, 2}, {0, 3}, {0, 4}, {0, 5}, {0, 6}, {0, 7}}, running)
 133
 134         require.Eventually(t, func() bool {
 135                 return summarizer.Active() == buildMax
 136         }, 10*time.Second, time.Millisecond)
 137
 138         close(summarizer.gate)
 139         require.NoError(t, memory.Wait(t.Context()))
 140         require.Equal(t, buildMax, summarizer.peak)
 141
 142         for level, count := range []int{12, 6, 3, 1} {
 143                 for index := range count {
 144                         require.True(t, treeIsBuilt(t, memory.tree, level, index), "node (%d, %d)", level, index)
 145                 }
 146         }
 147 }
 148
 149 func TestMemorySettleWaitsForThePendingLines(t *testing.T) {
 150         summarizer := newFakeSummarizer()
 151         summarizer.gate = make(chan struct{})
 152         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
 153         require.NoError(t, err)
 154
 155         defer func() {
 156                 require.NoError(t, memory.Close())
 157         }()
 158
 159         _, err = memory.Append("echo", memoryLongText(0))
 160         require.NoError(t, err)
 161         require.Contains(t, memory.View(), viewPending)
 162
 163         settled := make(chan error, 1)
 164         go func() {
 165                 settled <- memory.Settle(t.Context())
 166         }()
 167
 168         select {
 169         case err = <-settled:
 170                 require.Fail(t, "Settle returned before the build", "error: %v", err)
 171         case <-time.After(50 * time.Millisecond):
 172         }
 173
 174         close(summarizer.gate)
 175         require.NoError(t, <-settled)
 176         require.NotContains(t, memory.View(), viewPending)
 177 }
 178
 179 func TestMemorySettleFailsWhileABuildFailed(t *testing.T) {
 180         summarizer := newFakeSummarizer()
 181         summarizer.SetFailing(true)
 182         logs := &SyncBuffer{}
 183         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, log.New(logs, "", 0))
 184         require.NoError(t, err)
 185
 186         defer func() {
 187                 require.NoError(t, memory.Close())
 188         }()
 189
 190         _, err = memory.Append("echo", memoryLongText(0))
 191         require.NoError(t, err)
 192         require.ErrorIs(t, memory.Settle(t.Context()), errMemoryUnsettled)
 193         require.Equal(t, "a summary of the chat failed, and it is tried again at the next message: node (0, 0): fake failure\n", logs.String())
 194
 195         summarizer.SetFailing(false)
 196         _, err = memory.Append("user", "retry")
 197         require.NoError(t, err)
 198         require.NoError(t, memory.Settle(t.Context()))
 199 }
 200
 201 func TestMemorySettleEndsWithItsContext(t *testing.T) {
 202         summarizer := newFakeSummarizer()
 203         summarizer.gate = make(chan struct{})
 204         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
 205         require.NoError(t, err)
 206
 207         defer func() {
 208                 require.NoError(t, memory.Close())
 209         }()
 210
 211         _, err = memory.Append("echo", memoryLongText(0))
 212         require.NoError(t, err)
 213
 214         ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond)
 215         defer cancel()
 216
 217         require.ErrorIs(t, memory.Settle(ctx), context.DeadlineExceeded)
 218 }
 219
 220 func TestMemoryStartsALeafAfterFewerThanEightUnbuilt(t *testing.T) {
 221         summarizer := newFakeSummarizer()
 222         summarizer.SetFailing(true)
 223         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
 224         require.NoError(t, err)
 225
 226         defer func() {
 227                 require.NoError(t, memory.Close())
 228         }()
 229
 230         for n := range 9 {
 231                 _, err = memory.Append("echo", memoryLongText(n))
 232                 require.NoError(t, err)
 233         }
 234
 235         require.NoError(t, memory.Wait(t.Context()))
 236         for _, task := range summarizer.Tasks() {
 237                 require.Less(t, task.Index, 8)
 238         }
 239
 240         require.False(t, treeIsBuilt(t, memory.tree, 0, 8))
 241         summarizer.SetFailing(false)
 242
 243         _, err = memory.Append("user", "retry")
 244         require.NoError(t, err)
 245         require.NoError(t, memory.Wait(t.Context()))
 246         require.True(t, treeIsBuilt(t, memory.tree, 0, 8))
 247         require.True(t, treeIsBuilt(t, memory.tree, 3, 0))
 248 }
 249
 250 func TestMemoryMergesAfterBothHalves(t *testing.T) {
 251         summarizer := newFakeSummarizer()
 252         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
 253         require.NoError(t, err)
 254
 255         defer func() {
 256                 require.NoError(t, memory.Close())
 257         }()
 258
 259         _, err = memory.Append("echo", memoryLongText(0))
 260         require.NoError(t, err)
 261         require.NoError(t, memory.Wait(t.Context()))
 262
 263         _, err = memory.Append("echo", memoryLongText(1))
 264         require.NoError(t, err)
 265         require.NoError(t, memory.Wait(t.Context()))
 266         require.Len(t, summarizer.Tasks(), 3)
 267
 268         left := treeReadNode(t, memory.tree, 0, 0)
 269         right := treeReadNode(t, memory.tree, 0, 1)
 270         merge := summarizer.Tasks()[2]
 271         require.Equal(t, []string{left.Text, right.Text}, merge.Input)
 272         require.Equal(t, "<chat>\n0+1|"+left.Text+"\n1+1|"+right.Text+"\n</chat>", merge.Context)
 273         require.True(t, treeIsBuilt(t, memory.tree, 1, 0))
 274 }
 275
 276 func TestMemoryRestartRestoresTheQueue(t *testing.T) {
 277         root := t.TempDir()
 278         clock := newTestClock()
 279         summarizer := newFakeSummarizer()
 280         summarizer.SetFailing(true)
 281         memory, err := openMemory(t.Context(), root, clock.Now, summarizer, discardLog())
 282         require.NoError(t, err)
 283
 284         for n := range 4 {
 285                 _, err = memory.Append("echo", memoryLongText(n))
 286                 require.NoError(t, err)
 287         }
 288
 289         require.NoError(t, memory.Wait(t.Context()))
 290         require.NoError(t, memory.Close())
 291
 292         memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
 293         require.NoError(t, err)
 294
 295         defer func() {
 296                 require.NoError(t, memory.Close())
 297         }()
 298
 299         require.NoError(t, memory.Wait(t.Context()))
 300         require.True(t, treeIsBuilt(t, memory.tree, 2, 0))
 301 }
 302
 303 func TestMemoryBuildsFreeParents(t *testing.T) {
 304         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer(), discardLog())
 305         require.NoError(t, err)
 306
 307         defer func() {
 308                 require.NoError(t, memory.Close())
 309         }()
 310
 311         for _, text := range []string{"a", "b", "c", "d"} {
 312                 _, err = memory.Append("user", text)
 313                 require.NoError(t, err)
 314         }
 315
 316         node := treeReadNode(t, memory.tree, 2, 0)
 317         require.Equal(t, "user: a\nuser: b\nuser: c\nuser: d", node.Text)
 318 }
 319
 320 func TestMemoryKeepsALongMessageOutOfTheView(t *testing.T) {
 321         summarizer := newFakeSummarizer()
 322         summarizer.SetFailing(true)
 323         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
 324         require.NoError(t, err)
 325
 326         defer func() {
 327                 require.NoError(t, memory.Close())
 328         }()
 329
 330         _, err = memory.Append("echo", memoryLongText(0))
 331         require.NoError(t, err)
 332         require.NoError(t, memory.Wait(t.Context()))
 333         require.Equal(t, "<chat>\n0+1|"+viewPending+"\n</chat>", memory.View())
 334
 335         summarizer.SetFailing(false)
 336
 337         _, err = memory.Append("user", "again")
 338         require.NoError(t, err)
 339         require.NoError(t, memory.Wait(t.Context()))
 340
 341         task, err := memory.task(NodeKey{0, 0})
 342         require.NoError(t, err)
 343         require.Equal(t, "<chat>\n0+1|"+fakeLine(task)+"\n1+1|user: again\n</chat>", memory.View())
 344 }
 345
 346 func TestMemoryCoversMessagesAfterTheSavedView(t *testing.T) {
 347         root := t.TempDir()
 348         clock := newTestClock()
 349         memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
 350         require.NoError(t, err)
 351
 352         _, err = memory.Append("user", "a")
 353         require.NoError(t, err)
 354
 355         _, err = memory.log.Append("reply", "b")
 356         require.NoError(t, err)
 357         require.NoError(t, memory.Close())
 358
 359         memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
 360         require.NoError(t, err)
 361
 362         defer func() {
 363                 require.NoError(t, memory.Close())
 364         }()
 365
 366         require.Equal(t, "<chat>\n0+1|user: a\n1+1|reply: b\n</chat>", memory.View())
 367 }
 368
 369 func TestMemoryRefusesToRebuildAMissingView(t *testing.T) {
 370         root := t.TempDir()
 371         clock := newTestClock()
 372         memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
 373         require.NoError(t, err)
 374
 375         _, err = memory.Append("user", "a")
 376         require.NoError(t, err)
 377         require.NoError(t, memory.Close())
 378         require.NoError(t, os.Remove(filepath.Join(root, "view.json")))
 379
 380         _, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
 381         require.ErrorIs(t, err, errMemoryNoView)
 382 }
 383
 384 // memoryLongText returns a text that is too long for its own node.
 385 func memoryLongText(n int) string {
 386         return fmt.Sprintf("long %d ", n) + strings.Repeat("x", nodeBytes)
 387 }