byre/internal/store/store.go

package store

import (
	"context"
	"crypto/rand"
	"database/sql"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"sync"
	"time"

	_ "modernc.org/sqlite"

	"g.fc.run/misc/byre/internal/api"
)

var ErrNotFound = errors.New("not found")

var migrations = []string{`
CREATE TABLE IF NOT EXISTS users (
	login TEXT PRIMARY KEY,
	display_name TEXT NOT NULL,
	git_name TEXT NOT NULL DEFAULT '',
	git_email TEXT NOT NULL DEFAULT '',
	created_at TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS projects (
	id TEXT PRIMARY KEY,
	name TEXT NOT NULL UNIQUE,
	repo_path TEXT NOT NULL UNIQUE,
	default_branch TEXT NOT NULL,
	remote TEXT NOT NULL,
	created_by TEXT NOT NULL REFERENCES users(login),
	created_at TEXT NOT NULL
);

CREATE TABLE IF NOT EXISTS tasks (
	id TEXT PRIMARY KEY,
	project_id TEXT NOT NULL REFERENCES projects(id),
	title TEXT NOT NULL,
	branch TEXT NOT NULL,
	worktree_path TEXT NOT NULL,
	base_ref TEXT NOT NULL,
	status TEXT NOT NULL,
	created_by TEXT NOT NULL REFERENCES users(login),
	created_at TEXT NOT NULL,
	completed_at TEXT
);

CREATE INDEX IF NOT EXISTS tasks_project ON tasks(project_id, created_at);

CREATE TABLE IF NOT EXISTS events (
	seq INTEGER PRIMARY KEY AUTOINCREMENT,
	kind TEXT NOT NULL,
	project_id TEXT,
	task_id TEXT,
	actor TEXT NOT NULL,
	payload TEXT NOT NULL,
	created_at TEXT NOT NULL
);
`, `
ALTER TABLE events ADD COLUMN agent_session_id TEXT;

CREATE TABLE agent_sessions (
	id TEXT PRIMARY KEY,
	task_id TEXT NOT NULL REFERENCES tasks(id),
	vendor TEXT NOT NULL,
	model TEXT NOT NULL,
	vendor_session_id TEXT NOT NULL,
	state TEXT NOT NULL,
	role TEXT NOT NULL,
	created_by TEXT NOT NULL REFERENCES users(login),
	created_at TEXT NOT NULL
);

CREATE INDEX agent_sessions_task ON agent_sessions(task_id, created_at);
CREATE INDEX agent_sessions_state ON agent_sessions(state);
`, `
CREATE TABLE inbox_items (
	id TEXT PRIMARY KEY,
	project_id TEXT NOT NULL REFERENCES projects(id),
	task_id TEXT NOT NULL REFERENCES tasks(id),
	agent_session_id TEXT NOT NULL REFERENCES agent_sessions(id),
	thread_id TEXT NOT NULL,
	kind TEXT NOT NULL,
	priority INTEGER NOT NULL,
	blocking INTEGER NOT NULL,
	payload TEXT NOT NULL,
	state TEXT NOT NULL,
	snoozed_until TEXT,
	created_at TEXT NOT NULL,
	resolved_by TEXT REFERENCES users(login),
	resolved_at TEXT
);

CREATE INDEX inbox_items_open ON inbox_items(state, priority, created_at);
CREATE INDEX inbox_items_session ON inbox_items(agent_session_id, state);

CREATE TABLE inbox_responses (
	id TEXT PRIMARY KEY,
	item_id TEXT NOT NULL REFERENCES inbox_items(id),
	user_login TEXT NOT NULL REFERENCES users(login),
	type TEXT NOT NULL,
	option TEXT,
	text TEXT,
	attachments TEXT NOT NULL,
	snooze_until TEXT,
	created_at TEXT NOT NULL,
	delivered_at TEXT
);

CREATE INDEX inbox_responses_item ON inbox_responses(item_id, created_at);
`, `
CREATE INDEX events_task ON events(task_id, seq);
`, `
CREATE TABLE turn_queue (
	session_id TEXT PRIMARY KEY REFERENCES agent_sessions(id),
	actor TEXT NOT NULL,
	text TEXT NOT NULL,
	created_at TEXT NOT NULL
);
`}

