byre/internal/agents/mergein.go

package agents

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"slices"
	"strings"
	"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)
	}
	title, err := cleanMergeTitle(title)
	if err != nil {
		return api.Task{}, err
	}

	m.mergeMu.Lock()
	defer m.mergeMu.Unlock()

	if m.mergeHolders[project.Id] == 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 m.mergeHolders[project.Id] {
	case "":
		err := m.startMerge(ctx, 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 or already waiting to merge", ErrConflict, task.Id)
		}
		if err != nil {
			return api.Task{}, err
		}
	}
	return m.store.GetTask(ctx, task.Id)
}

func (m *Manager) startMerge(ctx context.Context, task api.Task, project api.Project, user api.User) error {
	err := git.FetchFastForward(project.RepoPath, project.DefaultBranch)
	if err != nil {
		return err
	}

	err = m.store.SetMergingBy(ctx, task.Id, user.Login)
	if err != nil {
		return err
	}
	m.mergeHolders[project.Id] = task.Id

	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)
	_, err = m.sendMergeIn(ctx, task, user, prompt)
	if err != nil {
		m.takeMerge(task.Id)
		delete(m.mergeHolders, project.Id)
	}
	return err
}

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

	m.mergeMu.Lock()
	defer m.mergeMu.Unlock()
	for _, t := range merging {
		m.mergeHolders[t.ProjectId] = t.Id
	}
	return nil
}

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) {
	m.mergeMu.Lock()
	defer m.mergeMu.Unlock()
	if m.mergeHolders[projectID] == taskID {
		delete(m.mergeHolders, projectID)
	}
	m.startNextMerge(projectID)
}

func (m *Manager) nextMerge(projectID string) {
	m.mergeMu.Lock()
	defer m.mergeMu.Unlock()
	m.startNextMerge(projectID)
}

func (m *Manager) startNextMerge(projectID string) {
	if m.mergeHolders[projectID] != "" {
		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, 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, 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, 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)
	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
	}

	pushed := ""
	if commit != "" {
		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 = git.Push(project.RepoPath, commit, project.DefaultBranch, notes)
		if err != nil {
			return err
		}
		pushed = git.Remote + "/" + project.DefaultBranch
		err = git.SetNotes(project.RepoPath, notes)
		if err != nil {
			return 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 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 err
	}

	_, err = m.store.MergeTask(ctx, taskID, commit, pushed, login)
	return err
}

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) Discard(ctx context.Context, task api.Task, project api.Project, actor string) (api.Task, error) {
	err := m.RemoveTask(ctx, task, project, actor)
	if err != nil {
		return api.Task{}, err
	}
	t, err := m.store.DiscardTask(ctx, task.Id, actor)
	if err != nil {
		return api.Task{}, err
	}
	go m.releaseMerge(project.Id, task.Id)
	return t, nil
}

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
}