Run workflow stages with finish_stage handoffs and lead signals

2026-10-15 22:11:25 -07:00
parent 0c698cd3c4
commit 2eccd23a98
15 changed files with 665 additions and 256 deletions
@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"log/slog"
 	"strings"

 	"g.fc.run/misc/byre/internal/api"
@@ -48,7 +49,9 @@ func (m *Manager) TaskWorkflow(ctx context.Context, taskID string) (api.TaskWork
 	return api.TaskWorkflow{Name: frozen.Name, Definition: string(file)}, nil
 }

-func (m *Manager) freeze(ctx context.Context, name string) (*store.FrozenWorkflow, error) {
+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
 	}
@@ -56,11 +59,232 @@ func (m *Manager) freeze(ctx context.Context, name string) (*store.FrozenWorkflo
 	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), " +
-	"vendor (claude or codex: gpt-* models run on codex, every other model on claude), model and prompt (the stage agent's instructions). " +
+	"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 {