byre/internal/store/review.go

package store

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

	"g.fc.run/ai/byre/internal/api"
	"g.fc.run/ai/byre/internal/workflow"
)

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

func (s *Store) CreateReview(ctx context.Context, r api.Review, definitionID string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		return insertReview(ctx, tx, emit, r, definitionID)
	})
}

func insertReview(ctx context.Context, tx *sql.Tx, emit emitFunc, r api.Review, definitionID string) error {
	finishedAt := ""
	if r.FinishedAt != nil {
		finishedAt = formatTime(*r.FinishedAt)
	}
	_, err := tx.ExecContext(ctx, `
		INSERT INTO reviews (id, task_id, snapshot_n, stage, definition_id, agent_session_id, vendor, model, state, error, created_at, finished_at)
		VALUES (?, ?, ?, ?, nullif(?, ''), nullif(?, ''), ?, ?, ?, ?, ?, ?)
	`, r.Id, r.TaskId, r.SnapshotN, r.Stage, definitionID, r.AgentSessionId, r.Vendor, r.Model, string(r.State), r.Error, 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")
}

type Reviewer struct {
	Session api.AgentSession
	Review  api.Review
}

func (s *Store) StartReviewers(ctx context.Context, projectID, definitionID string, reviewers []Reviewer, failed []api.Review) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		for _, r := range failed {
			err := insertReview(ctx, tx, emit, r, definitionID)
			if err != nil {
				return err
			}
		}
		for _, r := range reviewers {
			r.Session.State = api.AgentSessionStateRunning
			err := insertSession(ctx, tx, emit, r.Session, projectID)
			if err != nil {
				return err
			}
			r.Review.State = api.ReviewStateRunning
			r.Review.AgentSessionId = r.Session.Id
			err = insertReview(ctx, tx, emit, r.Review, definitionID)
			if err != nil {
				return err
			}
		}
		return nil
	})
}

func (s *Store) FinishSettling(ctx context.Context, taskID string, n int) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		_, err := tx.ExecContext(ctx, `UPDATE snapshots SET settle_state = 'done' WHERE task_id = ? AND n = ?`, taskID, n)
		return err
	})
}

type SnapshotRef struct {
	TaskID string
	N      int
}

type StageStart struct {
	TaskID       string
	SnapshotN    int
	DefinitionID string
	Definition   api.WorkflowDefinition
	Stage        string
	Inputs       []api.Review
}

type EndedReviewer struct {
	ReviewID  string
	SessionID string
	ProjectID string
	RepoPath  string
}

const currentReview = `(SELECT max(sn.n) FROM snapshots sn WHERE sn.task_id = tasks.id AND sn.review = 1)`

func endReviews(ctx context.Context, tx *sql.Tx, emit emitFunc) ([]EndedReviewer, error) {
	type running struct {
		EndedReviewer
		n       int
		current int
		active  bool
		state   string
	}
	rows, err := tx.QueryContext(ctx, `
		SELECT reviews.id, coalesce(reviews.agent_session_id, ''), tasks.project_id, (SELECT p.repo_path FROM projects p WHERE p.id = tasks.project_id),
			reviews.snapshot_n, coalesce(`+currentReview+`, 0), tasks.status = 'active',
			coalesce((SELECT s.state FROM agent_sessions s WHERE s.id = reviews.agent_session_id), '')
		FROM reviews JOIN tasks ON tasks.id = reviews.task_id
		WHERE reviews.state = 'running'`)
	if err != nil {
		return nil, err
	}
	found := []running{}
	for rows.Next() {
		r := running{}
		err = rows.Scan(&r.ReviewID, &r.SessionID, &r.ProjectID, &r.RepoPath, &r.n, &r.current, &r.active, &r.state)
		if err != nil {
			rows.Close()
			return nil, err
		}
		found = append(found, r)
	}
	err = errors.Join(rows.Err(), rows.Close())
	if err != nil {
		return nil, err
	}

	ended := []EndedReviewer{}
	for _, r := range found {
		reason := ""
		switch {
		case !r.active:
			reason = "the task was closed"
		case r.n < r.current:
			reason = fmt.Sprintf("superseded by snapshot %d", r.current)
		case r.state != string(api.AgentSessionStateRunning):
			reason = "the reviewer stopped before submitting a review"
		default:
			continue
		}
		err = finishReview(ctx, tx, emit, r.ReviewID, api.ReviewStateFailed, "", reason)
		if err != nil {
			return nil, err
		}
		if r.state != "" && r.state != string(api.AgentSessionStateCompleted) {
			err = transitionSessionTx(ctx, tx, emit, r.SessionID, nil, api.AgentSessionStateCompleted, "byre", nil)
			if err != nil {
				return nil, err
			}
		}
		ended = append(ended, r.EndedReviewer)
	}
	return ended, nil
}

