byre/internal/agents/subtasks.go

package agents

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"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) subtaskTools(l *live) []mcp.Tool {
	tool := func(name, description, schema string, run func(ctx context.Context, creator 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) {
				sess, err := m.store.GetSession(ctx, l.id)
				if err != nil {
					return "", err
				}
				creator, err := m.store.GetUser(ctx, sess.CreatedBy)
				if err != nil {
					return "", err
				}
				return run(ctx, creator, raw)
			},
		}
	}
	taskOnly := `{"type": "object", "properties": {"task_id": {"type": "string"}}, "required": ["task_id"]}`

	return []mcp.Tool{
		tool("create_subtask",
			"Create a subtask of this task: a byre task with its own branch and worktree, started from this task's HEAD (your committed work; byre commits at the end of each turn), "+
				"with an agent on it that gets prompt as its first message. Use it to run a piece of work in parallel or to get another model's view. "+
				"The new agent sees nothing of your context, so the prompt must be a complete brief. "+
				"Returns at once with the subtask's id. The final message of each of its turns reaches you as a message from byre; "+
				"manage it with list_subtasks, read_subtask, message_subtask, interrupt_subtask, merge_subtask and discard_subtask. "+
				"This task can't be merged or discarded while it has active subtasks.",
			`{"type": "object", "properties": {
				"title": {"type": "string", "description": "Short task title, shown in the task list."},
				"prompt": {"type": "string", "description": "The new agent's first message: a complete, self-contained brief."},
				"model": {"type": "string", "description": "Model for the new agent: gpt-* models run on Codex, every other model on Claude Code. Defaults to this session's model."}
			}, "required": ["title", "prompt"]}`,
			func(ctx context.Context, creator api.User, raw json.RawMessage) (string, error) {
				args := struct {
					Title  string `json:"title"`
					Prompt string `json:"prompt"`
					Model  string `json:"model"`
				}{}
				err := json.Unmarshal(raw, &args)
				if err != nil {
					return "", fmt.Errorf("malformed create_subtask call: %w", err)
				}
				if strings.TrimSpace(args.Title) == "" || strings.TrimSpace(args.Prompt) == "" {
					return "", errors.New("title and prompt are required")
				}
				project, err := m.store.GetProject(ctx, l.projectID)
				if err != nil {
					return "", err
				}
				base, err := git.Head(l.dir)
				if err != nil {
					return "", err
				}
				task, err := m.createTask(ctx, project, api.Task{Title: args.Title, CreatedBy: creator.Login, ParentTaskId: &l.taskID}, nil, base, creator.Login)
				if err != nil {
					return "", err
				}
				_, err = m.create(ctx, task, creator, args.Prompt, m.modelOr(args.Model, l), agentActor, nil)
				if err != nil {
					return "", fmt.Errorf("subtask %s created, but its agent did not start: %w", task.Id, err)
				}
				return fmt.Sprintf("Subtask %s %q created from %s, branch %s, worktree %s, and its agent started.", task.Id, task.Title, task.BaseRef, task.Branch, task.WorktreePath), nil
			}),

		tool("list_subtasks",
			"List this task's active subtasks: id, title, the agent's state and model, branch and worktree.",
			`{"type": "object", "properties": {}}`,
			func(ctx context.Context, creator api.User, raw json.RawMessage) (string, error) {
				subs, err := m.store.ActiveSubtasks(ctx, l.taskID)
				if err != nil {
					return "", err
				}
				if len(subs) == 0 {
					return "No active subtasks.", nil
				}
				b := &strings.Builder{}
				for _, t := range subs {
					agent := "no agent"
					sess, found, err := m.liveSession(ctx, t.Id)
					if err != nil {
						return "", err
					}
					if found {
						agent = fmt.Sprintf("agent %s on %s", sess.State, sess.Model)
					}
					fmt.Fprintf(b, "- %s %q: %s; branch %s; worktree %s", t.Id, t.Title, agent, t.Branch, t.WorktreePath)
					if t.MergingIntoParent {
						b.WriteString("; waiting for you to merge it")
					}
					b.WriteString("\n")
				}
				return b.String(), nil
			}),

		tool("read_subtask",
			"Read a subtask'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, creator 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
				}
				task, err := m.subtask(ctx, l, args.TaskID, false)
				if err != nil {
					return "", err
				}
				return m.transcriptText(ctx, task.Id, args.AfterSeq, args.Limit)
			}),

		tool("message_subtask",
			"Send a message to a subtask's agent. It starts a turn when the agent is idle, starts the agent when the subtask has none, "+
				"and reaches it mid-turn when it is working.",
			`{"type": "object", "properties": {
				"task_id": {"type": "string"},
				"text": {"type": "string"}
			}, "required": ["task_id", "text"]}`,
			func(ctx context.Context, creator api.User, raw json.RawMessage) (string, error) {
				args := struct {
					TaskID string `json:"task_id"`
					Text   string `json:"text"`
				}{}
				err := json.Unmarshal(raw, &args)
				if err != nil {
					return "", err
				}
				if strings.TrimSpace(args.Text) == "" {
					return "", errors.New("text is empty")
				}
				task, err := m.subtask(ctx, l, args.TaskID, true)
				if err != nil {
					return "", err
				}
				err = m.messageTask(ctx, task, creator, agentActor, args.Text)
				if err != nil {
					return "", err
				}
				return fmt.Sprintf("Sent to %s's agent.", task.Id), nil
			}),

		tool("interrupt_subtask", "Interrupt the turn a subtask's agent is running.", taskOnly,
			func(ctx context.Context, creator api.User, raw json.RawMessage) (string, error) {
				task, err := m.subtaskArg(ctx, l, raw)
				if err != nil {
					return "", err
				}
				sess, found, err := m.liveSession(ctx, task.Id)
				if err != nil {
					return "", err
				}
				if !found {
					return "", fmt.Errorf("%w: subtask %s has no agent", ErrConflict, task.Id)
				}
				_, err = m.Interrupt(ctx, sess.Id)
				if err != nil {
					return "", err
				}
				return "Interrupted.", nil
			}),

		tool("merge_subtask",
			"Merge a subtask. target parent: byre sends you a message asking you to merge the subtask's branch into this task's branch, "+
				"and closes the subtask once its branch is in this branch's history at the end of one of your turns. "+
				"target main: the subtask's agent merges the default branch in, then byre squashes it onto the default branch, pushes it and closes it, as the human's Merge does. "+
				"Refused while the subtask's agent is mid-turn.",
			`{"type": "object", "properties": {
				"task_id": {"type": "string"},
				"target": {"type": "string", "enum": ["parent", "main"]}
			}, "required": ["task_id", "target"]}`,
			func(ctx context.Context, creator api.User, raw json.RawMessage) (string, error) {
				args := struct {
					TaskID string `json:"task_id"`
					Target string `json:"target"`
				}{}
				err := json.Unmarshal(raw, &args)
				if err != nil {
					return "", err
				}
				task, err := m.subtask(ctx, l, args.TaskID, true)
				if err != nil {
					return "", err
				}
				switch args.Target {
				case "parent":
					_, err = m.MergeIntoParent(ctx, task, creator)
					if err != nil {
						return "", err
					}
					return fmt.Sprintf("Merge of %s into this task requested: byre's message asking you to merge %s follows once this turn ends.", task.Id, task.Branch), nil
				case "main":
					project, err := m.store.GetProject(ctx, task.ProjectId)
					if err != nil {
						return "", err
					}
					task, err = m.Merge(ctx, task, project, creator, task.MergeTitle)
					if err != nil {
						return "", err
					}
					if task.MergeQueuedAt != nil {
						return fmt.Sprintf("Merge of %s held until the merge in progress in the project is done.", task.Id), nil
					}
					return fmt.Sprintf("Merge of %s started.", task.Id), nil
				}
				return "", fmt.Errorf("%w: target must be parent or main, not %q", ErrInvalid, args.Target)
			}),

		tool("discard_subtask",
			"Discard a subtask: its agent gets one turn to write the task summary, then byre stops it and deletes its worktree and branch without merging. This destroys its unmerged work.",
			taskOnly,
			func(ctx context.Context, creator api.User, raw json.RawMessage) (string, error) {
				task, err := m.subtaskArg(ctx, l, raw)
				if err != nil {
					return "", err
				}
				_, err = m.Discard(ctx, task, creator.Login)
				if err != nil {
					return "", err
				}
				return fmt.Sprintf("Discard of %s started: its agent writes the task summary, then byre closes it.", task.Id), nil
			}),
	}
}

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

