byre/internal/agents/mergein.go

package agents

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"os"
	"strings"
	"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_titles.", 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
	}

	if title != task.MergeTitle {
		err = m.storeMergeTitle(ctx, task.Id, title, user.Login)
		if err != nil {
			return api.Task{}, err
		}
		task.MergeTitle = title
	}

	held, err := m.store.RequestMerge(ctx, task.Id, user.Login)
	switch {
	case 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)
	case errors.Is(err, store.ErrMerging):
		return api.Task{}, fmt.Errorf("%w: task %s is already merging", ErrConflict, task.Id)
	case err != nil:
		return api.Task{}, err
	}
	if held {
		m.headStartMerge(ctx, task, user, project.DefaultBranch)
	}
	m.clearMergeAfterReview(ctx, task.Id)
	return m.store.GetTask(ctx, task.Id)
}

func mergeTurnPrompt(project api.Project) string {
	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
}

func (m *Manager) startClaimedMerge(ctx context.Context, c store.ClaimedMerge) {
	err := m.beginClaimedMerge(ctx, c)
	if err == nil {
		return
	}
	slog.Error("start claimed merge", "task", c.TaskID, "err", err)
	_, terr := m.store.TakeMergingBy(ctx, c.TaskID)
	if terr != nil {
		slog.Error("drop claimed merge", "task", c.TaskID, "err", terr)
	}
	m.fileError(ctx, c.TaskID, "Starting the merge failed: "+err.Error())
}

func (m *Manager) beginClaimedMerge(ctx context.Context, c store.ClaimedMerge) error {
	task, err := m.store.GetTask(ctx, c.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, c.By)
	if err != nil {
		return err
	}
	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", task.Id, "by", c.By)
	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, err := m.leadSession(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)
}

func (m *Manager) deliverMerge(ctx context.Context, sess api.AgentSession, task api.Task, project api.Project, prompt string) error {
	_, compacting, err := m.recordMessage(ctx, sess, task, "byre", prompt, "", "", nil)
	if err != nil {
		return err
	}
	l, err := m.ensureLive(ctx, sess, task)
	if err != nil {
		return err
	}
	err = m.syncModel(ctx, l, sess.Model)
	if err != nil {
		return err
	}
	m.beginTurn(l, compacting)
	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 m.send(ctx, l, prompt, nil)
}

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(ctx, 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) fileError(ctx context.Context, taskID, msg string) {
	task, err := m.store.GetTask(ctx, taskID)
	if err != nil {
		slog.Error("load task to file error on", "task", taskID, "error", msg, "err", err)
		return
	}
	sessions, err := m.store.ListSessions(ctx, taskID)
	if err != nil {
		slog.Error("list sessions to file error on", "task", taskID, "error", msg, "err", err)
		return
	}
	m.fileTaskError(ctx, task, sessions, msg)
}

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) finalizeMerge(f store.FinalizingMerge) {
	ctx := context.Background()
	unlock := m.lockTask(f.Task.Id)
	err := m.landMerge(ctx, f)
	unlock()
	if err == nil {
		return
	}

	slog.Error("finish merge", "task", f.Task.Id, "err", err)
	unfinished := &unfinishedMergeError{}
	if errors.As(err, &unfinished) {
		m.fileError(ctx, f.Task.Id, "Merging the task failed: "+err.Error()+".")
		return
	}
	cerr := m.store.ClearFinalizing(ctx, f.Task.Id)
	if cerr != nil {
		slog.Error("clear merge finalizing", "task", f.Task.Id, "err", cerr)
	}
	m.fileError(ctx, f.Task.Id, "Merging the task failed: "+err.Error())
	m.deliverIdle(ctx, f.Task.Id)
}

func (m *Manager) deliverIdle(ctx context.Context, taskID string) {
	sessions, err := m.store.ListSessions(ctx, taskID)
	if err != nil {
		slog.Error("list sessions to deliver to", "task", taskID, "err", err)
		return
	}
	for _, s := range sessions {
		if s.Role != api.Implementer || s.State != api.AgentSessionStateIdle || m.interrupted(s.Id) {
			continue
		}
		err = m.Deliver(ctx, s.Id)
		if err != nil {
			slog.Error("deliver inbox responses", "session", s.Id, "err", err)
		}
	}
}

func (m *Manager) landMerge(ctx context.Context, f store.FinalizingMerge) error {
	task := f.Task
	project, err := m.store.GetProject(ctx, task.ProjectId)
	if err != nil {
		return err
	}
	if f.Unchanged {
		return m.closeUnchanged(ctx, task, project, f.By)
	}

	commit, notes := f.Commit, f.Notes
	if commit != "" {
		landed, err := pushLanded(ctx, project, commit)
		if err != nil {
			return err
		}
		if landed {
			return m.closePushed(ctx, task, project, commit, notes, f.By)
		}
	}
	if commit == "" {
		user, err := m.store.GetUser(ctx, f.By)
		if err != nil {
			return err
		}
		commit, notes, err = m.prepareMainCommit(ctx, task, project, user)
		if err != nil {
			return err
		}
		if commit == "" {
			err = m.store.SetFinalizingUnchanged(ctx, task.Id)
			if err != nil {
				return err
			}
			return m.closeUnchanged(ctx, task, project, f.By)
		}
		err = m.store.SetFinalizingCommit(ctx, task.Id, 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", task.Id, "commit", commit, "err", err)
	}
	return m.closePushed(ctx, task, project, commit, notes, f.By)
}

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 := advanceMainCheckout(project, commit, notes)
	if err != nil {
		return &unfinishedMergeError{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) 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
	}
	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(ctx, project.RepoPath, task.WorktreePath, task.Branch)
}

type setTitlesArgs struct {
	Title      string `json:"title"`
	MergeTitle string `json:"merge_title"`
}

func (m *Manager) setTitlesTool(l *live) mcp.Tool {
	return mcp.Tool{
		Name: "set_titles",
		Description: "Set this task's two titles. title names the task in byre's task list and switcher: very short (two to four words), distinct from the project's other tasks, and to the point, for example \"webhook retries\". " +
			"merge_title is the title of the single commit this task becomes when it is merged onto the default branch: one line in the imperative mood that describes all the work on the branch, " +
			"for example \"Add retry with backoff to webhook delivery\". Each replaces the earlier one.",
		InputSchema: json.RawMessage(`{
			"type": "object",
			"properties": {
				"title": {"type": "string", "description": "The task title: one short line."},
				"merge_title": {"type": "string", "description": "The merge commit title: one line, at most 72 characters."}
			},
			"required": ["title", "merge_title"]
		}`),
		Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
			return m.setTitles(ctx, l, raw)
		},
	}
}

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

	title := strings.TrimSpace(args.Title)
	if title == "" || strings.ContainsAny(title, "\r\n") {
		return "", fmt.Errorf("%w: title must be one non-empty line", ErrInvalid)
	}
	mergeTitle, err := cleanMergeTitle(args.MergeTitle)
	if err != nil {
		return "", err
	}

	err = m.store.SetTitles(ctx, l.taskID, title, mergeTitle, agentActor)
	if errors.Is(err, store.ErrNotFound) {
		return "", fmt.Errorf("%w: this task is closed, or its merge was already confirmed with its title", ErrConflict)
	}
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("Task title set to %q, merge commit title to %q.", title, mergeTitle), 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)
	}
	mergingBy, err := m.store.TakeMergingBy(ctx, l.taskID)
	if err != nil {
		return "", err
	}
	if mergingBy == "" {
		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
}