poimen-workflows/statemachine/orchestrator.go

package statemachine

import (
	"fmt"
	"io/ioutil"
	"path/filepath"
	"strings"
	"time"

	"go.temporal.io/sdk/workflow"
)

// OrchestratorWorkflow orchestrates multi-agent work on a target repository.
// Reconciliation loop: Plan → Dispatch → Review → Update → Repeat
func OrchestratorWorkflow(ctx workflow.Context, in OrchestratorInput) (OrchestratorOutput, error) {
	logger := workflow.GetLogger(ctx)
	output := OrchestratorOutput{
		MilestoneComplete: false,
		Done:              false,
		LastError:         "",
	}

	// Clone repo once at start
	activityOpts := workflow.ActivityOptions{
		StartToCloseTimeout:    10 * time.Minute,
		ScheduleToCloseTimeout: 15 * time.Minute,
	}
	ctxWithOpts := workflow.WithActivityOptions(ctx, activityOpts)

	if err := workflow.ExecuteActivity(ctxWithOpts, "CloneRepoActivity", map[string]interface{}{
		"RemoteURL":      in.RemoteURL,
		"TargetRepoPath": in.TargetRepoPath,
	}).Get(ctx, nil); err != nil {
		output.LastError = fmt.Sprintf("Clone failed: %v", err)
		return output, nil
	}

	// Prepare skills once
	if len(in.Config.Skills) > 0 {
		if err := workflow.ExecuteActivity(ctxWithOpts, "PrepareSkillsActivity", map[string]interface{}{
			"Skills":        in.Config.Skills,
			"StreamTimeout": 30 * time.Second,
			"Provider":      in.PiProvider,
		}).Get(ctx, nil); err != nil {
			logger.Info("skill preparation failed (continuing anyway)", "error", err)
		}
	}

	// Main reconciliation loop
	var paused bool
	var completedTasks int
	cycleCount := 0
	completedBranches := []string{}

	for cycleCount < in.MaxCyclesBeforeCAN {
		cycleCount++
		logger.Info("orchestrator cycle", "cycle", cycleCount)

		// Read board state
		bordState, _ := readTasksFromBoard(in.TargetRepoPath)

		// Call PlanningActivity to decide what to dispatch
		implOpts := workflow.ActivityOptions{
			StartToCloseTimeout:    30 * time.Minute,
			ScheduleToCloseTimeout: 35 * time.Minute,
		}
		implCtx := workflow.WithActivityOptions(ctx, implOpts)

		var planOutput map[string]interface{}
		if err := workflow.ExecuteActivity(implCtx, "PlanningActivity", map[string]interface{}{
			"Config":     in.Config,
			"BoardState": fmt.Sprintf("%v", bordState),
			"RepoPath":   in.TargetRepoPath,
			"Milestone":  in.Milestone,
		}).Get(ctx, &planOutput); err != nil {
			logger.Info("planning failed", "error", err)
			break
		}

		// Extract tasks to dispatch
		var tasksToDispatch []interface{}
		if tasks, ok := planOutput["TasksToDispatch"]; ok {
			tasksToDispatch = tasks.([]interface{})
		}

		if len(tasksToDispatch) == 0 {
			logger.Info("no tasks to dispatch, milestone complete")
			output.MilestoneComplete = true
			break
		}

		// Fan-out: Start TaskUnit workflows for each task
		logger.Info("dispatching task units", "count", len(tasksToDispatch))
		var childWFs []workflow.Future

		for _, taskIDRaw := range tasksToDispatch {
			taskID := taskIDRaw.(string)

			if paused {
				logger.Info("skipping dispatch due to pause signal", "task", taskID)
				continue
			}

			// Child workflow: TaskUnitWorkflow
			childOpts := workflow.ChildWorkflowOptions{
				WorkflowID: fmt.Sprintf("%s-%s-cycle%d", in.Milestone, taskID, cycleCount),
			}
			childCtx := workflow.WithChildOptions(ctx, childOpts)

			taskUnitInput := TaskUnitInput{
				TaskID:         taskID,
				RemoteURL:      in.RemoteURL,
				TargetRepoPath: in.TargetRepoPath,
				Milestone:      in.Milestone,
				Config:         in.Config,
				DryRun:         in.DryRun,
			}

			future := workflow.ExecuteChildWorkflow(childCtx, TaskUnitWorkflow, taskUnitInput)
			childWFs = append(childWFs, future)
		}

		// Fan-in: Wait for all TaskUnits to complete
		logger.Info("waiting for task units", "count", len(childWFs))
		for _, future := range childWFs {
			var taskOutput TaskUnitOutput
			if err := future.Get(ctx, &taskOutput); err != nil {
				logger.Info("task unit failed", "task", taskOutput.TaskID, "error", err)
			} else if taskOutput.Status == "success" {
				completedTasks++
				completedBranches = append(completedBranches, taskOutput.Branch)
				logger.Info("task unit succeeded", "task", taskOutput.TaskID)
			}
		}

		// Board update: Call PlanningActivity again to update board + commit
		if err := workflow.ExecuteActivity(implCtx, "PlanningActivity", map[string]interface{}{
			"Config":     in.Config,
			"BoardState": fmt.Sprintf("%v", bordState),
			"RepoPath":   in.TargetRepoPath,
			"Milestone":  in.Milestone,
		}).Get(ctx, nil); err != nil {
			logger.Info("board update failed", "error", err)
		}

		// Push completed branches
		if err := workflow.ExecuteActivity(ctxWithOpts, "GitPushActivity", map[string]interface{}{
			"RepoPath": in.TargetRepoPath,
		}).Get(ctx, nil); err != nil {
			logger.Info("push failed", "error", err)
		}

		// Squash merge completed branches when milestone ready
		if output.MilestoneComplete && len(completedBranches) > 0 {
			if err := workflow.ExecuteActivity(ctxWithOpts, "GitSquashMergeActivity", map[string]interface{}{
				"RepoPath": in.TargetRepoPath,
				"Branches": completedBranches,
				"Message":  fmt.Sprintf("%s: squash merge completed tasks", in.Milestone),
			}).Get(ctx, nil); err != nil {
				logger.Info("squash merge failed", "error", err)
			}
		}
	}

	// Use continue-as-new if hit cycle cap
	if cycleCount >= in.MaxCyclesBeforeCAN {
		logger.Info("cycle cap reached, continuing as new", "cycles", cycleCount)
		nextInput := in
		nextInput.CycleCount = cycleCount
		return output, workflow.NewContinueAsNewError(ctx, OrchestratorWorkflow, nextInput)
	}

	// Success
	output.Done = true
	output.LastError = fmt.Sprintf("Completed %d tasks in %d cycles", completedTasks, cycleCount)

	return output, nil
}

// readTasksFromBoard reads tasks from tasks/board.md
func readTasksFromBoard(repoPath string) ([]map[string]interface{}, error) {
	boardPath := filepath.Join(repoPath, "tasks", "board.md")

	content, err := ioutil.ReadFile(boardPath)
	if err != nil {
		return nil, err
	}

	lines := strings.Split(string(content), "\n")
	var tasks []map[string]interface{}

	for _, line := range lines {
		// Parse markdown table rows: | T1 | Description | [ ] | ...
		if strings.HasPrefix(strings.TrimSpace(line), "|") && !strings.Contains(line, "---|") && !strings.Contains(line, "ID") {
			parts := strings.Split(line, "|")
			if len(parts) >= 4 {
				id := strings.TrimSpace(parts[1])
				desc := strings.TrimSpace(parts[2])

				if id != "" && desc != "" {
					tasks = append(tasks, map[string]interface{}{
						"id":          id,
						"description": desc,
					})
				}
			}
		}
	}

	return tasks, nil
}