Agent Runtime State Machines¶

def start(self) -> None:
    """Transition from INITIALIZING to RUNNABLE.

    Call this after the agent has completed initialization and is
    ready to be scheduled for work.

    Precondition: state == INITIALIZING
    Raises: StateTransitionError if not in INITIALIZING state
    """

def schedule(self, task_id: str) -> None:
    """Transition from RUNNABLE to SCHEDULED.

    Called when the scheduler assigns a task to this agent.

    Args:
        task_id: The identifier of the task being assigned.

    Precondition: state == RUNNABLE
    Raises: StateTransitionError if not in RUNNABLE state
    """

def run(self) -> None:
    """Transition to RUNNING state.

    Valid from SCHEDULED, WAITING, WAITING_RESOURCES, or RESUMED states.

    Precondition: state in {SCHEDULED, WAITING, WAITING_RESOURCES, RESUMED}
    Raises: StateTransitionError if transition is not valid from current state
    """

def wait(self) -> None:
    """Transition from RUNNING to WAITING.

    Call this when the agent needs to wait for external input
    (e.g., waiting for a response from another agent or service).

    Precondition: state == RUNNING
    Raises: StateTransitionError if not in RUNNING state
    """

def wait_resources(self) -> None:
    """Transition from RUNNING to WAITING_RESOURCES.

    Call this when the agent needs to wait for resources to become
    available (e.g., memory, compute capacity, external service limits).

    Precondition: state == RUNNING
    Raises: StateTransitionError if not in RUNNING state
    """

def suspend(self) -> None:
    """Transition from RUNNING to SUSPENDED.

    Call this when the agent is being preempted and needs to
    suspend execution. The agent should save its state before
    calling this method.

    Precondition: state == RUNNING
    Raises: StateTransitionError if not in RUNNING state
    """

def resume(self) -> None:
    """Transition from SUSPENDED to RESUMED.

    Call this when resuming a suspended agent. After resuming,
    call run() to continue execution.

    Precondition: state == SUSPENDED
    Raises: StateTransitionError if not in SUSPENDED state
    """

def complete(self) -> None:
    """Transition from RUNNING to COMPLETED.

    Call this when the agent has successfully finished its task.
    COMPLETED is a terminal state.

    Precondition: state == RUNNING
    Raises: StateTransitionError if not in RUNNING state
    """

def fail(self, reason: str) -> None:
    """Transition to FAILED state.

    This transition is allowed from RUNNING, SHUTTING_DOWN, or
    RECOVERING (on recovery failure).

    Args:
        reason: Description of why the agent failed.

    Precondition: FAILED is a valid target from current state
    Raises: StateTransitionError if transition to FAILED is not valid
    """

def shutdown(self) -> None:
    """Transition from RUNNING to SHUTTING_DOWN.

    Call this to initiate a graceful shutdown. The agent should
    complete any critical operations and then transition to
    either COMPLETED or FAILED.

    Precondition: state == RUNNING
    Raises: StateTransitionError if not in RUNNING state
    """

def recover(self) -> None:
    """Transition from FAILED to RECOVERING.

    Call this to initiate recovery after a failure. Override
    on_recovery_start to implement recovery logic, then call
    either complete_recovery() or fail_recovery().

    Precondition: state == FAILED
    Raises: StateTransitionError if not in FAILED state
    """

def complete_recovery(self) -> None:
    """Complete recovery successfully, transitioning to RUNNABLE.

    Call this after recovery procedures have completed successfully.
    The agent will return to RUNNABLE state and can be scheduled again.

    Precondition: state == RECOVERING
    Raises: StateTransitionError if not in RECOVERING state
    """

def fail_recovery(self, reason: str) -> None:
    """Recovery failed, transitioning back to FAILED.

    Call this if recovery procedures fail. The agent will return
    to FAILED state.

    Args:
        reason: Description of why recovery failed.

    Precondition: state == RECOVERING
    Raises: StateTransitionError if not in RECOVERING state
    """

import time
import random

class RecoverableAgent(Agent):
    MAX_RETRIES = 5
    BASE_DELAY_S = 1.0

    def on_recovery_start(self) -> None:
        for attempt in range(self.MAX_RETRIES):
            try:
                self._reconnect_to_service()
                self.complete_recovery()
                return
            except ConnectionError:
                # Exponential backoff with jitter
                delay = self.BASE_DELAY_S * (2 ** attempt)
                jitter = random.uniform(0, delay * 0.1)
                time.sleep(delay + jitter)

        self.fail_recovery(f"Failed after {self.MAX_RETRIES} attempts")

