package main
import (
+ "context"
"fmt"
+ "log"
+ "maps"
"os"
"path/filepath"
+ "slices"
+ "strings"
"testing"
+ "time"
"github.com/stretchr/testify/require"
)
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 {
_, err = memory.Append("user", fmt.Sprint(n))
require.NoError(t, err)
}
view := memory.View()
+ context := memory.Context(0, 100)
saved, err := os.ReadFile(filepath.Join(root, "view.json"))
require.NoError(t, err)
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, view, memory.View())
+ require.Equal(t, context, memory.Context(0, 100))
loaded, err := os.ReadFile(filepath.Join(root, "view.json"))
require.NoError(t, err)
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, newFakeSummarizer(), discardLog())
+ require.NoError(t, err)
+
+ defer func() {
+ require.NoError(t, memory.Close())
+ }()
+
+ memory.chat.high = 200
+ memory.chat.low = 100
+
+ memory.cview.high = 80
+ memory.cview.low = 40
+ for n := range 64 {
+ _, err = memory.Append("user", fmt.Sprint(n%10))
+ require.NoError(t, err)
+
+ coarse := 0
+ for _, line := range memory.chat.lines {
+ outer := memory.cview.lines[coarse]
+ for nodeEnd(outer.level, outer.index) <= line.index<<line.level {
+ coarse++
+ outer = memory.cview.lines[coarse]
+ }
+
+ require.GreaterOrEqual(t, outer.level, line.level, "message %d", n)
+ }
+ }
+}
+
+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, log.New(logs, "", 0))
+ 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, "summary failed, retry at next message: node (0, 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())
- require.Equal(t, "<chat>\n0+1|"+pending+"\n</chat>", memory.View())
+ summarizer.SetFailing(false)
+
+ _, err = memory.Append("user", "again")
+ require.NoError(t, err)
+ require.NoError(t, memory.Wait(t.Context()))
+
+ 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)
+}
+
+// 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)
}
projects / onechat / blobdiff