func (m *Manager) subtask(ctx context.Context, l *live, taskID string, active bool) (api.Task, error) {
	task, err := m.store.GetTask(ctx, taskID)
	if err != nil {
		return api.Task{}, fmt.Errorf("task %s: %w", taskID, err)
	}
	if task.ParentTaskId == nil || *task.ParentTaskId != l.taskID {
		return api.Task{}, fmt.Errorf("%w: task %s is not a subtask of this task", ErrInvalid, taskID)
	}
	if active && task.Status != api.TaskStatusActive {
		return api.Task{}, fmt.Errorf("%w: subtask %s is %s", ErrConflict, task.Id, task.Status)
	}
	return task, nil
}

func (m *Manager) checkNoSubtasks(ctx context.Context, task api.Task) error {
	subs, err := m.store.ActiveSubtasks(ctx, task.Id)
	if err != nil {
		return err
	}
	if len(subs) == 0 {
		return nil
	}
	names := []string{}
	for _, t := range subs {
		names = append(names, fmt.Sprintf("%s %q", t.Id, t.Title))
	}
	return fmt.Errorf("%w: task %s has %d active subtasks; merge or discard them first: %s", ErrConflict, task.Id, len(subs), strings.Join(names, ", "))
}

func (m *Manager) messageTask(ctx context.Context, task api.Task, user api.User, actor, text string) error {
	sess, found, err := m.liveSession(ctx, task.Id)
	if err != nil {
		return err
	}
	if !found {
		_, err = m.create(ctx, task, user, text, "", actor, nil)
		return err
	}
	_, err = m.Message(ctx, sess.Id, actor, text, "")
	return err
}

