Claim a held merge as soon as the merge ahead of it finishes

internal/store/store.go

@@ -496,7 +496,7 @@ func (s *Store) TakeMergingBy(ctx context.Context, id string) (string, error) {
 		if login == "" {
 			return nil
 		}
-		_, err = tx.ExecContext(ctx, `UPDATE tasks SET merging_by = '' WHERE id = ?`, id)
+		_, err = tx.ExecContext(ctx, `UPDATE tasks SET merging_by = '', merge_after_turn = 0 WHERE id = ?`, id)
 		if err != nil {
 			return err
 		}
@@ -712,6 +712,72 @@ func (s *Store) TakeQueuedMerge(ctx context.Context, id string) (string, error)
 	return login, err
 }
 
+func (s *Store) ClaimQueuedMerge(ctx context.Context, id string) (string, bool, error) {
+	login := ""
+	claimed := false
+	err := s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
+		projectID := ""
+		err := tx.QueryRowContext(ctx, `SELECT project_id, merge_queued_by, EXISTS (SELECT 1 FROM agent_sessions s WHERE s.task_id = tasks.id AND s.role = ? AND s.state = ?) FROM tasks WHERE id = ?`,
+			string(api.Implementer), string(api.AgentSessionStateRunning), id).Scan(&projectID, &login, &claimed)
+		if errors.Is(err, sql.ErrNoRows) {
+			return ErrNotFound
+		}
+		if err != nil {
+			return err
+		}
+		if login == "" {
+			return nil
+		}
+		_, err = tx.ExecContext(ctx, `UPDATE tasks SET merge_queued_by = '', merge_queued_at = '' WHERE id = ?`, id)
+		if err != nil {
+			return err
+		}
+		err = emit(api.Event{Kind: "task.merge_unqueued", ProjectId: &projectID, TaskId: &id, Actor: "byre", Payload: map[string]any{"id": id}})
+		if err != nil || !claimed {
+			return err
+		}
+		_, err = tx.ExecContext(ctx, `UPDATE tasks SET merging_by = ?, merge_after_turn = 1 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 login, claimed, err
+}
+
+func (s *Store) TakeIdleMergeClaim(ctx context.Context, id string) (string, error) {
+	login := ""
+	err := s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
+		err := tx.QueryRowContext(ctx, `UPDATE tasks SET merge_after_turn = 0 WHERE id = ? AND merge_after_turn = 1
+			AND NOT EXISTS (SELECT 1 FROM agent_sessions s WHERE s.task_id = tasks.id AND s.role = ? AND s.state = ?) RETURNING merging_by`,
+			id, string(api.Implementer), string(api.AgentSessionStateRunning)).Scan(&login)
+		if errors.Is(err, sql.ErrNoRows) {
+			return nil
+		}
+		return err
+	})
+	return login, err
+}
+
+func (s *Store) MergeClaims(ctx context.Context) ([]string, error) {
+	rows, err := s.db.QueryContext(ctx, `SELECT id FROM tasks WHERE status = ? AND merge_after_turn = 1`, string(api.TaskStatusActive))
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	ids := []string{}
+	for rows.Next() {
+		id := ""
+		err = rows.Scan(&id)
+		if err != nil {
+			return nil, err
+		}
+		ids = append(ids, id)
+	}
+	return ids, rows.Err()
+}
+
 type MergeAfterReview struct {
 	By       string
 	ReviewID string