byre/internal/agents/workflows.go

package agents

import (
	"context"
	"encoding/json"
	"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"
	"g.fc.run/misc/byre/internal/workflow"
)

func (m *Manager) SaveWorkflow(ctx context.Context, name string, file []byte, actor string) (api.WorkflowSummary, error) {
	name, err := workflow.CleanName(name)
	if err != nil {
		return api.WorkflowSummary{}, err
	}
	def, err := workflow.Parse(file)
	if err != nil {
		return api.WorkflowSummary{}, err
	}
	err = m.store.PutWorkflow(ctx, name, def, actor)
	if err != nil {
		return api.WorkflowSummary{}, err
	}
	return api.WorkflowSummary{Name: name, Description: def.Description}, nil
}

func (m *Manager) WorkflowFile(ctx context.Context, name string) ([]byte, error) {
	def, err := m.store.GetWorkflow(ctx, name)
	if err != nil {
		return nil, err
	}
	return workflow.Format(def)
}

func (m *Manager) TaskWorkflow(ctx context.Context, taskID string) (api.TaskWorkflow, error) {
	frozen, err := m.store.TaskWorkflow(ctx, taskID)
	if err != nil {
		return api.TaskWorkflow{}, err
	}
	file, err := workflow.Format(frozen.Definition)
	if err != nil {
		return api.TaskWorkflow{}, err
	}
	return api.TaskWorkflow{Name: frozen.Name, Definition: string(file)}, nil
}

var errModelWithWorkflow = fmt.Errorf("%w: model can't be combined with workflow: each stage names its own model", ErrInvalid)

type parentStage struct {
	task   api.Task
	output string
	commit string
}

func (m *Manager) StartWorkflow(ctx context.Context, project api.Project, creator api.User, title, base, actor, name, brief string, todoItem *string) ([]api.Task, error) {
	if creator.GitName == "" || creator.GitEmail == "" {
		return nil, ErrNoGitIdentity
	}
	frozen, err := m.store.FreezeWorkflow(ctx, name)
	if err != nil {
		return nil, err
	}
	frozen.Brief = brief
	frozen.Title = title

	run := store.NewID()
	tasks := []api.Task{}
	for _, stage := range workflow.Starts(frozen.Definition) {
		from := base
		if workflow.FromMain(frozen.Definition, stage) {
			from = project.DefaultBranch
		}
		task, err := m.startStage(ctx, project, creator, frozen, run, stage, from, todoItem, nil, actor)
		if err != nil {
			return tasks, fmt.Errorf("workflow %q: stage %q did not start: %w", name, stage, err)
		}
		tasks = append(tasks, task)
	}
	return tasks, nil
}

func (m *Manager) startStage(ctx context.Context, project api.Project, creator api.User, frozen store.FrozenWorkflow, run, stage, base string, todoItem *string, parents []parentStage, actor string) (api.Task, error) {
	task, err := m.createTask(ctx, project, api.Task{
		Title:         fmt.Sprintf("%s: %s", frozen.Title, stage),
		MergeTitle:    frozen.Title,
		CreatedBy:     creator.Login,
		TodoItem:      todoItem,
		WorkflowRun:   &run,
		WorkflowStage: &stage,
	}, &frozen, base, actor)
	if err != nil {
		return api.Task{}, err
	}
	_, err = m.beginStage(ctx, task, project, creator, frozen, parents, actor)
	if err != nil {
		return task, fmt.Errorf("task %s created, but its agent did not start: %w", task.Id, err)
	}
	return task, nil
}

func (m *Manager) beginStage(ctx context.Context, task api.Task, project api.Project, creator api.User, frozen store.FrozenWorkflow, parents []parentStage, actor string) (api.AgentSession, error) {
	stage := *task.WorkflowStage
	after := []string{}
	for _, p := range parents {
		after = append(after, p.task.Id)
	}
	err := m.store.AppendEvent(ctx, api.Event{
		Kind:      "task.workflow_stage",
		ProjectId: &task.ProjectId,
		TaskId:    &task.Id,
		Actor:     actor,
		Payload:   map[string]any{"id": task.Id, "workflow": frozen.Name, "run": *task.WorkflowRun, "stage": stage, "after": after},
	})
	if err != nil {
		return api.AgentSession{}, err
	}
	return m.create(ctx, task, creator, stageMessage(frozen, stage, parents, project.DefaultBranch), workflow.Get(frozen.Definition, stage).Model, "byre", nil)
}

