Run reviews and nudges on the state-driven coordinator
This commit is contained in:
Ian Gulliver
+295
-41
@@ -5,9 +5,11 @@ import (
|
||||
"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`
|
||||
@@ -18,20 +20,280 @@ func (s *Store) CreateReview(ctx context.Context, r api.Review, definitionID str
|
||||
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, 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")
|
||||
return insertReview(ctx, tx, emit, r, definitionID, finishedAt)
|
||||
})
|
||||
}
|
||||
|
||||
func insertReview(ctx context.Context, tx *sql.Tx, emit emitFunc, r api.Review, definitionID, finishedAt string) error {
|
||||
_, 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, formatTime(*r.FinishedAt))
|
||||
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) ReviewDefinition(ctx context.Context, id string) (string, api.WorkflowDefinition, error) {
|
||||
defID := sql.NullString{}
|
||||
err := s.db.QueryRowContext(ctx, `SELECT definition_id FROM reviews WHERE id = ?`, id).Scan(&defID)
|
||||
@@ -70,29 +332,17 @@ func (s *Store) ListReviews(ctx context.Context, taskID string) ([]api.Review, e
|
||||
}
|
||||
|
||||
func (s *Store) SnapshotReviews(ctx context.Context, taskID string, n int) ([]api.Review, error) {
|
||||
rows, err := s.db.QueryContext(ctx, `SELECT `+reviewColumns+` FROM reviews WHERE task_id = ? AND snapshot_n = ? ORDER BY created_at`, taskID, n)
|
||||
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) 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) LatestDoneStages(ctx context.Context) ([]string, error) {
|
||||
return queryIDs(ctx, s.db, `
|
||||
SELECT r.id FROM reviews r JOIN tasks t ON t.id = r.task_id
|
||||
WHERE t.status = ? AND r.state = ? AND r.definition_id IS NOT NULL
|
||||
AND r.snapshot_n = (SELECT max(snapshot_n) FROM reviews WHERE task_id = r.task_id)
|
||||
ORDER BY r.created_at`, string(api.TaskStatusActive), string(api.ReviewStateDone))
|
||||
}
|
||||
|
||||
func (s *Store) FinishReview(ctx context.Context, id string, state api.ReviewState, result map[string]any, errText string) error {
|
||||
text := ""
|
||||
if result != nil {
|
||||
@@ -103,22 +353,26 @@ func (s *Store) FinishReview(ctx context.Context, id string, state api.ReviewSta
|
||||
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")
|
||||
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, `
|
||||
|
||||
Reference in New Issue
Block a user