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, or cancel its queued turn.", 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. "+
				"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
				}
				_, err = m.Merge(ctx, task, project, user)
				if err != nil {
					return "", err
				}
				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, inbox (the task's open items), todo (the project's TODO.md); with only project_id, only todo shows anything. "+
				"Omit both task_id and project_id to clear the pane. Refused while the human is working in the right pane.",
			`{"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", "inbox", "todo"]}},
				"tab": {"type": "string", "enum": ["transcript", "changes", "files", "inbox", "todo"]}
			}, "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
				}
				m.mu.Lock()
				focused := m.paneFocused[user.Login]
				m.mu.Unlock()
				if focused {
					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 at logical points, 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 (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)
			}
			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
	}

	b := &strings.Builder{}
	for _, ev := range events {
		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+"/", "") }
	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)
	}
	return ""
}