Vectrum
Sign In
102 lines · 2.6 KB · Go
package agent

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

	"harness/internal/llm"
)

// Session is a conversation saved to disk so --resume can pick it up.
type Session struct {
	ID       string        `json:"id"`
	Root     string        `json:"root"`
	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 = ".harness/sessions"

// Save writes the session for id, overwriting any previous copy.
func (a *Agent) Save(dir, id, root, model string) error {
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return err
	}
	usage, turns := a.state.Usage()
	sess := Session{
		ID:       id,
		Root:     root,
		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])
}

// 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
}