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"`
}

// stateResult holds the result of executing a state
type stateResult struct {
	result    interface{}
	nextState string
	isDone    bool
	err       error
}

// stateMachine manages workflow state execution
type stateMachine struct {
	spec       *routing.WorkflowSpec
	stateIndex map[string]*routing.State
	execCtx    *routing.ExecutionContext
	output     *RoutingWorkflowOutput
	current    string
}

// newStateMachine creates a state machine from spec
func newStateMachine(spec *routing.WorkflowSpec) *stateMachine {
	idx := make(map[string]*routing.State)
	for i := range spec.States {
		idx[spec.States[i].Name] = &spec.States[i]
	}
	return &stateMachine{
		spec:       spec,
		stateIndex: idx,
		execCtx: &routing.ExecutionContext{
			Input:       spec.Input,
			StepResults: make(map[string]interface{}),
		},
		output: &RoutingWorkflowOutput{
			Status:      "FAILED",
			StepResults: make(map[string]interface{}),
		},
		current: spec.States[0].Name,
	}
}

func (m *stateMachine) currentState() *routing.State {
	return m.stateIndex[m.current]
}

func (m *stateMachine) recordResult(name string, result interface{}) {
	wrapped := map[string]interface{}{"output": result}
	m.execCtx.StepResults[name] = wrapped
	m.output.StepResults[name] = result
}

func (m *stateMachine) complete(result interface{}) RoutingWorkflowOutput {
	m.output.Status = "COMPLETED"
	m.output.FinalOutput = result
	return *m.output
}

func (m *stateMachine) fail(errMsg string) RoutingWorkflowOutput {
	m.output.Error = errMsg
	return *m.output
}

// executeTask runs a Task state
func executeTask(ctx workflow.Context, state *routing.State, execCtx *routing.ExecutionContext, logger log.Logger) stateResult {
	result, nextState, err := executeTaskState(ctx, state, execCtx, logger)
	if err != nil {
		if nextState != "" {
			return stateResult{nextState: nextState} // Caught, continue to error handler
		}
		return stateResult{err: err}
	}
	return stateResult{result: result, nextState: state.Next, isDone: state.End}
}

// executePass runs a Pass state
func executePass(state *routing.State) stateResult {
	return stateResult{result: state.Result, nextState: state.Next, isDone: state.End}
}

// executeFail runs a Fail state
func executeFail(state *routing.State) stateResult {
	return stateResult{err: fmt.Errorf("%s: %s", state.Error, state.Cause)}
}

// RoutingWorkflow executes any WorkflowSpec generated by llm-router
func RoutingWorkflow(ctx workflow.Context, input RoutingWorkflowInput) (RoutingWorkflowOutput, error) {
	logger := workflow.GetLogger(ctx)

	if input.Spec == nil || len(input.Spec.States) == 0 {
		return RoutingWorkflowOutput{Status: "FAILED", Error: "empty workflow spec", StepResults: map[string]interface{}{}}, nil
	}

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

	m := newStateMachine(input.Spec)

	for {
		state := m.currentState()
		if state == nil {
			return m.fail(fmt.Sprintf("state not found: %s", m.current)), nil
		}

		logger.Info("executing state", "state", m.current, "type", state.Type)

		var res stateResult
		switch state.Type {
		case routing.StateTypeTask:
			res = executeTask(ctx, state, m.execCtx, logger)
		case routing.StateTypePass:
			res = executePass(state)
		case routing.StateTypeFail:
			res = executeFail(state)
		default:
			return m.fail(fmt.Sprintf("unknown state type: %s", state.Type)), nil
		}

		if res.err != nil {
			logger.Error("state failed", "state", m.current, "error", res.err)
			return m.fail(fmt.Sprintf("state %s failed: %v", m.current, res.err)), nil
		}

		if res.result != nil {
			m.recordResult(state.Name, res.result)
		}

		if res.isDone {
			logger.Info("RoutingWorkflow completed", "name", input.Spec.Name)
			return m.complete(res.result), nil
		}

		if res.nextState == "" {
			return m.fail("no next state and not end"), nil
		}
		m.current = res.nextState
	}
}

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