func stageMessage(frozen store.FrozenWorkflow, stage string, parents []parentStage, defaultBranch string) string {
	b := &strings.Builder{}
	fmt.Fprintf(b, "Workflow %q, stage %q\n\n%s", frozen.Name, stage, strings.TrimSpace(workflow.Get(frozen.Definition, stage).Prompt))
	if strings.TrimSpace(frozen.Brief) != "" {
		fmt.Fprintf(b, "\n\n## Brief\n\n%s", strings.TrimSpace(frozen.Brief))
	}
	for _, p := range parents {
		fmt.Fprintf(b, "\n\n## From stage %q (task %s, branch %s)\n\n%s", *p.task.WorkflowStage, p.task.Id, p.task.Branch, strings.TrimSpace(p.output))
	}
	switch {
	case len(parents) == 0:
	case workflow.FromMain(frozen.Definition, stage):
		fmt.Fprintf(b, "\n\nThis task starts from %s. The branch of each stage above holds that stage's work: `git merge` the ones this stage needs, or read them with `git show` and `git diff`.", defaultBranch)
	default:
		fmt.Fprintf(b, "\n\nThis task starts from the commit stage %q finished at, so its work is already in this worktree.", *parents[0].task.WorkflowStage)
	}
	return b.String()
}

func (m *Manager) advanceStage(l *live) {
	ctx := context.Background()
	_, err := m.nextStage(ctx, l.id)
	if err != nil {
		slog.Error("advance workflow stage", "session", l.id, "err", err)
		m.fileError(l, fmt.Sprintf("The workflow didn't move on from this stage: %v. byre tries again when it next starts.", err))
	}
	err = m.Deliver(ctx, l.id)
	if err != nil {
		slog.Error("deliver inbox responses", "session", l.id, "err", err)
	}
}

func (m *Manager) recoverStages(ctx context.Context) (map[string]bool, error) {
	ids, err := m.store.FinishedStages(ctx)
	if err != nil {
		return nil, err
	}
	started := map[string]bool{}
	for _, id := range ids {
		next, err := m.nextStage(ctx, id)
		if err != nil {
			slog.Error("advance workflow stage", "session", id, "err", err)
			continue
		}
		for _, n := range next {
			started[n] = true
		}
		slog.Info("advanced workflow stage", "session", id, "next", next)
	}
	return started, nil
}

func (m *Manager) nextStage(ctx context.Context, sessionID string) ([]string, error) {
	sess, err := m.store.GetSession(ctx, sessionID)
	if err != nil {
		return nil, err
	}
	task, err := m.store.GetTask(ctx, sess.TaskId)
	if err != nil {
		return nil, err
	}
	if task.Status != api.TaskStatusActive {
		return nil, m.store.ClearStage(ctx, sess.Id)
	}
	frozen, err := m.store.TaskWorkflow(ctx, task.Id)
	if err != nil {
		return nil, err
	}
	output, err := m.store.StageOutput(ctx, sess.Id)
	if err != nil {
		return nil, err
	}
	recorded, _, err := m.store.StageResult(ctx, task.Id)
	if err != nil {
		return nil, err
	}
	if recorded == "" {
		head, err := git.Head(task.WorktreePath)
		if err != nil {
			return nil, err
		}
		err = m.store.RecordStage(ctx, task.Id, output, head)
		if err != nil {
			return nil, err
		}
	}

	m.workflowMu.Lock()
	defer m.workflowMu.Unlock()

	next := workflow.Next(frozen.Definition, *task.WorkflowStage)
	if len(next) == 0 {
		err = m.discardStages(ctx, task)
		if err != nil {
			return nil, err
		}
		return nil, m.store.FinishWorkflow(ctx, task, sess.Id, frozen.Name, output)
	}

	started := []string{}
	for _, n := range next {
		id, err := m.advanceTo(ctx, task, frozen, n)
		if err != nil {
			return started, err
		}
		if id != "" {
			started = append(started, id)
		}
	}
	return started, m.store.ClearStage(ctx, sess.Id)
}

