byre/internal/store/merges.go

package store

import (
	"context"
	"database/sql"
	"errors"

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

const mergeSlotHolder = `o.project_id = tasks.project_id AND o.status = 'active' AND (o.merge_state IN ('claimed', 'turn', 'finalizing') OR o.pushing_to_main = 1)`

const implementerRunning = `EXISTS (SELECT 1 FROM agent_sessions s WHERE s.task_id = tasks.id AND s.role = 'implementer' AND s.state = 'running')`

func (s *Store) RequestMerge(ctx context.Context, id, login string) (bool, error) {
	held := false
	err := s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		projectID, status, discardState, mergeState := "", "", "", ""
		err := tx.QueryRowContext(ctx, `
			SELECT project_id, status, discard_state, merge_state,
				EXISTS (SELECT 1 FROM tasks o WHERE `+mergeSlotHolder+`)
				OR EXISTS (SELECT 1 FROM tasks o WHERE o.project_id = tasks.project_id AND o.status = 'active' AND o.merge_state = 'held')
			FROM tasks WHERE id = ?`, id).Scan(&projectID, &status, &discardState, &mergeState, &held)
		if errors.Is(err, sql.ErrNoRows) {
			return ErrNotFound
		}
		if err != nil {
			return err
		}
		switch {
		case status != string(api.TaskStatusActive) || discardState != "" || mergeState == "held":
			return ErrNotFound
		case mergeState != "":
			return ErrMerging
		}
		if held {
			_, err = tx.ExecContext(ctx, `UPDATE tasks SET merge_state = 'held', merge_by = ?, merge_queued_at = ? WHERE id = ?`, login, now(), id)
			if err != nil {
				return err
			}
			return emit(api.Event{Kind: "task.merge_queued", ProjectId: &projectID, TaskId: &id, Actor: login, Payload: map[string]any{"id": id}})
		}
		_, err = tx.ExecContext(ctx, `UPDATE tasks SET merge_state = 'claimed', merge_by = ? WHERE id = ?`, login, id)
		if err != nil {
			return err
		}
		return emit(api.Event{Kind: "task.merge_started", ProjectId: &projectID, TaskId: &id, Actor: login, Payload: map[string]any{"id": id}})
	})
	return held, err
}

type ClaimedMerge struct {
	TaskID string
	By     string
}

type FinalizingMerge struct {
	Task      api.Task
	By        string
	Commit    string
	Notes     string
	Unchanged bool
}

type idRow struct {
	id      string
	project string
	by      string
}

func queryIDRows(ctx context.Context, tx *sql.Tx, query string) ([]idRow, error) {
	rows, err := tx.QueryContext(ctx, query)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	found := []idRow{}
	for rows.Next() {
		r := idRow{}
		err = rows.Scan(&r.id, &r.project, &r.by)
		if err != nil {
			return nil, err
		}
		found = append(found, r)
	}
	return found, rows.Err()
}

func abandonMerges(ctx context.Context, tx *sql.Tx, emit emitFunc) ([]string, error) {
	found, err := queryIDRows(ctx, tx, `SELECT id, project_id, merge_by FROM tasks
		WHERE status = 'active' AND merge_state = 'turn' AND NOT `+implementerRunning)
	if err != nil {
		return nil, err
	}
	ids := []string{}
	for _, r := range found {
		_, err = tx.ExecContext(ctx, `UPDATE tasks SET merge_state = '', merge_by = '' WHERE id = ?`, r.id)
		if err != nil {
			return nil, err
		}
		err = emit(api.Event{Kind: "task.merge_ended", ProjectId: &r.project, TaskId: &r.id, Actor: "byre", Payload: map[string]any{"id": r.id}})
		if err != nil {
			return nil, err
		}
		ids = append(ids, r.id)
	}
	return ids, nil
}

func claimHeldMerges(ctx context.Context, tx *sql.Tx, emit emitFunc) error {
	found, err := queryIDRows(ctx, tx, `SELECT id, project_id, merge_by FROM tasks
		WHERE status = 'active' AND merge_state = 'held' AND discard_state = ''
			AND NOT EXISTS (SELECT 1 FROM tasks o WHERE `+mergeSlotHolder+`)
			AND NOT EXISTS (SELECT 1 FROM tasks o WHERE o.project_id = tasks.project_id AND o.status = 'active' AND o.merge_state = 'held'
				AND (o.merge_queued_at < tasks.merge_queued_at OR (o.merge_queued_at = tasks.merge_queued_at AND o.id < tasks.id)))`)
	if err != nil {
		return err
	}
	for _, r := range found {
		_, err = tx.ExecContext(ctx, `UPDATE tasks SET merge_state = 'claimed', merge_queued_at = '' WHERE id = ?`, r.id)
		if err != nil {
			return err
		}
		err = emit(api.Event{Kind: "task.merge_unqueued", ProjectId: &r.project, TaskId: &r.id, Actor: "byre", Payload: map[string]any{"id": r.id}})
		if err != nil {
			return err
		}
		err = emit(api.Event{Kind: "task.merge_started", ProjectId: &r.project, TaskId: &r.id, Actor: r.by, Payload: map[string]any{"id": r.id}})
		if err != nil {
			return err
		}
	}
	return nil
}

func startableMerges(ctx context.Context, tx *sql.Tx) ([]ClaimedMerge, error) {
	found, err := queryIDRows(ctx, tx, `SELECT id, project_id, merge_by FROM tasks
		WHERE status = 'active' AND merge_state = 'claimed' AND NOT `+implementerRunning)
	if err != nil {
		return nil, err
	}
	claims := []ClaimedMerge{}
	for _, r := range found {
		claims = append(claims, ClaimedMerge{TaskID: r.id, By: r.by})
	}
	return claims, nil
}

func finalizingMerges(ctx context.Context, tx *sql.Tx) ([]FinalizingMerge, error) {
	rows, err := tx.QueryContext(ctx, `SELECT `+taskColumns+`, merge_by, finalizing_commit, finalizing_notes, finalizing_unchanged FROM tasks WHERE status = 'active' AND merge_state = 'finalizing'`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	merges := []FinalizingMerge{}
	for rows.Next() {
		f := FinalizingMerge{}
		f.Task, err = scanTask(trailingScanner{rows, []any{&f.By, &f.Commit, &f.Notes, &f.Unchanged}})
		if err != nil {
			return nil, err
		}
		merges = append(merges, f)
	}
	return merges, rows.Err()
}

func (s *Store) ClaimMainPush(ctx context.Context, id string) (string, error) {
	holder := ""
	err := s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		err := tx.QueryRowContext(ctx, `SELECT coalesce((SELECT o.id FROM tasks o WHERE `+mergeSlotHolder+` LIMIT 1), '') FROM tasks WHERE id = ?`, id).Scan(&holder)
		if errors.Is(err, sql.ErrNoRows) {
			return ErrNotFound
		}
		if err != nil || holder != "" {
			return err
		}
		_, err = tx.ExecContext(ctx, `UPDATE tasks SET pushing_to_main = 1 WHERE id = ?`, id)
		return err
	})
	return holder, err
}

func (s *Store) EndMainPush(ctx context.Context, id string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		_, err := tx.ExecContext(ctx, `UPDATE tasks SET pushing_to_main = 0 WHERE id = ?`, id)
		return err
	})
}

func (s *Store) ClearMainPushes(ctx context.Context) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		_, err := tx.ExecContext(ctx, `UPDATE tasks SET pushing_to_main = 0 WHERE pushing_to_main = 1`)
		return err
	})
}