poimen-workflows/statemachine/orchestrator_recovery.go

package statemachine

import (
	"fmt"
	"time"

	"go.temporal.io/sdk/workflow"
	"github.com/rockliang/poimen/workflows/internal/recovery"
	"github.com/rockliang/poimen/workflows/internal/logging"
)

// OrchestratorWorkflowWithRecovery orchestrates multi-agent work with recovery capabilities
// It differs from the basic orchestrator by:
// 1. Using retry policies for all activities
// 2. Tracking workflow state via checkpoints
// 3. Using deadletter handling for permanently failed activities
// 4. Resuming from checkpoints after crashes
func OrchestratorWorkflowWithRecovery(ctx workflow.Context, in OrchestratorInput) (OrchestratorOutput, error) {
	output := OrchestratorOutput{
		MilestoneComplete: false,
		Done:              false,
		LastError:         "",
	}

	logger := logging.GetLogger()

	// Create activity options with retry policy
	retryPolicy := recovery.ActivityRetryPolicy()
	baseActivityOptions := workflow.ActivityOptions{
		StartToCloseTimeout:    10 * time.Minute,
		ScheduleToCloseTimeout: 15 * time.Minute,
		RetryPolicy:            retryPolicy.ToTemporalRetryPolicy(),
	}

	ctxWithOptions := workflow.WithActivityOptions(ctx, baseActivityOptions)

	// Step 1: Clone the repository with retry
	logger.Info("starting orchestrator workflow",
		logging.String("milestone", in.Milestone),
		logging.String("repo", in.TargetRepoPath))

	cloneErr := workflow.ExecuteActivity(
		ctxWithOptions,
		"CloneRepoActivity",
		map[string]interface{}{
			"RemoteURL":      in.RemoteURL,
			"TargetRepoPath": in.TargetRepoPath,
		},
	).Get(ctx, nil)

	if cloneErr != nil {
		logger.Error("clone failed",
			logging.Err(cloneErr),
			logging.String("repo", in.TargetRepoPath))
		output.LastError = fmt.Sprintf("Clone failed: %v", cloneErr)
		return output, nil
	}

	logger.Info("repository cloned",
		logging.String("repo", in.TargetRepoPath))

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

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

	logger.Info("tasks loaded",
		logging.Int("count", len(tasksToRun)))

	// Step 3: Process each task with recovery tracking
	completedTasks := 0
	failedTasks := []string{}

	// LLM activity uses longer timeout and more retries
	llmRetryPolicy := recovery.LLMActivityRetryPolicy()
	implOptions := workflow.ActivityOptions{
		StartToCloseTimeout:    30 * time.Minute,
		ScheduleToCloseTimeout: 35 * time.Minute,
		RetryPolicy:            llmRetryPolicy.ToTemporalRetryPolicy(),
	}
	implCtx := workflow.WithActivityOptions(ctx, implOptions)

	for taskIdx, task := range tasksToRun {
		taskID := task["id"].(string)
		taskDesc := task["description"].(string)

		logger.Info("processing task",
			logging.String("taskID", taskID),
			logging.Int("index", taskIdx+1),
			logging.Int("total", len(tasksToRun)))

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

		if wtErr != nil {
			logger.Error("worktree creation failed",
				logging.String("taskID", taskID),
				logging.Err(wtErr))
			failedTasks = append(failedTasks, taskID)
			continue
		}

		logger.Info("worktree created",
			logging.String("taskID", taskID),
			logging.String("path", worktreePath))

		// Call implementer
		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 {
			logger.Error("implementation failed",
				logging.String("taskID", taskID),
				logging.Err(implErr))
			failedTasks = append(failedTasks, taskID)
			continue
		}

		logger.Info("implementation succeeded",
			logging.String("taskID", taskID))

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

		if commitErr != nil {
			logger.Error("commit failed",
				logging.String("taskID", taskID),
				logging.Err(commitErr))
			failedTasks = append(failedTasks, taskID)
			continue
		}

		completedTasks++
		logger.Info("task completed",
			logging.String("taskID", taskID),
			logging.Int("completedCount", completedTasks))
	}

	// Step 4: Push to remote
	logger.Info("pushing changes to remote",
		logging.String("repo", in.TargetRepoPath))

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

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

	logger.Info("changes pushed to remote")

	// 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))
	}

	logger.Info("merging task branches",
		logging.Int("branchCount", len(branches)))

	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 {
		logger.Error("merge failed",
			logging.Err(mergeErr))
		output.LastError = fmt.Sprintf("Merge failed: %v", mergeErr)
		return output, nil
	}

	logger.Info("workflow completed",
		logging.Int("completed", completedTasks),
		logging.Int("failed", len(failedTasks)))

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

	return output, nil
}