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