Agent-Native/runtime/session/tree.go

package session

import (
	"bufio"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"strings"
	"time"

	"agent-native/runtime/llm"
)

// EntryType is the kind of one JSONL record.
type EntryType string

const (
	TypeHeader              EntryType = "header"
	TypeMessage             EntryType = "message"
	TypeModelChange         EntryType = "model_change"
	TypeThinkingLevelChange EntryType = "thinking_level_change"
	TypeCompaction          EntryType = "compaction"
	TypeLabel               EntryType = "label"
	TypeCustom              EntryType = "custom"
	TypeSessionInfo         EntryType = "session_info"
)

// Entry is one node in the session tree. Files written before the tree
// format (a bare llm.Message per line) are loaded as message entries.
type Entry struct {
	ID            string          `json:"id"`
	ParentID      string          `json:"parentId,omitempty"`
	Timestamp     string          `json:"timestamp"`
	Type          EntryType       `json:"type"`
	Message       *llm.Message    `json:"message,omitempty"`
	Model         string          `json:"model,omitempty"`
	ThinkingLevel string          `json:"thinkingLevel,omitempty"`
	Label         string          `json:"label,omitempty"`
	Name          string          `json:"name,omitempty"`
	CustomType    string          `json:"customType,omitempty"`
	Data          json.RawMessage `json:"data,omitempty"`
	Summary       string          `json:"summary,omitempty"`
	TailIDs       []string        `json:"tailIds,omitempty"`
	CWD           string          `json:"cwd,omitempty"`
}

// Entries returns every record in file order, including branches that are
// not on the active path.
func (s *Store) Entries() ([]Entry, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.readEntriesLocked()
}

// Leaf is the id of the current branch tip. Empty when the file has no entries.
func (s *Store) Leaf() (string, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	entries, err := s.readEntriesLocked()
	if err != nil {
		return "", err
	}
	return s.leafLocked(entries), nil
}

// SetLeaf moves the active branch to an existing entry.
func (s *Store) SetLeaf(id string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	entries, err := s.readEntriesLocked()
	if err != nil {
		return err
	}
	if !hasID(entries, id) {
		return fmt.Errorf("session: unknown entry %s", id)
	}
	return s.writeMetaKeyLocked("leaf", id)
}

// AppendEntry adds a non-message record (label, model change, custom, …)
// as a child of the current leaf and moves the leaf to it.
func (s *Store) AppendEntry(e Entry) (Entry, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.appendEntryLocked(e)
}

// Fork copies the tree into a new file and sets that file's leaf to atID.
// The source file is not modified. atID empty means the current leaf.
func (s *Store) Fork(dest, atID string) (*Store, error) {
	return s.copyTo(dest, atID, true)
}

// Clone copies the tree and keeps the same leaf.
func (s *Store) Clone(dest string) (*Store, error) {
	return s.copyTo(dest, "", false)
}

func (s *Store) copyTo(dest, atID string, moveLeaf bool) (*Store, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	entries, err := s.readEntriesLocked()
	if err != nil {
		return nil, err
	}
	if err := copyFile(s.path, dest); err != nil {
		return nil, err
	}
	out, err := Open(dest)
	if err != nil {
		return nil, err
	}
	meta := s.readMetaLocked()
	if meta == nil {
		meta = map[string]string{}
	}
	leaf := s.leafLocked(entries)
	if moveLeaf && atID != "" {
		if !hasID(entries, atID) {
			return nil, fmt.Errorf("session: unknown entry %s", atID)
		}
		leaf = atID
	}
	if leaf != "" {
		meta["leaf"] = leaf
	}
	if err := out.SaveMeta(meta); err != nil {
		return nil, err
	}
	return out, nil
}

// SetName records a human name for continue/resume.
func (s *Store) SetName(name string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if err := s.writeMetaKeyLocked("name", name); err != nil {
		return err
	}
	_, err := s.appendEntryLocked(Entry{Type: TypeSessionInfo, Name: name})
	return err
}

// Name returns the session name, or empty.
func (s *Store) Name() string {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.readMetaLocked()["name"]
}

func (s *Store) appendEntryLocked(e Entry) (Entry, error) {
	entries, err := s.readEntriesLocked()
	if err != nil {
		return Entry{}, err
	}
	if e.ID == "" {
		e.ID = newID()
	}
	if e.Timestamp == "" {
		e.Timestamp = time.Now().UTC().Format(time.RFC3339Nano)
	}
	if e.Type == "" {
		e.Type = TypeCustom
	}
	if e.ParentID == "" {
		e.ParentID = s.leafLocked(entries)
	}
	raw, err := json.Marshal(e)
	if err != nil {
		return Entry{}, err
	}
	f, err := os.OpenFile(s.path, os.O_APPEND|os.O_WRONLY, 0o644)
	if err != nil {
		return Entry{}, err
	}
	defer f.Close()
	if _, err := f.Write(append(raw, '\n')); err != nil {
		return Entry{}, err
	}
	if err := s.writeMetaKeyLocked("leaf", e.ID); err != nil {
		return Entry{}, err
	}
	return e, nil
}

