lvmh/daemon/store.go

package main

// store.go — SQLite persistence: sessions, events, settings, containers.

import (
	"database/sql"
	"encoding/json"
	"fmt"

	_ "modernc.org/sqlite"
)

// SessionInfo mirrors the hello/session_info payload snapshot in PROTOCOL.md.
type SessionInfo struct {
	ID        string  `json:"id"`
	Name      *string `json:"name"`
	Cwd       string  `json:"cwd"`
	Model     string  `json:"model"`
	Provider  string  `json:"provider"`
	Agent     bool    `json:"agent"`
	Repo      *string `json:"repo"`
	StartedAt int64   `json:"startedAt"`
}

// SessionRow is a persisted session with its bookkeeping columns.
type SessionRow struct {
	ID          string      `json:"id"`
	Info        SessionInfo `json:"-"`
	LastSeq     int64       `json:"-"`
	LastEventAt int64       `json:"lastEventAt"`
	OnlineDB    bool        `json:"-"`
}

// Event is a persisted agent event; Payload is the raw JSON payload object.
type Event struct {
	SessionID string
	Seq       int64
	TS        int64
	Type      string
	Payload   json.RawMessage
}

// ContainerRow tracks a daemon-spawned worker container per session.
type ContainerRow struct {
	SessionID   string
	ContainerID string
	Repo        string
}

// RepoImageRow maps a repo path to its registered custom worker image.
type RepoImageRow struct {
	Repo  string `json:"repo"`
	Image string `json:"image"`
	Built bool   `json:"built"`
}

const dsnParams string = "?_pragma=busy_timeout(5000)&_pragma=journal_mode(WAL)&_pragma=synchronous(NORMAL)"

const schema string = `
CREATE TABLE IF NOT EXISTS events(
  sessionId TEXT NOT NULL, seq INTEGER NOT NULL, ts INTEGER NOT NULL,
  type TEXT NOT NULL, payload TEXT NOT NULL, PRIMARY KEY(sessionId, seq));
CREATE TABLE IF NOT EXISTS sessions(
  id TEXT PRIMARY KEY, info TEXT NOT NULL,
  lastSeq INTEGER NOT NULL DEFAULT 0, lastEventAt INTEGER, online INTEGER DEFAULT 0);
CREATE TABLE IF NOT EXISTS settings(key TEXT PRIMARY KEY, value TEXT NOT NULL);
CREATE TABLE IF NOT EXISTS containers(
  sessionId TEXT PRIMARY KEY, containerId TEXT NOT NULL, repo TEXT NOT NULL);
CREATE TABLE IF NOT EXISTS repo_images(
  repo TEXT PRIMARY KEY, image TEXT NOT NULL);
CREATE INDEX IF NOT EXISTS idx_events_session ON events(sessionId, seq);
`

// OpenStore opens (creating if needed) the SQLite database at path with WAL.
func OpenStore(path string) (*Store, error) {
	db, err := sql.Open("sqlite", path+dsnParams)
	if err != nil {
		return nil, fmt.Errorf("open sqlite: %w", err)
	}
	// modernc sqlite serializes writes; a single conn avoids SQLITE_BUSY churn.
	db.SetMaxOpenConns(1)
	if _, err := db.Exec(schema); err != nil {
		db.Close()
		return nil, fmt.Errorf("apply schema: %w", err)
	}
	return &Store{db: db}, nil
}

// Store wraps the SQLite handle.
type Store struct {
	db *sql.DB
}

// Close closes the database.
func (s *Store) Close() error { return s.db.Close() }

// AppendEvent persists one event keyed (sessionId, seq). Duplicates are ignored.
func (s *Store) AppendEvent(e Event) error {
	payload := e.Payload
	if len(payload) == 0 {
		payload = json.RawMessage("{}")
	}
	_, err := s.db.Exec(
		`INSERT OR IGNORE INTO events(sessionId, seq, ts, type, payload) VALUES(?,?,?,?,?)`,
		e.SessionID, e.Seq, e.TS, e.Type, string(payload))
	return err
}

