byre/internal/agents/mergein.go

package agents

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"slices"
	"strings"
	"sync"
	"unicode/utf8"

	"g.fc.run/misc/byre/internal/api"
	"g.fc.run/misc/byre/internal/git"
	"g.fc.run/misc/byre/internal/mcp"
	"g.fc.run/misc/byre/internal/store"
)

const maxMergeTitle = 72

func mergeInPrompt(target string) string {
	return fmt.Sprintf("Merge %[1]s into this branch: run `git merge --no-ff --no-commit %[1]s`. "+
		"Resolve any conflicts keeping both sides' intent; regenerate generated files instead of hand-merging them. "+
		"Check that the build and tests pass. Leave the merge uncommitted: byre concludes it when the turn ends. "+
		"If resolving the conflicts took substantive changes rather than combining both sides as written, stop there: "+
		"end your turn with a description of those changes, so the review of the turn's snapshot covers them.", target)
}

func (m *Manager) MergeIn(ctx context.Context, task api.Task, user api.User, target string) (api.AgentSession, error) {
	prompt := mergeInPrompt(target) + fmt.Sprintf(" Finally, if the merge commit title (currently \"%s\") no longer describes all the work on this branch, set a better one with set_merge_title.", task.MergeTitle)
	return m.sendMergeIn(ctx, task, user, prompt)
}

func (m *Manager) sendMergeIn(ctx context.Context, task api.Task, user api.User, prompt string) (api.AgentSession, error) {
	sessions, err := m.store.ListSessions(ctx, task.Id)
	if err != nil {
		return api.AgentSession{}, err
	}

	for _, sess := range sessions {
		if sess.State != api.AgentSessionStateCompleted {
			return m.Send(ctx, sess.Id, "byre", prompt, "")
		}
	}

	return m.create(ctx, task, user, prompt, "", "byre", nil)
}

func (m *Manager) Merge(ctx context.Context, task api.Task, project api.Project, user api.User, title string) (api.Task, error) {
	if user.GitName == "" || user.GitEmail == "" {
		return api.Task{}, ErrNoGitIdentity
	}
	if task.MergeQueuedAt != nil {
		return api.Task{}, fmt.Errorf("%w: task %s is already waiting to merge", ErrConflict, task.Id)
	}
	if task.MergingIntoParent {
		return api.Task{}, fmt.Errorf("%w: task %s is merging into its parent", ErrConflict, task.Id)
	}
	err := m.checkNoSubtasks(ctx, task)
	if err != nil {
		return api.Task{}, err
	}
	title, err = cleanMergeTitle(title)
	if err != nil {
		return api.Task{}, err
	}

	pm := m.projectMerge(project.Id)
	pm.mu.Lock()
	defer pm.mu.Unlock()

	if pm.holder == task.Id {
		return api.Task{}, fmt.Errorf("%w: task %s is already merging", ErrConflict, task.Id)
	}
	if title != task.MergeTitle {
		err = m.storeMergeTitle(ctx, task.Id, title, user.Login)
		if err != nil {
			return api.Task{}, err
		}
		task.MergeTitle = title
	}

	switch pm.holder {
	case "":
		err := m.startMerge(ctx, pm, task, project, user)
		if err != nil {
			return api.Task{}, err
		}
	default:
		err := m.store.QueueMerge(ctx, task.Id, user.Login)
		if errors.Is(err, store.ErrNotFound) {
			return api.Task{}, fmt.Errorf("%w: task %s is closed, being discarded or already waiting to merge", ErrConflict, task.Id)
		}
		if err != nil {
			return api.Task{}, err
		}
	}
	m.clearMergeAfterReview(ctx, task.Id)
	return m.store.GetTask(ctx, task.Id)
}

type projectMerge struct {
	mu     sync.Mutex
	holder string
}

func (m *Manager) projectMerge(projectID string) *projectMerge {
	m.mergeMu.Lock()
	defer m.mergeMu.Unlock()
	pm := m.merges[projectID]
	if pm == nil {
		pm = &projectMerge{}
		m.merges[projectID] = pm
	}
	return pm
}

