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 }