Wait for the agent's turn instead of dropping a held merge

internal/agents/manager.go

@@ -495,7 +495,7 @@ func (m *Manager) failResume(ctx context.Context, sess api.AgentSession, task ap
 	case end.Discard == store.DiscardHeld:
 		go m.startDiscard(task.Id)
 	case end.MergeBy != "":
-		go m.startClaimedMerge(task.Id, end.MergeBy)
+		go m.advanceMergeClaim(task.Id)
 	case end.MergeWaiting:
 	case sess.Role == api.Implementer && m.leadsTask(ctx, task.Id, sess.Id):
 		m.abandonMerge(task.Id, task.ProjectId)
@@ -1031,7 +1031,7 @@ func (m *Manager) pump(l *live) {
 		return
 	}
 	if end.MergeBy != "" {
-		go m.startClaimedMerge(l.taskID, end.MergeBy)
+		go m.advanceMergeClaim(l.taskID)
 	}
 	if l.owner == "" {
 		if end.MergeBy == "" && !end.MergeWaiting && m.leadsTask(ctx, l.taskID, l.id) {
@@ -1337,7 +1337,7 @@ func (m *Manager) turnEnded(ctx context.Context, l *live, res endResult) error {
 		m.snapshotSoon(liveTarget(l), api.SnapshotTriggerTurnEnd, 0, review, false)
 	}
 	if end.MergeBy != "" {
-		go m.startClaimedMerge(l.taskID, end.MergeBy)
+		go m.advanceMergeClaim(l.taskID)
 		return nil
 	}
 

internal/agents/mergein.go

@@ -161,20 +161,12 @@ func (m *Manager) beginMergeTurn(ctx context.Context, task api.Task, project api
 	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, mergeTurnPrompt(project))
-	go m.launchMerge(sess, task, project, prompt)
 	return nil
 }
 
-func (m *Manager) advanceMergeClaim(taskID string) {
+func mergeTurnPrompt(project api.Project) string {
-	login, err := m.store.TakeIdleMergeClaim(context.Background(), taskID)
+	return 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
-	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) {
@@ -188,11 +180,10 @@ func (m *Manager) advanceMergeClaims(ctx context.Context) {
 	}
 }
 
-func (m *Manager) startClaimedMerge(taskID, login string) {
+func (m *Manager) advanceMergeClaim(taskID string) {
 	ctx := context.Background()
-	err := m.beginClaimedMerge(ctx, taskID, login)
+	err := m.beginClaimedMerge(ctx, taskID)
 	if err == nil {
-		slog.Info("started claimed merge", "task", taskID, "by", login)
 		return
 	}
 	slog.Error("start claimed merge", "task", taskID, "err", err)
@@ -210,7 +201,11 @@ func (m *Manager) startClaimedMerge(taskID, login string) {
 	m.fileTaskError(ctx, task, sessions, "Starting the held merge failed: "+err.Error())
 }
 
-func (m *Manager) beginClaimedMerge(ctx context.Context, taskID, login string) error {
+func (m *Manager) beginClaimedMerge(ctx context.Context, taskID string) error {
+	login, err := m.store.MergeClaim(ctx, taskID)
+	if err != nil || login == "" {
+		return err
+	}
 	task, err := m.store.GetTask(ctx, taskID)
 	if err != nil {
 		return err
@@ -223,21 +218,35 @@ func (m *Manager) beginClaimedMerge(ctx context.Context, taskID, login string) e
 	if err != nil {
 		return err
 	}
-	return m.beginMergeTurn(ctx, task, project, user)
+	if user.GitName == "" || user.GitEmail == "" {
+		return ErrNoGitIdentity
+	}
+	sess, err := m.leadSession(ctx, task, user)
+	if err != nil {
+		return err
+	}
+	reserved, err := m.store.ReserveMergeTurn(ctx, task.Id, sess.Id)
+	if err != nil || !reserved {
+		return err
+	}
+	slog.Info("started claimed merge", "task", taskID, "by", login)
+	go m.launchMerge(sess, task, project, mergeTurnPrompt(project))
+	return nil
+}
+
+func (m *Manager) leadSession(ctx context.Context, task api.Task, user api.User) (api.AgentSession, error) {
+	sess, found, err := m.liveSession(ctx, task.Id)
+	if err != nil || found {
+		return sess, err
+	}
+	return m.newSession(ctx, task, user, sess.Member, "")
 }
 
 func (m *Manager) mergeSession(ctx context.Context, task api.Task, user api.User) (api.AgentSession, error) {
-	sess, found, err := m.liveSession(ctx, task.Id)
+	sess, err := m.leadSession(ctx, task, user)
 	if err != nil {
 		return api.AgentSession{}, err
 	}
-	if !found {
-		sess, err = m.newSession(ctx, task, user, sess.Member, "")
-		if err != nil {
-			return api.AgentSession{}, err
-		}
-	}
-
 	err = transition(ctx, m.store, sess.Id, api.AgentSessionStateIdle, api.AgentSessionStateRunning, "byre")
 	return sess, err
 }
@@ -427,7 +436,7 @@ func (m *Manager) startNextMerge(pm *projectMerge, projectID string) {
 	}
 
 	for _, task := range held {
-		login, claimed, err := m.store.ClaimQueuedMerge(ctx, task.Id)
+		login, err := m.store.ClaimQueuedMerge(ctx, task.Id)
 		if err != nil {
 			slog.Error("take held merge", "task", task.Id, "err", err)
 			continue
@@ -435,42 +444,13 @@ func (m *Manager) startNextMerge(pm *projectMerge, projectID string) {
 		if login == "" {
 			continue
 		}
-		if claimed {
+		pm.holder = task.Id
-			pm.holder = task.Id
+		slog.Info("claimed held merge", "task", task.Id, "by", login)
-			slog.Info("claimed held merge until the agent's turn ends", "task", task.Id, "by", login)
+		go m.advanceMergeClaim(task.Id)
-			return
+		return
-		}
-
-		err = m.startHeldMerge(ctx, pm, task, login)
-		if err == nil {
-			slog.Info("started held merge", "task", task.Id, "by", login)
-			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())
 	}
 }
 
-func (m *Manager) startHeldMerge(ctx context.Context, pm *projectMerge, task api.Task, login string) error {
-	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
-	}
-	if user.GitName == "" || user.GitEmail == "" {
-		return ErrNoGitIdentity
-	}
-	return m.startMerge(ctx, pm, task, project, user)
-}
-
 func (m *Manager) fileTaskError(ctx context.Context, task api.Task, sessions []api.AgentSession, msg string) {
 	if len(sessions) == 0 {
 		slog.Error("no session to file task error on", "task", task.Id, "error", msg)

internal/store/discard.go

@@ -184,7 +184,7 @@ func (s *Store) EndTurn(ctx context.Context, sessionID, actor string, payload ma
 		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)
+			err = tx.QueryRowContext(ctx, `SELECT merging_by FROM tasks WHERE id = ?`, 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(&end.FinalizeBy)
 			if errors.Is(err, sql.ErrNoRows) {

internal/store/store.go

@@ -712,13 +712,11 @@ 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) {
+func (s *Store) ClaimQueuedMerge(ctx context.Context, id string) (string, 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 = ?`,
+		err := tx.QueryRowContext(ctx, `SELECT project_id, merge_queued_by FROM tasks WHERE id = ?`, id).Scan(&projectID, &login)
-			string(api.Implementer), string(api.AgentSessionStateRunning), id).Scan(&projectID, &login, &claimed)
 		if errors.Is(err, sql.ErrNoRows) {
 			return ErrNotFound
 		}
@@ -733,32 +731,59 @@ func (s *Store) ClaimQueuedMerge(ctx context.Context, id string) (string, bool,
 			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}})
+		res, err := tx.ExecContext(ctx, `UPDATE tasks SET merging_by = ?, merge_after_turn = 1 WHERE id = ? AND discard_state = '' AND finalizing = ''`, login, id)
-	})
+		if err != nil {
-	return login, claimed, err
+			return err
-}
+		}
-
+		n, err := res.RowsAffected()
-func (s *Store) TakeIdleMergeClaim(ctx context.Context, id string) (string, error) {
+		if err != nil {
-	login := ""
+			return err
-	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
+		if n == 0 {
-			AND NOT EXISTS (SELECT 1 FROM agent_sessions s WHERE s.task_id = tasks.id AND s.role = ? AND s.state = ?) RETURNING merging_by`,
+			login = ""
-			id, string(api.Implementer), string(api.AgentSessionStateRunning)).Scan(&login)
-		if errors.Is(err, sql.ErrNoRows) {
 			return 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, err
 }
 
+func (s *Store) MergeClaim(ctx context.Context, id string) (string, error) {
+	login := ""
+	err := s.db.QueryRowContext(ctx, `SELECT merging_by FROM tasks WHERE id = ? AND merge_after_turn = 1`, id).Scan(&login)
+	if errors.Is(err, sql.ErrNoRows) {
+		return "", nil
+	}
+	return login, err
+}
+
+func (s *Store) ReserveMergeTurn(ctx context.Context, id, sessionID string) (bool, error) {
+	reserved := false
+	err := s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
+		res, err := tx.ExecContext(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 = ?)`,
+			id, string(api.Implementer), string(api.AgentSessionStateRunning))
+		if err != nil {
+			return err
+		}
+		n, err := res.RowsAffected()
+		if err != nil || n == 0 {
+			return err
+		}
+		idle := api.AgentSessionStateIdle
+		err = transitionSessionTx(ctx, tx, emit, sessionID, &idle, api.AgentSessionStateRunning, "byre", nil)
+		if err != nil {
+			return err
+		}
+		reserved = true
+		return nil
+	})
+	return reserved, 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 {
@@ -1416,46 +1441,50 @@ func (s *Store) TransitionSession(ctx context.Context, id string, from, to api.A
 
 func (s *Store) transitionSession(ctx context.Context, id string, from *api.AgentSessionState, state api.AgentSessionState, actor string, payload map[string]any) error {
 	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
-		taskID := ""
+		return transitionSessionTx(ctx, tx, emit, id, from, state, actor, payload)
-		projectID := ""
-		owner := ""
-		current := ""
-		role := ""
-		discard := ""
-		finalizing := ""
-		err := tx.QueryRowContext(ctx, `
-			SELECT coalesce(t.id, ''), coalesce(t.project_id, ''), s.owner, s.state, s.role, coalesce(t.discard_state, ''), coalesce(t.finalizing, '') FROM agent_sessions s LEFT JOIN tasks t ON t.id = s.task_id WHERE s.id = ?
-		`, id).Scan(&taskID, &projectID, &owner, &current, &role, &discard, &finalizing)
-		if errors.Is(err, sql.ErrNoRows) {
-			return ErrNotFound
-		}
-		if err != nil {
-			return err
-		}
-
-		if from != nil && current != string(*from) {
-			return fmt.Errorf("%w: session %s is %s, not %s", ErrStateChanged, id, current, *from)
-		}
-		if state == api.AgentSessionStateRunning && role == string(api.Implementer) && discard != "" {
-			return fmt.Errorf("%w: task %s is being discarded", ErrStateChanged, taskID)
-		}
-		if state == api.AgentSessionStateRunning && role == string(api.Implementer) && finalizing != "" {
-			return fmt.Errorf("%w: task %s is finishing its merge", ErrStateChanged, taskID)
-		}
-
-		_, err = tx.ExecContext(ctx, `UPDATE agent_sessions SET state = ? WHERE id = ?`, string(state), id)
-		if err != nil {
-			return err
-		}
-
-		if payload == nil {
-			payload = map[string]any{}
-		}
-		payload["state"] = string(state)
-		return emit(sessionEvent("session.state", projectID, taskID, owner, id, actor, payload))
 	})
 }
 
+func transitionSessionTx(ctx context.Context, tx *sql.Tx, emit emitFunc, id string, from *api.AgentSessionState, state api.AgentSessionState, actor string, payload map[string]any) error {
+	taskID := ""
+	projectID := ""
+	owner := ""
+	current := ""
+	role := ""
+	discard := ""
+	finalizing := ""
+	err := tx.QueryRowContext(ctx, `
+		SELECT coalesce(t.id, ''), coalesce(t.project_id, ''), s.owner, s.state, s.role, coalesce(t.discard_state, ''), coalesce(t.finalizing, '') FROM agent_sessions s LEFT JOIN tasks t ON t.id = s.task_id WHERE s.id = ?
+	`, id).Scan(&taskID, &projectID, &owner, &current, &role, &discard, &finalizing)
+	if errors.Is(err, sql.ErrNoRows) {
+		return ErrNotFound
+	}
+	if err != nil {
+		return err
+	}
+
+	if from != nil && current != string(*from) {
+		return fmt.Errorf("%w: session %s is %s, not %s", ErrStateChanged, id, current, *from)
+	}
+	if state == api.AgentSessionStateRunning && role == string(api.Implementer) && discard != "" {
+		return fmt.Errorf("%w: task %s is being discarded", ErrStateChanged, taskID)
+	}
+	if state == api.AgentSessionStateRunning && role == string(api.Implementer) && finalizing != "" {
+		return fmt.Errorf("%w: task %s is finishing its merge", ErrStateChanged, taskID)
+	}
+
+	_, err = tx.ExecContext(ctx, `UPDATE agent_sessions SET state = ? WHERE id = ?`, string(state), id)
+	if err != nil {
+		return err
+	}
+
+	if payload == nil {
+		payload = map[string]any{}
+	}
+	payload["state"] = string(state)
+	return emit(sessionEvent("session.state", projectID, taskID, owner, id, actor, payload))
+}
+
 func (s *Store) SessionsWithSubagents(ctx context.Context) ([]api.AgentSession, error) {
 	return s.querySessions(ctx, `SELECT `+sessionColumns+` FROM agent_sessions WHERE subagents > 0`)
 }