byre/internal/store/review.go

package store

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"time"

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

const reviewColumns = `id, task_id, snapshot_n, coalesce(agent_session_id, ''), vendor, model, state, result, error, created_at, finished_at, finding_states`

func (s *Store) CreateReview(ctx context.Context, r api.Review) error {
	finishedAt := ""
	if r.FinishedAt != nil {
		finishedAt = formatTime(*r.FinishedAt)
	}
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		_, err := tx.ExecContext(ctx, `
			INSERT INTO reviews (id, task_id, snapshot_n, agent_session_id, vendor, model, state, created_at, finished_at)
			VALUES (?, ?, ?, nullif(?, ''), ?, ?, ?, ?, ?)
		`, r.Id, r.TaskId, r.SnapshotN, r.AgentSessionId, r.Vendor, r.Model, string(r.State), formatTime(r.CreatedAt), finishedAt)
		if err != nil {
			return err
		}
		if r.State == api.ReviewStateRunning {
			return emitReview(ctx, tx, emit, "review.started", r.Id, "byre")
		}
		return emitReview(ctx, tx, emit, "review.finished", r.Id, "byre")
	})
}

func (s *Store) GetReview(ctx context.Context, id string) (api.Review, error) {
	r, err := scanReview(s.db.QueryRowContext(ctx, `SELECT `+reviewColumns+` FROM reviews WHERE id = ?`, id))
	if errors.Is(err, sql.ErrNoRows) {
		return r, ErrNotFound
	}
	return r, err
}

func (s *Store) ReviewBySession(ctx context.Context, sessionID string) (api.Review, error) {
	r, err := scanReview(s.db.QueryRowContext(ctx, `SELECT `+reviewColumns+` FROM reviews WHERE agent_session_id = ?`, sessionID))
	if errors.Is(err, sql.ErrNoRows) {
		return r, ErrNotFound
	}
	return r, err
}

func (s *Store) ListReviews(ctx context.Context, taskID string) ([]api.Review, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT `+reviewColumns+` FROM reviews WHERE task_id = ? ORDER BY snapshot_n`, taskID)
	if err != nil {
		return nil, err
	}
	return scanReviews(rows)
}

func (s *Store) RunningReviews(ctx context.Context) ([]api.Review, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT `+reviewColumns+` FROM reviews WHERE state = ? ORDER BY created_at`, string(api.ReviewStateRunning))
	if err != nil {
		return nil, err
	}
	return scanReviews(rows)
}

func (s *Store) FinishReview(ctx context.Context, id string, state api.ReviewState, result map[string]any, errText string) error {
	text := ""
	if result != nil {
		b, err := json.Marshal(result)
		if err != nil {
			return err
		}
		text = string(b)
	}
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		res, err := tx.ExecContext(ctx, `UPDATE reviews SET state = ?, result = ?, error = ?, finished_at = ? WHERE id = ? AND state = ?`,
			string(state), text, errText, now(), id, string(api.ReviewStateRunning))
		if err != nil {
			return err
		}
		n, err := res.RowsAffected()
		if err != nil {
			return err
		}
		if n == 0 {
			return ErrNotFound
		}
		return emitReview(ctx, tx, emit, "review.finished", id, "byre")
	})
}

func (s *Store) SetFindingState(ctx context.Context, id string, index int, state api.FindingState, actor string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		res, err := tx.ExecContext(ctx, `
			UPDATE reviews SET finding_states = json_set(finding_states, '$."' || ? || '"', ?)
			WHERE id = ? AND state = ? AND ? < json_array_length(nullif(result, ''), '$.findings')
		`, index, string(state), id, string(api.ReviewStateDone), index)
		if err != nil {
			return err
		}
		n, err := res.RowsAffected()
		if err != nil {
			return err
		}
		if n == 0 {
			return ErrNotFound
		}
		return emitReview(ctx, tx, emit, "review.finding", id, actor)
	})
}

func (s *Store) ReviewerSessionIDs(ctx context.Context, taskID string) (map[string]bool, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT id FROM agent_sessions WHERE task_id = ? AND role = ?`, taskID, string(api.Reviewer))
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	ids := map[string]bool{}
	for rows.Next() {
		id := ""
		err = rows.Scan(&id)
		if err != nil {
			return nil, err
		}
		ids[id] = true
	}
	return ids, rows.Err()
}

func (s *Store) FirstTaskMessage(ctx context.Context, taskID string) (string, error) {
	text := ""
	err := s.db.QueryRowContext(ctx, `
		SELECT coalesce(json_extract(payload, '$.text'), '') FROM events
		WHERE task_id = ? AND kind = 'user.message'
		AND agent_session_id IN (SELECT id FROM agent_sessions WHERE task_id = ? AND role = ?)
		ORDER BY seq LIMIT 1
	`, taskID, taskID, string(api.Implementer)).Scan(&text)
	if errors.Is(err, sql.ErrNoRows) {
		return "", nil
	}
	return text, err
}

func emitReview(ctx context.Context, tx *sql.Tx, emit emitFunc, kind, id, actor string) error {
	r, err := scanReview(tx.QueryRowContext(ctx, `SELECT `+reviewColumns+` FROM reviews WHERE id = ?`, id))
	if err != nil {
		return err
	}
	projectID := ""
	err = tx.QueryRowContext(ctx, `SELECT project_id FROM tasks WHERE id = ?`, r.TaskId).Scan(&projectID)
	if err != nil {
		return err
	}
	return emit(api.Event{Kind: kind, ProjectId: &projectID, TaskId: &r.TaskId, Actor: actor, Payload: ToPayload(r)})
}

func questionReviews(ctx context.Context, q querier, item api.InboxItem) ([]api.Review, error) {
	if item.Kind != api.Question {
		return []api.Review{}, nil
	}
	rows, err := q.QueryContext(ctx, `
		SELECT `+reviewColumns+` FROM reviews WHERE task_id = ? AND snapshot_n = (
			SELECT min(n) FROM snapshots WHERE task_id = ? AND created_at >= ?
		)
	`, item.TaskId, item.TaskId, formatTime(item.CreatedAt))
	if err != nil {
		return nil, err
	}
	return scanReviews(rows)
}

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

	reviews := []api.Review{}
	for rows.Next() {
		r, err := scanReview(rows)
		if err != nil {
			return nil, err
		}
		reviews = append(reviews, r)
	}
	return reviews, rows.Err()
}

func scanReview(row scanner) (api.Review, error) {
	r := api.Review{}
	state := ""
	result := ""
	createdAt := ""
	finishedAt := ""
	findingStates := ""
	err := row.Scan(&r.Id, &r.TaskId, &r.SnapshotN, &r.AgentSessionId, &r.Vendor, &r.Model, &state, &result, &r.Error, &createdAt, &finishedAt, &findingStates)
	if err != nil {
		return r, err
	}
	r.State = api.ReviewState(state)
	r.FindingStates = map[string]api.FindingState{}
	err = json.Unmarshal([]byte(findingStates), &r.FindingStates)
	if err != nil {
		return r, err
	}
	if result != "" {
		res := map[string]any{}
		err = json.Unmarshal([]byte(result), &res)
		if err != nil {
			return r, err
		}
		r.Result = &res
	}
	r.CreatedAt, err = parseTime(createdAt)
	if err != nil {
		return r, err
	}
	if finishedAt != "" {
		var t time.Time
		t, err = parseTime(finishedAt)
		r.FinishedAt = &t
	}
	return r, err
}