func startableReviews(ctx context.Context, tx *sql.Tx) ([]SnapshotRef, error) {
	rows, err := tx.QueryContext(ctx, `
		SELECT c.id, c.n FROM (SELECT tasks.id, `+currentReview+` AS n FROM tasks WHERE status = 'active') c
		WHERE c.n IS NOT NULL AND NOT EXISTS (SELECT 1 FROM reviews r WHERE r.task_id = c.id AND r.snapshot_n >= c.n)`)
	if err != nil {
		return nil, err
	}
	return scanSnapshotRefs(rows)
}

func scanSnapshotRefs(rows *sql.Rows) ([]SnapshotRef, error) {
	defer rows.Close()

	refs := []SnapshotRef{}
	for rows.Next() {
		r := SnapshotRef{}
		err := rows.Scan(&r.TaskID, &r.N)
		if err != nil {
			return nil, err
		}
		refs = append(refs, r)
	}
	return refs, rows.Err()
}

func pendingStages(ctx context.Context, tx *sql.Tx) ([]StageStart, error) {
	rows, err := tx.QueryContext(ctx, `SELECT `+reviewColumns+`, coalesce(definition_id, '') FROM reviews
		WHERE (task_id, snapshot_n) IN (SELECT tasks.id, `+currentReview+` FROM tasks WHERE status = 'active')`)
	if err != nil {
		return nil, err
	}
	type defined struct {
		review api.Review
		defID  string
	}
	found := []defined{}
	for rows.Next() {
		d := defined{}
		d.review, err = scanReview(trailingScanner{rows, []any{&d.defID}})
		if err != nil {
			rows.Close()
			return nil, err
		}
		found = append(found, d)
	}
	err = errors.Join(rows.Err(), rows.Close())
	if err != nil {
		return nil, err
	}

	byTask := map[string]map[string]api.Review{}
	for _, d := range found {
		if byTask[d.review.TaskId] == nil {
			byTask[d.review.TaskId] = map[string]api.Review{}
		}
		byTask[d.review.TaskId][d.review.Stage] = d.review
	}
	defs := map[string]api.WorkflowDefinition{}
	starts := []StageStart{}
	seen := map[string]bool{}
	for _, d := range found {
		if d.review.State != api.ReviewStateDone || d.defID == "" {
			continue
		}
		def, loaded := defs[d.defID]
		if !loaded {
			def, err = loadDefinition(ctx, tx, d.defID)
			if err != nil {
				return nil, err
			}
			defs[d.defID] = def
		}
		byStage := byTask[d.review.TaskId]
		for _, next := range workflow.Next(def, d.review.Stage) {
			key := d.review.TaskId + "\x00" + next
			if _, started := byStage[next]; started || seen[key] {
				continue
			}
			inputs := []api.Review{}
			after := workflow.After(def, next)
			for _, a := range after {
				in := byStage[a]
				if in.State != api.ReviewStateDone {
					break
				}
				inputs = append(inputs, in)
			}
			if len(inputs) != len(after) {
				continue
			}
			seen[key] = true
			starts = append(starts, StageStart{TaskID: d.review.TaskId, SnapshotN: d.review.SnapshotN, DefinitionID: d.defID, Definition: def, Stage: next, Inputs: inputs})
		}
	}
	return starts, nil
}

