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

internal/agents/compact.go

@@ -71,6 +71,7 @@ func (m *Manager) continueFromHandoff(sessionID, handoff string) {
 	err := m.handoff(ctx, sessionID, handoff)
 	if err != nil {
 		slog.Error("continue from handoff", "session", sessionID, "err", err)
+		m.advanceSessionMergeClaim(ctx, sessionID)
 	}
 }
 

internal/agents/manager.go

@@ -470,6 +470,7 @@ func (m *Manager) Recover(ctx context.Context) error {
 
 	go m.deliverAllLeads()
 	go m.deliverAllQueued(context.Background())
+	go m.advanceMergeClaims(context.Background())
 	go m.startHeldMerges(context.Background())
 	go m.recoverMergeAfterReviews(context.Background())
 	go m.auditAll(context.Background())
@@ -484,15 +485,18 @@ func (m *Manager) failResume(ctx context.Context, sess api.AgentSession, task ap
 	if derr != nil {
 		return derr
 	}
-	next, _, serr := m.store.EndTurn(ctx, sess.Id, "byre", map[string]any{"error": err.Error()}, nil, false, false)
+	end, serr := m.store.EndTurn(ctx, sess.Id, "byre", map[string]any{"error": err.Error()}, nil, false, false, false)
 	if serr != nil {
 		return serr
 	}
 	switch {
-	case next == store.DiscardClosing:
+	case end.Discard == store.DiscardClosing:
 		go m.finishDiscard(task.Id)
-	case next == store.DiscardHeld:
+	case end.Discard == store.DiscardHeld:
 		go m.startDiscard(task.Id)
+	case end.MergeBy != "":
+		go m.startClaimedMerge(task.Id, end.MergeBy)
+	case end.MergeWaiting:
 	case sess.Role == api.Implementer && m.leadsTask(ctx, task.Id, sess.Id):
 		m.abandonMerge(task.Id, task.ProjectId)
 	}
@@ -679,6 +683,18 @@ func (m *Manager) failStart(ctx context.Context, sessionID, actor string, err er
 	if serr != nil {
 		slog.Error("record failed start", "session", sessionID, "err", serr)
 	}
+	go m.advanceSessionMergeClaim(context.Background(), sessionID)
+}
+
+func (m *Manager) advanceSessionMergeClaim(ctx context.Context, sessionID string) {
+	sess, err := m.store.GetSession(ctx, sessionID)
+	if err != nil {
+		slog.Error("load session for merge claim", "session", sessionID, "err", err)
+		return
+	}
+	if sess.TaskId != "" {
+		m.advanceMergeClaim(sess.TaskId)
+	}
 }
 
 func (m *Manager) syncModel(ctx context.Context, l *live, model string) error {
@@ -998,7 +1014,7 @@ func (m *Manager) pump(l *live) {
 		}
 	}
 	m.endMergeTurn(l.id, true)
-	next, _, err := m.store.EndTurn(ctx, l.id, agentActor, map[string]any{"error": "agent process exited mid-turn", "stopped": true}, nil, false, false)
+	end, err := m.store.EndTurn(ctx, l.id, agentActor, map[string]any{"error": "agent process exited mid-turn", "stopped": true}, nil, false, false, false)
 	if err != nil {
 		slog.Error("record session exit", "session", l.id, "err", err)
 	}
@@ -1006,7 +1022,7 @@ func (m *Manager) pump(l *live) {
 		m.finishReviewer(l, "the reviewer's process exited mid-turn")
 		return
 	}
-	switch next {
+	switch end.Discard {
 	case store.DiscardClosing:
 		m.finishDiscard(l.taskID)
 		return
@@ -1014,8 +1030,11 @@ func (m *Manager) pump(l *live) {
 		m.startDiscard(l.taskID)
 		return
 	}
+	if end.MergeBy != "" {
+		go m.startClaimedMerge(l.taskID, end.MergeBy)
+	}
 	if l.owner == "" {
-		if m.leadsTask(ctx, l.taskID, l.id) {
+		if end.MergeBy == "" && !end.MergeWaiting && m.leadsTask(ctx, l.taskID, l.id) {
 			m.abandonMerge(l.taskID, l.projectID)
 		}
 		m.recordError(l, "The agent process exited mid-turn. Send a message to resume it.")
@@ -1310,15 +1329,19 @@ func (m *Manager) turnEnded(ctx context.Context, l *live, res endResult) error {
 		state = map[string]any{"stopped": true}
 	}
 	turn := m.event(l, agentActor, "turn.ended", payload)
-	next, finalizeBy, err := m.store.EndTurn(ctx, l.id, agentActor, state, &turn, review && !held, finished && !handoff)
+	end, err := m.store.EndTurn(ctx, l.id, agentActor, state, &turn, review && !held, finished && !handoff, handoff)
 	if err != nil {
 		return err
 	}
 	if snapshot && !held {
 		m.snapshotSoon(liveTarget(l), api.SnapshotTriggerTurnEnd, 0, review, false)
 	}
+	if end.MergeBy != "" {
+		go m.startClaimedMerge(l.taskID, end.MergeBy)
+		return nil
+	}
 
-	switch next {
+	switch end.Discard {
 	case store.DiscardClosing:
 		m.logUndelivered(ctx, l.id, "turn ended in discard")
 		go m.finishDiscard(l.taskID)
@@ -1337,14 +1360,14 @@ func (m *Manager) turnEnded(ctx context.Context, l *live, res endResult) error {
 		go m.continueFromHandoff(l.id, res.result)
 		return nil
 	}
-	if finalizeBy != "" {
+	if end.FinalizeBy != "" {
 		m.logUndelivered(ctx, l.id, "turn ended in merge")
 		keep = owner
-		go m.finishMerge(l, finalizeBy)
+		go m.finishMerge(l, end.FinalizeBy)
 		return nil
 	}
 	mergeBy := ""
-	if m.leadsTask(ctx, l.taskID, l.id) {
+	if !end.MergeWaiting && m.leadsTask(ctx, l.taskID, l.id) {
 		mergeBy = m.takeMerge(l.taskID)
 	}
 	switch {

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)
+		login, claimed, err := m.store.ClaimQueuedMerge(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)
 		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())
 	}
 }

internal/store/discard.go

@@ -115,9 +115,15 @@ func (s *Store) StartDiscardTurn(ctx context.Context, taskID, sessionID, actor s
 	})
 }
 
-func (s *Store) EndTurn(ctx context.Context, sessionID, actor string, payload map[string]any, turn *api.Event, reviewRequest, finalizeMerge bool) (DiscardState, string, error) {
+type TurnEnd struct {
-	next := DiscardNone
+	Discard      DiscardState
-	finalizing := ""
+	FinalizeBy   string
+	MergeBy      string
+	MergeWaiting bool
+}
+
+func (s *Store) EndTurn(ctx context.Context, sessionID, actor string, payload map[string]any, turn *api.Event, reviewRequest, finalizeMerge, continues bool) (TurnEnd, error) {
+	end := TurnEnd{Discard: DiscardNone}
 	err := s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
 		taskID := ""
 		projectID := ""
@@ -126,10 +132,14 @@ func (s *Store) EndTurn(ctx context.Context, sessionID, actor string, payload ma
 		state := DiscardNone
 		discardSession := ""
 		lead := false
+		mergeAfterTurn := false
+		othersRunning := false
 		err := tx.QueryRowContext(ctx, `
 			SELECT coalesce(t.id, ''), coalesce(t.project_id, ''), s.owner, s.role, coalesce(t.discard_state, ''), coalesce(t.discard_session, ''),
-				s.member = coalesce((SELECT f.member FROM agent_sessions f WHERE f.task_id = s.task_id AND f.role = s.role ORDER BY f.created_at LIMIT 1), '')
+				s.member = coalesce((SELECT f.member FROM agent_sessions f WHERE f.task_id = s.task_id AND f.role = s.role ORDER BY f.created_at LIMIT 1), ''),
-			FROM agent_sessions s LEFT JOIN tasks t ON t.id = s.task_id WHERE s.id = ?`, sessionID).Scan(&taskID, &projectID, &owner, &role, &state, &discardSession, &lead)
+				coalesce(t.merge_after_turn, 0),
+				EXISTS (SELECT 1 FROM agent_sessions o WHERE o.task_id = s.task_id AND o.role = s.role AND o.state = ? AND o.id != s.id)
+			FROM agent_sessions s LEFT JOIN tasks t ON t.id = s.task_id WHERE s.id = ?`, string(api.AgentSessionStateRunning), sessionID).Scan(&taskID, &projectID, &owner, &role, &state, &discardSession, &lead, &mergeAfterTurn, &othersRunning)
 		if errors.Is(err, sql.ErrNoRows) {
 			return ErrNotFound
 		}
@@ -167,19 +177,23 @@ func (s *Store) EndTurn(ctx context.Context, sessionID, actor string, payload ma
 		}
 		switch {
 		case state == DiscardHeld:
-			next = DiscardHeld
+			end.Discard = DiscardHeld
 		case state == DiscardRunning && discardSession == sessionID:
-			next = DiscardClosing
+			end.Discard = DiscardClosing
 			_, err = tx.ExecContext(ctx, `UPDATE tasks SET discard_state = ? WHERE id = ?`, string(DiscardClosing), taskID)
+		case mergeAfterTurn && (othersRunning || continues):
+			end.MergeWaiting = true
+		case mergeAfterTurn:
+			err = tx.QueryRowContext(ctx, `UPDATE tasks SET merge_after_turn = 0 WHERE id = ? RETURNING merging_by`, taskID).Scan(&end.MergeBy)
 		case state == DiscardNone && finalizeMerge && lead:
-			err = tx.QueryRowContext(ctx, `UPDATE tasks SET finalizing = merging_by, merging_by = '' WHERE id = ? AND merging_by != '' RETURNING finalizing`, taskID).Scan(&finalizing)
+			err = tx.QueryRowContext(ctx, `UPDATE tasks SET finalizing = merging_by, merging_by = '' WHERE id = ? AND merging_by != '' RETURNING finalizing`, taskID).Scan(&end.FinalizeBy)
 			if errors.Is(err, sql.ErrNoRows) {
 				err = nil
 			}
 		}
 		return err
 	})
-	return next, finalizing, err
+	return end, err
 }
 
 func (s *Store) CloseDiscard(ctx context.Context, taskID string) error {

internal/store/migrations/089-add-task-merge-after-turn.sql

@@ -0,0 +1 @@
+ALTER TABLE tasks ADD COLUMN merge_after_turn INTEGER NOT NULL DEFAULT 0;

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