func (m *Manager) startMerge(ctx context.Context, pm *projectMerge, task api.Task, project api.Project, user api.User) error {
	err := m.store.SetMergingBy(ctx, task.Id, user.Login)
	if errors.Is(err, store.ErrNotFound) {
		return fmt.Errorf("%w: task %s is being discarded", ErrConflict, task.Id)
	}
	if err != nil {
		return err
	}
	pm.holder = task.Id

	sess, err := m.mergeSession(ctx, task, user)
	if err != nil {
		m.takeMerge(task.Id)
		pm.holder = ""
		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) mergeSession(ctx context.Context, task api.Task, user api.User) (api.AgentSession, error) {
	sessions, err := m.store.ListSessions(ctx, task.Id)
	if err != nil {
		return api.AgentSession{}, err
	}

	sess := api.AgentSession{}
	for _, s := range sessions {
		if s.State != api.AgentSessionStateCompleted {
			sess = s
		}
	}
	if sess.Id == "" {
		sess, err = m.newSession(ctx, task, user, "")
		if err != nil {
			return api.AgentSession{}, err
		}
	}

	err = transition(ctx, m.store, sess.Id, api.AgentSessionStateIdle, api.AgentSessionStateRunning, "byre")
	return sess, err
}

func (m *Manager) launchMerge(sess api.AgentSession, task api.Task, project api.Project, prompt string) {
	ctx := context.Background()
	err := m.deliverMerge(ctx, sess, task, project, prompt)
	if err == nil {
		return
	}
	slog.Error("start merge turn", "task", task.Id, "err", err)
	m.failStart(ctx, sess.Id, "byre", err)
	m.abandonMerge(task.Id, project.Id)
}

func (m *Manager) deliverMerge(ctx context.Context, sess api.AgentSession, task api.Task, project api.Project, prompt string) error {
	l, err := m.ensureLive(ctx, sess, task)
	if err != nil {
		return err
	}
	err = m.syncModel(ctx, l, sess.Model)
	if err != nil {
		return err
	}
	_, err = m.prepareMessage(ctx, l, "byre", prompt, nil)
	if err != nil {
		return err
	}
	m.mu.Lock()
	l.fetching = true
	m.mu.Unlock()

	err = git.FetchFastForward(ctx, project.RepoPath, project.DefaultBranch)

	m.mu.Lock()
	l.fetching = false
	interrupted := l.interrupted
	m.mu.Unlock()

	if err != nil {
		return err
	}
	if interrupted {
		return errors.New("interrupted before the merge started")
	}
	return l.sess.send(ctx, prompt, nil)
}

func (m *Manager) recoverMerges(ctx context.Context) error {
	merging, err := m.store.MergingTasks(ctx)
	if err != nil {
		return err
	}

	for _, t := range merging {
		pm := m.projectMerge(t.ProjectId)
		pm.mu.Lock()
		pm.holder = t.Id
		pm.mu.Unlock()
	}
	return nil
}

func (m *Manager) recoverFinalizing(ctx context.Context) error {
	finalizing, err := m.store.FinalizingMerges(ctx)
	if err != nil {
		return err
	}

	for _, f := range finalizing {
		ferr := m.resumeFinalizing(ctx, f)
		if ferr == nil {
			slog.Info("finished merge at startup", "task", f.Task.Id, "commit", f.Commit)
			continue
		}

		slog.Error("finish merge at startup", "task", f.Task.Id, "err", ferr)
		unfinished := &unfinishedMergeError{}
		if !errors.As(ferr, &unfinished) {
			err = m.store.ClearFinalizing(ctx, f.Task.Id)
			if err != nil {
				return err
			}
		}
		sessions, err := m.store.ListSessions(ctx, f.Task.Id)
		if err != nil {
			return err
		}
		m.fileTaskError(ctx, f.Task, sessions, "byre restarted while finishing this task's merge: "+ferr.Error()+".")
	}
	return nil
}

func (m *Manager) resumeFinalizing(ctx context.Context, f store.FinalizingMerge) error {
	project, err := m.store.GetProject(ctx, f.Task.ProjectId)
	if err != nil {
		return err
	}
	if f.Unchanged {
		return m.closeUnchanged(ctx, f.Task, project, f.By)
	}
	if f.Commit == "" {
		return errors.New("nothing was pushed. Merge again to retry")
	}

	landed, err := pushLanded(ctx, project, f.Commit)
	if err != nil {
		return err
	}
	if !landed {
		return fmt.Errorf("its squash commit %s is not on %s/%s, so the push didn't land. Merge again to retry", f.Commit, git.Remote, project.DefaultBranch)
	}
	return m.closePushed(ctx, f.Task, project, f.Commit, f.Notes, f.By)
}

func pushLanded(ctx context.Context, project api.Project, commit string) (bool, error) {
	target := git.Remote + "/" + project.DefaultBranch
	err := git.Fetch(ctx, project.RepoPath)
	if err != nil {
		return false, &unfinishedMergeError{fmt.Errorf("checking whether %s reached %s failed: %w", commit, target, err)}
	}
	landed, err := git.IsAncestor(project.RepoPath, commit, target)
	if err != nil {
		return false, &unfinishedMergeError{fmt.Errorf("checking whether %s reached %s failed: %w", commit, target, err)}
	}
	return landed, nil
}

type unfinishedMergeError struct {
	err error
}

func (e *unfinishedMergeError) Error() string {
	return e.err.Error() + "; byre finishes the merge when it next starts"
}

func (e *unfinishedMergeError) Unwrap() error {
	return e.err
}

func (m *Manager) startHeldMerges(ctx context.Context) {
	projects, err := m.store.ProjectsWithQueuedMerges(ctx)
	if err != nil {
		slog.Error("list projects with held merges", "err", err)
		return
	}
	for _, p := range projects {
		m.nextMerge(p)
	}
}

func (m *Manager) abandonMerge(taskID, projectID string) {
	if m.takeMerge(taskID) == "" {
		return
	}
	go m.releaseMerge(projectID, taskID)
}

func (m *Manager) releaseMerge(projectID, taskID string) {
	pm := m.projectMerge(projectID)
	pm.mu.Lock()
	defer pm.mu.Unlock()
	if pm.holder == taskID {
		pm.holder = ""
	}
	m.startNextMerge(pm, projectID)
}

func (m *Manager) nextMerge(projectID string) {
	pm := m.projectMerge(projectID)
	pm.mu.Lock()
	defer pm.mu.Unlock()
	m.startNextMerge(pm, projectID)
}

func (m *Manager) startNextMerge(pm *projectMerge, projectID string) {
	if pm.holder != "" {
		return
	}

	ctx := context.Background()
	held, err := m.store.QueuedMerges(ctx, projectID)
	if err != nil {
		slog.Error("list held merges", "project", projectID, "err", err)
		return
	}

	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)
		if err != nil {
			slog.Error("take held merge", "task", task.Id, "err", err)
			continue
		}
		if login == "" {
			continue
		}

		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)
		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)
		return
	}
	last := sessions[len(sessions)-1]
	m.fileError(&live{id: last.Id, taskID: task.Id, projectID: task.ProjectId}, msg)
}