func settleReviews(ctx context.Context, tx *sql.Tx, pending []StageStart) ([]SnapshotRef, error) {
	waiting := map[SnapshotRef]bool{}
	for _, p := range pending {
		waiting[SnapshotRef{TaskID: p.TaskID, N: p.SnapshotN}] = true
		_, err := tx.ExecContext(ctx, `UPDATE snapshots SET settle_state = '' WHERE task_id = ? AND n = ?`, p.TaskID, p.SnapshotN)
		if err != nil {
			return nil, err
		}
	}
	rows, err := tx.QueryContext(ctx, `
		SELECT sn.task_id, sn.n FROM tasks JOIN snapshots sn ON sn.task_id = tasks.id
		WHERE tasks.status = 'active' AND sn.settle_state = ''
			AND EXISTS (SELECT 1 FROM reviews r WHERE r.task_id = sn.task_id AND r.snapshot_n = sn.n)
			AND NOT EXISTS (SELECT 1 FROM reviews r WHERE r.task_id = sn.task_id AND r.snapshot_n = sn.n AND r.state = 'running')`)
	if err != nil {
		return nil, err
	}
	settled, err := scanSnapshotRefs(rows)
	if err != nil {
		return nil, err
	}
	for _, ref := range settled {
		if waiting[ref] {
			continue
		}
		_, err = tx.ExecContext(ctx, `UPDATE snapshots SET settle_state = 'running' WHERE task_id = ? AND n = ?`, ref.TaskID, ref.N)
		if err != nil {
			return nil, err
		}
	}
	rows, err = tx.QueryContext(ctx, `
		SELECT sn.task_id, sn.n FROM tasks JOIN snapshots sn ON sn.task_id = tasks.id
		WHERE tasks.status = 'active' AND sn.settle_state = 'running'`)
	if err != nil {
		return nil, err
	}
	return scanSnapshotRefs(rows)
}

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, created_at`, taskID)
	if err != nil {
		return nil, err
	}
	return scanReviews(rows)
}

func (s *Store) SnapshotReviews(ctx context.Context, taskID string, n int) ([]api.Review, error) {
	return snapshotReviews(ctx, s.db, taskID, n)
}

func snapshotReviews(ctx context.Context, q querier, taskID string, n int) ([]api.Review, error) {
	rows, err := q.QueryContext(ctx, `SELECT `+reviewColumns+` FROM reviews WHERE task_id = ? AND snapshot_n = ? ORDER BY created_at`, taskID, n)
	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 {
		return finishReview(ctx, tx, emit, id, state, text, errText)
	})
}

func finishReview(ctx context.Context, tx *sql.Tx, emit emitFunc, id string, state api.ReviewState, result, errText string) error {
	res, err := tx.ExecContext(ctx, `UPDATE reviews SET state = ?, result = ?, error = ?, finished_at = ? WHERE id = ? AND state = ?`,
		string(state), result, 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) SetFindingResponse(ctx context.Context, id string, index int, resp api.FindingResponse, actor string) error {
	b, err := json.Marshal(resp)
	if err != nil {
		return err
	}
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		res, err := tx.ExecContext(ctx, `
			UPDATE reviews SET
				finding_responses = json_set(finding_responses, '$."' || ?1 || '"', json(?2)),
				finding_states = CASE WHEN json_extract(finding_states, '$."' || ?1 || '"') = ?3
					THEN json_remove(finding_states, '$."' || ?1 || '"') ELSE finding_states END
			WHERE id = ?4 AND state = ?5 AND ?1 < json_array_length(nullif(result, ''), '$.findings')
		`, index, string(b), string(api.FindingFixed), id, string(api.ReviewStateDone))
		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 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 := ""
	findingResponses := ""
	err := row.Scan(&r.Id, &r.TaskId, &r.SnapshotN, &r.Stage, &r.AgentSessionId, &r.Vendor, &r.Model, &state, &result, &r.Error, &createdAt, &finishedAt, &findingStates, &findingResponses)
	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
	}
	r.FindingResponses = map[string]api.FindingResponse{}
	err = json.Unmarshal([]byte(findingResponses), &r.FindingResponses)
	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
}