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

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