Run workflow stages with finish_stage handoffs and lead signals

internal/store/workflow.go

@@ -13,6 +13,7 @@ import (
 type FrozenWorkflow struct {
 	Name       string
 	Definition api.WorkflowDefinition
+	Brief      string
 }
 
 func (s *Store) ListWorkflows(ctx context.Context) ([]api.WorkflowSummary, error) {
@@ -89,19 +90,104 @@ func (s *Store) FreezeWorkflow(ctx context.Context, name string) (FrozenWorkflow
 }
 
 func (s *Store) TaskWorkflow(ctx context.Context, taskID string) (FrozenWorkflow, error) {
-	name := ""
+	f := FrozenWorkflow{}
 	def := ""
-	err := s.db.QueryRowContext(ctx, `SELECT workflow_name, workflow_definition FROM tasks WHERE id = ?`, taskID).Scan(&name, &def)
+	err := s.db.QueryRowContext(ctx, `SELECT workflow_name, workflow_definition, workflow_brief FROM tasks WHERE id = ?`, taskID).Scan(&f.Name, &def, &f.Brief)
 	if errors.Is(err, sql.ErrNoRows) {
 		return FrozenWorkflow{}, fmt.Errorf("%w: no task %s", ErrNotFound, taskID)
 	}
 	if err != nil {
 		return FrozenWorkflow{}, err
 	}
-	if name == "" {
+	if f.Name == "" {
 		return FrozenWorkflow{}, fmt.Errorf("%w: task %s has no workflow", ErrNotFound, taskID)
 	}
-	f := FrozenWorkflow{Name: name}
 	err = json.Unmarshal([]byte(def), &f.Definition)
 	return f, err
 }
+
+func (s *Store) SetStageOutput(ctx context.Context, sessionID, output string) error {
+	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
+		_, err := tx.ExecContext(ctx, `UPDATE agent_sessions SET stage_output = ? WHERE id = ?`, output, sessionID)
+		return err
+	})
+}
+
+func (s *Store) FinishStage(ctx context.Context, sessionID string) (bool, error) {
+	finished := false
+	err := s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
+		res, err := tx.ExecContext(ctx, `UPDATE agent_sessions SET stage_finished = 1 WHERE id = ? AND stage_output != ''`, sessionID)
+		if err != nil {
+			return err
+		}
+		n, err := res.RowsAffected()
+		finished = n > 0
+		return err
+	})
+	return finished, err
+}
+
+func (s *Store) DropStageOutput(ctx context.Context, sessionID string) error {
+	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
+		_, err := tx.ExecContext(ctx, `UPDATE agent_sessions SET stage_output = '' WHERE id = ? AND stage_finished = 0`, sessionID)
+		return err
+	})
+}
+
+func (s *Store) StageOutput(ctx context.Context, sessionID string) (string, error) {
+	output := ""
+	err := s.db.QueryRowContext(ctx, `SELECT stage_output FROM agent_sessions WHERE id = ?`, sessionID).Scan(&output)
+	return output, err
+}
+
+func (s *Store) FinishedStages(ctx context.Context) ([]string, error) {
+	rows, err := s.db.QueryContext(ctx, `SELECT id FROM agent_sessions WHERE stage_finished = 1 ORDER BY created_at`)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	ids := []string{}
+	for rows.Next() {
+		id := ""
+		err = rows.Scan(&id)
+		if err != nil {
+			return nil, err
+		}
+		ids = append(ids, id)
+	}
+	return ids, rows.Err()
+}
+
+func (s *Store) ClearStage(ctx context.Context, sessionID string) error {
+	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
+		return clearStage(ctx, tx, sessionID)
+	})
+}
+
+func clearStage(ctx context.Context, tx *sql.Tx, sessionID string) error {
+	_, err := tx.ExecContext(ctx, `UPDATE agent_sessions SET stage_output = '', stage_finished = 0 WHERE id = ?`, sessionID)
+	return err
+}
+
+func (s *Store) FinishWorkflow(ctx context.Context, task api.Task, sessionID, workflowName, output string) error {
+	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
+		err := clearStage(ctx, tx, sessionID)
+		if err != nil {
+			return err
+		}
+		return emit(api.Event{
+			Kind:      "task.workflow_finished",
+			ProjectId: &task.ProjectId,
+			TaskId:    &task.Id,
+			Actor:     "byre",
+			Payload:   map[string]any{"id": task.Id, "workflow": workflowName, "output": output},
+		})
+	})
+}
+
+func (s *Store) WorkflowFinished(ctx context.Context, taskID string) (bool, error) {
+	finished := false
+	err := s.db.QueryRowContext(ctx, `SELECT EXISTS (SELECT 1 FROM events WHERE task_id = ? AND kind = 'task.workflow_finished')`, taskID).Scan(&finished)
+	return finished, err
+}