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

internal/agents/mergein.go

@@ -7,7 +7,6 @@ import (
 	"fmt"
 	"log/slog"
 	"os"
-	"slices"
 	"strings"
 	"sync"
 	"unicode/utf8"
@@ -149,18 +148,84 @@ func (m *Manager) startMerge(ctx context.Context, pm *projectMerge, task api.Tas
 	}
 	pm.holder = task.Id
 
-	sess, err := m.mergeSession(ctx, task, user)
+	err = m.beginMergeTurn(ctx, task, project, user)
 	if err != nil {
 		m.takeMerge(task.Id)
 		pm.holder = ""
+	}
+	return err
+}
+
+func (m *Manager) beginMergeTurn(ctx context.Context, task api.Task, project api.Project, user api.User) error {
+	sess, err := m.mergeSession(ctx, task, user)
+	if err != nil {
 		return err
 	}
-
 	prompt := mergeInPrompt(project.DefaultBranch) + fmt.Sprintf(" When this turn ends, byre squashes this branch onto %s and pushes it; if you stop, call stop_merge first so it doesn't. ", project.DefaultBranch) + summaryRequest
 	go m.launchMerge(sess, task, project, prompt)
 	return nil
 }
 
+func (m *Manager) advanceMergeClaim(taskID string) {
+	login, err := m.store.TakeIdleMergeClaim(context.Background(), taskID)
+	if err != nil {
+		slog.Error("take idle merge claim", "task", taskID, "err", err)
+		return
+	}
+	if login != "" {
+		m.startClaimedMerge(taskID, login)
+	}
+}
+
+func (m *Manager) advanceMergeClaims(ctx context.Context) {
+	ids, err := m.store.MergeClaims(ctx)
+	if err != nil {
+		slog.Error("list merge claims", "err", err)
+		return
+	}
+	for _, id := range ids {
+		m.advanceMergeClaim(id)
+	}
+}
+
+func (m *Manager) startClaimedMerge(taskID, login string) {
+	ctx := context.Background()
+	err := m.beginClaimedMerge(ctx, taskID, login)
+	if err == nil {
+		slog.Info("started claimed merge", "task", taskID, "by", login)
+		return
+	}
+	slog.Error("start claimed merge", "task", taskID, "err", err)
+	task, terr := m.store.GetTask(ctx, taskID)
+	if terr != nil {
+		slog.Error("load task for claimed merge", "task", taskID, "err", terr)
+		return
+	}
+	m.abandonMerge(taskID, task.ProjectId)
+	sessions, serr := m.store.ListSessions(ctx, taskID)
+	if serr != nil {
+		slog.Error("list sessions for claimed merge", "task", taskID, "err", serr)
+		return
+	}
+	m.fileTaskError(ctx, task, sessions, "Starting the held merge failed: "+err.Error())
+}
+
+func (m *Manager) beginClaimedMerge(ctx context.Context, taskID, login string) error {
+	task, err := m.store.GetTask(ctx, taskID)
+	if err != nil {
+		return err
+	}
+	project, err := m.store.GetProject(ctx, task.ProjectId)
+	if err != nil {
+		return err
+	}
+	user, err := m.store.GetUser(ctx, login)
+	if err != nil {
+		return err
+	}
+	return m.beginMergeTurn(ctx, task, project, user)
+}
+
 func (m *Manager) mergeSession(ctx context.Context, task api.Task, user api.User) (api.AgentSession, error) {
 	sess, found, err := m.liveSession(ctx, task.Id)
 	if err != nil {
@@ -362,18 +427,7 @@ func (m *Manager) startNextMerge(pm *projectMerge, projectID string) {
 	}
 
 	for _, task := range held {
-		sessions, err := m.store.ListSessions(ctx, task.Id)
-		if err != nil {
-			slog.Error("list sessions for held merge", "task", task.Id, "err", err)
-			continue
-		}
-		if slices.ContainsFunc(sessions, func(s api.AgentSession) bool {
-			return s.State == api.AgentSessionStateRunning
-		}) {
-			continue
-		}
-
-		login, err := m.store.TakeQueuedMerge(ctx, task.Id)
+		login, claimed, err := m.store.ClaimQueuedMerge(ctx, task.Id)
 		if err != nil {
 			slog.Error("take held merge", "task", task.Id, "err", err)
 			continue
@@ -381,6 +435,11 @@ func (m *Manager) startNextMerge(pm *projectMerge, projectID string) {
 		if login == "" {
 			continue
 		}
+		if claimed {
+			pm.holder = task.Id
+			slog.Info("claimed held merge until the agent's turn ends", "task", task.Id, "by", login)
+			return
+		}
 
 		err = m.startHeldMerge(ctx, pm, task, login)
 		if err == nil {
@@ -388,6 +447,11 @@ func (m *Manager) startNextMerge(pm *projectMerge, projectID string) {
 			return
 		}
 		slog.Error("start held merge", "task", task.Id, "err", err)
+		sessions, serr := m.store.ListSessions(ctx, task.Id)
+		if serr != nil {
+			slog.Error("list sessions for held merge", "task", task.Id, "err", serr)
+			continue
+		}
 		m.fileTaskError(ctx, task, sessions, "Starting the held merge failed: "+err.Error())
 	}
 }