func (m *Manager) takeMerge(taskID string) string {
	login, err := m.store.TakeMergingBy(context.Background(), taskID)
	if err != nil {
		slog.Error("take merge intent", "task", taskID, "err", err)
	}
	return login
}

func (m *Manager) finishMerge(l *live, login string) {
	defer m.releaseMerge(l.projectID, l.taskID)

	ctx := context.Background()
	err := m.squashAndClose(ctx, l.taskID, login)
	if err == nil {
		return
	}

	slog.Error("finish merge", "task", l.taskID, "err", err)
	unfinished := &unfinishedMergeError{}
	if errors.As(err, &unfinished) {
		m.fileError(l, "Merging the task failed: "+err.Error()+".")
		return
	}
	cerr := m.store.ClearFinalizing(ctx, l.taskID)
	if cerr != nil {
		slog.Error("clear merge finalizing", "task", l.taskID, "err", cerr)
	}
	m.fileError(l, "Merging the task failed: "+err.Error())
	err = m.Deliver(ctx, l.id)
	if err != nil {
		slog.Error("deliver inbox responses", "session", l.id, "err", err)
	}
}

func (m *Manager) squashAndClose(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
	}

	coAuthors, err := m.CoAuthors(ctx, taskID, login)
	if err != nil {
		return err
	}

	env := []string{
		"GIT_AUTHOR_NAME=" + user.GitName,
		"GIT_AUTHOR_EMAIL=" + user.GitEmail,
		"GIT_COMMITTER_NAME=" + user.GitName,
		"GIT_COMMITTER_EMAIL=" + user.GitEmail,
	}
	commit, err := git.SquashCommit(project.RepoPath, task.WorktreePath, project.DefaultBranch, squashMessage(task.MergeTitle, coAuthors), env)
	if err != nil {
		return err
	}

	if commit == "" {
		err = m.store.SetFinalizingUnchanged(ctx, taskID)
		if err != nil {
			return err
		}
		return m.closeUnchanged(ctx, task, project, login)
	}

	note, err := m.mergeNote(ctx, task, login)
	if err != nil {
		return err
	}
	notes, err := git.NoteCommit(project.RepoPath, commit, note, env)
	if err != nil {
		return err
	}
	err = m.store.SetFinalizingCommit(ctx, taskID, commit, notes)
	if err != nil {
		return err
	}
	err = git.Push(ctx, project.RepoPath, commit, project.DefaultBranch, notes)
	if errors.Is(err, git.ErrRemoteMoved) {
		return err
	}
	if err != nil {
		landed, lerr := pushLanded(ctx, project, commit)
		if lerr != nil {
			return &unfinishedMergeError{fmt.Errorf("%w, and %w", err, errors.Unwrap(lerr))}
		}
		if !landed {
			return err
		}
		slog.Warn("push reported failure but landed", "task", taskID, "commit", commit, "err", err)
	}
	return m.closePushed(ctx, task, project, commit, notes, login)
}

