byre/internal/workflow/workflow.go

package workflow

import (
	"errors"
	"fmt"
	"strings"
	"unicode"
	"unicode/utf8"

	"g.fc.run/ai/byre/internal/api"
)

var ErrInvalid = errors.New("invalid workflow")

const maxNameRunes = 100

func VendorOf(model string) api.AgentSessionVendor {
	switch {
	case strings.HasPrefix(model, "gpt-"):
		return api.Codex
	case strings.HasPrefix(model, "gemini-"):
		return api.Antigravity
	}
	return api.Claude
}

func CleanName(name string) (string, error) {
	name = strings.TrimSpace(name)
	problem := nameProblem(name)
	if problem != "" {
		return "", fmt.Errorf("%w: workflow name %s", ErrInvalid, problem)
	}
	return name, nil
}

func nameProblem(name string) string {
	switch {
	case strings.TrimSpace(name) == "":
		return "is empty"
	case strings.TrimSpace(name) != name:
		return fmt.Sprintf("%q has leading or trailing spaces", name)
	case utf8.RuneCountInString(name) > maxNameRunes:
		return fmt.Sprintf("%q is longer than %d characters", name, maxNameRunes)
	case strings.ContainsFunc(name, unicode.IsControl):
		return fmt.Sprintf("%q isn't a single line", name)
	}
	return ""
}

func Validate(def api.WorkflowDefinition) error {
	problems := []string{}
	if len(def.Stages) == 0 {
		problems = append(problems, "it has no stages")
	}

	seen := map[string]bool{}
	for i, s := range def.Stages {
		where := fmt.Sprintf("stage %d (%q)", i+1, s.Name)
		p := nameProblem(s.Name)
		if p != "" {
			problems = append(problems, fmt.Sprintf("%s: name %s", where, p))
		}
		if seen[s.Name] {
			problems = append(problems, fmt.Sprintf("%s: another stage has the same name", where))
		}
		seen[s.Name] = true

		switch {
		case s.Agents != nil && strings.TrimSpace(s.Model) != "":
			problems = append(problems, fmt.Sprintf("%s: model can't be given with agents, which name their own", where))
		case s.Agents != nil:
			for _, p := range TeamProblems(*s.Agents) {
				problems = append(problems, fmt.Sprintf("%s: %s", where, p))
			}
		case strings.TrimSpace(s.Model) == "":
			problems = append(problems, fmt.Sprintf("%s: model is required", where))
		}

		if strings.TrimSpace(s.Prompt) == "" {
			problems = append(problems, fmt.Sprintf("%s: prompt is empty", where))
		}

		if s.Base != nil && !s.Base.Valid() {
			problems = append(problems, fmt.Sprintf("%s: base must be parent or main, not %q", where, *s.Base))
		}
		if s.Base != nil && len(After(def, s.Name)) > 1 {
			problems = append(problems, fmt.Sprintf("%s: base can't be given on a stage with more than one after stage, which always branches from the default branch", where))
		}
		if s.After == nil {
			continue
		}
		listed := map[string]bool{}
		for _, a := range *s.After {
			switch {
			case a == s.Name:
				problems = append(problems, fmt.Sprintf("%s: after lists the stage itself", where))
			case Index(def, a) < 0:
				problems = append(problems, fmt.Sprintf("%s: after lists %q, which isn't a stage", where, a))
			case listed[a]:
				problems = append(problems, fmt.Sprintf("%s: after lists %q twice", where, a))
			}
			listed[a] = true
		}
	}

	if len(problems) == 0 {
		problems = append(problems, graphProblems(def)...)
	}

	if len(problems) > 0 {
		return fmt.Errorf("%w: %s", ErrInvalid, strings.Join(problems, "; "))
	}
	return nil
}

func TeamProblems(members []api.TeamMember) []string {
	problems := []string{}
	if len(members) < 2 {
		problems = append(problems, "a team needs at least two agents")
	}
	seen := map[string]bool{}
	for i, mem := range members {
		where := fmt.Sprintf("agent %d (%q)", i+1, mem.Name)
		p := nameProblem(mem.Name)
		if p != "" {
			problems = append(problems, fmt.Sprintf("%s: name %s", where, p))
		}
		if seen[mem.Name] {
			problems = append(problems, fmt.Sprintf("%s: another agent has the same name", where))
		}
		seen[mem.Name] = true
		if strings.TrimSpace(mem.Model) == "" {
			problems = append(problems, fmt.Sprintf("%s: model is required", where))
		}
		if strings.TrimSpace(mem.Prompt) == "" {
			problems = append(problems, fmt.Sprintf("%s: prompt is empty", where))
		}
	}
	return problems
}

func graphProblems(def api.WorkflowDefinition) []string {
	problems := []string{}
	finals := []string{}
	for _, s := range def.Stages {
		if len(Next(def, s.Name)) == 0 {
			finals = append(finals, fmt.Sprintf("%q", s.Name))
		}
	}
	if len(finals) != 1 {
		problems = append(problems, fmt.Sprintf("it must have exactly one final stage (one no other stage comes after), but has %d: %s", len(finals), strings.Join(finals, ", ")))
	}

	state := map[string]int{}
	var visit func(name string, path []string)
	visit = func(name string, path []string) {
		switch state[name] {
		case 1:
			problems = append(problems, fmt.Sprintf("its stages form a cycle: %s", strings.Join(append(path, fmt.Sprintf("%q", name)), " after ")))
			return
		case 2:
			return
		}
		state[name] = 1
		for _, a := range After(def, name) {
			visit(a, append(path, fmt.Sprintf("%q", name)))
		}
		state[name] = 2
	}
	for _, s := range def.Stages {
		visit(s.Name, nil)
	}
	return problems
}

func Index(def api.WorkflowDefinition, name string) int {
	for i, s := range def.Stages {
		if s.Name == name {
			return i
		}
	}
	return -1
}

func Get(def api.WorkflowDefinition, name string) api.Stage {
	return def.Stages[Index(def, name)]
}

func After(def api.WorkflowDefinition, name string) []string {
	i := Index(def, name)
	s := def.Stages[i]
	if s.After != nil {
		return *s.After
	}
	if i == 0 {
		return []string{}
	}
	return []string{def.Stages[i-1].Name}
}

func Next(def api.WorkflowDefinition, name string) []string {
	next := []string{}
	for _, s := range def.Stages {
		for _, a := range After(def, s.Name) {
			if a == name {
				next = append(next, s.Name)
			}
		}
	}
	return next
}

func Starts(def api.WorkflowDefinition) []string {
	starts := []string{}
	for _, s := range def.Stages {
		if len(After(def, s.Name)) == 0 {
			starts = append(starts, s.Name)
		}
	}
	return starts
}

func FromMain(def api.WorkflowDefinition, name string) bool {
	s := Get(def, name)
	return len(After(def, name)) > 1 || (s.Base != nil && *s.Base == api.StageBaseMain)
}