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

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