script-build-and-framework-tmp/agent/agent/orchestration/state_machine.py

"""Legal task transitions for explicit validation."""

from __future__ import annotations

from typing import Dict, FrozenSet

from .models import TaskStatus


class InvalidTaskTransition(ValueError):
    pass


LEGAL_TRANSITIONS: Dict[TaskStatus, FrozenSet[TaskStatus]] = {
    TaskStatus.PENDING: frozenset({TaskStatus.RUNNING, TaskStatus.WAITING_CHILDREN, TaskStatus.BLOCKED, TaskStatus.CANCELLED, TaskStatus.SUPERSEDED}),
    TaskStatus.RUNNING: frozenset({TaskStatus.AWAITING_VALIDATION, TaskStatus.NEEDS_REPLAN, TaskStatus.BLOCKED, TaskStatus.CANCELLED, TaskStatus.SUPERSEDED}),
    TaskStatus.AWAITING_VALIDATION: frozenset({TaskStatus.VALIDATING, TaskStatus.NEEDS_REPLAN, TaskStatus.BLOCKED, TaskStatus.CANCELLED, TaskStatus.SUPERSEDED}),
    TaskStatus.VALIDATING: frozenset({TaskStatus.AWAITING_DECISION, TaskStatus.NEEDS_REPLAN, TaskStatus.BLOCKED, TaskStatus.CANCELLED, TaskStatus.SUPERSEDED}),
    TaskStatus.AWAITING_DECISION: frozenset({TaskStatus.COMPLETED, TaskStatus.PENDING, TaskStatus.WAITING_CHILDREN, TaskStatus.BLOCKED, TaskStatus.CANCELLED, TaskStatus.SUPERSEDED}),
    TaskStatus.NEEDS_REPLAN: frozenset({TaskStatus.PENDING, TaskStatus.RUNNING, TaskStatus.VALIDATING, TaskStatus.WAITING_CHILDREN, TaskStatus.BLOCKED, TaskStatus.CANCELLED, TaskStatus.SUPERSEDED}),
    TaskStatus.WAITING_CHILDREN: frozenset({TaskStatus.NEEDS_REPLAN, TaskStatus.BLOCKED, TaskStatus.CANCELLED, TaskStatus.SUPERSEDED}),
    TaskStatus.BLOCKED: frozenset({TaskStatus.NEEDS_REPLAN, TaskStatus.CANCELLED, TaskStatus.SUPERSEDED}),
    TaskStatus.COMPLETED: frozenset(),
    TaskStatus.CANCELLED: frozenset(),
    TaskStatus.SUPERSEDED: frozenset(),
}


def transition(current: TaskStatus, target: TaskStatus) -> TaskStatus:
    if target not in LEGAL_TRANSITIONS[current]:
        raise InvalidTaskTransition(f"Illegal task transition: {current.value} -> {target.value}")
    return target


__all__ = ["InvalidTaskTransition", "LEGAL_TRANSITIONS", "transition"]