Run workflow stages with finish_stage handoffs and lead signals

internal/agents/manager.go

@@ -74,6 +74,8 @@ type live struct {
 	projectID   string
 	owner       string
 	reviewID    string
+	stage       int
+	stages      int
 	stopping    bool
 	interrupted bool
 	asked       bool
@@ -126,6 +128,18 @@ func (m *Manager) Create(ctx context.Context, task api.Task, creator api.User, p
 }
 
 func (m *Manager) create(ctx context.Context, task api.Task, creator api.User, prompt, model, actor string, images []string) (api.AgentSession, error) {
+	existing, err := m.store.ListSessions(ctx, task.Id)
+	if err != nil {
+		return api.AgentSession{}, err
+	}
+	stage := 0
+	if len(existing) > 0 && existing[len(existing)-1].WorkflowStage != nil {
+		stage = *existing[len(existing)-1].WorkflowStage
+	}
+	return m.createAt(ctx, task, creator, prompt, model, actor, images, stage)
+}
+
+func (m *Manager) createAt(ctx context.Context, task api.Task, creator api.User, prompt, model, actor string, images []string, stage int) (api.AgentSession, error) {
 	if creator.GitName == "" || creator.GitEmail == "" {
 		return api.AgentSession{}, ErrNoGitIdentity
 	}
@@ -157,6 +171,9 @@ func (m *Manager) create(ctx context.Context, task api.Task, creator api.User, p
 		CreatedBy:       creator.Login,
 		CreatedAt:       time.Now().UTC(),
 	}
+	if stage > 0 {
+		sess.WorkflowStage = &stage
+	}
 
 	err = m.store.CreateSession(ctx, sess, task.ProjectId)
 	if err != nil {
@@ -265,12 +282,20 @@ func (m *Manager) Recover(ctx context.Context) error {
 		return err
 	}
 
+	staged, err := m.recoverStages(ctx)
+	if err != nil {
+		return err
+	}
+
 	sessions, err := m.store.SessionsInState(ctx, api.AgentSessionStateRunning)
 	if err != nil {
 		return err
 	}
 
 	for _, sess := range sessions {
+		if staged[sess.Id] {
+			continue
+		}
 		task, err := m.taskFor(ctx, sess)
 		if err != nil {
 			return err
@@ -280,9 +305,13 @@ func (m *Manager) Recover(ctx context.Context) error {
 		if err != nil {
 			slog.Error("resume session", "session", sess.Id, "err", err)
 			m.logUndelivered(ctx, sess.Id, "resume failed")
-			err = m.store.SetSessionState(ctx, sess.Id, api.AgentSessionStateIdle, "byre", map[string]any{"error": err.Error()})
-			if err != nil {
-				return err
+			derr := m.store.DropStageOutput(ctx, sess.Id)
+			if derr != nil {
+				return derr
+			}
+			serr := m.store.SetSessionState(ctx, sess.Id, api.AgentSessionStateIdle, "byre", map[string]any{"error": err.Error()})
+			if serr != nil {
+				return serr
 			}
 			if sess.Role == api.Implementer {
 				m.abandonMerge(task.Id, task.ProjectId)
@@ -518,6 +547,16 @@ func (m *Manager) ensureLive(ctx context.Context, sess api.AgentSession, task ap
 		return nil, err
 	}
 	tools := []mcp.Tool{m.askTool(l), m.cancelQuestionTool(l), m.createTaskTool(l), m.setMergeTitleTool(l), m.stopMergeTool(l), m.exposePortTool(l)}
+	if sess.Role == api.Implementer && sess.WorkflowStage != nil {
+		frozen, err := m.store.TaskWorkflow(ctx, task.Id)
+		if err != nil {
+			return nil, err
+		}
+		l.stage = *sess.WorkflowStage
+		l.stages = len(frozen.Definition.Stages)
+		g += stageGuidance(frozen, l.stage)
+		tools = append(tools, m.finishStageTool(l))
+	}
 	env := portEnv(task)
 	stdio := []mcp.Stdio{}
 	if sess.Role == api.Reviewer {
@@ -588,6 +627,12 @@ func (m *Manager) pump(l *live) {
 	}
 
 	m.logUndelivered(ctx, l.id, "agent process exited mid-turn")
+	if l.stage > 0 {
+		err = m.store.DropStageOutput(ctx, l.id)
+		if err != nil {
+			slog.Error("drop stage output", "session", l.id, "err", err)
+		}
+	}
 	err = m.store.SetSessionState(ctx, l.id, api.AgentSessionStateIdle, agentActor, map[string]any{"error": "agent process exited mid-turn"})
 	if err != nil {
 		slog.Error("record session exit", "session", l.id, "err", err)
@@ -788,6 +833,18 @@ func (m *Manager) turnEnded(ctx context.Context, l *live, res endResult) error {
 		m.fileError(l, failure)
 	}
 
+	stageFinished := false
+	switch {
+	case l.stage == 0:
+	case finished && !handoff && mergingBy == "":
+		stageFinished, err = m.store.FinishStage(ctx, l.id)
+	default:
+		err = m.store.DropStageOutput(ctx, l.id)
+	}
+	if err != nil {
+		return err
+	}
+
 	err = m.store.SetSessionState(ctx, l.id, api.AgentSessionStateIdle, agentActor, nil)
 	if err != nil {
 		return err
@@ -809,6 +866,12 @@ func (m *Manager) turnEnded(ctx context.Context, l *live, res endResult) error {
 	default:
 		go m.nextMerge(l.projectID)
 	}
+	if stageFinished {
+		go m.advanceStage(l)
+		if l.stage < l.stages {
+			return nil
+		}
+	}
 	err = m.Deliver(ctx, l.id)
 	if err != nil {
 		slog.Error("deliver inbox responses", "session", l.id, "err", err)