byre/internal/agents/leadtools.go

package agents

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"slices"
	"strings"

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

func (m *Manager) leadTools(l *live) []mcp.Tool {
	tool := func(name, description, schema string, run func(ctx context.Context, user api.User, raw json.RawMessage) (string, error)) mcp.Tool {
		return mcp.Tool{
			Name:        name,
			Description: description,
			InputSchema: json.RawMessage(schema),
			Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
				user, err := m.store.GetUser(ctx, l.owner)
				if err != nil {
					return "", err
				}
				return run(ctx, user, raw)
			},
		}
	}
	taskOnly := `{"type": "object", "properties": {"task_id": {"type": "string"}}, "required": ["task_id"]}`

	return []mcp.Tool{
		tool("list_tasks",
			"List every project and its active tasks: id, title, worktree path, the agent's state and model, and open inbox items.",
			`{"type": "object", "properties": {}}`,
			func(ctx context.Context, user api.User, raw json.RawMessage) (string, error) {
				return m.listTasksText(ctx)
			}),

		tool("read_transcript",
			"Read a task's transcript: messages, tool calls (briefly), tool results (cut short), turn ends and inbox filings, each with its event seq. "+
				"Without after_seq it returns the last limit events; with it, the events after that seq.",
			`{"type": "object", "properties": {
				"task_id": {"type": "string"},
				"after_seq": {"type": "integer"},
				"limit": {"type": "integer", "description": "At most this many events, default 100, at most 500."}
			}, "required": ["task_id"]}`,
			func(ctx context.Context, user api.User, raw json.RawMessage) (string, error) {
				args := struct {
					TaskID   string `json:"task_id"`
					AfterSeq *int64 `json:"after_seq"`
					Limit    int    `json:"limit"`
				}{}
				err := json.Unmarshal(raw, &args)
				if err != nil {
					return "", err
				}
				return m.transcriptText(ctx, args.TaskID, args.AfterSeq, args.Limit)
			}),

		tool("read_inbox",
			"Read open inbox items as JSON, with each question's full context, choices and responses. Give task_id for one task's items.",
			`{"type": "object", "properties": {"task_id": {"type": "string"}}}`,
			func(ctx context.Context, user api.User, raw json.RawMessage) (string, error) {
				args := struct {
					TaskID string `json:"task_id"`
				}{}
				err := json.Unmarshal(raw, &args)
				if err != nil {
					return "", err
				}
				items, err := m.store.ListInbox(ctx, store.InboxFilter{State: api.Open, TaskID: args.TaskID})
				if err != nil {
					return "", err
				}
				b, err := json.MarshalIndent(items, "", "  ")
				return string(b), err
			}),

		tool("answer_question",
			"Respond to an open inbox item on the human's behalf. Only when the answer is clear or the human told you it. "+
				"type option: options holds the chosen choices' exact text (one for pick-one; pick-any may add text). "+
				"type other: free text to a question or an error. type dismiss: clears a report or an error.",
			`{"type": "object", "properties": {
				"item_id": {"type": "string"},
				"type": {"type": "string", "enum": ["option", "other", "dismiss"]},
				"options": {"type": "array", "items": {"type": "string"}},
				"text": {"type": "string"}
			}, "required": ["item_id", "type"]}`,
			func(ctx context.Context, user api.User, raw json.RawMessage) (string, error) {
				c := api.InboxResponseCreate{}
				err := json.Unmarshal(raw, &c)
				if err != nil {
					return "", err
				}
				m.mu.Lock()
				m.leadAnswered[c.ItemId] = true
				m.mu.Unlock()
				_, err = m.Respond(ctx, user.Login, []api.InboxResponseCreate{c})
				if err != nil {
					m.mu.Lock()
					delete(m.leadAnswered, c.ItemId)
					m.mu.Unlock()
					return "", err
				}
				return fmt.Sprintf("Responded to %s. The agent gets it once it is idle and has no other blocking question waiting.", c.ItemId), nil
			}),

		tool("send_message",
			"Send a message to a task's agent; it starts a turn. Starts the agent when the task has none. The agent must be idle: "+
				"if it is mid-turn, wait for its report or interrupt it first.",
			`{"type": "object", "properties": {
				"task_id": {"type": "string"},
				"text": {"type": "string"},
				"model": {"type": "string", "description": "Switch the agent to this model from this turn on."}
			}, "required": ["task_id", "text"]}`,
			func(ctx context.Context, user api.User, raw json.RawMessage) (string, error) {
				args := struct {
					TaskID string `json:"task_id"`
					Text   string `json:"text"`
					Model  string `json:"model"`
				}{}
				err := json.Unmarshal(raw, &args)
				if err != nil {
					return "", err
				}
				if strings.TrimSpace(args.Text) == "" {
					return "", errors.New("text is empty")
				}
				task, _, err := m.activeTask(ctx, args.TaskID)
				if err != nil {
					return "", err
				}
				sess, found, err := m.liveSession(ctx, task.Id)
				if err != nil {
					return "", err
				}
				if !found {
					_, err = m.create(ctx, task, user, args.Text, args.Model, leadActor(user.Login), nil)
					if err != nil {
						return "", err
					}
					return fmt.Sprintf("Started an agent on %s with your message.", task.Id), nil
				}
				_, err = m.Send(ctx, sess.Id, leadActor(user.Login), args.Text, args.Model)
				if err != nil {
					return "", err
				}
				return fmt.Sprintf("Sent to %s's agent.", task.Id), nil
			}),

		tool("create_task",
			"Create a task with its own branch and worktree, and start an agent on it with prompt as its first message. "+
				"The new agent sees nothing else, so the prompt must be a complete brief. It starts from the project's default branch, "+
				"or from from_task_id's committed HEAD.",
			`{"type": "object", "properties": {
				"project_id": {"type": "string"},
				"title": {"type": "string"},
				"prompt": {"type": "string"},
				"from_task_id": {"type": "string"},
				"model": {"type": "string"}
			}, "required": ["project_id", "title", "prompt"]}`,
			func(ctx context.Context, user api.User, raw json.RawMessage) (string, error) {
				args := struct {
					ProjectID  string `json:"project_id"`
					Title      string `json:"title"`
					Prompt     string `json:"prompt"`
					FromTaskID string `json:"from_task_id"`
					Model      string `json:"model"`
				}{}
				err := json.Unmarshal(raw, &args)
				if err != nil {
					return "", err
				}
				if strings.TrimSpace(args.Title) == "" || strings.TrimSpace(args.Prompt) == "" {
					return "", errors.New("title and prompt are required")
				}
				if user.GitName == "" || user.GitEmail == "" {
					return "", ErrNoGitIdentity
				}
				project, err := m.store.GetProject(ctx, args.ProjectID)
				if err != nil {
					return "", err
				}
				base := project.DefaultBranch
				if args.FromTaskID != "" {
					from, _, err := m.activeTask(ctx, args.FromTaskID)
					if err != nil {
						return "", err
					}
					base, err = git.Head(from.WorktreePath)
					if err != nil {
						return "", err
					}
				}
				task, err := m.CreateTask(ctx, project, args.Title, user.Login, base, leadActor(user.Login))
				if err != nil {
					return "", err
				}
				_, err = m.create(ctx, task, user, args.Prompt, m.modelOr(args.Model, l), leadActor(user.Login), nil)
				if err != nil {
					return "", fmt.Errorf("task %s created, but its agent did not start: %w", task.Id, err)
				}
				return fmt.Sprintf("Task %s %q created from %s and its agent started.", task.Id, task.Title, task.BaseRef), nil
			}),

		tool("start_todo",
			"Start a task on an item of a project's TODO.md (its title, exactly as list_tasks' project or the file on the default branch has it).",
			`{"type": "object", "properties": {
				"project_id": {"type": "string"},
				"title": {"type": "string"},
				"model": {"type": "string"}
			}, "required": ["project_id", "title"]}`,
			func(ctx context.Context, user api.User, raw json.RawMessage) (string, error) {
				args := struct {
					ProjectID string `json:"project_id"`
					Title     string `json:"title"`
					Model     string `json:"model"`
				}{}
				err := json.Unmarshal(raw, &args)
				if err != nil {
					return "", err
				}
				project, err := m.store.GetProject(ctx, args.ProjectID)
				if err != nil {
					return "", err
				}
				task, err := m.StartTodo(ctx, project, user, args.Title, m.modelOr(args.Model, l), leadActor(user.Login))
				if err != nil {
					return "", err
				}
				return fmt.Sprintf("Task %s %q started.", task.Id, task.Title), nil
			}),

		tool("interrupt", "Interrupt the turn a task's agent is running.", taskOnly,
			func(ctx context.Context, user api.User, raw json.RawMessage) (string, error) {
				sess, err := m.taskSessionArg(ctx, raw)
				if err != nil {
					return "", err
				}
				_, err = m.Interrupt(ctx, sess.Id)
				if err != nil {
					return "", err
				}
				return "Interrupted.", nil
			}),

		tool("compact_task", "Have a task's idle agent write a handoff and continue in a fresh session.", taskOnly,
			func(ctx context.Context, user api.User, raw json.RawMessage) (string, error) {
				sess, err := m.taskSessionArg(ctx, raw)
				if err != nil {
					return "", err
				}
				_, err = m.Compact(ctx, sess.Id)
				if err != nil {
					return "", err
				}
				return "Compaction started.", nil
			}),

		tool("merge_task",
			"Merge a task as the human: its agent merges the default branch in, then byre squashes it into one commit, pushes it to origin and closes the task. "+
				"While another task in the project is merging, byre holds the merge and starts it when that one is done. "+
				"Only when the human asks. You get a signal when it is merged; failures are filed in the task's inbox.",
			taskOnly,
			func(ctx context.Context, user api.User, raw json.RawMessage) (string, error) {
				task, project, err := m.taskArg(ctx, raw)
				if err != nil {
					return "", err
				}
				task, err = m.Merge(ctx, task, project, user, task.MergeTitle)
				if err != nil {
					return "", err
				}
				if task.MergeQueuedAt != nil {
					return fmt.Sprintf("Merge of %s held until the merge in progress in its project is done.", task.Id), nil
				}
				return fmt.Sprintf("Merge of %s started.", task.Id), nil
			}),

		tool("discard_task",
			"Discard a task: stop its agent and delete its worktree and branch without merging. This destroys its unmerged work. Only when the human asks.",
			taskOnly,
			func(ctx context.Context, user api.User, raw json.RawMessage) (string, error) {
				task, project, err := m.taskArg(ctx, raw)
				if err != nil {
					return "", err
				}
				_, err = m.Discard(ctx, task, project, user.Login)
				if err != nil {
					return "", err
				}
				return fmt.Sprintf("Discarded %s.", task.Id), nil
			}),

		tool("show",
			"Point the human's right pane at a task, or at a project with no task: tabs are the tabs to put first, in order (the rest follow), and tab is the one to select. "+
				"Tabs: transcript, changes, files, terminal (a shell in the task's worktree), preview (the dev servers the task's agent exposed with expose_port), inbox (the task's open items), reviews (the task's reviews by reviewer agents), project (the project's dashboard: its active tasks and its TODO.md items); with only project_id, only project shows anything. "+
				"Omit both task_id and project_id to clear the pane. "+
				"If the pane already shows that task (or project) on that tab, only the tab order changes and you are told so, even while the human is working in the pane. "+
				"Otherwise refused while the human is working in the right pane, unless it is on the project tab or what it shows is finished (a merged or discarded task, or an inbox tab with nothing left).",
			`{"type": "object", "properties": {
				"task_id": {"type": "string"},
				"project_id": {"type": "string", "description": "Ignored when task_id is given."},
				"tabs": {"type": "array", "items": {"type": "string", "enum": ["transcript", "changes", "files", "terminal", "preview", "inbox", "reviews", "project"]}},
				"tab": {"type": "string", "enum": ["transcript", "changes", "files", "terminal", "preview", "inbox", "reviews", "project"]}
			}, "required": ["tabs", "tab"]}`,
			func(ctx context.Context, user api.User, raw json.RawMessage) (string, error) {
				f := api.Focus{}
				err := json.Unmarshal(raw, &f)
				if err != nil {
					return "", err
				}
				current, err := m.store.GetFocus(ctx, user.Login)
				if err != nil {
					return "", err
				}
				if sameView(current, f) {
					if slices.Equal(current.Tabs, f.Tabs) {
						return "Already shown: the right pane shows that now, so nothing changed.", nil
					}
					err = m.SetFocus(ctx, user.Login, f, leadActor(user.Login))
					if err != nil {
						return "", err
					}
					return "Already shown on that tab; the tab order was updated.", nil
				}
				m.mu.Lock()
				focused := m.paneFocused[user.Login]
				m.mu.Unlock()
				if focused && current.Tab != api.TabProject {
					finished, err := m.paneFinished(ctx, current)
					if err != nil {
						return "", err
					}
					if finished == "" {
						return "", fmt.Errorf("%w: the human is working in the right pane, so it was left as it is; tell them in chat instead", ErrConflict)
					}
				}
				err = m.SetFocus(ctx, user.Login, f, leadActor(user.Login))
				if err != nil {
					return "", err
				}
				return "Shown.", nil
			}),

		tool("compact_self",
			"Ask byre to compact you when this turn ends: you write a handoff, and a fresh lead session continues from it. Call it only once your context is past half its window, at a logical point, then end the turn with a short final message.",
			`{"type": "object", "properties": {}}`,
			func(ctx context.Context, user api.User, raw json.RawMessage) (string, error) {
				m.mu.Lock()
				l.wantCompact = true
				m.mu.Unlock()
				return "byre will ask for your handoff once this turn ends.", nil
			}),
	}
}