func (m *Manager) MergeIntoParent(ctx context.Context, task api.Task, user api.User) (api.Task, error) {
	if task.ParentTaskId == nil {
		return api.Task{}, fmt.Errorf("%w: task %s is not a subtask", ErrInvalid, task.Id)
	}
	err := m.checkNoSubtasks(ctx, task)
	if err != nil {
		return api.Task{}, err
	}
	parent, _, err := m.activeTask(ctx, *task.ParentTaskId)
	if err != nil {
		return api.Task{}, err
	}
	sess, found, err := m.liveSession(ctx, task.Id)
	if err != nil {
		return api.Task{}, err
	}
	if found && sess.State == api.AgentSessionStateRunning {
		return api.Task{}, fmt.Errorf("%w: subtask %s's agent is mid-turn; wait for the turn to end or interrupt it", ErrConflict, task.Id)
	}
	mergingBy, err := m.store.MergingBy(ctx, task.Id)
	if err != nil {
		return api.Task{}, err
	}
	if mergingBy != "" || task.MergeQueuedAt != nil {
		return api.Task{}, fmt.Errorf("%w: subtask %s is merging into the default branch", ErrConflict, task.Id)
	}

	err = m.store.SetParentMerge(ctx, task.Id, user.Login)
	switch {
	case errors.Is(err, store.ErrNotFound):
		return api.Task{}, fmt.Errorf("%w: subtask %s is closed", ErrConflict, task.Id)
	case errors.Is(err, store.ErrDiscardPending), errors.Is(err, store.ErrParentMergeRequested):
		return api.Task{}, fmt.Errorf("%w: subtask %s: %w", ErrConflict, task.Id, err)
	case err != nil:
		return api.Task{}, err
	}

	prompt := mergeInPrompt(task.Branch) + fmt.Sprintf(" This merges in your subtask %q (%s); byre closes the subtask once its branch is in this branch's history at the end of a turn.", task.Title, task.Id)
	err = m.messageTask(ctx, parent, user, "byre", prompt)
	if err != nil {
		cerr := m.store.ClearParentMerge(ctx, task.Id)
		if cerr != nil {
			slog.Error("clear parent merge", "task", task.Id, "err", cerr)
		}
		return api.Task{}, err
	}
	return m.store.GetTask(ctx, task.Id)
}

func (m *Manager) closeMergedSubtasks(parentID, dir string) {
	ctx := context.Background()
	subs, err := m.store.ActiveSubtasks(ctx, parentID)
	if err != nil {
		slog.Error("list subtasks", "task", parentID, "err", err)
		return
	}
	for _, t := range subs {
		if !t.MergingIntoParent {
			continue
		}
		err = m.closeMergedSubtask(ctx, t, parentID, dir)
		if err != nil {
			slog.Error("close merged subtask", "task", t.Id, "parent", parentID, "err", err)
		}
	}
}

func (m *Manager) closeMergedSubtask(ctx context.Context, t api.Task, parentID, dir string) error {
	merged, err := git.IsAncestor(dir, t.Branch, "HEAD")
	if err != nil || !merged {
		return err
	}
	err = m.checkNoSubtasks(ctx, t)
	if err != nil {
		return err
	}
	login, err := m.store.ParentMergeBy(ctx, t.Id)
	if err != nil {
		return err
	}
	head, err := git.Head(dir)
	if err != nil {
		return err
	}
	project, err := m.store.GetProject(ctx, t.ProjectId)
	if err != nil {
		return err
	}
	err = m.RemoveTask(ctx, t, project, login)
	if err != nil {
		return err
	}
	_, err = m.store.MergeTaskIntoParent(ctx, t.Id, parentID, head, login)
	if err != nil {
		return err
	}
	slog.Info("subtask merged into parent", "task", t.Id, "parent", parentID, "commit", head)
	return nil
}

func (m *Manager) reportToParent(taskID, result string) {
	ctx := context.Background()
	task, err := m.store.GetTask(ctx, taskID)
	if err != nil {
		slog.Error("load subtask to report", "task", taskID, "err", err)
		return
	}
	if task.ParentTaskId == nil || task.MergingIntoParent {
		return
	}
	parent, err := m.store.GetTask(ctx, *task.ParentTaskId)
	if err != nil {
		slog.Error("load parent to report to", "task", taskID, "err", err)
		return
	}
	if parent.Status != api.TaskStatusActive {
		return
	}
	user, err := m.store.GetUser(ctx, task.CreatedBy)
	if err != nil {
		slog.Error("load subtask creator", "task", taskID, "err", err)
		return
	}
	text := fmt.Sprintf("Your subtask %q (%s, branch %s) finished a turn:\n\n%s", task.Title, task.Id, task.Branch, quote(result))
	err = m.messageTask(ctx, parent, user, "byre", text)
	if err != nil {
		slog.Error("report subtask turn to parent", "task", taskID, "parent", parent.Id, "err", err)
	}
}