byre/internal/agents/mergemain.go

package agents

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"

	"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"
)

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

func (m *Manager) mergeToMainTool(l *live) mcp.Tool {
	return mcp.Tool{
		Name: "merge_to_main",
		Description: "Land this task's work so far on the default branch now and keep working: byre commits the worktree, squashes the branch onto the default branch as one commit titled title, pushes it, " +
			"and merges that commit back into this branch, so the task stays open and its later Merge carries only what comes after. " +
			"Refused when the branch conflicts with the default branch: merge the default branch in, resolve the conflicts, and call it again.",
		InputSchema: json.RawMessage(`{
			"type": "object",
			"properties": {
				"title": {"type": "string", "description": "The commit's title: one line in the imperative mood, at most 72 characters, describing the work landing now."}
			},
			"required": ["title"]
		}`),
		Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
			return m.mergeToMain(ctx, l, raw)
		},
	}
}

func (m *Manager) mergeToMain(ctx context.Context, l *live, raw json.RawMessage) (string, error) {
	args := mergeToMainArgs{}
	err := json.Unmarshal(raw, &args)
	if err != nil {
		return "", fmt.Errorf("malformed merge_to_main call: %w", err)
	}
	title, err := cleanMergeTitle(args.Title)
	if err != nil {
		return "", err
	}
	if m.ownsTask(l.id) {
		return "", fmt.Errorf("%w: this turn is merging into the worktree; end it first", ErrConflict)
	}

	sess, err := m.store.GetSession(ctx, l.id)
	if err != nil {
		return "", err
	}
	login := sess.CreatedBy
	if by := m.landingBy(l.taskID, l.id); by != "" {
		login = by
	}
	user, err := m.store.GetUser(ctx, login)
	if err != nil {
		return "", err
	}
	if user.GitName == "" || user.GitEmail == "" {
		return "", ErrNoGitIdentity
	}

	unlock := m.lockTask(l.taskID)
	defer unlock()

	task, err := m.store.GetTask(ctx, l.taskID)
	if err != nil {
		return "", err
	}
	if task.Status != api.TaskStatusActive || task.Merging || task.MergeQueuedAt != nil || task.MergingIntoParent || task.Discarding {
		return "", fmt.Errorf("%w: task %s is %s", ErrConflict, task.Id, task.State)
	}
	project, err := m.store.GetProject(ctx, task.ProjectId)
	if err != nil {
		return "", err
	}

	pm := m.projectMerge(project.Id)
	pm.mu.Lock()
	holder := pm.holder
	if holder == "" {
		pm.holder = task.Id
	}
	pm.mu.Unlock()
	if holder != "" {
		return "", fmt.Errorf("%w: task %s is merging into %s; call again once it is done", ErrConflict, holder, project.DefaultBranch)
	}
	defer m.releaseMerge(project.Id, task.Id)

	commit, err := m.pushToMain(ctx, l, task, project, user, title)
	if err != nil {
		return "", err
	}
	if commit == "" {
		return fmt.Sprintf("Nothing to push: %s already has all of this branch's work, so byre merged %s into this branch. Keep working; the task stays open.", project.DefaultBranch, project.DefaultBranch), nil
	}

	pushed := git.Remote + "/" + project.DefaultBranch
	err = m.emit(ctx, l, agentActor, "task.merged_to_main", map[string]any{"id": task.Id, "commit": commit, "pushed": pushed, "merge_title": title})
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("Pushed %s %q to %s and merged it back into this branch. Keep working; the task stays open.", commit, title, pushed), nil
}