def safe_point(self) -> bool:
    """Return True if preemption is requested and caller should yield.

    Call this at safe points in your agent's execution loop to
    check if preemption has been requested.

    Returns:
        True if preemption is requested, False otherwise.
    """

def process_items(self, items: list) -> None:
    for item in items:
        # Check for preemption at each iteration
        if self.safe_point():
            self.suspend()
            return

        self.process_single_item(item)

@contextlib.contextmanager
def non_preemptible(self, *, deadline_ms: Optional[int] = None) -> ContextManager[None]:
    """Context manager for critical sections that should not be preempted.

    While inside this context, safe_point() will return False even if
    preemption has been requested. Note that the scheduler may still
    enforce a hard kill externally after the deadline.

    Args:
        deadline_ms: Optional deadline in milliseconds for the critical
            section. The scheduler may force preemption after this time.

    Yields:
        None
    """

def commit_transaction(self) -> None:
    # This section must complete atomically
    with self.non_preemptible(deadline_ms=5000):
        self.db.begin_transaction()
        self.db.write(self.pending_changes)
        self.db.commit()

class MyAgent(Agent):
    def on_startup(self) -> None:
        """Called during agent startup initialization."""
        self.logger.info("Agent starting up")
        self._init_connections()

    def on_shutdown(self) -> None:
        """Called during agent shutdown."""
        self.logger.info("Agent shutting down")
        self._close_connections()

    def on_scheduled(self, task_id: str) -> None:
        """Called when the agent is scheduled with a task."""
        self.logger.info(f"Assigned task: {task_id}")
        self._prepare_for_task(task_id)

    def on_state_change(self, old_state: int, new_state: int) -> None:
        """Called whenever the agent state changes."""
        self.logger.debug(
            f"State: {AgentState.name(old_state)} -> {AgentState.name(new_state)}"
        )

    def on_preempt_request(self, reason: str) -> None:
        """Called when a preemption is requested."""
        self.logger.warning(f"Preemption requested: {reason}")

    def on_suspend(self) -> None:
        """Called when the agent is suspended."""
        self._save_checkpoint()

    def on_resume(self) -> None:
        """Called when the agent resumes from suspension."""
        self._load_checkpoint()

    def on_recovery_start(self) -> None:
        """Called when recovery begins after a failure."""
        self.logger.info("Starting recovery")

    def on_recovery_complete(self, success: bool) -> None:
        """Called when recovery completes."""
        if success:
            self.logger.info("Recovery successful")
        else:
            self.logger.error("Recovery failed")

from sw4rm.runtime.agent import Agent, AgentState, StateTransitionError

agent = Agent("agent-1", "Worker")
# Agent starts in INITIALIZING

try:
    # Invalid: cannot schedule without first calling start()
    agent.schedule("task-123")
except StateTransitionError as e:
    print(f"Current state: {e.current_state}")  # 0 (INITIALIZING)
    print(f"Target state: {e.target_state}")    # 2 (SCHEDULED)
    print(str(e))  # "Invalid state transition: INITIALIZING -> SCHEDULED"

@property
def state(self) -> int:
    """Return the current agent state (read-only)."""

@property
def state_name(self) -> str:
    """Return the human-readable name of the current state."""

def is_terminal(self) -> bool:
    """Check if the agent is in a terminal state (COMPLETED)."""

def can_transition_to(self, target_state: int) -> bool:
    """Check if a transition to the target state is valid."""

from sw4rm.runtime.agent import Agent, AgentState

class WorkerAgent(Agent):
    def __init__(self, agent_id: str):
        super().__init__(agent_id, "Worker")
        self.work_queue = []

    def on_startup(self) -> None:
        print(f"[{self.agent_id}] Starting up...")

    def on_scheduled(self, task_id: str) -> None:
        print(f"[{self.agent_id}] Scheduled for task: {task_id}")

    def on_suspend(self) -> None:
        print(f"[{self.agent_id}] Suspending, saving state...")

    def on_resume(self) -> None:
        print(f"[{self.agent_id}] Resuming from checkpoint...")

    def process_work(self) -> None:
        """Main work loop with preemption support."""
        for item in self.work_queue:
            if self.safe_point():
                print(f"[{self.agent_id}] Yielding to preemption")
                self.suspend()
                return

            self._process_item(item)

        self.complete()


# Usage
agent = WorkerAgent("worker-1")
assert agent.state == AgentState.INITIALIZING

agent.start()
assert agent.state == AgentState.RUNNABLE

agent.schedule("task-abc")
assert agent.state == AgentState.SCHEDULED

agent.run()
assert agent.state == AgentState.RUNNING

# Agent processes work and completes
agent.complete()
assert agent.state == AgentState.COMPLETED
assert agent.is_terminal()