byre/internal/agents/team.go

package agents

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

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

const boardActor = "board"

func teamLead(sessions []api.AgentSession) string {
	if len(sessions) == 0 {
		return ""
	}
	return sessions[0].Member
}

func memberSession(sessions []api.AgentSession, member string) (api.AgentSession, bool) {
	for i := len(sessions) - 1; i >= 0; i-- {
		s := sessions[i]
		if s.Member == member && s.State != api.AgentSessionStateCompleted {
			return s, true
		}
	}
	return api.AgentSession{}, false
}

func teamSessions(sessions []api.AgentSession) []api.AgentSession {
	team := []api.AgentSession{}
	seen := map[string]bool{}
	for _, s := range sessions {
		if s.Member == "" || seen[s.Member] {
			continue
		}
		seen[s.Member] = true
		live, found := memberSession(sessions, s.Member)
		if !found {
			live = latestSession(sessions, s.Member)
		}
		team = append(team, live)
	}
	return team
}

func latestSession(sessions []api.AgentSession, member string) api.AgentSession {
	latest := api.AgentSession{}
	for _, s := range sessions {
		if s.Member == member {
			latest = s
		}
	}
	return latest
}

func (m *Manager) leadsTask(ctx context.Context, taskID, sessionID string) bool {
	sess, found, err := m.liveSession(ctx, taskID)
	if err != nil {
		slog.Error("find task's first agent", "task", taskID, "err", err)
		return false
	}
	return found && sess.Id == sessionID
}

func checkTeam(members []api.TeamMember) error {
	problems := workflow.TeamProblems(members)
	if len(problems) > 0 {
		return fmt.Errorf("%w: %s", ErrInvalid, strings.Join(problems, "; "))
	}
	return nil
}

func (m *Manager) startTeam(ctx context.Context, task api.Task, creator api.User, shared string, members []api.TeamMember, actor string) error {
	err := checkTeam(members)
	if err != nil {
		return err
	}
	if creator.GitName == "" || creator.GitEmail == "" {
		return ErrNoGitIdentity
	}
	sessions := []api.AgentSession{}
	for _, mem := range members {
		sess, err := m.newSession(ctx, task, creator, mem.Name, mem.Model)
		if err != nil {
			return err
		}
		sessions = append(sessions, sess)
	}
	failed := []string{}
	for i, sess := range sessions {
		_, err = m.turn(ctx, sess, task, actor, teamPrompt(shared, members[i]), nil)
		if err != nil {
			failed = append(failed, fmt.Sprintf("%q: %v", members[i].Name, err))
		}
	}
	if len(failed) > 0 {
		return fmt.Errorf("these agents did not start, so the team is incomplete; discard the task and create it again: %s", strings.Join(failed, "; "))
	}
	return nil
}

func teamPrompt(shared string, mem api.TeamMember) string {
	return fmt.Sprintf("%s\n\n## Your part: %s\n\n%s", strings.TrimSpace(shared), mem.Name, strings.TrimSpace(mem.Prompt))
}

func teamGuidance(sessions []api.AgentSession, member string, stage bool) string {
	team := teamSessions(sessions)
	names := []string{}
	for _, s := range team {
		names = append(names, fmt.Sprintf("%q (%s)", s.Member, s.Model))
	}
	lead := teamLead(sessions)
	g := fmt.Sprintf("\n- This task is worked by a team of agents in this one worktree: %s. You are %q. "+
		"The team talks through a shared board: post_to_board sends a post to every other agent (a working agent gets it after its next tool call; an idle one starts a turn with it), and read_board reads the posts so far. "+
		"The human and the lead post there too. Post only when you add something: a finding, a proposal, an objection, a decision, a request. Don't post acknowledgements or thanks, since every post wakes every other agent. "+
		"Your final message goes to the human, not to the team. The other agents change the same files: re-read a file before you edit it, and say on the board when you take on a file or an area.",
		strings.Join(names, ", "), member)
	if !stage {
		return g
	}
	if member == lead {
		return g + " You are the team's first agent, so you call finish_stage once the team's work on this stage is done; the others tell you on the board when their parts are."
	}
	return g + fmt.Sprintf(" %q, the team's first agent, finishes this stage with finish_stage; tell it on the board when your part is done.", lead)
}

type poster struct {
	actor   string
	member  string
	model   string
	session string
}

func (m *Manager) Post(ctx context.Context, task api.Task, actor, text string) error {
	return m.post(ctx, task, poster{actor: actor}, text)
}