func (m *Manager) SetFocus(ctx context.Context, login string, f api.Focus, actor string) error {
	for _, t := range append(slices.Clone(f.Tabs), f.Tab) {
		if !t.Valid() {
			return fmt.Errorf("%w: unknown tab %q", ErrInvalid, t)
		}
	}
	if f.TaskId != nil && *f.TaskId == "" {
		f.TaskId = nil
	}
	if f.ProjectId != nil && *f.ProjectId == "" {
		f.ProjectId = nil
	}
	if f.TaskId != nil {
		task, _, err := m.activeTask(ctx, *f.TaskId)
		if err != nil {
			return err
		}
		f.ProjectId = &task.ProjectId
	} else if f.ProjectId != nil {
		_, err := m.store.GetProject(ctx, *f.ProjectId)
		if err != nil {
			return fmt.Errorf("project %s: %w", *f.ProjectId, err)
		}
	}
	return m.store.SetFocus(ctx, login, f, actor)
}

func sameView(a, b api.Focus) bool {
	if a.Tab != b.Tab {
		return false
	}
	at, bt := deref(a.TaskId), deref(b.TaskId)
	if at != "" || bt != "" {
		return at == bt
	}
	return deref(a.ProjectId) == deref(b.ProjectId)
}

const inboxEmptied = "has nothing left in its inbox"

