Vectrum
Sign In
230 lines · 6.7 KB · Go
package agent

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"codey/internal/llm"
	"codey/internal/statedir"
)

// Session is a conversation saved to disk so --resume can pick it up.
type Session struct {
	ID   string `json:"id"`
	Root string `json:"root"`
	// Worktree is the tree /worktree moved the work into, and is empty for a
	// session that never left Root. The session file itself stays filed under
	// Root — one home for every conversation, so the main checkout's chooser
	// still lists the ones whose work wandered off — and this is what sends
	// the work back there when such a session is reopened.
	Worktree string        `json:"worktree,omitempty"`
	Model    string        `json:"model"`
	Created  time.Time     `json:"created"`
	Updated  time.Time     `json:"updated"`
	Turns    int           `json:"turns"`
	Usage    llm.Usage     `json:"usage"`
	Messages []llm.Message `json:"messages"`
}

// SessionDir is where sessions live, relative to the working root.
const SessionDir = statedir.Dir + "/sessions"

// Save writes the session for id, overwriting any previous copy. root is the
// directory the session is filed under and worktree the tree the work has been
// moved into, which is usually empty and never replaces root.
func (a *Agent) Save(dir, id, root, worktree, model string) error {
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return err
	}
	usage, turns := a.state.Usage()
	sess := Session{
		ID:       id,
		Root:     root,
		Worktree: worktree,
		Model:    model,
		Created:  time.Now().UTC(),
		Updated:  time.Now().UTC(),
		Turns:    turns,
		Usage:    usage,
		Messages: a.state.History(),
	}
	if prev, err := LoadSession(filepath.Join(dir, id+".json")); err == nil {
		sess.Created = prev.Created
	}
	b, err := json.MarshalIndent(sess, "", "  ")
	if err != nil {
		return err
	}
	return os.WriteFile(filepath.Join(dir, id+".json"), b, 0o644)
}

// LoadSession reads one session file.
func LoadSession(path string) (Session, error) {
	var s Session
	b, err := os.ReadFile(path)
	if err != nil {
		return s, err
	}
	if err := json.Unmarshal(b, &s); err != nil {
		return s, fmt.Errorf("session %s: %w", path, err)
	}
	return s, nil
}

// ResolveSession finds a session by id, or the most recently updated one when
// id is "latest" or empty.
func ResolveSession(dir, id string) (Session, error) {
	if id != "" && id != "latest" {
		return LoadSession(filepath.Join(dir, id+".json"))
	}
	entries, err := os.ReadDir(dir)
	if err != nil {
		return Session{}, fmt.Errorf("no sessions in %s", dir)
	}
	var files []string
	for _, e := range entries {
		if !e.IsDir() && filepath.Ext(e.Name()) == ".json" {
			files = append(files, filepath.Join(dir, e.Name()))
		}
	}
	if len(files) == 0 {
		return Session{}, fmt.Errorf("no sessions in %s", dir)
	}
	sort.Slice(files, func(i, j int) bool {
		a, _ := os.Stat(files[i])
		b, _ := os.Stat(files[j])
		return a.ModTime().After(b.ModTime())
	})
	return LoadSession(files[0])
}

// BelongsTo reports whether a session was recorded in this working root. The
// session directory already lives inside the root, so this is mostly belt and
// braces — but a file copied, symlinked or resumed from elsewhere would
// otherwise let one checkout reopen another's conversation, which is exactly
// what a directory-scoped chooser must not do. A session with no root recorded
// is taken to belong where it is filed: it predates the field, and the
// directory it sits in is the only evidence there is.
func (s Session) BelongsTo(root string) bool {
	if root == "" || s.Root == "" {
		return true
	}
	return filepath.Clean(s.Root) == filepath.Clean(root)
}

// titleMax is how much of the opening prompt a session is named by. Long
// enough to tell two mornings' work apart, short enough to leave the picker
// room for the date beside it.
const titleMax = 64

// Title names the session by what was first asked of it, which is what someone
// scanning a list of old sessions is actually looking for. Sessions carry no
// title of their own: one derived from the first prompt costs nothing to
// maintain and cannot go stale.
func (s Session) Title() string {
	for _, msg := range s.Messages {
		if msg.Role != "user" {
			continue
		}
		for _, b := range msg.Content {
			// A user turn can also be tool results being handed back; those
			// are the conversation's plumbing, not anything anyone said.
			if b.Type != "text" && b.Type != "" {
				continue
			}
			if t := firstLine(b.Text); t != "" {
				return t
			}
		}
	}
	return ""
}

// firstLine is the first non-blank line of s, trimmed and capped.
func firstLine(s string) string {
	for _, line := range strings.Split(s, "\n") {
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		r := []rune(line)
		if len(r) > titleMax {
			return strings.TrimRight(string(r[:titleMax-1]), " ") + "…"
		}
		return line
	}
	return ""
}

// SessionInfo is what a chooser needs to know about a saved session: enough to
// describe it in one row, without the whole conversation coming along.
type SessionInfo struct {
	ID       string
	Root     string
	Model    string
	Title    string
	Created  time.Time
	Updated  time.Time
	Turns    int
	Messages int
	Usage    llm.Usage
}

// ListSessions describes the sessions in dir that belong to root, newest
// first. A directory that is not there yet is not an error — it means no
// session has been saved, which is a thing a chooser has to say anyway. A file
// that will not parse is skipped rather than failing the list: one bad session
// must not hide the good ones.
func ListSessions(dir, root string) ([]SessionInfo, error) {
	entries, err := os.ReadDir(dir)
	if errors.Is(err, os.ErrNotExist) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	var out []SessionInfo
	for _, e := range entries {
		if e.IsDir() || filepath.Ext(e.Name()) != ".json" {
			continue
		}
		s, err := LoadSession(filepath.Join(dir, e.Name()))
		if err != nil || !s.BelongsTo(root) {
			continue
		}
		updated := s.Updated
		if updated.IsZero() {
			if fi, err := e.Info(); err == nil {
				updated = fi.ModTime()
			}
		}
		out = append(out, SessionInfo{
			ID:       s.ID,
			Root:     s.Root,
			Model:    s.Model,
			Title:    s.Title(),
			Created:  s.Created,
			Updated:  updated,
			Turns:    s.Turns,
			Messages: len(s.Messages),
			Usage:    s.Usage,
		})
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Updated.After(out[j].Updated) })
	return out, nil
}

// Restore seeds the agent's state from a saved session.
func (a *Agent) Restore(s Session) {
	a.state.Replace(s.Messages)
	a.state.RestoreUsage(s.Usage, s.Turns)
	// The last request's size is unknown after a restore; the next turn's
	// response re-establishes it before compaction is consulted again.
	a.lastInputTokens = 0
}