--- /dev/null
+package main
+
+import (
+ "context"
+ "encoding/json"
+ "errors"
+ "fmt"
+ "os"
+ "path/filepath"
+ "sync"
+ "time"
+)
+
+// errNoView tells that a chat has messages but no saved view.
+var errNoView = errors.New("the chat has messages but no saved view")
+
+// Memory is the memory of one chat: the message log, the summary tree, and the
+// view. It saves the view to view.json after each change and loads it at start.
+// It never builds the view again from the log, because a new view differs from
+// the live one and so loses the prompt cache.
+type Memory struct {
+ mutex sync.Mutex
+ root string
+ log *Log
+ tree *Tree
+ chat *View
+}
+
+// viewFile is the content of view.json.
+type viewFile struct {
+ Chat viewState `json:"chat"`
+}
+
+// openMemory opens the chat in root. The clock gives the dates of new messages
+// and the days of new nodes.
+func openMemory(ctx context.Context, root string, clock func() time.Time) (*Memory, error) {
+ log, err := openLog(ctx, root, clock)
+ if err != nil {
+ return nil, err
+ }
+
+ m := &Memory{
+ root: root,
+ log: log,
+ }
+
+ m.tree, err = openTree(ctx, root, log.Len(), clock)
+ if err != nil {
+ return nil, errors.Join(err, m.Close())
+ }
+
+ err = m.load()
+ if err != nil {
+ return nil, errors.Join(fmt.Errorf("load view: %w", err), m.Close())
+ }
+
+ return m, nil
+}
+
+// load reads the saved view. A crash can stop an append after the log write and
+// before the view save, so load continues the view with the messages that it
+// does not cover yet.
+func (m *Memory) load() error {
+ data, err := os.ReadFile(m.viewPath())
+ if errors.Is(err, os.ErrNotExist) && m.log.Len() == 0 {
+ m.chat = newView(chatHigh, chatLow)
+ return nil
+ }
+
+ if errors.Is(err, os.ErrNotExist) {
+ return errNoView
+ }
+
+ if err != nil {
+ return fmt.Errorf("read view: %w", err)
+ }
+
+ var file viewFile
+ err = json.Unmarshal(data, &file)
+ if err != nil {
+ return fmt.Errorf("decode view: %w", err)
+ }
+
+ m.chat, err = restoreView(file.Chat, chatHigh, chatLow, m.tree)
+ if err != nil {
+ return err
+ }
+
+ if m.chat.end > m.log.Len() {
+ return fmt.Errorf("view covers %d messages, but the log has %d", m.chat.end, m.log.Len())
+ }
+
+ if m.chat.end == m.log.Len() {
+ return nil
+ }
+
+ for m.chat.end < m.log.Len() {
+ err = m.addMessage(m.log.Message(m.chat.end))
+ if err != nil {
+ return err
+ }
+ }
+
+ return m.save()
+}
+
+// Close closes the files of the chat and releases its lock.
+func (m *Memory) Close() error {
+ m.mutex.Lock()
+ defer m.mutex.Unlock()
+
+ var err error
+ if m.tree != nil {
+ err = m.tree.Close()
+ }
+
+ return errors.Join(err, m.log.Close())
+}
+
+// Append logs a new message and adds it to the tree and the view.
+func (m *Memory) Append(kind, text string) (Message, error) {
+ m.mutex.Lock()
+ defer m.mutex.Unlock()
+
+ message, err := m.log.Append(kind, text)
+ if err != nil {
+ return Message{}, err
+ }
+
+ err = m.addMessage(message)
+ if err != nil {
+ return Message{}, err
+ }
+
+ err = m.save()
+ if err != nil {
+ return Message{}, err
+ }
+
+ return message, nil
+}
+
+// addMessage builds the leaf of message when it needs no model call and appends
+// it to the view.
+func (m *Memory) addMessage(message Message) error {
+ _, built := m.tree.Node(0, message.ID)
+ leaf, free := freeLeaf(message)
+ if free && !built {
+ err := m.build(leaf)
+ if err != nil {
+ return err
+ }
+ }
+
+ m.chat.Append(m.tree)
+
+ return nil
+}
+
+// build adds node to the tree, and then each parent above it that needs no
+// model call.
+func (m *Memory) build(node Node) error {
+ for {
+ err := m.tree.Add(node)
+ if err != nil {
+ return err
+ }
+
+ sibling, built := m.tree.Node(node.Level, node.Index^1)
+ if !built {
+ return nil
+ }
+
+ left, right := node, sibling
+ if node.Index%2 != 0 {
+ left, right = sibling, node
+ }
+
+ parent, free := freeMerge(left, right)
+ if !free {
+ return nil
+ }
+
+ node = parent
+ }
+}
+
+// save writes the views to view.json.
+func (m *Memory) save() error {
+ data, err := json.Marshal(viewFile{
+ Chat: m.chat.State(),
+ })
+ if err != nil {
+ return fmt.Errorf("encode view: %w", err)
+ }
+
+ err = writeFileAtomic(m.viewPath(), append(data, '\n'))
+ if err != nil {
+ return fmt.Errorf("save view: %w", err)
+ }
+
+ return nil
+}
+
+func (m *Memory) viewPath() string {
+ return filepath.Join(m.root, "view.json")
+}
+
+// View returns the chat view as the model reads it.
+func (m *Memory) View() string {
+ m.mutex.Lock()
+ defer m.mutex.Unlock()
+
+ return m.chat.Render(m.tree)
+}
+
+// writeFileAtomic replaces the file at path with data. A reader or a crash sees
+// the old file or the new file, never a part of one.
+func writeFileAtomic(path string, data []byte) error {
+ dir := filepath.Dir(path)
+ file, err := os.CreateTemp(dir, "."+filepath.Base(path)+".*")
+ if err != nil {
+ return fmt.Errorf("create temporary file: %w", err)
+ }
+
+ _, err = file.Write(data)
+ if err == nil {
+ err = file.Sync()
+ }
+
+ err = errors.Join(err, file.Close())
+ if err == nil {
+ err = os.Rename(file.Name(), path)
+ }
+
+ if err != nil {
+ return errors.Join(fmt.Errorf("replace file: %w", err), os.Remove(file.Name()))
+ }
+
+ return syncDir(dir)
+}
+
+// syncDir flushes the entries of dir to the disk.
+func syncDir(dir string) (err error) {
+ handle, err := os.Open(dir)
+ if err != nil {
+ return fmt.Errorf("open directory: %w", err)
+ }
+
+ defer func() {
+ err = errors.Join(err, handle.Close())
+ }()
+
+ err = handle.Sync()
+ if err != nil {
+ return fmt.Errorf("flush directory: %w", err)
+ }
+
+ return nil
+}
--- /dev/null
+package main
+
+import (
+ "fmt"
+ "os"
+ "path/filepath"
+ "testing"
+
+ "github.com/stretchr/testify/require"
+)
+
+func TestMemoryRestartLoadsTheSameView(t *testing.T) {
+ root := t.TempDir()
+ clock := newTestClock()
+ memory, err := openMemory(t.Context(), root, clock.Now)
+ require.NoError(t, err)
+
+ memory.chat.high = 200
+ memory.chat.low = 100
+ for n := range 100 {
+ _, err = memory.Append("user", fmt.Sprint(n))
+ require.NoError(t, err)
+ }
+
+ view := memory.View()
+ 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)
+ require.NoError(t, err)
+
+ defer func() {
+ require.NoError(t, memory.Close())
+ }()
+
+ require.Equal(t, view, memory.View())
+
+ loaded, err := os.ReadFile(filepath.Join(root, "view.json"))
+ require.NoError(t, err)
+ require.Equal(t, saved, loaded)
+}
+
+func TestMemoryBuildsFreeParents(t *testing.T) {
+ memory, err := openMemory(t.Context(), t.TempDir(), newTestClock().Now)
+ 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, built := memory.tree.Node(2, 0)
+ require.True(t, built)
+ 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)
+ require.NoError(t, err)
+
+ defer func() {
+ require.NoError(t, memory.Close())
+ }()
+
+ _, err = memory.Append("echo", string(make([]byte, nodeBytes)))
+ require.NoError(t, err)
+
+ require.Equal(t, "<chat>\n0+1|"+pending+"\n</chat>", memory.View())
+}
+
+func TestMemoryCoversMessagesAfterTheSavedView(t *testing.T) {
+ root := t.TempDir()
+ clock := newTestClock()
+ memory, err := openMemory(t.Context(), root, clock.Now)
+ 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)
+ require.NoError(t, err)
+
+ defer func() {
+ require.NoError(t, memory.Close())
+ }()
+
+ require.Equal(t, "<chat>\n0+1|user: a\n1+1|reply: b\n</chat>", memory.View())
+}
+
+func TestMemoryRefusesToRebuildAMissingView(t *testing.T) {
+ root := t.TempDir()
+ clock := newTestClock()
+ memory, err := openMemory(t.Context(), root, clock.Now)
+ 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)
+ require.ErrorIs(t, err, errNoView)
+}
--- /dev/null
+package main
+
+import (
+ "fmt"
+ "math"
+ "math/bits"
+ "strconv"
+ "strings"
+)
+
+const (
+ // chatHigh and chatLow are the limits of the chat view in bytes.
+ chatHigh = 128000
+ chatLow = 64000
+
+ // pending is the text of a line that is not built yet.
+ pending = "(not summarized yet: zoom it)"
+)
+
+// nodeReader gives the built nodes of a tree.
+type nodeReader interface {
+ Node(level, index int) (Node, bool)
+}
+
+// View is a list of tree nodes that covers messages 0 to end-1 without a gap,
+// oldest first. Each new message appends its leaf as a line. Once the size of
+// the view goes above high, one batch merges the most due sibling pairs until
+// the size is at most low. A batch merges only pairs whose parent is built. If
+// the batch cannot get to low, it continues at each new message until it does.
+//
+// The size of a line is the size of its node text, or the size of the pending
+// text when the node is not built.
+//
+// For the sibling pair (l, i) and (l, i+1), due is (T - last) / 2^l, with T the
+// number of messages and last the ID of the last message of the pair. That is
+// how long ago the pair ended, measured in its own line size. With a budget of
+// the length of Taelin's rollback push list, this order makes the same merges
+// as the push. The oldest pair wins a tie.
+type View struct {
+ high int
+ low int
+ lines []nodeKey
+ end int
+ merging bool
+}
+
+// viewState is the saved form of a View. The limits are not part of it.
+type viewState struct {
+ Lines [][2]int `json:"lines"`
+ Merging bool `json:"merging"`
+}
+
+// newView returns an empty view with the limits high and low.
+func newView(high, low int) *View {
+ return &View{
+ high: high,
+ low: low,
+ }
+}
+
+// restoreView returns the view of state with the limits high and low. Each line
+// must continue the one before it, and each line above level 0 must be built.
+func restoreView(state viewState, high, low int, tree nodeReader) (*View, error) {
+ v := newView(high, low)
+ v.merging = state.Merging
+ for _, pair := range state.Lines {
+ line := nodeKey{pair[0], pair[1]}
+ if line.level < 0 || line.level >= bits.UintSize-2 || line.index < 0 || line.index<<line.level != v.end {
+ return nil, fmt.Errorf("view line (%d, %d) does not start at message %d", line.level, line.index, v.end)
+ }
+
+ _, built := tree.Node(line.level, line.index)
+ if line.level > 0 && !built {
+ return nil, fmt.Errorf("view line (%d, %d) is not built", line.level, line.index)
+ }
+
+ v.lines = append(v.lines, line)
+ v.end = nodeEnd(line.level, line.index)
+ }
+
+ return v, nil
+}
+
+// State returns the saved form of the view.
+func (v *View) State() viewState {
+ state := viewState{
+ Lines: make([][2]int, len(v.lines)),
+ Merging: v.merging,
+ }
+
+ for n, line := range v.lines {
+ state.Lines[n] = [2]int{line.level, line.index}
+ }
+
+ return state
+}
+
+// Append adds the leaf of the next message as a line and merges when the view
+// is too large.
+func (v *View) Append(tree nodeReader) {
+ v.lines = append(v.lines, nodeKey{0, v.end})
+ v.end++
+ v.fit(tree)
+}
+
+func (v *View) fit(tree nodeReader) {
+ size := v.Size(tree)
+ if size > v.high {
+ v.merging = true
+ }
+
+ for v.merging && size > v.low {
+ best := v.mostDue(tree)
+ if best < 0 {
+ return
+ }
+
+ left, right := v.lines[best], v.lines[best+1]
+ parent := nodeKey{left.level + 1, left.index / 2}
+ size += lineSize(tree, parent) - lineSize(tree, left) - lineSize(tree, right)
+ v.lines[best] = parent
+ v.lines = append(v.lines[:best+1], v.lines[best+2:]...)
+ }
+
+ v.merging = false
+}
+
+// mostDue returns the position of the left line of the most due sibling pair
+// whose parent is built, or -1 when there is no such pair.
+func (v *View) mostDue(tree nodeReader) int {
+ best := -1
+ bestDue := 0.0
+ for n := 0; n+1 < len(v.lines); n++ {
+ left, right := v.lines[n], v.lines[n+1]
+ if left.index%2 != 0 || right.level != left.level || right.index != left.index+1 {
+ continue
+ }
+
+ _, built := tree.Node(left.level+1, left.index/2)
+ if !built {
+ continue
+ }
+
+ last := nodeEnd(right.level, right.index) - 1
+ due := math.Ldexp(float64(v.end-last), -left.level)
+ if best < 0 || due > bestDue {
+ best = n
+ bestDue = due
+ }
+ }
+
+ return best
+}
+
+// Size returns the sum of the line sizes.
+func (v *View) Size(tree nodeReader) int {
+ size := 0
+ for _, line := range v.lines {
+ size += lineSize(tree, line)
+ }
+
+ return size
+}
+
+// Render returns the view as the model reads it: one line "id+n|text" for each
+// node inside <chat> tags, with the newlines of the text as spaces.
+func (v *View) Render(tree nodeReader) string {
+ var text strings.Builder
+ text.WriteString("<chat>\n")
+ for _, line := range v.lines {
+ renderLine(&text, tree, line)
+ }
+
+ text.WriteString("</chat>")
+
+ return text.String()
+}
+
+// lineSize returns the size of a view line.
+func lineSize(tree nodeReader, line nodeKey) int {
+ node, built := tree.Node(line.level, line.index)
+ if !built {
+ return len(pending)
+ }
+
+ return node.Size
+}
+
+// renderLine writes one view line, "id+n|text", and a newline.
+func renderLine(text *strings.Builder, tree nodeReader, line nodeKey) {
+ body := pending
+ node, built := tree.Node(line.level, line.index)
+ if built {
+ body = strings.ReplaceAll(node.Text, "\n", " ")
+ }
+
+ text.WriteString(strconv.Itoa(line.index << line.level))
+ text.WriteByte('+')
+ text.WriteString(strconv.Itoa(1 << line.level))
+ text.WriteByte('|')
+ text.WriteString(body)
+ text.WriteByte('\n')
+}
--- /dev/null
+package main
+
+import (
+ "slices"
+ "strings"
+ "testing"
+
+ "github.com/stretchr/testify/require"
+)
+
+func TestViewMergesInTheOrderOfTaelinsPush(t *testing.T) {
+ tree := builtTree{}
+ view := newView(0, 0)
+ var states *pushEntry
+ for tick := range 20001 {
+ states = push(tick, states)
+ starts := states.starts()
+ view.high = len(starts)
+ view.low = len(starts)
+ view.Append(tree)
+
+ lines := make([]int, len(view.lines))
+ for n, line := range view.lines {
+ lines[n] = line.index << line.level
+ }
+
+ require.Equal(t, starts, lines, "t=%d", tick)
+ }
+}
+
+func TestViewMergesInOneBatchFromHighToLow(t *testing.T) {
+ tree := builtTree{}
+ view := newView(10, 5)
+ sizes := make([]int, 0, 20)
+ for range 20 {
+ view.Append(tree)
+ sizes = append(sizes, view.Size(tree))
+ }
+
+ require.Equal(t, []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 5, 6, 7, 8, 9, 10, 5, 6, 7, 8}, sizes)
+}
+
+func TestViewMergesOnlyBuiltParents(t *testing.T) {
+ tree := partTree{
+ sizes: map[nodeKey]int{
+ {0, 0}: 1,
+ {0, 1}: 1,
+ {0, 2}: 1,
+ {0, 3}: 1,
+ {0, 4}: 1,
+ {0, 5}: 1,
+ },
+ }
+
+ view := newView(3, 2)
+ for range 4 {
+ view.Append(tree)
+ }
+
+ require.Equal(t, []nodeKey{{0, 0}, {0, 1}, {0, 2}, {0, 3}}, view.lines)
+ require.True(t, view.merging)
+
+ tree.sizes[nodeKey{1, 0}] = 1
+ view.Append(tree)
+ require.Equal(t, []nodeKey{{1, 0}, {0, 2}, {0, 3}, {0, 4}}, view.lines)
+ require.True(t, view.merging)
+
+ tree.sizes[nodeKey{1, 1}] = 1
+ tree.sizes[nodeKey{1, 2}] = 1
+ tree.sizes[nodeKey{2, 0}] = 1
+ view.Append(tree)
+ require.Equal(t, []nodeKey{{2, 0}, {1, 2}}, view.lines)
+ require.False(t, view.merging)
+}
+
+func TestViewRendersPendingLinesWithoutTheirMessage(t *testing.T) {
+ tree := partTree{
+ sizes: map[nodeKey]int{
+ {0, 0}: 3,
+ {0, 1}: 3,
+ {1, 0}: 7,
+ },
+ }
+
+ view, err := restoreView(viewState{
+ Lines: [][2]int{{1, 0}, {0, 2}},
+ }, chatHigh, chatLow, tree)
+ require.NoError(t, err)
+ require.Equal(t, "<chat>\n0+2|aaa aaa\n2+1|"+pending+"\n</chat>", view.Render(tree))
+ require.Equal(t, 7+len(pending), view.Size(tree))
+}
+
+func TestRestoreViewRejectsAGap(t *testing.T) {
+ _, err := restoreView(viewState{
+ Lines: [][2]int{{0, 0}, {0, 2}},
+ }, chatHigh, chatLow, builtTree{})
+ require.ErrorContains(t, err, "does not start at message 1")
+}
+
+// pushEntry is one entry of Taelin's rollback state list, newest first. Under
+// push alone, life stays 0, so the entry leaves it out.
+type pushEntry struct {
+ keep bool
+ state int
+ older *pushEntry
+}
+
+// push is the push of rollback_state_list.js.
+func push(state int, states *pushEntry) *pushEntry {
+ if states == nil {
+ return &pushEntry{
+ state: state,
+ }
+ }
+
+ if !states.keep {
+ return &pushEntry{
+ keep: true,
+ state: states.state,
+ older: states.older,
+ }
+ }
+
+ return &pushEntry{
+ state: state,
+ older: push(states.state, states.older),
+ }
+}
+
+// starts returns the states of the list, oldest first.
+func (e *pushEntry) starts() []int {
+ var starts []int
+ for entry := e; entry != nil; entry = entry.older {
+ starts = append(starts, entry.state)
+ }
+
+ slices.Reverse(starts)
+
+ return starts
+}
+
+// builtTree is a tree with each node built, with 1 byte of text.
+type builtTree struct{}
+
+func (builtTree) Node(level, index int) (Node, bool) {
+ return testNode(level, index, "a"), true
+}
+
+// partTree is a tree with the built nodes of its sizes.
+type partTree struct {
+ sizes map[nodeKey]int
+}
+
+func (p partTree) Node(level, index int) (Node, bool) {
+ size, built := p.sizes[nodeKey{level, index}]
+ if !built {
+ return Node{}, false
+ }
+
+ text := strings.Repeat("a", size)
+ if size == 7 {
+ text = "aaa\naaa"
+ }
+
+ return testNode(level, index, text), true
+}
projects / onechat / commitdiff