func (m *Manager) advanceTo(ctx context.Context, from api.Task, frozen store.FrozenWorkflow, stage string) (string, error) {
	tasks, err := m.store.RunTasks(ctx, *from.WorkflowRun)
	if err != nil {
		return "", err
	}
	parents := []parentStage{}
	for _, a := range workflow.After(frozen.Definition, stage) {
		p, found, err := m.finishedStage(ctx, tasks, a)
		if err != nil || !found {
			return "", err
		}
		parents = append(parents, p)
	}
	project, err := m.store.GetProject(ctx, from.ProjectId)
	if err != nil {
		return "", err
	}
	creator, err := m.store.GetUser(ctx, from.CreatedBy)
	if err != nil {
		return "", err
	}

	existing, found := latestStage(tasks, stage)
	if found {
		if existing.Status != api.TaskStatusActive {
			return "", nil
		}
		return m.resumeStage(ctx, existing, project, creator, frozen, parents)
	}

	base := parents[0].commit
	if workflow.FromMain(frozen.Definition, stage) {
		base = project.DefaultBranch
	}
	task, err := m.startStage(ctx, project, creator, frozen, *from.WorkflowRun, stage, base, from.TodoItem, parents, "byre")
	if err != nil {
		return "", err
	}
	sessions, err := m.store.ListSessions(ctx, task.Id)
	if err != nil || len(sessions) == 0 {
		return "", err
	}
	return sessions[0].Id, nil
}

func latestStage(tasks []api.Task, stage string) (api.Task, bool) {
	for i := len(tasks) - 1; i >= 0; i-- {
		if *tasks[i].WorkflowStage == stage {
			return tasks[i], true
		}
	}
	return api.Task{}, false
}

func (m *Manager) finishedStage(ctx context.Context, tasks []api.Task, stage string) (parentStage, bool, error) {
	for i := len(tasks) - 1; i >= 0; i-- {
		if *tasks[i].WorkflowStage != stage {
			continue
		}
		output, commit, err := m.store.StageResult(ctx, tasks[i].Id)
		if err != nil {
			return parentStage{}, false, err
		}
		if output != "" {
			return parentStage{task: tasks[i], output: output, commit: commit}, true, nil
		}
	}
	return parentStage{}, false, nil
}

func (m *Manager) resumeStage(ctx context.Context, task api.Task, project api.Project, creator api.User, frozen store.FrozenWorkflow, parents []parentStage) (string, error) {
	sessions, err := m.store.ListSessions(ctx, task.Id)
	if err != nil {
		return "", err
	}
	if len(sessions) == 0 {
		sess, err := m.beginStage(ctx, task, project, creator, frozen, parents, "byre")
		return sess.Id, err
	}
	sess := sessions[0]
	started, err := m.store.SessionStarted(ctx, sess.Id)
	if err != nil || started || sess.State == api.AgentSessionStateCompleted {
		return "", err
	}
	if sess.State == api.AgentSessionStateRunning {
		err = m.store.SetSessionState(ctx, sess.Id, api.AgentSessionStateIdle, "byre", nil)
		if err != nil {
			return "", err
		}
	}
	_, err = m.turn(ctx, sess, task, "byre", stageMessage(frozen, *task.WorkflowStage, parents, project.DefaultBranch), nil)
	return sess.Id, err
}

func (m *Manager) discardStages(ctx context.Context, final api.Task) error {
	tasks, err := m.store.RunTasks(ctx, *final.WorkflowRun)
	if err != nil {
		return err
	}
	project, err := m.store.GetProject(ctx, final.ProjectId)
	if err != nil {
		return err
	}
	for _, t := range tasks {
		if t.Id == final.Id || t.Status != api.TaskStatusActive {
			continue
		}
		_, err = m.Discard(ctx, t, project, "byre")
		if err != nil {
			return fmt.Errorf("discard stage task %s: %w", t.Id, err)
		}
	}
	return nil
}

