package main
import (
+ "context"
+ "errors"
"fmt"
+ "maps"
+ "math/rand/v2"
"os"
"path/filepath"
+ "runtime"
+ "slices"
+ "strings"
"testing"
+ "time"
"github.com/stretchr/testify/require"
)
+const (
+ // replayMessages is the length of the synthetic log.
+ replayMessages = 100000
+
+ // replayShortMessages is the length of the synthetic log with -short.
+ // It is long enough for batches of the chat view.
+ replayShortMessages = 5000
+)
+
func TestMemoryRestartLoadsTheSameView(t *testing.T) {
root := t.TempDir()
clock := newTestClock()
- memory, err := openMemory(t.Context(), root, clock.Now)
+ memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
require.NoError(t, err)
memory.chat.high = 200
memory.chat.low = 100
+
memory.cview.high = 80
memory.cview.low = 40
for n := range 100 {
require.Less(t, len(memory.chat.lines), 100)
require.NoError(t, memory.Close())
- memory, err = openMemory(t.Context(), root, clock.Now)
+ memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
require.NoError(t, err)
defer func() {
require.Equal(t, saved, loaded)
}
+func TestMemoryOpensAfterACrashBeforeTheViewSave(t *testing.T) {
+ root := t.TempDir()
+ clock := newTestClock()
+ memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
+ require.NoError(t, err)
+
+ _, err = memory.Append("user", "first")
+ require.NoError(t, err)
+
+ saved, err := os.ReadFile(filepath.Join(root, "view.json"))
+ require.NoError(t, err)
+
+ _, err = memory.Append("user", "second")
+ require.NoError(t, err)
+
+ view := memory.View()
+ require.NoError(t, memory.Close())
+
+ // The log and the tree hold message 1, but the view does not.
+ require.NoError(t, os.WriteFile(filepath.Join(root, "view.json"), saved, 0o600))
+
+ memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
+ require.NoError(t, err)
+ require.Equal(t, view, memory.View())
+ require.NoError(t, memory.Close())
+}
+
func TestMemoryKeepsTheCompactionViewCoarser(t *testing.T) {
- memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now)
+ memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer(), discardLog())
require.NoError(t, err)
defer func() {
memory.chat.high = 200
memory.chat.low = 100
+
memory.cview.high = 80
memory.cview.low = 40
for n := range 64 {
}
}
+func TestMemoryRunsAtMostEightBuilds(t *testing.T) {
+ summarizer := newFakeSummarizer()
+ summarizer.gate = make(chan struct{})
+ memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
+ require.NoError(t, err)
+
+ defer func() {
+ require.NoError(t, memory.Close())
+ }()
+
+ for n := range 12 {
+ _, err = memory.Append("echo", memoryLongText(n))
+ require.NoError(t, err)
+ }
+
+ memory.mutex.Lock()
+
+ running := slices.Collect(maps.Keys(memory.queue.running))
+ memory.mutex.Unlock()
+ slices.SortFunc(running, func(a, b NodeKey) int {
+ return a.index - b.index
+ })
+
+ require.Equal(t, []NodeKey{{0, 0}, {0, 1}, {0, 2}, {0, 3}, {0, 4}, {0, 5}, {0, 6}, {0, 7}}, running)
+
+ require.Eventually(t, func() bool {
+ return summarizer.Active() == buildMax
+ }, 10*time.Second, time.Millisecond)
+
+ close(summarizer.gate)
+ require.NoError(t, memory.Wait(t.Context()))
+ require.Equal(t, buildMax, summarizer.peak)
+
+ for level, count := range []int{12, 6, 3, 1} {
+ for index := range count {
+ require.True(t, treeIsBuilt(t, memory.tree, level, index), "node (%d, %d)", level, index)
+ }
+ }
+}
+
+func TestMemorySettleWaitsForThePendingLines(t *testing.T) {
+ summarizer := newFakeSummarizer()
+ summarizer.gate = make(chan struct{})
+ memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
+ require.NoError(t, err)
+
+ defer func() {
+ require.NoError(t, memory.Close())
+ }()
+
+ _, err = memory.Append("echo", memoryLongText(0))
+ require.NoError(t, err)
+ require.Contains(t, memory.View(), viewPending)
+
+ settled := make(chan error, 1)
+ go func() {
+ settled <- memory.Settle(t.Context())
+ }()
+
+ select {
+ case err = <-settled:
+ require.Fail(t, "Settle returned before the build", "error: %v", err)
+ case <-time.After(50 * time.Millisecond):
+ }
+
+ close(summarizer.gate)
+ require.NoError(t, <-settled)
+ require.NotContains(t, memory.View(), viewPending)
+}
+
+func TestMemorySettleFailsWhileABuildFailed(t *testing.T) {
+ summarizer := newFakeSummarizer()
+ summarizer.SetFailing(true)
+
+ logs := &SyncBuffer{}
+ memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, daemonTestLogger(logs))
+ require.NoError(t, err)
+
+ defer func() {
+ require.NoError(t, memory.Close())
+ }()
+
+ _, err = memory.Append("echo", memoryLongText(0))
+ require.NoError(t, err)
+ require.ErrorIs(t, memory.Settle(t.Context()), errMemoryUnsettled)
+ require.Equal(t, "<4>summary failed level=0 index=0: fake failure\n", logs.String())
+
+ summarizer.SetFailing(false)
+
+ _, err = memory.Append("user", "retry")
+ require.NoError(t, err)
+ require.NoError(t, memory.Settle(t.Context()))
+}
+
+func TestMemorySettleEndsWithItsContext(t *testing.T) {
+ summarizer := newFakeSummarizer()
+ summarizer.gate = make(chan struct{})
+ memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
+ require.NoError(t, err)
+
+ defer func() {
+ require.NoError(t, memory.Close())
+ }()
+
+ _, err = memory.Append("echo", memoryLongText(0))
+ require.NoError(t, err)
+
+ ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond)
+ defer cancel()
+
+ require.ErrorIs(t, memory.Settle(ctx), context.DeadlineExceeded)
+}
+
+func TestMemoryStartsALeafAfterFewerThanEightUnbuilt(t *testing.T) {
+ summarizer := newFakeSummarizer()
+ summarizer.SetFailing(true)
+
+ memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
+ require.NoError(t, err)
+
+ defer func() {
+ require.NoError(t, memory.Close())
+ }()
+
+ for n := range 9 {
+ _, err = memory.Append("echo", memoryLongText(n))
+ require.NoError(t, err)
+ }
+
+ require.NoError(t, memory.Wait(t.Context()))
+
+ for _, task := range summarizer.Tasks() {
+ require.Less(t, task.Index, 8)
+ }
+
+ require.False(t, treeIsBuilt(t, memory.tree, 0, 8))
+ summarizer.SetFailing(false)
+
+ _, err = memory.Append("user", "retry")
+ require.NoError(t, err)
+ require.NoError(t, memory.Wait(t.Context()))
+ require.True(t, treeIsBuilt(t, memory.tree, 0, 8))
+ require.True(t, treeIsBuilt(t, memory.tree, 3, 0))
+}
+
+func TestMemoryMergesAfterBothHalves(t *testing.T) {
+ summarizer := newFakeSummarizer()
+ memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
+ require.NoError(t, err)
+
+ defer func() {
+ require.NoError(t, memory.Close())
+ }()
+
+ _, err = memory.Append("echo", memoryLongText(0))
+ require.NoError(t, err)
+ require.NoError(t, memory.Wait(t.Context()))
+
+ _, err = memory.Append("echo", memoryLongText(1))
+ require.NoError(t, err)
+ require.NoError(t, memory.Wait(t.Context()))
+ require.Len(t, summarizer.Tasks(), 3)
+
+ left := treeReadNode(t, memory.tree, 0, 0)
+ right := treeReadNode(t, memory.tree, 0, 1)
+ merge := summarizer.Tasks()[2]
+ require.Equal(t, []string{left.Text, right.Text}, merge.Input)
+ require.Equal(t, "<chat>\n0+1|"+left.Text+"\n1+1|"+right.Text+"\n</chat>", merge.Context)
+ require.True(t, treeIsBuilt(t, memory.tree, 1, 0))
+}
+
+func TestMemoryRestartRestoresTheQueue(t *testing.T) {
+ root := t.TempDir()
+ clock := newTestClock()
+ summarizer := newFakeSummarizer()
+ summarizer.SetFailing(true)
+
+ memory, err := openMemory(t.Context(), root, clock.Now, summarizer, discardLog())
+ require.NoError(t, err)
+
+ for n := range 4 {
+ _, err = memory.Append("echo", memoryLongText(n))
+ require.NoError(t, err)
+ }
+
+ require.NoError(t, memory.Wait(t.Context()))
+ require.NoError(t, memory.Close())
+
+ memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
+ require.NoError(t, err)
+
+ defer func() {
+ require.NoError(t, memory.Close())
+ }()
+
+ require.NoError(t, memory.Wait(t.Context()))
+ require.True(t, treeIsBuilt(t, memory.tree, 2, 0))
+}
+
func TestMemoryBuildsFreeParents(t *testing.T) {
- memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now)
+ memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer(), discardLog())
require.NoError(t, err)
defer func() {
require.NoError(t, err)
}
- node, built := memory.tree.Node(2, 0)
- require.True(t, built)
+ node := treeReadNode(t, memory.tree, 2, 0)
require.Equal(t, "user: a\nuser: b\nuser: c\nuser: d", node.Text)
}
func TestMemoryKeepsALongMessageOutOfTheView(t *testing.T) {
- memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now)
+ summarizer := newFakeSummarizer()
+ summarizer.SetFailing(true)
+
+ memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, summarizer, discardLog())
require.NoError(t, err)
defer func() {
require.NoError(t, memory.Close())
}()
- _, err = memory.Append("echo", string(make([]byte, nodeBytes)))
+ _, err = memory.Append("echo", memoryLongText(0))
+ require.NoError(t, err)
+ require.NoError(t, memory.Wait(t.Context()))
+ require.Equal(t, "<chat>\n0+1|"+viewPending+"\n</chat>", memory.View())
+
+ summarizer.SetFailing(false)
+
+ _, err = memory.Append("user", "again")
require.NoError(t, err)
+ require.NoError(t, memory.Wait(t.Context()))
- require.Equal(t, "<chat>\n0+1|"+pending+"\n</chat>", memory.View())
+ task, err := memory.task(NodeKey{0, 0})
+ require.NoError(t, err)
+ require.Equal(t, "<chat>\n0+1|"+fakeLine(task)+"\n1+1|user: again\n</chat>", memory.View())
}
func TestMemoryCoversMessagesAfterTheSavedView(t *testing.T) {
root := t.TempDir()
clock := newTestClock()
- memory, err := openMemory(t.Context(), root, clock.Now)
+ memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
require.NoError(t, err)
_, err = memory.Append("user", "a")
require.NoError(t, err)
require.NoError(t, memory.Close())
- memory, err = openMemory(t.Context(), root, clock.Now)
+ memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
require.NoError(t, err)
defer func() {
func TestMemoryRefusesToRebuildAMissingView(t *testing.T) {
root := t.TempDir()
clock := newTestClock()
- memory, err := openMemory(t.Context(), root, clock.Now)
+ memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
require.NoError(t, err)
_, err = memory.Append("user", "a")
require.NoError(t, memory.Close())
require.NoError(t, os.Remove(filepath.Join(root, "view.json")))
- _, err = openMemory(t.Context(), root, clock.Now)
- require.ErrorIs(t, err, errNoView)
+ _, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
+ require.ErrorIs(t, err, errMemoryNoView)
+}
+
+func TestQueuedAppendResumesWithoutADuplicate(t *testing.T) {
+ dir := t.TempDir()
+ memory, err := openMemory(t.Context(), dir, newTestClock().Now, newFakeSummarizer(), discardLog())
+ require.NoError(t, err)
+
+ var saved *int
+
+ first, err := memory.AppendOnce("user", "Record the gate size.", nil, func(id int) error {
+ saved = &id
+
+ require.Equal(t, 0, memory.log.Len())
+
+ return nil
+ })
+ require.NoError(t, err)
+ require.NoError(t, memory.Close())
+
+ memory, err = openMemory(t.Context(), dir, newTestClock().Now, newFakeSummarizer(), discardLog())
+ require.NoError(t, err)
+
+ defer func() {
+ require.NoError(t, memory.Close())
+ }()
+
+ second, err := memory.AppendOnce("user", first.Text, saved, func(int) error {
+ require.FailNow(t, "a recorded message must not be prepared again")
+
+ return nil
+ })
+ require.NoError(t, err)
+ require.Equal(t, first, second)
+ require.Equal(t, 1, memory.Len())
+}
+
+func TestQueuedAppendRetriesAnUnwrittenPreparedMessage(t *testing.T) {
+ memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer(), discardLog())
+ require.NoError(t, err)
+
+ defer func() {
+ require.NoError(t, memory.Close())
+ }()
+
+ failure := errors.New("fake queue write failure")
+
+ var saved *int
+
+ _, err = memory.AppendOnce("work", "[Garden] Check the gate.", nil, func(id int) error {
+ saved = &id
+
+ return failure
+ })
+ require.ErrorIs(t, err, failure)
+ require.Equal(t, 0, memory.Len())
+
+ message, err := memory.AppendOnce("work", "[Garden] Check the gate.", saved, func(id int) error {
+ require.Equal(t, *saved, id)
+
+ return nil
+ })
+ require.NoError(t, err)
+ require.Equal(t, 0, message.ID)
+ require.Equal(t, 1, memory.Len())
+}
+
+func TestQueuedAppendRefusesAnInvalidReceipt(t *testing.T) {
+ memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now, newFakeSummarizer(), discardLog())
+ require.NoError(t, err)
+
+ defer func() {
+ require.NoError(t, memory.Close())
+ }()
+
+ _, err = memory.Append("user", "Keep the original.")
+ require.NoError(t, err)
+
+ testCases := []struct {
+ name string
+ id int
+ cause string
+ }{
+ {
+ name: "receipt points to different text",
+ id: 0,
+ cause: "does not match",
+ },
+ {
+ name: "receipt skips a log position",
+ id: 2,
+ cause: "beyond the log",
+ },
+ {
+ name: "receipt has a negative position",
+ id: -1,
+ cause: "read queued message",
+ },
+ }
+
+ for _, tc := range testCases {
+ t.Run(tc.name, func(t *testing.T) {
+ _, err := memory.AppendOnce("user", "Other text.", &tc.id, func(int) error {
+ require.FailNow(t, "an invalid receipt must not be prepared")
+
+ return nil
+ })
+ require.ErrorContains(t, err, tc.cause)
+ require.Equal(t, 1, memory.Len())
+ })
+ }
+}
+
+// TestReplaySyntheticLog runs a synthetic log through the memory core with the
+// fake summarizer. After each message it waits for the builds, as a turn does.
+// It checks that each view stays within its sawtooth, that no view line holds a
+// whole long message or a part of an unbuilt one, and that a restart in the
+// middle and at the end loads the same views byte for byte. With -short, the
+// log is shorter, and the test still requires a batch of the chat view.
+func TestReplaySyntheticLog(t *testing.T) {
+ start := time.Now()
+ messages := replayMessages
+ if testing.Short() {
+ messages = replayShortMessages
+ }
+
+ root := t.TempDir()
+ clock := newTestClock()
+ random := rand.New(rand.NewPCG(1, 2))
+ memory, err := openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
+ require.NoError(t, err)
+
+ defer func() {
+ require.NoError(t, memory.Close())
+ }()
+
+ batches := 0
+
+ low := viewChatHigh
+ long := make([]bool, messages)
+ for id := range messages {
+ if id == messages/2 {
+ memory = replayRequireRestart(t, memory, root, clock)
+ }
+
+ clock.now = clock.now.Add(time.Minute)
+ kind, text := replayMessage(random, id)
+ long[id] = len(kind)+len(": ")+len(text) > nodeBytes
+ lines := replayViewLines(memory)
+ _, err = memory.Append(kind, text)
+ require.NoError(t, err)
+
+ size, merged := replayRequireViews(t, memory, lines, long, id%1000 == 0)
+ if merged {
+ batches++
+ low = min(low, size)
+ }
+
+ require.NoError(t, memory.Wait(t.Context()))
+ }
+
+ memory = replayRequireRestart(t, memory, root, clock)
+ require.Equal(t, messages, memory.log.Len())
+ require.Positive(t, batches)
+
+ require.NoError(t, memory.Wait(t.Context()))
+ t.Logf("%d messages, %d nodes, %d batches in %s", messages, memory.tree.Count(), batches, time.Since(start).Round(time.Second))
+ t.Logf("chat view: %d lines, %d bytes, lowest after a batch %d bytes", len(memory.chat.lines), memory.chat.Size(), low)
+ t.Logf("compaction view: %d lines, %d bytes", len(memory.cview.lines), memory.cview.Size())
+
+ runtime.GC()
+
+ var stats runtime.MemStats
+
+ runtime.ReadMemStats(&stats)
+ t.Logf("heap in use: %d KB", stats.HeapInuse/1024)
+}
+
+// memoryLongText returns a text that is too long for its own node.
+func memoryLongText(n int) string {
+ return fmt.Sprintf("long %d ", n) + strings.Repeat("x", nodeBytes)
+}
+
+// replayRequireRestart closes memory, opens it again, and checks that the views
+// and view.json stay the same byte for byte.
+func replayRequireRestart(t *testing.T, memory *Memory, root string, clock *TestClock) *Memory {
+ t.Helper()
+
+ require.NoError(t, memory.Wait(t.Context()))
+
+ view := memory.View()
+ context := memory.Context(0, memory.log.Len())
+ saved, err := os.ReadFile(filepath.Join(root, "view.json"))
+ require.NoError(t, err)
+ require.NoError(t, memory.Close())
+
+ memory, err = openMemory(t.Context(), root, clock.Now, newFakeSummarizer(), discardLog())
+ require.NoError(t, err)
+ require.NoError(t, memory.Wait(t.Context()))
+ require.Equal(t, view, memory.View())
+ require.Equal(t, context, memory.Context(0, memory.log.Len()))
+
+ loaded, err := os.ReadFile(filepath.Join(root, "view.json"))
+ require.NoError(t, err)
+ require.Equal(t, saved, loaded)
+
+ return memory
+}
+
+// replayMessage returns the kind and the text of message id. Most messages are
+// tool calls and their output. Users write short lines. One output in a hundred
+// is over 30 KB.
+func replayMessage(random *rand.Rand, id int) (string, string) {
+ roll := random.IntN(100)
+ switch {
+ case roll < 5:
+ return "user", replayText(id, 10+random.IntN(200))
+ case roll < 10:
+ return "reply", replayText(id, 100+random.IntN(1500))
+ case roll < 12:
+ return "note", replayText(id, 50+random.IntN(600))
+ case roll < 13:
+ return "work", replayText(id, 300+random.IntN(3000))
+ case roll < 56:
+ return "tool", replayText(id, 60+random.IntN(400))
+ case roll < 57:
+ return "echo", replayText(id, 30000+random.IntN(30000))
+ default:
+ return "echo", replayText(id, 20+random.IntN(4000))
+ }
+}
+
+// replayViewLines returns the number of lines in the chat view of memory.
+func replayViewLines(memory *Memory) int {
+ memory.mutex.Lock()
+ defer memory.mutex.Unlock()
+
+ return len(memory.chat.lines)
+}
+
+// replayRequireViews checks the views after a new message, and with full
+// compares them with the tree. Each view is at most its high limit, and a batch
+// merges the chat view to its low limit. The leaf of the new message can be
+// built after the append, so the sizes count it as the pending mark, as the
+// append did. It returns the size of the chat view and if a batch merged it,
+// from the number of lines before the message.
+func replayRequireViews(t *testing.T, memory *Memory, lines int, long []bool, full bool) (int, bool) {
+ t.Helper()
+
+ memory.mutex.Lock()
+ defer memory.mutex.Unlock()
+
+ size := replayAppendSize(memory.chat)
+ require.LessOrEqual(t, size, viewChatHigh)
+ require.LessOrEqual(t, replayAppendSize(memory.cview), viewCompactionHigh)
+
+ merged := len(memory.chat.lines) < lines
+ if merged {
+ require.LessOrEqual(t, size, viewChatLow)
+ }
+
+ replayRequireOnlySummaries(t, memory, long, full)
+
+ return size, merged
+}
+
+// replayText returns a text of about size bytes that starts with its message ID
+// and holds some non-ASCII characters.
+func replayText(id, size int) string {
+ head := fmt.Sprintf("m%d ", id)
+ words := "ação data line café "
+
+ return head + strings.Repeat(words, max(1, (size-len(head))/len(words)))
+}
+
+// replayAppendSize returns the size of view with its last line as the pending
+// mark.
+func replayAppendSize(view *View) int {
+ last := view.lines[len(view.lines)-1]
+ return view.Size() - view.lineSize(last) + len(viewPending)
+}
+
+// replayRequireOnlySummaries checks each line of the chat view as the model
+// reads it. A built line holds its node, a long message has a summary of the
+// fake summarizer and not its text, and an unbuilt line holds only the pending
+// mark. With full, it also checks that the view keeps the node of each built
+// line. The caller holds the lock of memory.
+func replayRequireOnlySummaries(t *testing.T, memory *Memory, long []bool, full bool) {
+ t.Helper()
+
+ rendered := strings.Split(memory.chat.Render(), "\n")
+ require.Len(t, rendered, len(memory.chat.lines)+2)
+
+ for n, line := range memory.chat.lines {
+ head := fmt.Sprintf("%d+%d|", line.index<<line.level, 1<<line.level)
+ node, built := memory.chat.nodes[line]
+ if full {
+ require.Equal(t, treeIsBuilt(t, memory.tree, line.level, line.index), built)
+ }
+
+ if !built {
+ require.Equal(t, 0, line.level)
+ require.Equal(t, head+viewPending, rendered[n+1])
+
+ continue
+ }
+
+ require.Equal(t, head+strings.ReplaceAll(node.Text, "\n", " "), rendered[n+1])
+
+ if line.level == 0 && long[line.index] {
+ require.True(t, strings.HasPrefix(node.Text, "summary "), "message %d", line.index)
+ }
+ }
}
projects / onechat / blobdiff