poimen-workflows/statemachine/routing_workflow.go

package statemachine

import (
	"fmt"
	"time"

	"github.com/rockliang/poimen/workflows/internal/routing"
	"go.temporal.io/sdk/log"
	"go.temporal.io/sdk/temporal"
	"go.temporal.io/sdk/workflow"
)

// RoutingWorkflowInput is input for the routing workflow
type RoutingWorkflowInput struct {
	Spec *routing.WorkflowSpec `json:"spec"`
}

// RoutingWorkflowOutput is output from the routing workflow
type RoutingWorkflowOutput struct {
	Status      string                 `json:"status"` // "COMPLETED", "FAILED"
	FinalOutput interface{}            `json:"finalOutput,omitempty"`
	StepResults map[string]interface{} `json:"stepResults"`
	Error       string                 `json:"error,omitempty"`
}

// RoutingWorkflow executes any WorkflowSpec generated by llm-router
func RoutingWorkflow(ctx workflow.Context, input RoutingWorkflowInput) (RoutingWorkflowOutput, error) {
	logger := workflow.GetLogger(ctx)
	
	output := RoutingWorkflowOutput{
		Status:      "FAILED",
		StepResults: make(map[string]interface{}),
	}

	if input.Spec == nil || len(input.Spec.States) == 0 {
		output.Error = "empty workflow spec"
		return output, nil
	}

	logger.Info("RoutingWorkflow started", "name", input.Spec.Name, "stateCount", len(input.Spec.States))

	// Build execution context
	execCtx := &routing.ExecutionContext{
		Input:       input.Spec.Input,
		StepResults: make(map[string]interface{}),
	}

	// Build state index for fast lookup
	stateIndex := make(map[string]*routing.State)
	for i := range input.Spec.States {
		stateIndex[input.Spec.States[i].Name] = &input.Spec.States[i]
	}

	// Find first state (first in array)
	currentStateName := input.Spec.States[0].Name

	// State machine loop
	for {
		state, ok := stateIndex[currentStateName]
		if !ok {
			output.Error = fmt.Sprintf("state not found: %s", currentStateName)
			return output, nil
		}

		logger.Info("executing state", "state", currentStateName, "type", state.Type)

		switch state.Type {
		case routing.StateTypeTask:
			result, nextState, err := executeTaskState(ctx, state, execCtx, logger)
			if err != nil {
				// Check for catch clause
				if nextState != "" {
					currentStateName = nextState
					continue
				}
				output.Error = fmt.Sprintf("state %s failed: %v", currentStateName, err)
				return output, nil
			}
			// Wrap result in output key for JSONPath compatibility (e.g., ${Clone.output.path})
			wrappedResult := map[string]interface{}{"output": result}
			execCtx.StepResults[state.Name] = wrappedResult
			output.StepResults[state.Name] = result // Keep original for output

			if state.End {
				output.Status = "COMPLETED"
				output.FinalOutput = result
				logger.Info("RoutingWorkflow completed", "name", input.Spec.Name)
				return output, nil
			}
			currentStateName = state.Next

		case routing.StateTypePass:
			execCtx.StepResults[state.Name] = state.Result
			output.StepResults[state.Name] = state.Result

			if state.End {
				output.Status = "COMPLETED"
				output.FinalOutput = state.Result
				return output, nil
			}
			currentStateName = state.Next

		case routing.StateTypeFail:
			output.Error = fmt.Sprintf("%s: %s", state.Error, state.Cause)
			logger.Error("RoutingWorkflow failed at Fail state", "state", currentStateName, "error", state.Error)
			return output, nil

		default:
			output.Error = fmt.Sprintf("unknown state type: %s", state.Type)
			return output, nil
		}

		// Safety check
		if currentStateName == "" {
			output.Error = "no next state and not end"
			return output, nil
		}
	}
}

// executeTaskState executes a Task state with retry policy
func executeTaskState(ctx workflow.Context, state *routing.State, execCtx *routing.ExecutionContext, logger log.Logger) (interface{}, string, error) {
	// Parse timeout
	timeout := 5 * time.Minute
	if state.Timeout != "" {
		if parsed, err := time.ParseDuration(state.Timeout); err == nil {
			timeout = parsed
		}
	}

	// Build activity options
	activityOpts := workflow.ActivityOptions{
		StartToCloseTimeout:    timeout,
		ScheduleToCloseTimeout: timeout + 5*time.Minute,
	}

	// Add retry policy if specified
	if state.Retry != nil {
		initialInterval := time.Second
		if state.Retry.InitialInterval != "" {
			if parsed, err := time.ParseDuration(state.Retry.InitialInterval); err == nil {
				initialInterval = parsed
			}
		}
		maxInterval := 30 * time.Second
		if state.Retry.MaxInterval != "" {
			if parsed, err := time.ParseDuration(state.Retry.MaxInterval); err == nil {
				maxInterval = parsed
			}
		}

		activityOpts.RetryPolicy = &temporal.RetryPolicy{
			InitialInterval:    initialInterval,
			BackoffCoefficient: state.Retry.BackoffRate,
			MaximumInterval:    maxInterval,
			MaximumAttempts:    state.Retry.MaxAttempts,
		}
	}

	actCtx := workflow.WithActivityOptions(ctx, activityOpts)

	// Resolve parameters using JSONPath
	resolver := routing.NewJSONPathResolver(execCtx.Input, execCtx.StepResults)
	resolvedParams, err := resolver.ResolvePaths(state.Parameters)
	if err != nil {
		return nil, "", fmt.Errorf("failed to resolve parameters: %w", err)
	}

	logger.Info("executing activity", "activity", state.Resource, "params", resolvedParams)

	// Execute activity
	var result interface{}
	err = workflow.ExecuteActivity(actCtx, state.Resource, resolvedParams).Get(ctx, &result)
	
	if err != nil {
		logger.Error("activity failed", "activity", state.Resource, "error", err)
		
		// Check for catch clauses
		for _, catch := range state.Catch {
			if matchesError(err, catch.ErrorEquals) {
				logger.Info("error caught", "handler", catch.Next)
				return nil, catch.Next, err
			}
		}
		return nil, "", err
	}

	logger.Info("activity completed", "activity", state.Resource)
	return result, "", nil
}

// matchesError checks if error matches any of the error types
func matchesError(err error, errorEquals []string) bool {
	errStr := err.Error()
	for _, errType := range errorEquals {
		switch errType {
		case "ActivityError":
			return true // Match all activity errors
		case "TimeoutError":
			if temporal.IsTimeoutError(err) {
				return true
			}
		case "ApplicationError":
			if temporal.IsApplicationError(err) {
				return true
			}
		default:
			// Match by error string contains
			if contains(errStr, errType) {
				return true
			}
		}
	}
	return false
}

func contains(s, substr string) bool {
	return len(s) >= len(substr) && (s == substr || len(s) > 0 && containsHelper(s, substr))
}

func containsHelper(s, substr string) bool {
	for i := 0; i <= len(s)-len(substr); i++ {
		if s[i:i+len(substr)] == substr {
			return true
		}
	}
	return false
}