func (m *Manager) pushToMain(ctx context.Context, l *live, task api.Task, project api.Project, user api.User, title string) (string, error) {
	err := git.FetchFastForward(ctx, project.RepoPath, project.DefaultBranch)
	if err != nil {
		return "", err
	}
	_, err = git.CommitAll(l.dir, title, l.gitName, l.gitEmail)
	if err != nil {
		return "", err
	}

	coAuthors, err := m.CoAuthors(ctx, task.Id, user.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, l.dir, project.DefaultBranch, squashMessage(title, coAuthors), env, m.signer(ctx, user))
	conflict := &git.ConflictError{}
	if errors.As(err, &conflict) {
		return "", fmt.Errorf("%w; run `git merge --no-ff --no-commit %s`, resolve them, and call merge_to_main again", err, project.DefaultBranch)
	}
	if err != nil {
		return "", err
	}
	if commit == "" {
		err = git.Merge(l.dir, project.DefaultBranch, "Merge "+project.DefaultBranch, l.gitName, l.gitEmail)
		return "", err
	}

	task.MergeTitle = title
	note, err := m.mergeNote(ctx, task, user.Login)
	if err != nil {
		return "", err
	}
	notes, err := git.NoteCommit(project.RepoPath, commit, note, env)
	if err != nil {
		return "", err
	}
	err = git.Push(ctx, project.RepoPath, commit, project.DefaultBranch, notes)
	if errors.Is(err, git.ErrRemoteMoved) {
		return "", fmt.Errorf("%w; call merge_to_main again", err)
	}
	if err != nil {
		return "", err
	}

	pushed := git.Remote + "/" + project.DefaultBranch
	err = git.AdvanceNotes(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 = git.Merge(l.dir, commit, fmt.Sprintf("Merge %s from %s", commit[:7], project.DefaultBranch), l.gitName, l.gitEmail)
	if err != nil {
		return "", fmt.Errorf("pushed %s to %s, but merging it back into this branch failed: %w", commit, pushed, err)
	}
	return commit, nil
}

func landPrompt(target, title string) string {
	return fmt.Sprintf("The human asked to merge this task's work so far into %[1]s and keep the task open (Merge & continue); byre holds the task's other agents until this turn ends. "+
		"First merge %[1]s in with `git merge --no-ff --no-commit %[1]s`, resolve any conflicts keeping both sides' intent, regenerate generated files instead of hand-merging them, and check that the build and tests pass. "+
		"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 covers them before anything lands. "+
		"Otherwise call merge_to_main with title %[2]q, which commits the merge and lands the work on %[1]s, then end your turn with one line.", target, title)
}

func (m *Manager) MergeToMain(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
	}
	title, err := cleanMergeTitle(title)
	if err != nil {
		return api.Task{}, err
	}
	if task.Merging || task.MergeQueuedAt != nil || task.MergingIntoParent || task.Discarding {
		return api.Task{}, fmt.Errorf("%w: task %s is %s", ErrConflict, task.Id, task.State)
	}
	sess, found, err := m.liveSession(ctx, task.Id)
	if err != nil {
		return api.Task{}, err
	}
	if !found {
		return api.Task{}, fmt.Errorf("%w: task %s has no agent", ErrConflict, task.Id)
	}

	unlock := m.lockDelivery(sess.Id)
	defer unlock()

	m.mu.Lock()
	holder := m.landing[task.Id]
	if holder.session == "" {
		m.landing[task.Id] = landing{session: sess.Id, login: user.Login}
	}
	m.mu.Unlock()
	if holder.session != "" {
		return api.Task{}, fmt.Errorf("%w: task %s is already merging and continuing", ErrConflict, task.Id)
	}
	sessions, err := m.store.ListSessions(ctx, task.Id)
	if err == nil {
		for _, s := range sessions {
			if s.Role == api.Implementer && s.State == api.AgentSessionStateRunning {
				err = fmt.Errorf("%w: an agent on this task is mid-turn; merge when they're all idle", ErrConflict)
			}
		}
	}
	if err == nil {
		m.clearMergeAfterReview(ctx, task.Id)
		_, err = m.Send(ctx, sess.Id, "byre", landPrompt(project.DefaultBranch, title), "")
	}
	if err != nil {
		m.endLanding(task.Id, sess.Id)
		return api.Task{}, err
	}

	err = m.store.ClearReady(ctx, task.Id, user.Login)
	if err != nil && !errors.Is(err, store.ErrNotFound) {
		return api.Task{}, err
	}
	return m.store.GetTask(ctx, task.Id)
}

type landing struct {
	session string
	login   string
}

func (m *Manager) heldForLanding(taskID, sessionID string) bool {
	m.mu.Lock()
	defer m.mu.Unlock()
	holder := m.landing[taskID].session
	return holder != "" && holder != sessionID
}

func (m *Manager) landingBy(taskID, sessionID string) string {
	m.mu.Lock()
	defer m.mu.Unlock()
	if m.landing[taskID].session != sessionID {
		return ""
	}
	return m.landing[taskID].login
}

func (m *Manager) endLanding(taskID, sessionID string) {
	m.mu.Lock()
	held := m.landing[taskID].session == sessionID
	if held {
		delete(m.landing, taskID)
	}
	m.mu.Unlock()
	if held {
		go m.releaseTeam(taskID, sessionID)
	}
}

func (m *Manager) releaseTeam(taskID, landingID string) {
	ctx := context.Background()
	sessions, err := m.store.ListSessions(ctx, taskID)
	if err != nil {
		slog.Error("list sessions after merge and continue", "task", taskID, "err", err)
		return
	}
	for _, s := range sessions {
		if s.Id == landingID || 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 after merge and continue", "session", s.Id, "err", err)
		}
	}
}