Workflow State Machines

from enum import Enum

class AgentState(Enum):
    IDLE = "idle"
    PLANNING = "planning"
    OBSERVING = "observing"
    PROPOSING = "proposing"
    AWAITING_APPROVAL = "awaiting_approval"
    EXECUTING = "executing"
    RETRYING = "retrying"
    COMPENSATING = "compensating"
    COMPLETED = "completed"
    FAILED = "failed"
    TIMED_OUT = "timed_out"

class AgentWorkflow:
    def __init__(self, agent_id):
        self.agent_id = agent_id
        self.state = AgentState.IDLE
        self.steps = []
        self.current_step = 0
        self.context = {}

    def transition(self, new_state, reason=""):
        """Transition to a new state with audit logging."""
        old_state = self.state
        self.state = new_state
        log_transition(self.agent_id, old_state, new_state, reason)

    def save(self):
        """Persist state to durable storage."""
        save_workflow_state(self.agent_id, {
            "state": self.state.value,
            "steps": self.steps,
            "current_step": self.current_step,
            "context": self.context
        })

VALID_TRANSITIONS = {
    AgentState.IDLE: [AgentState.PLANNING],
    AgentState.PLANNING: [AgentState.OBSERVING, AgentState.PROPOSING],
    AgentState.OBSERVING: [AgentState.PROPOSING, AgentState.PLANNING],
    AgentState.PROPOSING: [AgentState.AWAITING_APPROVAL, AgentState.EXECUTING],
    AgentState.AWAITING_APPROVAL: [AgentState.EXECUTING, AgentState.IDLE],
    AgentState.EXECUTING: [AgentState.COMPLETED, AgentState.RETRYING, AgentState.FAILED],
    AgentState.RETRYING: [AgentState.EXECUTING, AgentState.COMPENSATING, AgentState.FAILED],
    AgentState.COMPENSATING: [AgentState.IDLE, AgentState.FAILED],
    AgentState.COMPLETED: [AgentState.IDLE],
    AgentState.FAILED: [AgentState.IDLE],
    AgentState.TIMED_OUT: [AgentState.COMPENSATING, AgentState.IDLE],
}

def can_transition(from_state, to_state):
    """Check if a state transition is valid."""
    return to_state in VALID_TRANSITIONS.get(from_state, [])

class RetryManager:
    def __init__(self, max_retries=3, base_delay=5):
        self.max_retries = max_retries
        self.base_delay = base_delay

    def should_retry(self, error, attempt):
        """Determine if we should retry."""
        if attempt >= self.max_retries:
            return False

        # Don't retry on permanent errors
        permanent_errors = ["PermissionError", "ValidationError"]
        if any(e in str(error) for e in permanent_errors):
            return False

        return True

    def get_delay(self, attempt):
        """Get retry delay with exponential backoff."""
        return self.base_delay * (2 ** attempt)

# In the workflow
def execute_with_retry(workflow, step):
    """Execute a step with retry logic."""
    retry_mgr = RetryManager()

    for attempt in range(retry_mgr.max_retries + 1):
        try:
            result = execute_step(step)
            return result
        except Exception as e:
            if not retry_mgr.should_retry(e, attempt):
                workflow.transition(AgentState.FAILED, str(e))
                raise
            delay = retry_mgr.get_delay(attempt)
            workflow.transition(AgentState.RETRYING, f"Attempt {attempt + 1} failed: {e}")
            time.sleep(delay)

class CompensatingAction:
    def __init__(self, action, compensation):
        self.action = action
        self.compensation = compensation

# Define compensations for each action
COMPENSATIONS = {
    "kubectl_scale": CompensatingAction(
        action=lambda **kw: kubectl_scale(**kw),
        compensation=lambda original, **kw: kubectl_scale(
            deployment=original["deployment"],
            replicas=original["original_replicas"],
            namespace=original["namespace"]
        )
    ),
    "config_patch": CompensatingAction(
        action=lambda **kw: apply_config_patch(**kw),
        compensation=lambda original, **kw: apply_config_patch(
            resource=original["resource"],
            patch=original["original_config"],
            namespace=original["namespace"]
        )
    ),
}

def execute_with_compensation(steps):
    """Execute steps with automatic compensation on failure."""
    completed = []
    for step in steps:
        try:
            result = step["action"](**step["arguments"])
            completed.append({"step": step, "result": result})
        except Exception as e:
            # Compensate completed steps in reverse order
            for completed_step in reversed(completed):
                comp = COMPENSATIONS[completed_step["step"]["name"]]
                comp.compensation(completed_step["result"])
            raise