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(), discardLog())
  21         require.NoError(t, err)
  22
  23         memory.chat.high = 200
  24         memory.chat.low = 100
  25
  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
  93         memory.cview.high = 80
  94         memory.cview.low = 40
  95         for n := range 64 {
  96                 _, err = memory.Append("user", fmt.Sprint(n%10))
  97                 require.NoError(t, err)
  98
  99                 coarse := 0
 100                 for _, line := range memory.chat.lines {
 101                         outer := memory.cview.lines[coarse]
 102                         for nodeEnd(outer.level, outer.index) <= line.index<<line.level {
 103                                 coarse++
 104                                 outer = memory.cview.lines[coarse]
 105                         }
 106
 107                         require.GreaterOrEqual(t, outer.level, line.level, "message %d", n)
 108                 }
 109         }
 110 }
 111
 112 func TestMemoryRunsAtMostEightBuilds(t *testing.T) {
 113         summarizer := newFakeSummarizer()
 114         summarizer.gate = make(chan struct{})
 115         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
 116         require.NoError(t, err)
 117
 118         defer func() {
 119                 require.NoError(t, memory.Close())
 120         }()
 121
 122         for n := range 12 {
 123                 _, err = memory.Append("echo", memoryLongText(n))
 124                 require.NoError(t, err)
 125         }
 126
 127         memory.mutex.Lock()
 128
 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() == buildMax
 138         }, 10*time.Second, time.Millisecond)
 139
 140         close(summarizer.gate)
 141         require.NoError(t, memory.Wait(t.Context()))
 142         require.Equal(t, buildMax, summarizer.peak)
 143
 144         for level, count := range []int{12, 6, 3, 1} {
 145                 for index := range count {
 146                         require.True(t, treeIsBuilt(t, memory.tree, level, index), "node (%d, %d)", level, index)
 147                 }
 148         }
 149 }
 150
 151 func TestMemorySettleWaitsForThePendingLines(t *testing.T) {
 152         summarizer := newFakeSummarizer()
 153         summarizer.gate = make(chan struct{})
 154         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
 155         require.NoError(t, err)
 156
 157         defer func() {
 158                 require.NoError(t, memory.Close())
 159         }()
 160
 161         _, err = memory.Append("echo", memoryLongText(0))
 162         require.NoError(t, err)
 163         require.Contains(t, memory.View(), viewPending)
 164
 165         settled := make(chan error, 1)
 166         go func() {
 167                 settled <- memory.Settle(t.Context())
 168         }()
 169
 170         select {
 171         case err = <-settled:
 172                 require.Fail(t, "Settle returned before the build", "error: %v", err)
 173         case <-time.After(50 * time.Millisecond):
 174         }
 175
 176         close(summarizer.gate)
 177         require.NoError(t, <-settled)
 178         require.NotContains(t, memory.View(), viewPending)
 179 }
 180
 181 func TestMemorySettleFailsWhileABuildFailed(t *testing.T) {
 182         summarizer := newFakeSummarizer()
 183         summarizer.SetFailing(true)
 184
 185         logs := &SyncBuffer{}
 186         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, daemonTestLogger(logs))
 187         require.NoError(t, err)
 188
 189         defer func() {
 190                 require.NoError(t, memory.Close())
 191         }()
 192
 193         _, err = memory.Append("echo", memoryLongText(0))
 194         require.NoError(t, err)
 195         require.ErrorIs(t, memory.Settle(t.Context()), errMemoryUnsettled)
 196         require.Equal(t, "<4>summary failed level=0 index=0: fake failure\n", logs.String())
 197
 198         summarizer.SetFailing(false)
 199
 200         _, err = memory.Append("user", "retry")
 201         require.NoError(t, err)
 202         require.NoError(t, memory.Settle(t.Context()))
 203 }
 204
 205 func TestMemorySettleEndsWithItsContext(t *testing.T) {
 206         summarizer := newFakeSummarizer()
 207         summarizer.gate = make(chan struct{})
 208         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
 209         require.NoError(t, err)
 210
 211         defer func() {
 212                 require.NoError(t, memory.Close())
 213         }()
 214
 215         _, err = memory.Append("echo", memoryLongText(0))
 216         require.NoError(t, err)
 217
 218         ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond)
 219         defer cancel()
 220
 221         require.ErrorIs(t, memory.Settle(ctx), context.DeadlineExceeded)
 222 }
 223
 224 func TestMemoryStartsALeafAfterFewerThanEightUnbuilt(t *testing.T) {
 225         summarizer := newFakeSummarizer()
 226         summarizer.SetFailing(true)
 227
 228         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
 229         require.NoError(t, err)
 230
 231         defer func() {
 232                 require.NoError(t, memory.Close())
 233         }()
 234
 235         for n := range 9 {
 236                 _, err = memory.Append("echo", memoryLongText(n))
 237                 require.NoError(t, err)
 238         }
 239
 240         require.NoError(t, memory.Wait(t.Context()))
 241
 242         for _, task := range summarizer.Tasks() {
 243                 require.Less(t, task.Index, 8)
 244         }
 245
 246         require.False(t, treeIsBuilt(t, memory.tree, 0, 8))
 247         summarizer.SetFailing(false)
 248
 249         _, err = memory.Append("user", "retry")
 250         require.NoError(t, err)
 251         require.NoError(t, memory.Wait(t.Context()))
 252         require.True(t, treeIsBuilt(t, memory.tree, 0, 8))
 253         require.True(t, treeIsBuilt(t, memory.tree, 3, 0))
 254 }
 255
 256 func TestMemoryMergesAfterBothHalves(t *testing.T) {
 257         summarizer := newFakeSummarizer()
 258         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
 259         require.NoError(t, err)
 260
 261         defer func() {
 262                 require.NoError(t, memory.Close())
 263         }()
 264
 265         _, err = memory.Append("echo", memoryLongText(0))
 266         require.NoError(t, err)
 267         require.NoError(t, memory.Wait(t.Context()))
 268
 269         _, err = memory.Append("echo", memoryLongText(1))
 270         require.NoError(t, err)
 271         require.NoError(t, memory.Wait(t.Context()))
 272         require.Len(t, summarizer.Tasks(), 3)
 273
 274         left := treeReadNode(t, memory.tree, 0, 0)
 275         right := treeReadNode(t, memory.tree, 0, 1)
 276         merge := summarizer.Tasks()[2]
 277         require.Equal(t, []string{left.Text, right.Text}, merge.Input)
 278         require.Equal(t, "<chat>\n0+1|"+left.Text+"\n1+1|"+right.Text+"\n</chat>", merge.Context)
 279         require.True(t, treeIsBuilt(t, memory.tree, 1, 0))
 280 }
 281
 282 func TestMemoryRestartRestoresTheQueue(t *testing.T) {
 283         root := t.TempDir()
 284         clock := newTestClock()
 285         summarizer := newFakeSummarizer()
 286         summarizer.SetFailing(true)
 287
 288         memory, err := openMemory(t.Context(), root, clock.Now, summarizer, discardLog())
 289         require.NoError(t, err)
 290
 291         for n := range 4 {
 292                 _, err = memory.Append("echo", memoryLongText(n))
 293                 require.NoError(t, err)
 294         }
 295
 296         require.NoError(t, memory.Wait(t.Context()))
 297         require.NoError(t, memory.Close())
 298
 299         memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
 300         require.NoError(t, err)
 301
 302         defer func() {
 303                 require.NoError(t, memory.Close())
 304         }()
 305
 306         require.NoError(t, memory.Wait(t.Context()))
 307         require.True(t, treeIsBuilt(t, memory.tree, 2, 0))
 308 }
 309
 310 func TestMemoryBuildsFreeParents(t *testing.T) {
 311         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer(), discardLog())
 312         require.NoError(t, err)
 313
 314         defer func() {
 315                 require.NoError(t, memory.Close())
 316         }()
 317
 318         for _, text := range []string{"a", "b", "c", "d"} {
 319                 _, err = memory.Append("user", text)
 320                 require.NoError(t, err)
 321         }
 322
 323         node := treeReadNode(t, memory.tree, 2, 0)
 324         require.Equal(t, "user: a\nuser: b\nuser: c\nuser: d", node.Text)
 325 }
 326
 327 func TestMemoryKeepsALongMessageOutOfTheView(t *testing.T) {
 328         summarizer := newFakeSummarizer()
 329         summarizer.SetFailing(true)
 330
 331         memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
 332         require.NoError(t, err)
 333
 334         defer func() {
 335                 require.NoError(t, memory.Close())
 336         }()
 337
 338         _, err = memory.Append("echo", memoryLongText(0))
 339         require.NoError(t, err)
 340         require.NoError(t, memory.Wait(t.Context()))
 341         require.Equal(t, "<chat>\n0+1|"+viewPending+"\n</chat>", memory.View())
 342
 343         summarizer.SetFailing(false)
 344
 345         _, err = memory.Append("user", "again")
 346         require.NoError(t, err)
 347         require.NoError(t, memory.Wait(t.Context()))
 348
 349         task, err := memory.task(NodeKey{0, 0})
 350         require.NoError(t, err)
 351         require.Equal(t, "<chat>\n0+1|"+fakeLine(task)+"\n1+1|user: again\n</chat>", memory.View())
 352 }
 353
 354 func TestMemoryCoversMessagesAfterTheSavedView(t *testing.T) {
 355         root := t.TempDir()
 356         clock := newTestClock()
 357         memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
 358         require.NoError(t, err)
 359
 360         _, err = memory.Append("user", "a")
 361         require.NoError(t, err)
 362
 363         _, err = memory.log.Append("reply", "b")
 364         require.NoError(t, err)
 365         require.NoError(t, memory.Close())
 366
 367         memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
 368         require.NoError(t, err)
 369
 370         defer func() {
 371                 require.NoError(t, memory.Close())
 372         }()
 373
 374         require.Equal(t, "<chat>\n0+1|user: a\n1+1|reply: b\n</chat>", memory.View())
 375 }
 376
 377 func TestMemoryRefusesToRebuildAMissingView(t *testing.T) {
 378         root := t.TempDir()
 379         clock := newTestClock()
 380         memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
 381         require.NoError(t, err)
 382
 383         _, err = memory.Append("user", "a")
 384         require.NoError(t, err)
 385         require.NoError(t, memory.Close())
 386         require.NoError(t, os.Remove(filepath.Join(root, "view.json")))
 387
 388         _, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
 389         require.ErrorIs(t, err, errMemoryNoView)
 390 }
 391
 392 // memoryLongText returns a text that is too long for its own node.
 393 func memoryLongText(n int) string {
 394         return fmt.Sprintf("long %d ", n) + strings.Repeat("x", nodeBytes)
 395 }