diff --git a/internal/agents/compact.go b/internal/agents/compact.go index 288b8f3..4857b50 100644 --- a/internal/agents/compact.go +++ b/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) } } diff --git a/internal/agents/manager.go b/internal/agents/manager.go index 038661a..936c11c 100644 --- a/internal/agents/manager.go +++ b/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 { diff --git a/internal/agents/mergein.go b/internal/agents/mergein.go index eec6c02..1c9e786 100644 --- a/internal/agents/mergein.go +++ b/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()) } } diff --git a/internal/store/discard.go b/internal/store/discard.go index e99ef9b..fce73e7 100644 --- a/internal/store/discard.go +++ b/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) { - next := DiscardNone - finalizing := "" +type TurnEnd struct { + Discard DiscardState + 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), '') - 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) + 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), ''), + 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 { diff --git a/internal/store/migrations/089-add-task-merge-after-turn.sql b/internal/store/migrations/089-add-task-merge-after-turn.sql new file mode 100644 index 0000000..5660995 --- /dev/null +++ b/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; diff --git a/internal/store/store.go b/internal/store/store.go index c33c245..93cf085 100644 --- a/internal/store/store.go +++ b/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