type Store struct {
	db      *sql.DB
	path    string
	writeMu sync.Mutex
	onEvent func(api.Event)
}

func Open(path string, onEvent func(api.Event)) (*Store, error) {
	dsn := "file:" + path + "?_pragma=journal_mode(WAL)&_pragma=busy_timeout(5000)&_pragma=foreign_keys(1)&_txlock=immediate"
	db, err := sql.Open("sqlite", dsn)
	if err != nil {
		return nil, err
	}

	err = migrate(db)
	if err != nil {
		db.Close()
		return nil, err
	}

	return &Store{db: db, path: path, onEvent: onEvent}, nil
}

func migrate(db *sql.DB) error {
	version := 0
	err := db.QueryRow(`PRAGMA user_version`).Scan(&version)
	if err != nil {
		return err
	}

	for i := version; i < len(migrations); i++ {
		tx, err := db.Begin()
		if err != nil {
			return err
		}

		_, err = tx.Exec(migrations[i])
		if err != nil {
			tx.Rollback()
			return fmt.Errorf("migration %d: %w", i, err)
		}

		_, err = tx.Exec(fmt.Sprintf(`PRAGMA user_version = %d`, i+1))
		if err != nil {
			tx.Rollback()
			return fmt.Errorf("migration %d: %w", i, err)
		}

		err = tx.Commit()
		if err != nil {
			return fmt.Errorf("migration %d: %w", i, err)
		}
	}

	return nil
}

func (s *Store) Close() error {
	return s.db.Close()
}

func (s *Store) Path() string {
	return s.path
}

func NewID() string {
	b := make([]byte, 6)
	rand.Read(b)
	return hex.EncodeToString(b)
}

func (s *Store) EnsureUser(ctx context.Context, login, displayName string) (api.User, error) {
	_, err := s.db.ExecContext(ctx, `
		INSERT INTO users (login, display_name, created_at) VALUES (?, ?, ?)
		ON CONFLICT (login) DO UPDATE SET display_name = excluded.display_name
	`, login, displayName, now())
	if err != nil {
		return api.User{}, err
	}
	return s.GetUser(ctx, login)
}

func (s *Store) GetUser(ctx context.Context, login string) (api.User, error) {
	u := api.User{}
	err := s.db.QueryRowContext(ctx, `SELECT login, display_name, git_name, git_email FROM users WHERE login = ?`, login).
		Scan(&u.Login, &u.DisplayName, &u.GitName, &u.GitEmail)
	if errors.Is(err, sql.ErrNoRows) {
		return u, ErrNotFound
	}
	return u, err
}

func (s *Store) UpdateUserGit(ctx context.Context, login, gitName, gitEmail string) (api.User, error) {
	err := s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		_, err := tx.ExecContext(ctx, `UPDATE users SET git_name = ?, git_email = ? WHERE login = ?`, gitName, gitEmail, login)
		if err != nil {
			return err
		}
		return emit(api.Event{Kind: "user.updated", Actor: login, Payload: map[string]any{"login": login}})
	})
	if err != nil {
		return api.User{}, err
	}
	return s.GetUser(ctx, login)
}

func (s *Store) CreateProject(ctx context.Context, p api.Project) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		_, err := tx.ExecContext(ctx, `
			INSERT INTO projects (id, name, repo_path, default_branch, remote, created_by, created_at)
			VALUES (?, ?, ?, ?, ?, ?, ?)
		`, p.Id, p.Name, p.RepoPath, p.DefaultBranch, p.Remote, p.CreatedBy, formatTime(p.CreatedAt))
		if err != nil {
			return err
		}
		return emit(api.Event{Kind: "project.created", ProjectId: &p.Id, Actor: p.CreatedBy, Payload: ToPayload(p)})
	})
}

