poimen-workflows/statemachine/taskunit.go

package statemachine

import (
	"fmt"
	"time"

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

// TaskUnitWorkflow executes a single task with retry loops, timeout escalation, and lessons injection.
func TaskUnitWorkflow(ctx workflow.Context, in TaskUnitInput) (TaskUnitOutput, error) {
	// Initialize output
	output := TaskUnitOutput{
		TaskID:  in.TaskID,
		Verdict: "fail",
	}

	// 1. Add worktree for isolated work
	var worktreePath string
	wtErr := workflow.ExecuteActivity(
		ctx,
		"GitWorktreeAddActivity",
		map[string]interface{}{
			"RepoPath": in.TargetRepoPath,
			"TaskID":   in.TaskID,
		},
	).Get(ctx, &worktreePath)
	if wtErr != nil {
		output.Critique = fmt.Sprintf("Failed to create worktree: %v", wtErr)
		return output, nil
	}

	// 2. Retry loop with separate timeout and judge attempt tracking
	timeoutAttempt := 1
	for judgeAttempt := 1; judgeAttempt <= in.MaxJudgeRetries; judgeAttempt++ {
		// Calculate timeouts for this attempt
		baseTimeout := in.BaseTimeout * time.Duration(timeoutAttempt)
		heartbeatTimeout := baseTimeout / 4

		// Prepare activity options with escalating timeout
		ao := workflow.ActivityOptions{
			ScheduleToCloseTimeout: baseTimeout,
			StartToCloseTimeout:    baseTimeout,
			HeartbeatTimeout:       heartbeatTimeout,
			RetryPolicy: &temporal.RetryPolicy{
				MaximumAttempts: 1, // We manage retries in this loop
			},
		}
		ctxWithOptions := workflow.WithActivityOptions(ctx, ao)

		// Call implementer activity
		var implOutput map[string]interface{}
		implErr := workflow.ExecuteActivity(
			ctxWithOptions,
			"ImplementerActivity",
			map[string]interface{}{
				"TaskID":       in.TaskID,
				"WorktreePath": worktreePath,
				"Prompt":       in.ImplementerSpec,
			},
		).Get(ctx, &implOutput)

		// Check if it's a timeout error
		if implErr != nil && isStartToCloseTimeout(implErr) {
			// Timeout: escalate and retry without consuming judge attempt
			timeoutAttempt++
			judgeAttempt-- // Don't consume a judge retry on timeout
			continue
		}
		if implErr != nil {
			output.Critique = fmt.Sprintf("Implementer failed: %v", implErr)
			return output, nil
		}

		// Call judge activity
		judgeTimeout := time.Minute * 5
		judgeCtx := workflow.WithActivityOptions(ctx, workflow.ActivityOptions{
			ScheduleToCloseTimeout: judgeTimeout,
			StartToCloseTimeout:    judgeTimeout,
		})
		var judgeOutput map[string]interface{}
		judgeErr := workflow.ExecuteActivity(
			judgeCtx,
			"JudgeActivity",
			map[string]interface{}{
				"TaskID":       in.TaskID,
				"WorktreePath": worktreePath,
				"Prompt":       in.JudgeSpec,
			},
		).Get(ctx, &judgeOutput)
		if judgeErr != nil {
			output.Critique = fmt.Sprintf("Judge error: %v", judgeErr)
			return output, nil
		}

		// Check judge verdict
		verdict := ""
		if judgeOutput != nil {
			if v, ok := judgeOutput["Verdict"].(string); ok {
				verdict = v
			}
		}

		if verdict == "pass" {
			// Commit in worktree
			commitErr := workflow.ExecuteActivity(
				ctx,
				"GitCommitActivity",
				map[string]interface{}{
					"WorktreePath": worktreePath,
					"Message":      fmt.Sprintf("%s: implementation", in.TaskID),
				},
			).Get(ctx, nil)
			if commitErr != nil {
				output.Critique = fmt.Sprintf("Commit failed: %v", commitErr)
				return output, nil
			}

			// Success!
			output.Verdict = "pass"
			output.Branch = "task/" + in.TaskID
			return output, nil
		}

		// Judge failed: update lessons and retry
		critique := ""
		if judgeOutput != nil {
			if c, ok := judgeOutput["Critique"].(string); ok {
				critique = c
			}
		}

		updateErr := workflow.ExecuteActivity(
			ctx,
			"UpdateLessonsActivity",
			map[string]interface{}{
				"TargetRepoPath": in.TargetRepoPath,
				"TaskID":         in.TaskID,
				"Attempt":        judgeAttempt,
				"Critique":       critique,
			},
		).Get(ctx, nil)
		if updateErr != nil {
			output.Critique = fmt.Sprintf("Failed to update lessons: %v", updateErr)
			return output, nil
		}

		// Continue to next judge attempt with lessons injected
	}

	// Retries exhausted
	output.Verdict = "fail"
	output.Critique = fmt.Sprintf("Exhausted %d judge retries", in.MaxJudgeRetries)
	return output, nil
}

// isStartToCloseTimeout checks if an error is a StartToCloseTimeout error
func isStartToCloseTimeout(err error) bool {
	if err == nil {
		return false
	}
	return fmt.Sprint(err) == "context deadline exceeded"
}