func (m *Manager) paneFinished(ctx context.Context, f api.Focus) (string, error) {
	if f.TaskId == nil {
		return "", nil
	}
	task, err := m.store.GetTask(ctx, *f.TaskId)
	if err != nil {
		return "", err
	}
	if task.Status != api.TaskStatusActive {
		return "is " + string(task.Status), nil
	}
	if f.Tab != api.TabInbox {
		return "", nil
	}
	items, err := m.store.ListInbox(ctx, store.InboxFilter{State: api.Open, TaskID: task.Id})
	if err != nil {
		return "", err
	}
	if len(items) > 0 {
		return "", nil
	}
	return inboxEmptied, nil
}

func (m *Manager) SetPaneFocus(login string, focused bool) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.paneFocused[login] = focused
}

func (m *Manager) modelOr(model string, l *live) string {
	if model != "" {
		return model
	}
	return l.model
}

func (m *Manager) activeTask(ctx context.Context, taskID string) (api.Task, api.Project, error) {
	task, err := m.store.GetTask(ctx, taskID)
	if err != nil {
		return api.Task{}, api.Project{}, fmt.Errorf("task %s: %w", taskID, err)
	}
	if task.Status != api.TaskStatusActive {
		return api.Task{}, api.Project{}, fmt.Errorf("%w: task %s is %s", ErrConflict, task.Id, task.Status)
	}
	project, err := m.store.GetProject(ctx, task.ProjectId)
	return task, project, err
}