func (s *Store) readEntriesLocked() ([]Entry, error) {
	raw, err := os.ReadFile(s.path)
	if err != nil {
		return nil, err
	}
	text := string(raw)
	if text != "" && !strings.HasSuffix(text, "\n") {
		if idx := strings.LastIndexByte(text, '\n'); idx >= 0 {
			if !json.Valid([]byte(text[idx+1:])) {
				text = text[:idx+1]
			}
		} else if !json.Valid([]byte(text)) {
			text = ""
		}
	}
	var out []Entry
	sc := bufio.NewScanner(strings.NewReader(text))
	sc.Buffer(make([]byte, 0, 64*1024), 8*1024*1024)
	legacyParent := ""
	n := 0
	for sc.Scan() {
		line := sc.Bytes()
		if len(line) == 0 {
			continue
		}
		var probe map[string]json.RawMessage
		if err := json.Unmarshal(line, &probe); err != nil {
			return nil, err
		}
		if _, ok := probe["type"]; ok {
			var e Entry
			if err := json.Unmarshal(line, &e); err != nil {
				return nil, err
			}
			out = append(out, e)
			legacyParent = e.ID
			n++
			continue
		}
		var msg llm.Message
		if err := json.Unmarshal(line, &msg); err != nil {
			return nil, err
		}
		id := fmt.Sprintf("legacy-%d", n)
		out = append(out, Entry{
			ID: id, ParentID: legacyParent, Type: TypeMessage, Message: &msg,
		})
		legacyParent = id
		n++
	}
	return out, sc.Err()
}

func (s *Store) leafLocked(entries []Entry) string {
	if leaf := s.readMetaLocked()["leaf"]; leaf != "" && hasID(entries, leaf) {
		return leaf
	}
	if len(entries) == 0 {
		return ""
	}
	return entries[len(entries)-1].ID
}

func visibleMessages(entries []Entry, leaf string) []llm.Message {
	path := pathTo(entries, leaf)
	var last *Entry
	for i := range path {
		if path[i].Type == TypeCompaction {
			last = &path[i]
		}
	}
	if last == nil {
		var msgs []llm.Message
		for _, e := range path {
			if e.Type == TypeMessage && e.Message != nil {
				msgs = append(msgs, *e.Message)
			}
		}
		return msgs
	}
	byID := map[string]Entry{}
	for _, e := range entries {
		byID[e.ID] = e
	}
	var msgs []llm.Message
	for _, e := range path {
		if e.ID == last.ID {
			break
		}
		if e.Type == TypeMessage && e.Message != nil && e.Message.Role == llm.RoleSystem {
			msgs = append(msgs, *e.Message)
		}
	}
	msgs = append(msgs, llm.Message{
		Role:    llm.RoleUser,
		Content: "[conversation summary]\n" + last.Summary + "\n[recent messages follow]",
	})
	for _, id := range last.TailIDs {
		e, ok := byID[id]
		if ok && e.Message != nil {
			msgs = append(msgs, *e.Message)
		}
	}
	after := false
	for _, e := range path {
		if e.ID == last.ID {
			after = true
			continue
		}
		if after && e.Type == TypeMessage && e.Message != nil {
			msgs = append(msgs, *e.Message)
		}
	}
	return msgs
}

func pathTo(entries []Entry, leaf string) []Entry {
	byID := map[string]Entry{}
	for _, e := range entries {
		byID[e.ID] = e
	}
	var rev []Entry
	seen := map[string]bool{}
	for id := leaf; id != "" && !seen[id]; {
		e, ok := byID[id]
		if !ok {
			break
		}
		seen[id] = true
		rev = append(rev, e)
		id = e.ParentID
	}
	for i, j := 0, len(rev)-1; i < j; i, j = i+1, j-1 {
		rev[i], rev[j] = rev[j], rev[i]
	}
	return rev
}

func hasID(entries []Entry, id string) bool {
	for _, e := range entries {
		if e.ID == id {
			return true
		}
	}
	return false
}

func (s *Store) readMetaLocked() map[string]string {
	raw, err := os.ReadFile(s.path + ".meta.json")
	if err != nil {
		return nil
	}
	var meta map[string]string
	if err := json.Unmarshal(raw, &meta); err != nil {
		return nil
	}
	return meta
}

func (s *Store) writeMetaKeyLocked(key, value string) error {
	meta := s.readMetaLocked()
	if meta == nil {
		meta = map[string]string{}
	}
	meta[key] = value
	raw, err := json.MarshalIndent(meta, "", "  ")
	if err != nil {
		return err
	}
	tmp := s.path + ".meta.json.tmp"
	if err := os.WriteFile(tmp, raw, 0o644); err != nil {
		return err
	}
	return os.Rename(tmp, s.path+".meta.json")
}

func newID() string {
	var b [8]byte
	if _, err := rand.Read(b[:]); err != nil {
		return fmt.Sprintf("%d", time.Now().UnixNano())
	}
	return hex.EncodeToString(b[:])
}

func copyFile(src, dst string) error {
	in, err := os.Open(src)
	if err != nil {
		return err
	}
	defer in.Close()
	if err := os.MkdirAll(dirOf(dst), 0o755); err != nil {
		return err
	}
	out, err := os.OpenFile(dst, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o644)
	if err != nil {
		return err
	}
	defer out.Close()
	_, err = io.Copy(out, in)
	return err
}

func dirOf(path string) string {
	i := strings.LastIndex(path, "/")
	if i < 0 {
		return "."
	}
	return path[:i]
}