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.
func OrchestratorWorkflow(ctx workflow.Context, in OrchestratorInput) (OrchestratorOutput, error) {
	output := OrchestratorOutput{
		MilestoneComplete: false,
		Done:              false,
		LastError:         "",
	}

	// Step 1: Clone the repository
	activityOptions := workflow.ActivityOptions{
		StartToCloseTimeout: 10 * time.Minute,
		ScheduleToCloseTimeout: 15 * time.Minute,
	}
	ctxWithOptions := workflow.WithActivityOptions(ctx, activityOptions)
	
	cloneErr := workflow.ExecuteActivity(
		ctxWithOptions,
		"CloneRepoActivity",
		map[string]interface{}{
			"RemoteURL":      in.RemoteURL,
			"TargetRepoPath": in.TargetRepoPath,
		},
	).Get(ctx, nil)

	if cloneErr != nil {
		output.LastError = fmt.Sprintf("Clone failed: %v", cloneErr)
		return output, nil
	}

	// Step 2: Read tasks from board.md
	tasksToRun, err := readTasksFromBoard(in.TargetRepoPath)
	if err != nil {
		output.LastError = fmt.Sprintf("Failed to read tasks: %v", err)
		return output, nil
	}

	if len(tasksToRun) == 0 {
		output.LastError = "No tasks found in board.md"
		return output, nil
	}

	// Step 3: Process each task
	completedTasks := 0
	for _, task := range tasksToRun {
		taskID := task["id"].(string)
		taskDesc := task["description"].(string)

			// taskID will be used for worktree and branch

		// Add worktree
		var worktreePath string
		wtErr := workflow.ExecuteActivity(
			ctxWithOptions,
			"GitWorktreeAddActivity",
			map[string]interface{}{
				"RepoPath": in.TargetRepoPath,
				"TaskID":   taskID,
			},
		).Get(ctx, &worktreePath)

		if wtErr != nil {
			continue // Skip this task on error
		}

		// Call implementer to generate code (longer timeout for LLM calls)
		implOptions := workflow.ActivityOptions{
			StartToCloseTimeout: 30 * time.Minute,
			ScheduleToCloseTimeout: 35 * time.Minute,
		}
		implCtx := workflow.WithActivityOptions(ctx, implOptions)
		
		var implOutput map[string]interface{}
		implErr := workflow.ExecuteActivity(
			implCtx,
			"ImplementerActivity",
			map[string]interface{}{
				"TaskID":       taskID,
				"Description":  taskDesc,
				"WorktreePath": worktreePath,
				"Prompt": PromptSpec{
					TemplateRef: "implementer/default.tmpl",
					Model: ModelSpec{
						ModelID: in.Config.RolePrompts["implementer"].Model.ModelID,
					},
				},
			},
		).Get(ctx, &implOutput)

		if implErr != nil {
			continue
		}

		// Commit changes
		commitErr := workflow.ExecuteActivity(
			ctxWithOptions,
			"GitCommitActivity",
			map[string]interface{}{
				"WorktreePath": worktreePath,
				"Message":      fmt.Sprintf("%s: implementation", taskID),
			},
		).Get(ctx, nil)

		if commitErr == nil {
			completedTasks++
		}
	}

	// Step 4: Push to remote
	pushErr := workflow.ExecuteActivity(
		ctxWithOptions,
		"GitPushActivity",
		map[string]interface{}{
			"RepoPath": in.TargetRepoPath,
		},
	).Get(ctx, nil)

	if pushErr != nil {
		output.LastError = fmt.Sprintf("Push failed: %v", pushErr)
		return output, nil
	}

	// Step 5: Squash merge all task branches
	branches := make([]string, len(tasksToRun))
	for i, task := range tasksToRun {
		branches[i] = fmt.Sprintf("task/%s", task["id"].(string))
	}

	mergeErr := workflow.ExecuteActivity(
		ctxWithOptions,
		"GitSquashMergeActivity",
		map[string]interface{}{
			"RepoPath": in.TargetRepoPath,
			"Branches": branches,
			"Message":  fmt.Sprintf("%s: squash merge all tasks", in.Milestone),
		},
	).Get(ctx, nil)

	if mergeErr != nil {
		output.LastError = fmt.Sprintf("Merge failed: %v", mergeErr)
		return output, nil
	}

	// Success!
	output.MilestoneComplete = true
	output.Done = true
	output.LastError = fmt.Sprintf("Completed %d tasks successfully", completedTasks)

	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
}

// isPiStreamTimeout checks if an error is a 504 stream timeout from Pi command
func isPiStreamTimeout(err error) bool {
	if err == nil {
		return false
	}
	return strings.Contains(err.Error(), "PiStreamTimeout")
}