func (m *Manager) liveSession(ctx context.Context, taskID string) (api.AgentSession, bool, error) {
	sessions, err := m.store.ListSessions(ctx, taskID)
	if err != nil {
		return api.AgentSession{}, false, err
	}
	for _, s := range sessions {
		if s.State != api.AgentSessionStateCompleted {
			return s, true, nil
		}
	}
	return api.AgentSession{}, false, nil
}

func (m *Manager) taskArg(ctx context.Context, raw json.RawMessage) (api.Task, api.Project, error) {
	args := struct {
		TaskID string `json:"task_id"`
	}{}
	err := json.Unmarshal(raw, &args)
	if err != nil {
		return api.Task{}, api.Project{}, err
	}
	return m.activeTask(ctx, args.TaskID)
}

func (m *Manager) taskSessionArg(ctx context.Context, raw json.RawMessage) (api.AgentSession, error) {
	task, _, err := m.taskArg(ctx, raw)
	if err != nil {
		return api.AgentSession{}, err
	}
	sess, found, err := m.liveSession(ctx, task.Id)
	if err != nil {
		return api.AgentSession{}, err
	}
	if !found {
		return api.AgentSession{}, fmt.Errorf("task %s has no agent", task.Id)
	}
	return sess, nil
}

func (m *Manager) listTasksText(ctx context.Context) (string, error) {
	projects, err := m.store.ListProjects(ctx)
	if err != nil {
		return "", err
	}
	items, err := m.store.ListInbox(ctx, store.InboxFilter{State: api.Open})
	if err != nil {
		return "", err
	}

	b := &strings.Builder{}
	for _, p := range projects {
		fmt.Fprintf(b, "## %s (project %s, repo %s, default branch %s)\n", p.Name, p.Id, p.RepoPath, p.DefaultBranch)
		tasks, err := m.store.ListTasks(ctx, p.Id)
		if err != nil {
			return "", err
		}
		listed := false
		for _, t := range tasks {
			if t.Status != api.TaskStatusActive {
				continue
			}
			listed = true
			agent := "no agent"
			sess, found, err := m.liveSession(ctx, t.Id)
			if err != nil {
				return "", err
			}
			if found {
				agent = fmt.Sprintf("agent %s", sess.State)
				if sess.Model != "" {
					agent += " on " + sess.Model
				}
			}
			blocked, fyi := 0, 0
			for _, i := range items {
				if i.TaskId != t.Id {
					continue
				}
				if i.Priority == api.Blocked {
					blocked++
				} else {
					fyi++
				}
			}
			fmt.Fprintf(b, "- %s %q: %s; inbox %d blocked, %d fyi; worktree %s", t.Id, t.Title, agent, blocked, fyi, t.WorktreePath)
			if t.TodoItem != nil {
				fmt.Fprintf(b, "; from TODO item %q", *t.TodoItem)
			}
			if t.ForkedFrom != nil {
				fmt.Fprintf(b, "; forked from %s at snapshot %d", *t.ForkedFrom, *t.ForkedAt)
			}
			b.WriteString("\n")
		}
		if !listed {
			b.WriteString("- no active tasks\n")
		}
	}
	return b.String(), nil
}

