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

func (m *Manager) freeze(ctx context.Context, name, brief string) (*store.FrozenWorkflow, error) {
	if name == "" {
		return nil, nil
	}
	frozen, err := m.store.FreezeWorkflow(ctx, name)
	if err != nil {
		return nil, err
	}
	frozen.Brief = brief
	return &frozen, nil
}

func stageMessage(frozen store.FrozenWorkflow, n int, previous string) string {
	stages := frozen.Definition.Stages
	stage := stages[n-1]

	b := &strings.Builder{}
	fmt.Fprintf(b, "Workflow %q, stage %d of %d: %s\n\n%s", frozen.Name, n, len(stages), stage.Name, strings.TrimSpace(stage.Prompt))
	if strings.TrimSpace(frozen.Brief) != "" {
		fmt.Fprintf(b, "\n\n## Brief\n\n%s", strings.TrimSpace(frozen.Brief))
	}
	if n > 1 {
		fmt.Fprintf(b, "\n\n## From the previous stage (%s)\n\n%s", stages[n-2].Name, strings.TrimSpace(previous))
	}
	return b.String()
}

func (m *Manager) startStage(ctx context.Context, task api.Task, creator api.User, frozen store.FrozenWorkflow, n int, previous, actor string) (api.AgentSession, error) {
	stage := frozen.Definition.Stages[n-1]
	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, "stage": n, "stages": len(frozen.Definition.Stages), "name": stage.Name},
	})
	if err != nil {
		return api.AgentSession{}, err
	}
	return m.createAt(ctx, task, creator, stageMessage(frozen, n, previous), stage.Model, "byre", nil, n)
}

func (m *Manager) startWorkflow(ctx context.Context, task api.Task, creator api.User, actor string) error {
	frozen, err := m.store.TaskWorkflow(ctx, task.Id)
	if err == nil {
		_, err = m.startStage(ctx, task, creator, frozen, 1, "", actor)
	}
	if err != nil {
		return fmt.Errorf("task %s created, but its workflow did not start: %w", task.Id, err)
	}
	return nil
}

func (m *Manager) advanceStage(l *live) {
	_, err := m.nextStage(context.Background(), 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))
	}
}

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
		}
		if next != "" {
			started[next] = 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 "", err
	}
	task, err := m.store.GetTask(ctx, sess.TaskId)
	if err != nil {
		return "", err
	}
	if task.Status != api.TaskStatusActive {
		return "", m.store.ClearStage(ctx, sess.Id)
	}
	frozen, err := m.store.TaskWorkflow(ctx, task.Id)
	if err != nil {
		return "", err
	}
	output, err := m.store.StageOutput(ctx, sess.Id)
	if err != nil {
		return "", err
	}

	next := *sess.WorkflowStage + 1
	if next > len(frozen.Definition.Stages) {
		return "", m.store.FinishWorkflow(ctx, task, sess.Id, frozen.Name, output)
	}

	creator, err := m.store.GetUser(ctx, sess.CreatedBy)
	if err != nil {
		return "", err
	}
	if sess.State != api.AgentSessionStateCompleted {
		m.logUndelivered(ctx, sess.Id, "stage finished")
		err = m.retire(ctx, sess, "byre")
		if err != nil {
			return "", err
		}
	}

	successor, found, err := m.successor(ctx, sess, next)
	if err != nil {
		return "", err
	}
	sent := true
	switch {
	case !found:
		successor, err = m.startStage(ctx, task, creator, frozen, next, output, "byre")
	default:
		sent, err = m.resendStage(ctx, successor, task, frozen, next, output)
	}
	if err != nil {
		return "", err
	}
	err = m.store.ClearStage(ctx, sess.Id)
	if err != nil || !sent {
		return "", err
	}
	return successor.Id, nil
}

func (m *Manager) successor(ctx context.Context, sess api.AgentSession, n int) (api.AgentSession, bool, error) {
	sessions, err := m.store.ListSessions(ctx, sess.TaskId)
	if err != nil {
		return api.AgentSession{}, false, err
	}
	for _, s := range sessions {
		if s.CreatedAt.Before(sess.CreatedAt) || s.Id == sess.Id || s.WorkflowStage == nil || *s.WorkflowStage != n {
			continue
		}
		return s, true, nil
	}
	return api.AgentSession{}, false, nil
}

func (m *Manager) resendStage(ctx context.Context, successor api.AgentSession, task api.Task, frozen store.FrozenWorkflow, n int, previous string) (bool, error) {
	started, err := m.store.SessionStarted(ctx, successor.Id)
	if err != nil || started || successor.State == api.AgentSessionStateCompleted {
		return false, err
	}
	if successor.State == api.AgentSessionStateRunning {
		err = m.store.SetSessionState(ctx, successor.Id, api.AgentSessionStateIdle, "byre", nil)
		if err != nil {
			return false, err
		}
	}
	_, err = m.turn(ctx, successor, task, "byre", stageMessage(frozen, n, previous), nil)
	return err == nil, err
}

func (m *Manager) workflowStatus(ctx context.Context, taskID string) (string, error) {
	finished, err := m.store.WorkflowFinished(ctx, taskID)
	if err != nil || finished {
		return "finished", err
	}
	frozen, err := m.store.TaskWorkflow(ctx, taskID)
	if err != nil {
		return "", err
	}
	sessions, err := m.store.ListSessions(ctx, taskID)
	if err != nil {
		return "", err
	}
	if len(sessions) == 0 || sessions[len(sessions)-1].WorkflowStage == nil {
		return "not started", nil
	}
	n := *sessions[len(sessions)-1].WorkflowStage
	stages := frozen.Definition.Stages
	return fmt.Sprintf("stage %d/%d %s", n, len(stages), stages[n-1].Name), nil
}

func (m *Manager) finishStageTool(l *live) mcp.Tool {
	return mcp.Tool{
		Name: "finish_stage",
		Description: "Finish your stage of this task's workflow. Call it once the stage's work is done; byre starts the next stage, in a fresh session, when this turn ends " +
			"(after the last 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 next stage starts from: everything it needs from this stage, since it sees nothing else of your context. For the last 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)
			}
			err = m.store.SetStageOutput(ctx, l.id, args.Output)
			if err != nil {
				return "", err
			}
			if l.stage == l.stages {
				return "This is the workflow's last stage: it finishes when this turn ends. End your turn now with your summary.", nil
			}
			return fmt.Sprintf("byre starts stage %d when this turn ends. End your turn now with your summary.", l.stage+1), nil
		},
	}
}

func stageGuidance(frozen store.FrozenWorkflow, n int) string {
	return fmt.Sprintf("\n- This task runs byre workflow %q, and you are its stage %d of %d, %q. When this stage's work is done, call byre's finish_stage tool with its output, then end your turn: "+
		"byre then starts the next stage in a fresh session, which gets the output and nothing else of your context. Until you call it, the task stays in your stage.",
		frozen.Name, n, len(frozen.Definition.Stages), frozen.Definition.Stages[n-1].Name)
}

const workflowFileHelp = "A workflow's file is YAML: an optional description, then stages, each with name (one line, unique in the workflow), " +
	"model (which also picks the vendor: gpt-* models run on Codex, every other model on Claude Code) and prompt (the stage agent's instructions). " +
	"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
			}),
	}
}