func (s *Store) ListProjects(ctx context.Context) ([]api.Project, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT `+projectColumns+` FROM projects ORDER BY name`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	projects := []api.Project{}
	for rows.Next() {
		p, err := scanProject(rows)
		if err != nil {
			return nil, err
		}
		projects = append(projects, p)
	}
	return projects, rows.Err()
}

func (s *Store) GetProject(ctx context.Context, id string) (api.Project, error) {
	p, err := scanProject(s.db.QueryRowContext(ctx, `SELECT `+projectColumns+` FROM projects WHERE id = ?`, id))
	if errors.Is(err, sql.ErrNoRows) {
		return p, ErrNotFound
	}
	return p, err
}

func (s *Store) CreateTask(ctx context.Context, t api.Task) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		_, err := tx.ExecContext(ctx, `
			INSERT INTO tasks (id, project_id, title, branch, worktree_path, base_ref, status, created_by, created_at)
			VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
		`, t.Id, t.ProjectId, t.Title, t.Branch, t.WorktreePath, t.BaseRef, string(t.Status), t.CreatedBy, formatTime(t.CreatedAt))
		if err != nil {
			return err
		}
		return emit(api.Event{Kind: "task.created", ProjectId: &t.ProjectId, TaskId: &t.Id, Actor: t.CreatedBy, Payload: ToPayload(t)})
	})
}

func (s *Store) ListTasks(ctx context.Context, projectID string) ([]api.Task, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT `+taskColumns+` FROM tasks WHERE project_id = ? ORDER BY created_at DESC`, projectID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	tasks := []api.Task{}
	for rows.Next() {
		t, err := scanTask(rows)
		if err != nil {
			return nil, err
		}
		tasks = append(tasks, t)
	}
	return tasks, rows.Err()
}

func (s *Store) GetTask(ctx context.Context, id string) (api.Task, error) {
	t, err := scanTask(s.db.QueryRowContext(ctx, `SELECT `+taskColumns+` FROM tasks WHERE id = ?`, id))
	if errors.Is(err, sql.ErrNoRows) {
		return t, ErrNotFound
	}
	return t, err
}

func (s *Store) CompleteTask(ctx context.Context, id, actor string) (api.Task, error) {
	err := s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		res, err := tx.ExecContext(ctx, `UPDATE tasks SET status = ?, completed_at = ? WHERE id = ? AND status = ?`,
			string(api.TaskStatusCompleted), now(), id, string(api.TaskStatusActive))
		if err != nil {
			return err
		}
		n, err := res.RowsAffected()
		if err != nil {
			return err
		}
		if n == 0 {
			return ErrNotFound
		}

		projectID := ""
		err = tx.QueryRowContext(ctx, `SELECT project_id FROM tasks WHERE id = ?`, id).Scan(&projectID)
		if err != nil {
			return err
		}
		return emit(api.Event{Kind: "task.completed", ProjectId: &projectID, TaskId: &id, Actor: actor, Payload: map[string]any{"id": id}})
	})
	if err != nil {
		return api.Task{}, err
	}
	return s.GetTask(ctx, id)
}

func (s *Store) EventsAfter(ctx context.Context, after int64, limit int) ([]api.Event, error) {
	rows, err := s.db.QueryContext(ctx, `
		SELECT seq, kind, project_id, task_id, agent_session_id, actor, payload, created_at
		FROM events WHERE seq > ? ORDER BY seq LIMIT ?
	`, after, limit)
	if err != nil {
		return nil, err
	}
	return scanEvents(rows)
}

func (s *Store) TaskEvents(ctx context.Context, taskID string, after int64, limit int) ([]api.Event, error) {
	rows, err := s.db.QueryContext(ctx, `
		SELECT seq, kind, project_id, task_id, agent_session_id, actor, payload, created_at
		FROM events WHERE task_id = ? AND seq > ? ORDER BY seq LIMIT ?
	`, taskID, after, limit)
	if err != nil {
		return nil, err
	}
	return scanEvents(rows)
}

func scanEvents(rows *sql.Rows) ([]api.Event, error) {
	defer rows.Close()

	events := []api.Event{}
	for rows.Next() {
		ev := api.Event{}
		payload := ""
		createdAt := ""
		err := rows.Scan(&ev.Seq, &ev.Kind, &ev.ProjectId, &ev.TaskId, &ev.AgentSessionId, &ev.Actor, &payload, &createdAt)
		if err != nil {
			return nil, err
		}
		err = json.Unmarshal([]byte(payload), &ev.Payload)
		if err != nil {
			return nil, err
		}
		ev.CreatedAt, err = parseTime(createdAt)
		if err != nil {
			return nil, err
		}
		events = append(events, ev)
	}
	return events, rows.Err()
}