func (m *Manager) transcriptText(ctx context.Context, taskID string, after *int64, limit int) (string, error) {
	task, err := m.store.GetTask(ctx, taskID)
	if err != nil {
		return "", fmt.Errorf("task %s: %w", taskID, err)
	}
	if limit <= 0 {
		limit = 100
	}
	limit = min(limit, 500)

	var events []api.Event
	if after == nil {
		events, err = m.store.TaskEventsTail(ctx, task.Id, limit)
	} else {
		events, err = m.store.TaskEvents(ctx, task.Id, *after, limit)
	}
	if err != nil {
		return "", err
	}
	reviewers, err := m.store.ReviewerSessionIDs(ctx, task.Id)
	if err != nil {
		return "", err
	}

	b := &strings.Builder{}
	for _, ev := range events {
		if ev.AgentSessionId != nil && reviewers[*ev.AgentSessionId] {
			continue
		}
		line := transcriptLine(ev, task.WorktreePath)
		if line != "" {
			fmt.Fprintf(b, "#%d %s\n", ev.Seq, line)
		}
	}
	if len(events) == limit && limit > 0 {
		fmt.Fprintf(b, "(more may follow: call again with after_seq %d)\n", events[len(events)-1].Seq)
	}
	if b.Len() == 0 {
		return "No events.", nil
	}
	return b.String(), nil
}

func transcriptLine(ev api.Event, root string) string {
	p := ev.Payload
	text := func(k string) string { return fmt.Sprint(p[k]) }
	rel := func(s string) string { return strings.ReplaceAll(s, root+"/", "") }
	if p["parent_tool_use_id"] != nil {
		switch ev.Kind {
		case "agent.message":
			return "subagent: " + text("text")
		case "agent.tool_call":
			return "subagent tool " + text("name")
		}
		return ""
	}
	switch ev.Kind {
	case "user.message":
		if ev.Actor == "byre" {
			return "byre: " + clip(text("text"), 300)
		}
		return ev.Actor + ": " + text("text")
	case "agent.message":
		return "agent: " + text("text")
	case "agent.tool_call":
		input, _ := json.Marshal(p["input"])
		return fmt.Sprintf("tool %s %s", text("name"), clip(rel(string(input)), 200))
	case "agent.tool_result":
		content, _ := json.Marshal(p["content"])
		return "result: " + clip(rel(string(content)), 300)
	case "turn.ended":
		if p["is_error"] == true {
			return fmt.Sprintf("turn ended with an error (%s)", text("subtype"))
		}
		return "turn ended"
	case "inbox.created":
		return fmt.Sprintf("filed %s %s", text("kind"), text("id"))
	case "session.created":
		return fmt.Sprintf("session started by %s", ev.Actor)
	case "session.state":
		if p["error"] != nil {
			return fmt.Sprintf("session %s: %s", text("state"), text("error"))
		}
	case "task.merged", "task.discarded":
		return fmt.Sprintf("%s by %s", ev.Kind, ev.Actor)
	case "review.finished":
		r := api.Review{}
		err := remarshal(p, &r)
		if err != nil {
			return "review finished: " + err.Error()
		}
		return reviewLine(r)
	}
	return ""
}