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