func (m *Manager) workflowStatus(ctx context.Context, task api.Task) (string, error) {
	finished, err := m.store.WorkflowFinished(ctx, *task.WorkflowRun)
	if err != nil {
		return "", err
	}
	output, _, err := m.store.StageResult(ctx, task.Id)
	if err != nil {
		return "", err
	}
	status := fmt.Sprintf("stage %q", *task.WorkflowStage)
	if output != "" {
		status += " (finished)"
	}
	if finished {
		status += ", workflow finished"
	}
	return status, nil
}

func (m *Manager) finishStageTool(l *live) mcp.Tool {
	return mcp.Tool{
		Name: "finish_stage",
		Description: "Finish this task's stage of its workflow. Call it once the stage's work is done; when this turn ends, byre hands the output to the stages that come after this one, " +
			"each of which starts in its own task once every stage it comes after has finished (after the final stage, the workflow is finished). Calling it again in the same turn replaces the output. " +
			"The call counts only for the turn that makes it: if that turn ends in an error or is interrupted, it is discarded, and you call it again once the stage's work is done.",
		InputSchema: json.RawMessage(`{
			"type": "object",
			"properties": {
				"output": {"type": "string", "description": "What the stages after this one start from: everything they need from this stage, since they see nothing else of your context. For the final stage, what the workflow produced."}
			},
			"required": ["output"]
		}`),
		Handler: func(ctx context.Context, raw json.RawMessage) (string, error) {
			args := struct {
				Output string `json:"output"`
			}{}
			err := json.Unmarshal(raw, &args)
			if err != nil {
				return "", fmt.Errorf("malformed finish_stage call: %w", err)
			}
			if strings.TrimSpace(args.Output) == "" {
				return "", fmt.Errorf("%w: output is required", ErrInvalid)
			}
			recorded, _, err := m.store.StageResult(ctx, l.taskID)
			if err != nil {
				return "", err
			}
			if recorded != "" {
				return "", fmt.Errorf("%w: stage %q already finished; its output went to the stages after it", ErrConflict, l.stage)
			}
			frozen, err := m.store.TaskWorkflow(ctx, l.taskID)
			if err != nil {
				return "", err
			}
			err = m.store.SetStageOutput(ctx, l.id, args.Output)
			if err != nil {
				return "", err
			}
			next := workflow.Next(frozen.Definition, l.stage)
			if len(next) == 0 {
				return "This is the workflow's final stage: it finishes when this turn ends, and byre discards the workflow's other stage tasks. End your turn now with your summary.", nil
			}
			return fmt.Sprintf("When this turn ends, byre hands the output to %s. End your turn now with your summary.", quoteList(next)), nil
		},
	}
}

func describeTasks(tasks []api.Task) string {
	parts := []string{}
	for _, t := range tasks {
		parts = append(parts, fmt.Sprintf("task %s %q from %s", t.Id, t.Title, t.BaseRef))
	}
	return strings.Join(parts, ", ")
}

func quoteList(names []string) string {
	quoted := []string{}
	for _, n := range names {
		quoted = append(quoted, fmt.Sprintf("%q", n))
	}
	return strings.Join(quoted, ", ")
}

func stageGuidance(frozen store.FrozenWorkflow, stage string) string {
	next := workflow.Next(frozen.Definition, stage)
	then := "it is the workflow's final stage, so the workflow then finishes and byre discards its other stage tasks"
	if len(next) > 0 {
		then = fmt.Sprintf("byre then hands the output to %s, each in its own task, which gets the output and nothing else of your context", quoteList(next))
	}
	return fmt.Sprintf("\n- This task runs stage %q of byre workflow %q. When this stage's work is done, call byre's finish_stage tool with its output, then end your turn: "+
		"%s. Until you call it, the workflow waits on this stage.", stage, frozen.Name, then)
}

