package main import ( "context" "fmt" "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, 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, newFakeSummarizer(), discardLog()) require.NoError(t, err) defer func() { require.NoError(t, memory.Close()) }() 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<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, "\n0+1|"+left.Text+"\n1+1|"+right.Text+"\n", 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, newFakeSummarizer(), discardLog()) require.NoError(t, err) defer func() { require.NoError(t, memory.Close()) }() for _, text := range []string{"a", "b", "c", "d"} { _, err = memory.Append("user", text) require.NoError(t, err) } 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) { 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", memoryLongText(0)) require.NoError(t, err) require.NoError(t, memory.Wait(t.Context())) require.Equal(t, "\n0+1|"+viewPending+"\n", 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, "\n0+1|"+fakeLine(task)+"\n1+1|user: again\n", memory.View()) } func TestMemoryCoversMessagesAfterTheSavedView(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", "a") require.NoError(t, err) _, err = memory.log.Append("reply", "b") require.NoError(t, err) 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.Equal(t, "\n0+1|user: a\n1+1|reply: b\n", memory.View()) } func TestMemoryRefusesToRebuildAMissingView(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", "a") require.NoError(t, err) require.NoError(t, memory.Close()) require.NoError(t, os.Remove(filepath.Join(root, "view.json"))) _, 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) }