// LastSeq returns the highest persisted seq for the session, 0 if none.
func (s *Store) LastSeq(sessionID string) (int64, error) {
	var seq sql.NullInt64
	err := s.db.QueryRow(`SELECT MAX(seq) FROM events WHERE sessionId=?`, sessionID).Scan(&seq)
	if err != nil {
		return 0, err
	}
	return seq.Int64, nil
}

// EventsAfter returns persisted events for the session with seq > after,
// ascending, at most limit rows.
func (s *Store) EventsAfter(sessionID string, after int64, limit int) ([]Event, error) {
	rows, err := s.db.Query(
		`SELECT sessionId, seq, ts, type, payload FROM events
		 WHERE sessionId=? AND seq>? ORDER BY seq LIMIT ?`,
		sessionID, after, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []Event
	for rows.Next() {
		var e Event
		var payload string
		if err := rows.Scan(&e.SessionID, &e.Seq, &e.TS, &e.Type, &payload); err != nil {
			return nil, err
		}
		e.Payload = json.RawMessage(payload)
		out = append(out, e)
	}
	return out, rows.Err()
}

// EventsLatest returns the newest limit events for the session, ascending.
func (s *Store) EventsLatest(sessionID string, limit int) ([]Event, error) {
	return s.eventsDesc(sessionID, ` WHERE sessionId=? ORDER BY seq DESC LIMIT ?`, limit, sessionID)
}

// EventsBefore returns the newest window of at most limit events with
// seq < before, ascending.
func (s *Store) EventsBefore(sessionID string, before int64, limit int) ([]Event, error) {
	return s.eventsDesc(sessionID, ` WHERE sessionId=? AND seq<? ORDER BY seq DESC LIMIT ?`,
		limit, sessionID, before)
}

// eventsDesc runs a newest-first query and returns the rows reversed
// (oldest-first), for newest-window reads.
func (s *Store) eventsDesc(sessionID, where string, limit int, args ...any) ([]Event, error) {
	rows, err := s.db.Query(
		`SELECT sessionId, seq, ts, type, payload FROM events`+where, append(args, limit)...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var desc []Event
	for rows.Next() {
		var e Event
		var payload string
		if err := rows.Scan(&e.SessionID, &e.Seq, &e.TS, &e.Type, &payload); err != nil {
			return nil, err
		}
		e.Payload = json.RawMessage(payload)
		desc = append(desc, e)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	out := make([]Event, len(desc))
	for i, e := range desc {
		out[len(desc)-1-i] = e
	}
	return out, nil
}

// UpsertSession inserts or refreshes the info snapshot for a session.
func (s *Store) UpsertSession(info SessionInfo) error {
	blob, err := json.Marshal(info)
	if err != nil {
		return err
	}
	_, err = s.db.Exec(
		`INSERT INTO sessions(id, info, lastSeq, lastEventAt, online) VALUES(?,?,0,NULL,0)
		 ON CONFLICT(id) DO UPDATE SET info=excluded.info`,
		info.ID, string(blob))
	return err
}

// TouchSession advances lastSeq (monotonic) and lastEventAt.
func (s *Store) TouchSession(sessionID string, lastSeq, lastEventAt int64) error {
	_, err := s.db.Exec(
		`UPDATE sessions SET lastSeq=MAX(lastSeq,?), lastEventAt=? WHERE id=?`,
		lastSeq, lastEventAt, sessionID)
	return err
}

// SetOnline records the persisted online flag (memory hub is authoritative).
func (s *Store) SetOnline(sessionID string, online bool) error {
	_, err := s.db.Exec(`UPDATE sessions SET online=? WHERE id=?`, online, sessionID)
	return err
}

// SetSessionName merges a new name into the persisted session info blob,
// leaving every other snapshot field untouched. Returns false when the
// session row does not exist.
func (s *Store) SetSessionName(sessionID, name string) (bool, error) {
	var infoBlob string
	err := s.db.QueryRow(`SELECT info FROM sessions WHERE id=?`, sessionID).Scan(&infoBlob)
	if err == sql.ErrNoRows {
		return false, nil
	}
	if err != nil {
		return false, err
	}
	var info SessionInfo
	if err := json.Unmarshal([]byte(infoBlob), &info); err != nil {
		return false, fmt.Errorf("decode session info %s: %w", sessionID, err)
	}
	info.Name = &name
	blob, err := json.Marshal(info)
	if err != nil {
		return true, err
	}
	_, err = s.db.Exec(`UPDATE sessions SET info=? WHERE id=?`, string(blob), sessionID)
	return true, err
}

// Sessions returns every known session row ordered by session id.
func (s *Store) Sessions() ([]SessionRow, error) {
	rows, err := s.db.Query(`SELECT id, info, lastSeq, lastEventAt, online FROM sessions ORDER BY id`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []SessionRow
	for rows.Next() {
		var row SessionRow
		var infoBlob string
		var lastEventAt sql.NullInt64
		var online bool
		if err := rows.Scan(&row.ID, &infoBlob, &row.LastSeq, &lastEventAt, &online); err != nil {
			return nil, err
		}
		if err := json.Unmarshal([]byte(infoBlob), &row.Info); err != nil {
			return nil, fmt.Errorf("decode session info %s: %w", row.ID, err)
		}
		row.LastEventAt = lastEventAt.Int64
		row.OnlineDB = online
		out = append(out, row)
	}
	return out, rows.Err()
}

// DeleteSessionEvents drops every persisted event for the session and
// returns the number of rows removed (0 when the session had none).
func (s *Store) DeleteSessionEvents(sessionID string) (int64, error) {
	res, err := s.db.Exec(`DELETE FROM events WHERE sessionId=?`, sessionID)
	if err != nil {
		return 0, err
	}
	return res.RowsAffected()
}

// DeleteSession removes a session entirely: its events, container row and
// the session row itself.
func (s *Store) DeleteSession(sessionID string) error {
	if _, err := s.db.Exec(`DELETE FROM events WHERE sessionId=?`, sessionID); err != nil {
		return err
	}
	if _, err := s.db.Exec(`DELETE FROM containers WHERE sessionId=?`, sessionID); err != nil {
		return err
	}
	_, err := s.db.Exec(`DELETE FROM sessions WHERE id=?`, sessionID)
	return err
}

// UsageStats aggregates token/cost usage over agent_end events.
type UsageStats struct {
	Turns        int64   `json:"turns"`
	InputTokens  int64   `json:"inputTokens"`
	OutputTokens int64   `json:"outputTokens"`
	TotalCost    float64 `json:"totalCost"`
}

// StatsTotals is UsageStats aggregated across all sessions, plus how many
// sessions contributed agent_end events.
type StatsTotals struct {
	UsageStats
	SessionsCount int64 `json:"sessionsCount"`
}

// agentEndUsage is the usage sub-object of an agent_end payload; absent
// fields decode as zero.
type agentEndUsage struct {
	Usage struct {
		InputTokens  int64   `json:"inputTokens"`
		OutputTokens int64   `json:"outputTokens"`
		TotalCost    float64 `json:"totalCost"`
	} `json:"usage"`
}

// SessionStats aggregates usage over a session's persisted agent_end events.
// Each agent_end row counts as one turn; rows without (or with partial)
// usage payloads contribute zeros.
func (s *Store) SessionStats(sessionID string) (UsageStats, error) {
	rows, err := s.db.Query(`SELECT sessionId, payload FROM events WHERE sessionId=? AND type=?`,
		sessionID, evAgentEnd)
	if err != nil {
		return UsageStats{}, err
	}
	defer rows.Close()
	totals, err := scanAgentEndUsage(rows)
	return totals.UsageStats, err
}

// StatsTotals aggregates usage over every persisted agent_end event.
func (s *Store) StatsTotals() (StatsTotals, error) {
	rows, err := s.db.Query(`SELECT sessionId, payload FROM events WHERE type=?`, evAgentEnd)
	if err != nil {
		return StatsTotals{}, err
	}
	defer rows.Close()
	return scanAgentEndUsage(rows)
}

// scanAgentEndUsage folds agent_end rows into totals; unparseable payloads
// still count as turns but add no usage.
func scanAgentEndUsage(rows *sql.Rows) (StatsTotals, error) {
	var out StatsTotals
	seen := make(map[string]struct{})
	for rows.Next() {
		var sessionID, payload string
		if err := rows.Scan(&sessionID, &payload); err != nil {
			return out, err
		}
		seen[sessionID] = struct{}{}
		out.Turns++
		var u agentEndUsage
		if json.Unmarshal([]byte(payload), &u) == nil {
			out.InputTokens += u.Usage.InputTokens
			out.OutputTokens += u.Usage.OutputTokens
			out.TotalCost += u.Usage.TotalCost
		}
	}
	out.SessionsCount = int64(len(seen))
	return out, rows.Err()
}

// GetSetting returns a settings value and whether it exists.
func (s *Store) GetSetting(key string) (string, bool, error) {
	var val string
	err := s.db.QueryRow(`SELECT value FROM settings WHERE key=?`, key).Scan(&val)
	if err == sql.ErrNoRows {
		return "", false, nil
	}
	if err != nil {
		return "", false, err
	}
	return val, true, nil
}

// SetSetting upserts a settings value.
func (s *Store) SetSetting(key, value string) error {
	_, err := s.db.Exec(`INSERT INTO settings(key, value) VALUES(?,?)
		ON CONFLICT(key) DO UPDATE SET value=excluded.value`, key, value)
	return err
}

// DeleteSetting removes a settings key.
func (s *Store) DeleteSetting(key string) error {
	_, err := s.db.Exec(`DELETE FROM settings WHERE key=?`, key)
	return err
}

// UpsertContainer records the container created for a session.
func (s *Store) UpsertContainer(sessionID, containerID, repo string) error {
	_, err := s.db.Exec(`INSERT INTO containers(sessionId, containerId, repo) VALUES(?,?,?)
		ON CONFLICT(sessionId) DO UPDATE SET containerId=excluded.containerId, repo=excluded.repo`,
		sessionID, containerID, repo)
	return err
}

// GetContainer returns the container row for a session, if any.
func (s *Store) GetContainer(sessionID string) (ContainerRow, bool, error) {
	var row ContainerRow
	err := s.db.QueryRow(`SELECT sessionId, containerId, repo FROM containers WHERE sessionId=?`,
		sessionID).Scan(&row.SessionID, &row.ContainerID, &row.Repo)
	if err == sql.ErrNoRows {
		return row, false, nil
	}
	if err != nil {
		return row, false, err
	}
	return row, true, nil
}

// DeleteContainer forgets the container row for a session.
func (s *Store) DeleteContainer(sessionID string) error {
	_, err := s.db.Exec(`DELETE FROM containers WHERE sessionId=?`, sessionID)
	return err
}

// GetRepoImage returns the custom image registered for a repo, if any.
func (s *Store) GetRepoImage(repo string) (string, bool, error) {
	var image string
	err := s.db.QueryRow(`SELECT image FROM repo_images WHERE repo=?`, repo).Scan(&image)
	if err == sql.ErrNoRows {
		return "", false, nil
	}
	if err != nil {
		return "", false, err
	}
	return image, true, nil
}

// SetRepoImage upserts the custom image for a repo.
func (s *Store) SetRepoImage(repo, image string) error {
	_, err := s.db.Exec(`INSERT INTO repo_images(repo, image) VALUES(?,?)
		ON CONFLICT(repo) DO UPDATE SET image=excluded.image`, repo, image)
	return err
}

// DeleteRepoImage removes a repo's custom image registration (idempotent).
func (s *Store) DeleteRepoImage(repo string) error {
	_, err := s.db.Exec(`DELETE FROM repo_images WHERE repo=?`, repo)
	return err
}

// ListRepoImages returns every repo→image registration, sorted by repo.
func (s *Store) ListRepoImages() ([]RepoImageRow, error) {
	rows, err := s.db.Query(`SELECT repo, image FROM repo_images ORDER BY repo`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []RepoImageRow
	for rows.Next() {
		var row RepoImageRow
		if err := rows.Scan(&row.Repo, &row.Image); err != nil {
			return nil, err
		}
		out = append(out, row)
	}
	return out, rows.Err()
}