func (m *Manager) closeUnchanged(ctx context.Context, task api.Task, project api.Project, login string) error {
	err := m.RemoveTask(ctx, task, project, login)
	if err != nil {
		return &unfinishedMergeError{fmt.Errorf("the merge changed nothing, but removing the task's worktree failed: %w", err)}
	}
	_, err = m.store.MergeTask(ctx, task.Id, "", "", login)
	if err != nil {
		return &unfinishedMergeError{fmt.Errorf("the merge changed nothing, but closing the task failed: %w", err)}
	}
	return nil
}

func (m *Manager) closePushed(ctx context.Context, task api.Task, project api.Project, commit, notes, login string) error {
	pushed := git.Remote + "/" + project.DefaultBranch
	err := git.AdvanceNotes(project.RepoPath, notes)
	if err != nil {
		return &unfinishedMergeError{fmt.Errorf("pushed %s to %s, but updating %s in the project checkout failed: %w", commit, pushed, git.NotesRef, err)}
	}
	err = git.FastForward(project.RepoPath, commit)
	if err != nil {
		return &unfinishedMergeError{fmt.Errorf("pushed %s to %s, but fast-forwarding the project checkout failed: %w", commit, pushed, err)}
	}
	err = m.RemoveTask(ctx, task, project, login)
	if err != nil {
		return &unfinishedMergeError{fmt.Errorf("pushed %s to %s, but removing the task's worktree failed: %w", commit, pushed, err)}
	}
	_, err = m.store.MergeTask(ctx, task.Id, commit, pushed, login)
	if err != nil {
		return &unfinishedMergeError{fmt.Errorf("pushed %s to %s, but closing the task failed: %w", commit, pushed, err)}
	}
	return nil
}

func (m *Manager) CoAuthors(ctx context.Context, taskID, login string) ([]api.Participant, error) {
	participants, err := m.store.Participants(ctx, taskID)
	if err != nil {
		return nil, err
	}

	coAuthors := []api.Participant{}
	for _, p := range participants {
		if p.Login != login {
			coAuthors = append(coAuthors, p)
		}
	}
	return coAuthors, nil
}

func squashMessage(title string, coAuthors []api.Participant) string {
	trailers := []string{}
	for _, p := range coAuthors {
		if p.GitName == "" || p.GitEmail == "" {
			continue
		}
		trailers = append(trailers, fmt.Sprintf("Co-authored-by: %s <%s>", p.GitName, p.GitEmail))
	}
	if len(trailers) == 0 {
		return title
	}
	return title + "\n\n" + strings.Join(trailers, "\n")
}

