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

   1 package main
   2
   3 import (
   4         "fmt"
   5         "math/rand/v2"
   6         "os"
   7         "path/filepath"
   8         "strings"
   9         "testing"
  10         "time"
  11
  12         "github.com/stretchr/testify/require"
  13 )
  14
  15 // syntheticMessages is the length of the synthetic log.
  16 const syntheticMessages = 100000
  17
  18 // TestReplaySyntheticLog runs a synthetic log through the memory core with the
  19 // fake summarizer. After each message it waits for the builds, as a turn does.
  20 // It checks that each view stays within its sawtooth, that no view line holds a
  21 // whole long message or a part of an unbuilt one, and that a restart in the
  22 // middle and at the end loads the same views byte for byte.
  23 func TestReplaySyntheticLog(t *testing.T) {
  24         start := time.Now()
  25         root := t.TempDir()
  26         clock := newTestClock()
  27         random := rand.New(rand.NewPCG(1, 2))
  28         memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
  29         require.NoError(t, err)
  30
  31         defer func() {
  32                 require.NoError(t, memory.Close())
  33         }()
  34
  35         batches := 0
  36         low := chatHigh
  37         for id := range syntheticMessages {
  38                 if id == syntheticMessages/2 {
  39                         memory = requireRestart(t, memory, root, clock)
  40                 }
  41
  42                 clock.now = clock.now.Add(time.Minute)
  43                 kind, text := syntheticMessage(random, id)
  44                 lines := viewLines(memory)
  45                 _, err = memory.Append(kind, text)
  46                 require.NoError(t, err)
  47
  48                 size, merged := requireViews(t, memory, lines)
  49                 if merged {
  50                         batches++
  51                         low = min(low, size)
  52                 }
  53
  54                 require.NoError(t, memory.Wait(t.Context()))
  55         }
  56
  57         memory = requireRestart(t, memory, root, clock)
  58         require.Equal(t, syntheticMessages, memory.log.Len())
  59         require.Greater(t, batches, 0)
  60
  61         require.NoError(t, memory.Wait(t.Context()))
  62         t.Logf("%d messages, %d nodes, %d batches in %s", syntheticMessages, len(memory.tree.nodes), batches, time.Since(start).Round(time.Second))
  63         t.Logf("chat view: %d lines, %d bytes, lowest after a batch %d bytes", len(memory.chat.lines), memory.chat.Size(memory.tree), low)
  64         t.Logf("compaction view: %d lines, %d bytes", len(memory.cview.lines), memory.cview.Size(memory.tree))
  65 }
  66
  67 // syntheticMessage returns the kind and the text of message id. Most messages
  68 // are tool calls and their output. Users write short lines. One output in a
  69 // hundred is over 30 KB.
  70 func syntheticMessage(random *rand.Rand, id int) (string, string) {
  71         roll := random.IntN(100)
  72         switch {
  73         case roll < 5:
  74                 return "user", syntheticText(id, 10+random.IntN(200))
  75         case roll < 10:
  76                 return "reply", syntheticText(id, 100+random.IntN(1500))
  77         case roll < 12:
  78                 return "note", syntheticText(id, 50+random.IntN(600))
  79         case roll < 13:
  80                 return "work", syntheticText(id, 300+random.IntN(3000))
  81         case roll < 56:
  82                 return "tool", syntheticText(id, 60+random.IntN(400))
  83         case roll < 57:
  84                 return "echo", syntheticText(id, 30000+random.IntN(30000))
  85         default:
  86                 return "echo", syntheticText(id, 20+random.IntN(4000))
  87         }
  88 }
  89
  90 // syntheticText returns a text of about size bytes that starts with its message
  91 // ID and holds some non-ASCII characters.
  92 func syntheticText(id, size int) string {
  93         head := fmt.Sprintf("m%d ", id)
  94         words := "ação data line café "
  95
  96         return head + strings.Repeat(words, max(1, (size-len(head))/len(words)))
  97 }
  98
  99 // requireViews checks the views after a new message. Each view is at most its
 100 // high limit, and a batch merges the chat view to its low limit. The leaf of
 101 // the new message can be built after the append, so the sizes count it as the
 102 // pending mark, as the append did. It returns the size of the chat view and if
 103 // a batch merged it, from the number of lines before the message.
 104 func requireViews(t *testing.T, memory *Memory, lines int) (int, bool) {
 105         t.Helper()
 106
 107         memory.mutex.Lock()
 108         defer memory.mutex.Unlock()
 109
 110         size := appendSize(memory.chat, memory.tree)
 111         require.LessOrEqual(t, size, chatHigh)
 112         require.LessOrEqual(t, appendSize(memory.cview, memory.tree), compactionHigh)
 113
 114         merged := len(memory.chat.lines) < lines
 115         if merged {
 116                 require.LessOrEqual(t, size, chatLow)
 117         }
 118
 119         requireOnlySummaries(t, memory)
 120
 121         return size, merged
 122 }
 123
 124 // requireOnlySummaries checks each line of the chat view as the model reads it.
 125 // A built line holds its node, a long message has a summary and not its text,
 126 // and an unbuilt line holds only the pending mark. The caller holds the lock of
 127 // memory.
 128 func requireOnlySummaries(t *testing.T, memory *Memory) {
 129         t.Helper()
 130
 131         rendered := strings.Split(memory.chat.Render(memory.tree), "\n")
 132         require.Len(t, rendered, len(memory.chat.lines)+2)
 133
 134         for n, line := range memory.chat.lines {
 135                 head := fmt.Sprintf("%d+%d|", line.index<<line.level, 1<<line.level)
 136                 node, built := memory.tree.Node(line.level, line.index)
 137                 if !built {
 138                         require.Equal(t, 0, line.level)
 139                         require.Equal(t, head+pending, rendered[n+1])
 140
 141                         continue
 142                 }
 143
 144                 require.Equal(t, head+strings.ReplaceAll(node.Text, "\n", " "), rendered[n+1])
 145                 if line.level == 0 {
 146                         source := leafSource(memory.log.Message(line.index))
 147                         if len(source) > nodeBytes {
 148                                 require.NotEqual(t, source, node.Text, "message %d", line.index)
 149                         }
 150                 }
 151         }
 152 }
 153
 154 // appendSize returns the size of view with its last line as the pending mark.
 155 func appendSize(view *View, tree nodeReader) int {
 156         last := view.lines[len(view.lines)-1]
 157         return view.Size(tree) - lineSize(tree, last) + len(pending)
 158 }
 159
 160 // viewLines returns the number of lines in the chat view of memory.
 161 func viewLines(memory *Memory) int {
 162         memory.mutex.Lock()
 163         defer memory.mutex.Unlock()
 164
 165         return len(memory.chat.lines)
 166 }
 167
 168 // requireRestart closes memory, opens it again, and checks that the views and
 169 // view.json stay the same byte for byte.
 170 func requireRestart(t *testing.T, memory *Memory, root string, clock *testClock) *Memory {
 171         t.Helper()
 172
 173         require.NoError(t, memory.Wait(t.Context()))
 174
 175         view := memory.View()
 176         context := memory.Context(0, memory.log.Len())
 177         saved, err := os.ReadFile(filepath.Join(root, "view.json"))
 178         require.NoError(t, err)
 179         require.NoError(t, memory.Close())
 180
 181         memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer())
 182         require.NoError(t, err)
 183         require.NoError(t, memory.Wait(t.Context()))
 184         require.Equal(t, view, memory.View())
 185         require.Equal(t, context, memory.Context(0, memory.log.Len()))
 186
 187         loaded, err := os.ReadFile(filepath.Join(root, "view.json"))
 188         require.NoError(t, err)
 189         require.Equal(t, saved, loaded)
 190
 191         return memory
 192 }