func (m *Manager) post(ctx context.Context, task api.Task, from poster, text string) error {
	if strings.TrimSpace(text) == "" {
		return fmt.Errorf("%w: text is empty", ErrInvalid)
	}
	if task.Status != api.TaskStatusActive {
		return fmt.Errorf("%w: task %s is %s", ErrConflict, task.Id, task.Status)
	}
	sessions, err := m.store.ListSessions(ctx, task.Id)
	if err != nil {
		return err
	}
	team := teamSessions(sessions)
	if len(team) == 0 {
		return fmt.Errorf("%w: task %s has no team, so no board", ErrInvalid, task.Id)
	}

	payload := map[string]any{"text": text}
	if from.member != "" {
		payload["member"] = from.member
		payload["model"] = from.model
	}
	ev := api.Event{Kind: "board.posted", ProjectId: &task.ProjectId, TaskId: &task.Id, Actor: from.actor, Payload: payload}
	if from.session != "" {
		ev.AgentSessionId = &from.session
	}
	err = m.store.AppendEvent(ctx, ev)
	if err != nil {
		return err
	}

	who := from.actor
	if from.member != "" {
		who = fmt.Sprintf("%q", from.member)
	}
	msg := fmt.Sprintf("%s posted on the board:\n\n%s", who, text)
	for _, s := range team {
		if s.Member == from.member {
			continue
		}
		m.deliverPost(context.WithoutCancel(ctx), task, s, msg)
	}
	return nil
}

func (m *Manager) deliverPost(ctx context.Context, task api.Task, s api.AgentSession, msg string) {
	started, err := m.store.SessionStarted(ctx, s.Id)
	if err == nil && started {
		_, _, err = m.Message(ctx, s.Id, boardActor, msg, "")
	}
	if err == nil && started {
		return
	}
	if err != nil && !errors.Is(err, ErrConflict) {
		slog.Error("deliver board post", "task", task.Id, "member", s.Member, "err", err)
	}
	err = m.store.AddTaskInput(ctx, store.TaskInput{TaskID: task.Id, Member: s.Member, Actor: boardActor, Text: msg}, task.ProjectId, s.Id)
	if err != nil {
		slog.Error("queue board post", "task", task.Id, "member", s.Member, "err", err)
	}
}

func (m *Manager) deliverTeam(ctx context.Context, l *live) {
	if l.member == "" {
		return
	}
	sessions, err := m.store.ListSessions(ctx, l.taskID)
	if err != nil {
		slog.Error("list team sessions", "task", l.taskID, "err", err)
		return
	}
	for _, s := range teamSessions(sessions) {
		if s.Id == l.id || s.State != api.AgentSessionStateIdle || m.interrupted(s.Id) {
			continue
		}
		err = m.Deliver(ctx, s.Id)
		if err != nil {
			slog.Error("deliver to team member", "session", s.Id, "err", err)
		}
	}
}

func (m *Manager) boardTools(l *live) []mcp.Tool {
	return []mcp.Tool{
		{
			Name:        "post_to_board",
			Description: "Post to this task's team board. Every other agent of the team gets the post: a working agent after its next tool call, an idle one as a new turn. The human sees the board in the task's Board tab.",
			InputSchema: json.RawMessage(`{"type": "object", "properties": {"text": {"type": "string", "description": "Markdown."}}, "required": ["text"]}`),
			Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
				args := struct {
					Text string `json:"text"`
				}{}
				err := json.Unmarshal(raw, &args)
				if err != nil {
					return "", fmt.Errorf("malformed post_to_board call: %w", err)
				}
				task, err := m.store.GetTask(ctx, l.taskID)
				if err != nil {
					return "", err
				}
				err = m.post(ctx, task, poster{actor: agentActor, member: l.member, model: l.model, session: l.id}, args.Text)
				if err != nil {
					return "", err
				}
				return "Posted to the board.", nil
			},
		},
		{
			Name:        "read_board",
			Description: "Read this task's team board: every post, oldest first, each with its seq and who posted it. With after_seq, only the posts after it.",
			InputSchema: json.RawMessage(`{"type": "object", "properties": {"after_seq": {"type": "integer"}}}`),
			Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
				args := struct {
					AfterSeq int64 `json:"after_seq"`
				}{}
				err := json.Unmarshal(raw, &args)
				if err != nil {
					return "", fmt.Errorf("malformed read_board call: %w", err)
				}
				posts, err := m.store.BoardPosts(ctx, l.taskID, args.AfterSeq)
				if err != nil {
					return "", err
				}
				if len(posts) == 0 {
					return "No posts.", nil
				}
				b := &strings.Builder{}
				for _, ev := range posts {
					fmt.Fprintf(b, "#%d %s:\n%v\n\n", ev.Seq, postAuthor(ev), ev.Payload["text"])
				}
				return strings.TrimSpace(b.String()), nil
			},
		},
	}
}

func postAuthor(ev api.Event) string {
	member, _ := ev.Payload["member"].(string)
	if member != "" {
		return fmt.Sprintf("%q (%v)", member, ev.Payload["model"])
	}
	return ev.Actor
}