const workflowFileHelp = "A workflow's file is YAML: an optional description, then stages, each with name (one line, unique in the workflow), " +
	"after (the names of the stages whose outputs it gets; it starts once all of them have finished; absent means the previous stage in the list, and for the first stage none; [] means it starts from the brief), " +
	"base (parent, the default: branch from the commit its one after stage finished at, or for a stage with none, from where the workflow was started; main: branch from the default branch; " +
	"refused on a stage with more than one after stage, which always branches from the default branch and is told the branches it can merge), " +
	"model (which also picks the vendor: gpt-* models run on Codex, every other model on Claude Code) and prompt (the stage agent's instructions). " +
	"Every stage runs as its own task. Exactly one stage is final (no other stage comes after it), and cycles are refused. " +
	"The workflow's name is not in the file."

func (m *Manager) workflowTools(tool func(name, description, schema string, run func(ctx context.Context, user api.User, raw json.RawMessage) (string, error)) mcp.Tool) []mcp.Tool {
	return []mcp.Tool{
		tool("list_workflows",
			"List the workflows on the server: each one's name and description.",
			`{"type": "object", "properties": {}}`,
			func(ctx context.Context, user api.User, raw json.RawMessage) (string, error) {
				workflows, err := m.store.ListWorkflows(ctx)
				if err != nil {
					return "", err
				}
				if len(workflows) == 0 {
					return "No workflows.", nil
				}
				b := &strings.Builder{}
				for _, w := range workflows {
					fmt.Fprintf(b, "- %q", w.Name)
					if w.Description != "" {
						fmt.Fprintf(b, ": %s", w.Description)
					}
					b.WriteString("\n")
				}
				return b.String(), nil
			}),

		tool("read_workflow",
			"Read a workflow's file: with name, the workflow on the server; with task_id, the copy that task froze when it was created. "+workflowFileHelp,
			`{"type": "object", "properties": {
				"name": {"type": "string"},
				"task_id": {"type": "string"}
			}}`,
			func(ctx context.Context, user api.User, raw json.RawMessage) (string, error) {
				args := struct {
					Name   string `json:"name"`
					TaskID string `json:"task_id"`
				}{}
				err := json.Unmarshal(raw, &args)
				if err != nil {
					return "", err
				}
				if (args.Name == "") == (args.TaskID == "") {
					return "", fmt.Errorf("%w: give exactly one of name and task_id", ErrInvalid)
				}
				if args.TaskID != "" {
					tw, err := m.TaskWorkflow(ctx, args.TaskID)
					if err != nil {
						return "", err
					}
					return fmt.Sprintf("Task %s froze workflow %q:\n\n%s", args.TaskID, tw.Name, tw.Definition), nil
				}
				file, err := m.WorkflowFile(ctx, args.Name)
				return string(file), err
			}),

		tool("save_workflow",
			"Create or replace a workflow on the server from its file. Tasks that already froze it keep their copy. "+
				"The name is one line of 1 to 100 characters, shown to the human on buttons and pickers. "+workflowFileHelp,
			`{"type": "object", "properties": {
				"name": {"type": "string"},
				"definition": {"type": "string", "description": "The workflow's file, as YAML."}
			}, "required": ["name", "definition"]}`,
			func(ctx context.Context, user api.User, raw json.RawMessage) (string, error) {
				args := struct {
					Name       string `json:"name"`
					Definition string `json:"definition"`
				}{}
				err := json.Unmarshal(raw, &args)
				if err != nil {
					return "", err
				}
				w, err := m.SaveWorkflow(ctx, args.Name, []byte(args.Definition), leadActor(user.Login))
				if err != nil {
					return "", err
				}
				return fmt.Sprintf("Saved workflow %q.", w.Name), nil
			}),

		tool("delete_workflow",
			"Delete a workflow from the server. Tasks that froze it keep their copy. Only when the human asks.",
			`{"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]}`,
			func(ctx context.Context, user api.User, raw json.RawMessage) (string, error) {
				args := struct {
					Name string `json:"name"`
				}{}
				err := json.Unmarshal(raw, &args)
				if err != nil {
					return "", err
				}
				err = m.store.DeleteWorkflow(ctx, args.Name, leadActor(user.Login))
				if err != nil {
					return "", err
				}
				return fmt.Sprintf("Deleted workflow %q.", args.Name), nil
			}),
	}
}