func (m *Manager) RemoveTask(ctx context.Context, task api.Task, project api.Project, actor string) error {
	err := m.StopTask(ctx, task.Id, actor)
	if err != nil {
		return err
	}
	m.cancelReviews(ctx, task, project)
	err = m.dropSnapshots(project.RepoPath, task.Id)
	if err != nil {
		return err
	}
	return git.RemoveWorktree(project.RepoPath, task.WorktreePath, task.Branch)
}

type setMergeTitleArgs struct {
	Title string `json:"title"`
}

func (m *Manager) setMergeTitleTool(l *live) mcp.Tool {
	return mcp.Tool{
		Name: "set_merge_title",
		Description: "Set the title of the single commit this task becomes when it is merged onto the default branch. " +
			"It starts as the task title. Write one line in the imperative mood that describes all the work on the branch, " +
			"for example \"Add retry with backoff to webhook delivery\". It replaces any earlier title.",
		InputSchema: json.RawMessage(`{
			"type": "object",
			"properties": {
				"title": {"type": "string", "description": "One line, at most 72 characters."}
			},
			"required": ["title"]
		}`),
		Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
			return m.setMergeTitle(ctx, l, raw)
		},
	}
}

func (m *Manager) setMergeTitle(ctx context.Context, l *live, raw json.RawMessage) (string, error) {
	args := setMergeTitleArgs{}
	err := json.Unmarshal(raw, &args)
	if err != nil {
		return "", fmt.Errorf("malformed set_merge_title call: %w", err)
	}

	title, err := cleanMergeTitle(args.Title)
	if err != nil {
		return "", err
	}

	err = m.storeMergeTitle(ctx, l.taskID, title, agentActor)
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("Merge commit title set to %q.", title), nil
}

type stopMergeArgs struct {
	Reason string `json:"reason"`
}

func (m *Manager) stopMergeTool(l *live) mcp.Tool {
	return mcp.Tool{
		Name: "stop_merge",
		Description: "Stop the merge byre is to make when this turn ends, so the task stays active and its snapshot is reviewed before anything is pushed. " +
			"Call it during a Merge turn when resolving conflicts took substantive changes, then end your turn describing them. The human merges again after the review.",
		InputSchema: json.RawMessage(`{
			"type": "object",
			"properties": {
				"reason": {"type": "string", "description": "The substantive changes the resolution made, in a sentence or two."}
			},
			"required": ["reason"]
		}`),
		Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
			return m.stopMerge(ctx, l, raw)
		},
	}
}

func (m *Manager) stopMerge(ctx context.Context, l *live, raw json.RawMessage) (string, error) {
	args := stopMergeArgs{}
	err := json.Unmarshal(raw, &args)
	if err != nil {
		return "", fmt.Errorf("malformed stop_merge call: %w", err)
	}
	reason := strings.TrimSpace(args.Reason)
	if reason == "" {
		return "", fmt.Errorf("%w: reason must say what the resolution changed", ErrInvalid)
	}

	if m.takeMerge(l.taskID) == "" {
		return "No merge is pending for this task; just end your turn.", nil
	}
	go m.releaseMerge(l.projectID, l.taskID)

	err = m.emit(ctx, l, agentActor, "task.merge_stopped", map[string]any{"id": l.taskID, "reason": reason})
	if err != nil {
		return "", err
	}
	return "Merge stopped: byre won't squash or push when this turn ends. End your turn now with a description of the changes.", nil
}

func cleanMergeTitle(raw string) (string, error) {
	title := strings.TrimSpace(raw)
	if title == "" || strings.ContainsAny(title, "\r\n") {
		return "", fmt.Errorf("%w: title must be one non-empty line", ErrInvalid)
	}
	if utf8.RuneCountInString(title) > maxMergeTitle {
		return "", fmt.Errorf("%w: title is %d characters; keep it to %d", ErrInvalid, utf8.RuneCountInString(title), maxMergeTitle)
	}
	return title, nil
}

func (m *Manager) storeMergeTitle(ctx context.Context, taskID, title, actor string) error {
	err := m.store.SetMergeTitle(ctx, taskID, title, actor)
	if errors.Is(err, store.ErrNotFound) {
		return fmt.Errorf("%w: this task is closed, or its merge was already confirmed with its title", ErrConflict)
	}
	return err
}