func (s *Store) AppendEvent(ctx context.Context, ev api.Event) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		return emit(ev)
	})
}

func (s *Store) CreateSession(ctx context.Context, sess api.AgentSession, projectID string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		_, err := tx.ExecContext(ctx, `
			INSERT INTO agent_sessions (id, task_id, vendor, model, vendor_session_id, state, role, created_by, created_at)
			VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
		`, sess.Id, sess.TaskId, string(sess.Vendor), sess.Model, sess.VendorSessionId, string(sess.State), string(sess.Role), sess.CreatedBy, formatTime(sess.CreatedAt))
		if err != nil {
			return err
		}
		return emit(api.Event{
			Kind:           "session.created",
			ProjectId:      &projectID,
			TaskId:         &sess.TaskId,
			AgentSessionId: &sess.Id,
			Actor:          sess.CreatedBy,
			Payload:        ToPayload(sess),
		})
	})
}

func (s *Store) GetSession(ctx context.Context, id string) (api.AgentSession, error) {
	sess, err := scanSession(s.db.QueryRowContext(ctx, `SELECT `+sessionColumns+` FROM agent_sessions WHERE id = ?`, id))
	if errors.Is(err, sql.ErrNoRows) {
		return sess, ErrNotFound
	}
	return sess, err
}

func (s *Store) ListSessions(ctx context.Context, taskID string) ([]api.AgentSession, error) {
	return s.querySessions(ctx, `SELECT `+sessionColumns+` FROM agent_sessions WHERE task_id = ? ORDER BY created_at`, taskID)
}

func (s *Store) SessionsInState(ctx context.Context, state api.AgentSessionState) ([]api.AgentSession, error) {
	return s.querySessions(ctx, `SELECT `+sessionColumns+` FROM agent_sessions WHERE state = ? ORDER BY created_at`, string(state))
}

func (s *Store) querySessions(ctx context.Context, query string, args ...any) ([]api.AgentSession, error) {
	rows, err := s.db.QueryContext(ctx, query, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	sessions := []api.AgentSession{}
	for rows.Next() {
		sess, err := scanSession(rows)
		if err != nil {
			return nil, err
		}
		sessions = append(sessions, sess)
	}
	return sessions, rows.Err()
}

var ErrStateChanged = errors.New("session state changed")

func (s *Store) SetSessionState(ctx context.Context, id string, state api.AgentSessionState, actor string, payload map[string]any) error {
	return s.transitionSession(ctx, id, nil, state, actor, payload)
}

func (s *Store) TransitionSession(ctx context.Context, id string, from, to api.AgentSessionState, actor string) error {
	return s.transitionSession(ctx, id, &from, to, actor, nil)
}

func (s *Store) transitionSession(ctx context.Context, id string, from *api.AgentSessionState, state api.AgentSessionState, actor string, payload map[string]any) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		taskID := ""
		projectID := ""
		current := ""
		err := tx.QueryRowContext(ctx, `
			SELECT t.id, t.project_id, s.state FROM agent_sessions s JOIN tasks t ON t.id = s.task_id WHERE s.id = ?
		`, id).Scan(&taskID, &projectID, &current)
		if errors.Is(err, sql.ErrNoRows) {
			return ErrNotFound
		}
		if err != nil {
			return err
		}

		if from != nil && current != string(*from) {
			return fmt.Errorf("%w: session %s is %s, not %s", ErrStateChanged, id, current, *from)
		}

		_, err = tx.ExecContext(ctx, `UPDATE agent_sessions SET state = ? WHERE id = ?`, string(state), id)
		if err != nil {
			return err
		}

		if payload == nil {
			payload = map[string]any{}
		}
		payload["state"] = string(state)
		return emit(api.Event{
			Kind:           "session.state",
			ProjectId:      &projectID,
			TaskId:         &taskID,
			AgentSessionId: &id,
			Actor:          actor,
			Payload:        payload,
		})
	})
}

type Counts struct {
	LastEventSeq int64
	Projects     int
	ActiveTasks  int
}

