byre/internal/store/coordinate.go

package store

import (
	"context"
	"database/sql"
	"time"

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

type Work struct {
	AbandonedMerges []string
	StartMerges     []ClaimedMerge
	FinalizeMerges  []FinalizingMerge
	StartDiscards   []string
	FinishDiscards  []string
	EndReviewers    []EndedReviewer
	StartReviews    []SnapshotRef
	StartStages     []StageStart
	SettleReviews   []SnapshotRef
	ReadyPrompts    []ReadyPrompt
	SelfCloses      []SelfClose
	AdvanceStages   []string
	LeadDeliveries  []LeadDelivery
	QueuedDelivery  []QueuedDelivery
	DueSchedules    []api.Schedule
	NextSchedule    *time.Time
}

func (s *Store) Coordinate(ctx context.Context) (Work, error) {
	work := Work{}
	err := s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		var err error
		work.AbandonedMerges, err = abandonMerges(ctx, tx, emit)
		if err != nil {
			return err
		}
		err = claimHeldMerges(ctx, tx, emit)
		if err != nil {
			return err
		}
		work.StartMerges, err = startableMerges(ctx, tx)
		if err != nil {
			return err
		}
		work.FinalizeMerges, err = finalizingMerges(ctx, tx)
		if err != nil {
			return err
		}
		_, err = tx.ExecContext(ctx, `UPDATE tasks SET discard_state = 'closing'
			WHERE status = 'active' AND discard_state = 'running'
				AND NOT EXISTS (SELECT 1 FROM agent_sessions s WHERE s.id = tasks.discard_session AND s.state = 'running')`)
		if err != nil {
			return err
		}
		work.StartDiscards, err = queryIDs(ctx, tx, `SELECT id FROM tasks WHERE status = 'active' AND discard_state = 'held' AND NOT `+implementerRunning)
		if err != nil {
			return err
		}
		work.FinishDiscards, err = queryIDs(ctx, tx, `SELECT id FROM tasks WHERE status = 'active' AND discard_state = 'closing'`)
		if err != nil {
			return err
		}
		work.EndReviewers, err = endReviews(ctx, tx, emit)
		if err != nil {
			return err
		}
		work.StartReviews, err = startableReviews(ctx, tx)
		if err != nil {
			return err
		}
		work.StartStages, err = pendingStages(ctx, tx)
		if err != nil {
			return err
		}
		work.SettleReviews, err = settleReviews(ctx, tx, work.StartStages)
		if err != nil {
			return err
		}
		work.ReadyPrompts, err = readyPrompts(ctx, tx, emit)
		if err != nil {
			return err
		}
		work.SelfCloses, err = selfCloses(ctx, tx)
		if err != nil {
			return err
		}
		work.AdvanceStages, err = finishedStages(ctx, tx)
		if err != nil {
			return err
		}
		work.LeadDeliveries, err = leadDeliveries(ctx, tx)
		if err != nil {
			return err
		}
		err = endLandings(ctx, tx)
		if err != nil {
			return err
		}
		work.QueuedDelivery, err = queuedDeliveries(ctx, tx)
		if err != nil {
			return err
		}
		work.DueSchedules, work.NextSchedule, err = dueSchedules(ctx, tx)
		return err
	})
	return work, err
}

type SelfClose struct {
	TaskID string
	Ready  string
	By     string
}

func selfCloses(ctx context.Context, tx *sql.Tx) ([]SelfClose, error) {
	rows, err := tx.QueryContext(ctx, `
		SELECT tasks.id, tasks.ready, r.created_by
		FROM tasks JOIN tasks r ON r.id = tasks.parent_task_id JOIN workflow_definitions d ON d.id = r.workflow_definition_id
		WHERE tasks.status = 'active' AND tasks.workflow_stage != '' AND tasks.discard_state = ''
			AND tasks.merging_by = '' AND tasks.finalizing = '' AND tasks.merge_queued_by = '' AND tasks.parent_merge_by = '' AND tasks.landing_session = ''
			AND ((tasks.ready IN ('complete', 'continue') AND d.auto_merge) OR (tasks.ready = 'discard' AND d.auto_discard))
			AND NOT `+implementerRunning+`
		ORDER BY tasks.created_at`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	closes := []SelfClose{}
	for rows.Next() {
		c := SelfClose{}
		err = rows.Scan(&c.TaskID, &c.Ready, &c.By)
		if err != nil {
			return nil, err
		}
		closes = append(closes, c)
	}
	return closes, rows.Err()
}