byre/internal/agents/mergein.go

package agents

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

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

const maxMergeTitle = 72

const mergeInMarker = " into this branch: run `git merge --no-ff --no-commit "

func isMergeInPrompt(actor, text string) bool {
	return actor == "byre" && strings.Contains(text, mergeInMarker)
}

func mergeInPrompt(target string) string {
	return fmt.Sprintf("Merge %[1]s"+mergeInMarker+"%[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) headStartMerge(ctx context.Context, task api.Task, user api.User, target string) {
	sess, err := m.mergeSession(ctx, task, user)
	if errors.Is(err, ErrConflict) {
		slog.Info("skip held merge head start", "task", task.Id, "err", err)
		return
	}
	if err != nil {
		slog.Error("held merge head start", "task", task.Id, "err", err)
		return
	}
	prompt := mergeInPrompt(target) + fmt.Sprintf(" This task's Merge is held while another task merges into %s; this turn gets a head start, and byre merges %s in again when the held merge starts. If you stop for substantive changes, call stop_merge first, so byre drops the held merge and nothing is pushed before the review.", target, target)
	go func() {
		_, err := m.launch(context.Background(), sess, task, "byre", prompt, nil)
		if err != nil {
			slog.Error("start held merge head start", "task", task.Id, "err", err)
		}
	}()
}

func (m *Manager) sendMergeIn(ctx context.Context, task api.Task, user api.User, prompt string) (api.AgentSession, error) {
	sess, found, err := m.liveSession(ctx, task.Id)
	if err != nil {
		return api.AgentSession{}, err
	}
	if found {
		return m.Send(ctx, sess.Id, "byre", prompt, "")
	}
	return m.create(ctx, task, user, sess.Member, 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.headStartMerge(ctx, task, user, project.DefaultBranch)
	}
	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) {
	sess, found, err := m.liveSession(ctx, task.Id)
	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
}

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.recordError(&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)
	m.endMergeTurn(l.id, false)
	if err == nil {
		return
	}

	slog.Error("finish merge", "task", l.taskID, "err", err)
	unfinished := &unfinishedMergeError{}
	if errors.As(err, &unfinished) {
		m.recordError(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.recordError(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, m.signer(ctx, user))
	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
	}
	err = os.RemoveAll(MediaDir(task.WorktreePath))
	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, or drop the held merge waiting for another task's, so the task stays active and its snapshot is reviewed before anything is pushed. " +
			"Call it during a Merge turn, or a held merge's head-start merge-in, 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.leadsTask(ctx, l.taskID, l.id) {
		return "", fmt.Errorf("%w: only the team's first agent runs the merge, so only it can stop it", ErrConflict)
	}
	switch {
	case m.takeMerge(l.taskID) != "":
		go m.releaseMerge(l.projectID, l.taskID)
	default:
		queuedBy, err := m.store.TakeQueuedMerge(ctx, l.taskID)
		if err != nil {
			return "", err
		}
		if queuedBy == "" {
			return "No merge is pending for this task; just end your turn.", nil
		}
	}

	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
}