func (s *Store) Counts(ctx context.Context) (Counts, error) {
	c := Counts{}
	err := s.db.QueryRowContext(ctx, `
		SELECT
			(SELECT COALESCE(MAX(seq), 0) FROM events),
			(SELECT COUNT(*) FROM projects),
			(SELECT COUNT(*) FROM tasks WHERE status = ?)
	`, string(api.TaskStatusActive)).Scan(&c.LastEventSeq, &c.Projects, &c.ActiveTasks)
	return c, err
}

type emitFunc func(api.Event) error

func (s *Store) write(ctx context.Context, fn func(tx *sql.Tx, emit emitFunc) error) error {
	s.writeMu.Lock()
	defer s.writeMu.Unlock()

	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	events := []api.Event{}
	emit := func(ev api.Event) error {
		ev.CreatedAt = time.Now().UTC()
		payload, err := json.Marshal(ev.Payload)
		if err != nil {
			return err
		}
		err = tx.QueryRowContext(ctx, `
			INSERT INTO events (kind, project_id, task_id, agent_session_id, actor, payload, created_at)
			VALUES (?, ?, ?, ?, ?, ?, ?) RETURNING seq
		`, ev.Kind, ev.ProjectId, ev.TaskId, ev.AgentSessionId, ev.Actor, string(payload), formatTime(ev.CreatedAt)).Scan(&ev.Seq)
		if err != nil {
			return err
		}
		events = append(events, ev)
		return nil
	}

	err = fn(tx, emit)
	if err != nil {
		return err
	}

	err = tx.Commit()
	if err != nil {
		return err
	}

	for _, ev := range events {
		s.onEvent(ev)
	}
	return nil
}

const projectColumns = `id, name, repo_path, default_branch, remote, created_by, created_at`

const taskColumns = `id, project_id, title, branch, worktree_path, base_ref, status, created_by, created_at, completed_at`

const sessionColumns = `id, task_id, vendor, model, vendor_session_id, state, role, created_by, created_at`

type scanner interface {
	Scan(dest ...any) error
}

func scanProject(row scanner) (api.Project, error) {
	p := api.Project{}
	createdAt := ""
	err := row.Scan(&p.Id, &p.Name, &p.RepoPath, &p.DefaultBranch, &p.Remote, &p.CreatedBy, &createdAt)
	if err != nil {
		return p, err
	}
	p.CreatedAt, err = parseTime(createdAt)
	return p, err
}

func scanTask(row scanner) (api.Task, error) {
	t := api.Task{}
	status := ""
	createdAt := ""
	completedAt := sql.NullString{}
	err := row.Scan(&t.Id, &t.ProjectId, &t.Title, &t.Branch, &t.WorktreePath, &t.BaseRef, &status, &t.CreatedBy, &createdAt, &completedAt)
	if err != nil {
		return t, err
	}
	t.Status = api.TaskStatus(status)
	t.CreatedAt, err = parseTime(createdAt)
	if err != nil {
		return t, err
	}
	if completedAt.Valid {
		ct, err := parseTime(completedAt.String)
		if err != nil {
			return t, err
		}
		t.CompletedAt = &ct
	}
	return t, nil
}

func scanSession(row scanner) (api.AgentSession, error) {
	sess := api.AgentSession{}
	vendor := ""
	state := ""
	role := ""
	createdAt := ""
	err := row.Scan(&sess.Id, &sess.TaskId, &vendor, &sess.Model, &sess.VendorSessionId, &state, &role, &sess.CreatedBy, &createdAt)
	if err != nil {
		return sess, err
	}
	sess.Vendor = api.AgentSessionVendor(vendor)
	sess.State = api.AgentSessionState(state)
	sess.Role = api.AgentSessionRole(role)
	sess.CreatedAt, err = parseTime(createdAt)
	return sess, err
}

func ToPayload(v any) map[string]any {
	b, err := json.Marshal(v)
	if err != nil {
		panic(err)
	}
	m := map[string]any{}
	err = json.Unmarshal(b, &m)
	if err != nil {
		panic(err)
	}
	return m
}

func now() string {
	return formatTime(time.Now())
}

const timeLayout = "2006-01-02T15:04:05.000000000Z07:00"

func formatTime(t time.Time) string {
	return t.UTC().Format(timeLayout)
}

func parseTime(s string) (time.Time, error) {
	return time.Parse(time.RFC3339Nano, s)
}